TraceQL: Attribute iterators collect matched array values (#3867)

* Test case for queries containing array attributes

* Move makeIterFn to definition of makeIterFunc

* Static.AsAnyValue() supports arrays

* it is fast but does it pass CI

* Add test for isMatchingOperand and isMatchingArrayElement

* clean up and add benchmark blockID and path

* Add tests for arrys AsAnyValue

* refactor arrays support in traceql

* make array support symmetric + tests

* rework attribute collector keep group

* make getelements use a VisitFunc

* cleanup symmetric array support

---------

Co-authored-by: Suraj Nath <9503187+electron0zero@users.noreply.github.com>

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## main / unreleased
 
+* [ENHANCEMENT] TraceQL: Attribute iterators collect matched array values [#3867](https://github.com/grafana/tempo/pull/3867) (@electron0zero, @stoewer)
 * [ENHANCEMENT] Add bytes and spans received to usage stats [#3983](https://github.com/grafana/tempo/pull/3983) (@joe-elliott)
 
 # v2.6.0-rc.0

pkg/traceql/ast.go

@@ -731,6 +731,52 @@ func (s Static) compare(o *Static) int {
 	}
 }
 
+type VisitFunc func(Static) bool // Return false to stop iteration
+
+// GetElements turns arrays into slice of Static elements to iterate over.
+func (s Static) GetElements(fn VisitFunc) error {
+	switch s.Type {
+	case TypeIntArray:
+		ints, _ := s.IntArray()
+		for _, n := range ints {
+			if !fn(NewStaticInt(n)) {
+				break // stop early if the callback returns false
+			}
+		}
+		return nil
+
+	case TypeFloatArray:
+		floats, _ := s.FloatArray()
+		for _, f := range floats {
+			if !fn(NewStaticFloat(f)) {
+				break
+			}
+		}
+		return nil
+
+	case TypeStringArray:
+		strs, _ := s.StringArray()
+		for _, str := range strs {
+			if !fn(NewStaticString(str)) {
+				break
+			}
+		}
+		return nil
+
+	case TypeBooleanArray:
+		bools, _ := s.BooleanArray()
+		for _, b := range bools {
+			if !fn(NewStaticBool(b)) {
+				break
+			}
+		}
+		return nil
+
+	default:
+		return fmt.Errorf("unsupported type")
+	}
+}
+
 func (s Static) Int() (int, bool) {
 	if s.Type != TypeInt {
 		return 0, false

pkg/traceql/ast_execute.go

@@ -381,7 +381,7 @@ func (o *BinaryOperation) execute(span Span) (Static, error) {
 	}
 
 	// if both sides are integers then do integer math, otherwise we can drop to the
-	// catch all below
+	// catch-all below
 	if lhsT == TypeInt && rhsT == TypeInt {
 		lhsN, _ := lhs.Int()
 		rhsN, _ := rhs.Int()
@@ -422,6 +422,39 @@ func (o *BinaryOperation) execute(span Span) (Static, error) {
 		}
 	}
 
+	if lhsT.isMatchingArrayElement(rhsT) {
+		// we only support boolean op in the arrays
+		if !o.Op.isBoolean() {
+			return NewStaticNil(), errors.ErrUnsupported
+		}
+
+		elemOp := &BinaryOperation{Op: o.Op, LHS: lhs, RHS: rhs}
+		arraySide := lhs
+		// to support symmetric operations
+		if rhsT.isArray() {
+			// for regex operations, TraceQL makes an assumption that RHS is the regex, and compiles it.
+			// we can support symmetric array operations by flipping the sides and executing the binary operation.
+			elemOp = &BinaryOperation{Op: getFlippedOp(o.Op), LHS: rhs, RHS: lhs}
+			arraySide = rhs
+		}
+
+		var res Static
+		err := arraySide.GetElements(func(elem Static) bool {
+			elemOp.LHS = elem
+			res, err = elemOp.execute(span)
+			if err != nil {
+				return false // stop iteration early if there's an error
+			}
+			match, ok := res.Bool()
+			return !(ok && match) // stop if a match is found
+		})
+		if err != nil {
+			return NewStaticNil(), err
+		}
+
+		return res, err
+	}
+
 	switch o.Op {
 	case OpAdd:
 		return NewStaticFloat(lhs.Float() + rhs.Float()), nil
@@ -452,6 +485,22 @@ func (o *BinaryOperation) execute(span Span) (Static, error) {
 	}
 }
 
+// getFlippedOp will return the flipped op, used when flipping the LHS and RHS of a BinaryOperation
+func getFlippedOp(op Operator) Operator {
+	switch op {
+	case OpGreater:
+		return OpLess
+	case OpGreaterEqual:
+		return OpLessEqual
+	case OpLess:
+		return OpGreater
+	case OpLessEqual:
+		return OpGreaterEqual
+	default:
+		return op
+	}
+}
+
 // why does this and the above exist?
 func binOp(op Operator, lhs, rhs Static) (bool, error) {
 	lhsT := lhs.impliedType()