TraceQL: Proposal to address mixed-type attribute querying limitations

tempodb/encoding/vparquet4/block_traceql.go

@@ -2047,7 +2047,7 @@ func createSpanIterator(makeIter makeIterFn, innerIterators []parquetquery.Itera
 		}
 	}
 
-	attrIter, err := createAttributeIterator(makeIter, genericConditions, DefinitionLevelResourceSpansILSSpanAttrs,
+	attrIter, err := createAttributeIterator2(makeIter, genericConditions, DefinitionLevelResourceSpansILSSpanAttrs,
 		columnPathSpanAttrKey, columnPathSpanAttrString, columnPathSpanAttrInt, columnPathSpanAttrDouble, columnPathSpanAttrBool, allConditions, selectAll)
 	if err != nil {
 		return nil, fmt.Errorf("creating span attribute iterator: %w", err)
@@ -2697,6 +2697,139 @@ func createAttributeIterator(makeIter makeIterFn, conditions []traceql.Condition
 	return nil, nil
 }
 
+func createAttributeIterator2(makeIter makeIterFn, conditions []traceql.Condition,
+	definitionLevel int,
+	keyPath, strPath, intPath, floatPath, boolPath string,
+	allConditions bool, selectAll bool,
+) (parquetquery.Iterator, error) {
+	if selectAll {
+		return parquetquery.NewLeftJoinIterator(definitionLevel,
+			[]parquetquery.Iterator{makeIter(keyPath, nil, "key")},
+			[]parquetquery.Iterator{
+				makeIter(strPath, nil, "string"),
+				makeIter(intPath, nil, "int"),
+				makeIter(floatPath, nil, "float"),
+				makeIter(boolPath, nil, "bool"),
+			},
+			&attributeCollector{},
+			parquetquery.WithPool(pqAttrPool))
+	}
+
+	var (
+		attrKeys        = []string{}
+		attrStringPreds = []parquetquery.Predicate{}
+		attrIntPreds    = []parquetquery.Predicate{}
+		attrFltPreds    = []parquetquery.Predicate{}
+		boolPreds       = []parquetquery.Predicate{}
+	)
+	for _, cond := range conditions {
+
+		attrKeys = append(attrKeys, cond.Attribute.Name)
+
+		if cond.Op == traceql.OpNone {
+			// This means we have to scan all values, we don't know what type
+			// to expect
+			attrStringPreds = append(attrStringPreds, nil)
+			attrIntPreds = append(attrIntPreds, nil)
+			attrFltPreds = append(attrFltPreds, nil)
+			boolPreds = append(boolPreds, nil)
+			continue
+		}
+
+		switch cond.Operands[0].Type {
+
+		case traceql.TypeString:
+			pred, err := createStringPredicate(cond.Op, cond.Operands)
+			if err != nil {
+				return nil, fmt.Errorf("creating attribute predicate: %w", err)
+			}
+			attrStringPreds = append(attrStringPreds, pred)
+
+		case traceql.TypeInt:
+			{
+				pred, err := createIntPredicate(cond.Op, cond.Operands)
+				if err != nil {
+					return nil, fmt.Errorf("creating attribute predicate: %w", err)
+				}
+				attrIntPreds = append(attrIntPreds, pred)
+			}
+
+			{
+				if i, ok := cond.Operands[0].Int(); ok {
+					operands := traceql.Operands{traceql.NewStaticFloat(float64(i))}
+					pred, err := createFloatPredicate(cond.Op, operands)
+					if err != nil {
+						return nil, fmt.Errorf("creating attribute predicate: %w", err)
+					}
+					attrFltPreds = append(attrFltPreds, pred)
+				}
+			}
+
+		case traceql.TypeFloat:
+			{
+				operands := traceql.Operands{traceql.NewStaticInt(int(cond.Operands[0].Float()))}
+				pred, err := createIntPredicate(cond.Op, operands)
+				if err != nil {
+					return nil, fmt.Errorf("creating attribute predicate: %w", err)
+				}
+				attrIntPreds = append(attrIntPreds, pred)
+			}
+
+			{
+				pred, err := createFloatPredicate(cond.Op, cond.Operands)
+				if err != nil {
+					return nil, fmt.Errorf("creating attribute predicate: %w", err)
+				}
+				attrFltPreds = append(attrFltPreds, pred)
+			}
+
+		case traceql.TypeBoolean:
+			pred, err := createBoolPredicate(cond.Op, cond.Operands)
+			if err != nil {
+				return nil, fmt.Errorf("creating attribute predicate: %w", err)
+			}
+			boolPreds = append(boolPreds, pred)
+		}
+	}
+
+	var valueIters []parquetquery.Iterator
+	if len(attrStringPreds) > 0 {
+		valueIters = append(valueIters, makeIter(strPath, orIfNeeded(attrStringPreds), "string"))
+	}
+	if len(attrIntPreds) > 0 {
+		valueIters = append(valueIters, makeIter(intPath, orIfNeeded(attrIntPreds), "int"))
+	}
+	if len(attrFltPreds) > 0 {
+		valueIters = append(valueIters, makeIter(floatPath, orIfNeeded(attrFltPreds), "float"))
+	}
+	if len(boolPreds) > 0 {
+		valueIters = append(valueIters, makeIter(boolPath, orIfNeeded(boolPreds), "bool"))
+	}
+
+	if len(valueIters) > 0 {
+		// LeftJoin means only look at rows where the key is what we want.
+		// Bring in any of the typed values as needed.
+
+		// if all conditions must be true we can use a simple join iterator to test the values one column at a time.
+		// len(valueIters) must be 1 to handle queries like `{ span.foo = "x" && span.bar > 1}`
+		if allConditions && len(valueIters) == 1 {
+			iters := append([]parquetquery.Iterator{makeIter(keyPath, parquetquery.NewStringInPredicate(attrKeys), "key")}, valueIters...)
+			return parquetquery.NewJoinIterator(definitionLevel,
+				iters,
+				&attributeCollector{},
+				parquetquery.WithPool(pqAttrPool)), nil
+		}
+
+		return parquetquery.NewLeftJoinIterator(definitionLevel,
+			[]parquetquery.Iterator{makeIter(keyPath, parquetquery.NewStringInPredicate(attrKeys), "key")},
+			valueIters,
+			&attributeCollector{},
+			parquetquery.WithPool(pqAttrPool))
+	}
+
+	return nil, nil
+}
+
 // This turns groups of span values into Span objects
 type spanCollector struct {
 	minAttributes int