TraceQL Perf: Reduce metadata retrieved by implementing two pass iteration.

pkg/parquetquery/iters.go

@@ -9,6 +9,7 @@ import (
 	"sync"
 	"sync/atomic"
 
+	"github.com/grafana/tempo/pkg/util"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pkg/errors"
 	pq "github.com/segmentio/parquet-go"
@@ -148,6 +149,15 @@ func (r *IteratorResult) AppendOtherValue(k string, v interface{}) {
 	}{k, v})
 }
 
+func (r *IteratorResult) OtherValueFromKey(k string) interface{} {
+	for _, e := range r.OtherEntries {
+		if e.Key == k {
+			return e.Value
+		}
+	}
+	return nil
+}
+
 // ToMap converts the unstructured list of data into a map containing an entry
 // for each column, and the lists of values.  The order of columns is
 // not preseved, but the order of values within each column is.
@@ -184,6 +194,8 @@ func (r *IteratorResult) Columns(buffer [][]pq.Value, names ...string) [][]pq.Va
 
 // iterator - Every iterator follows this interface and can be composed.
 type Iterator interface {
+	fmt.Stringer
+
 	// Next returns nil when done
 	Next() (*IteratorResult, error)
 
@@ -288,6 +300,10 @@ func NewColumnIterator(ctx context.Context, rgs []pq.RowGroup, column int, colum
 	return c
 }
 
+func (c *ColumnIterator) String() string {
+	return fmt.Sprintf("ColumnIterator: %s \n\t%s", c.colName, util.TabOut(c.filter))
+}
+
 func (c *ColumnIterator) iterate(ctx context.Context, readSize int) {
 	defer close(c.ch)
 
@@ -577,6 +593,14 @@ func NewJoinIterator(definitionLevel int, iters []Iterator, pred GroupPredicate)
 	return &j
 }
 
+func (j *JoinIterator) String() string {
+	var iters string
+	for _, iter := range j.iters {
+		iters += "\n\t" + util.TabOut(iter)
+	}
+	return fmt.Sprintf("JoinIterator: %d\t%s\n%s)", j.definitionLevel, iters, j.pred)
+}
+
 func (j *JoinIterator) Next() (*IteratorResult, error) {
 	// Here is the algorithm for joins:  On each pass of the iterators
 	// we remember which ones are pointing at the earliest rows. If all
@@ -739,6 +763,18 @@ func NewLeftJoinIterator(definitionLevel int, required, optional []Iterator, pre
 	return &j
 }
 
+func (j *LeftJoinIterator) String() string {
+	srequired := "required: "
+	for _, r := range j.required {
+		srequired += "\n\t\t" + util.TabOut(r)
+	}
+	soptional := "optional: "
+	for _, o := range j.optional {
+		soptional += "\n\t\t" + util.TabOut(o)
+	}
+	return fmt.Sprintf("LeftJoinIterator: %d\n\t%s\n\t%s\n\t%s", j.definitionLevel, srequired, soptional, j.pred)
+}
+
 func (j *LeftJoinIterator) Next() (*IteratorResult, error) {
 
 	// Here is the algorithm for joins:  On each pass of the iterators
@@ -926,6 +962,14 @@ func NewUnionIterator(definitionLevel int, iters []Iterator, pred GroupPredicate
 	return &j
 }
 
+func (u *UnionIterator) String() string {
+	var iters string
+	for _, iter := range u.iters {
+		iters += iter.String() + ", "
+	}
+	return fmt.Sprintf("UnionIterator(%s)", iters)
+}
+
 func (u *UnionIterator) Next() (*IteratorResult, error) {
 	// Here is the algorithm for unions:  On each pass of the iterators
 	// we remember which ones are pointing at the earliest same row. The
@@ -1040,6 +1084,8 @@ func (u *UnionIterator) Close() {
 }
 
 type GroupPredicate interface {
+	fmt.Stringer
+
 	KeepGroup(*IteratorResult) bool
 }
 
@@ -1068,6 +1114,18 @@ func NewKeyValueGroupPredicate(keys, values []string) *KeyValueGroupPredicate {
 	return p
 }
 
+func (a *KeyValueGroupPredicate) String() string {
+	var skeys []string
+	var svals []string
+	for _, k := range a.keys {
+		skeys = append(skeys, string(k))
+	}
+	for _, v := range a.vals {
+		svals = append(svals, string(v))
+	}
+	return fmt.Sprintf("KeyValueGroupPredicate{%v, %v}", skeys, svals)
+}
+
 // KeepGroup checks if the given group contains all of the requested
 // key/value pairs.
 func (a *KeyValueGroupPredicate) KeepGroup(group *IteratorResult) bool {

pkg/parquetquery/predicate_test.go

@@ -27,6 +27,7 @@ func newAlwaysFalsePredicate() *mockPredicate {
 	return &mockPredicate{ret: false}
 }
 
+func (p *mockPredicate) String() string                           { return "mockPredicate{}" }
 func (p *mockPredicate) KeepValue(parquet.Value) bool             { p.valCalled = true; return p.ret }
 func (p *mockPredicate) KeepPage(parquet.Page) bool               { p.pageCalled = true; return p.ret }
 func (p *mockPredicate) KeepColumnChunk(parquet.ColumnChunk) bool { p.chunkCalled = true; return p.ret }

pkg/parquetquery/predicates.go

@@ -2,6 +2,7 @@ package parquetquery
 
 import (
 	"bytes"
+	"fmt"
 	"regexp"
 	"strings"
 
@@ -11,6 +12,8 @@ import (
 // Predicate is a pushdown predicate that can be applied at
 // the chunk, page, and value levels.
 type Predicate interface {
+	fmt.Stringer
+
 	KeepColumnChunk(cc pq.ColumnChunk) bool
 	KeepPage(page pq.Page) bool
 	KeepValue(pq.Value) bool
@@ -36,6 +39,14 @@ func NewStringInPredicate(ss []string) Predicate {
 	return p
 }
 
+func (p *StringInPredicate) String() string {
+	var strings string
+	for _, s := range p.ss {
+		strings += fmt.Sprintf("%s, ", string(s))
+	}
+	return fmt.Sprintf("StringInPredicate{%s}", strings)
+}
+
 func (p *StringInPredicate) KeepColumnChunk(cc pq.ColumnChunk) bool {
 	p.helper.setNewRowGroup()
 
@@ -97,6 +108,14 @@ func NewRegexInPredicate(regs []string) (*RegexInPredicate, error) {
 	return p, nil
 }
 
+func (p *RegexInPredicate) String() string {
+	var strings string
+	for _, s := range p.regs {
+		strings += fmt.Sprintf("%s, ", s.String())
+	}
+	return fmt.Sprintf("RegexInPredicate{%s}", strings)
+}
+
 func (p *RegexInPredicate) keep(v *pq.Value) bool {
 	if v.IsNull() {
 		// Null
@@ -154,6 +173,10 @@ func NewSubstringPredicate(substring string) *SubstringPredicate {
 	}
 }
 
+func (p *SubstringPredicate) String() string {
+	return fmt.Sprintf("SubstringPredicate{%s}", p.substring)
+}
+
 func (p *SubstringPredicate) KeepColumnChunk(cc pq.ColumnChunk) bool {
 	p.helper.setNewRowGroup()
 
@@ -192,6 +215,10 @@ func NewIntBetweenPredicate(min, max int64) *IntBetweenPredicate {
 	return &IntBetweenPredicate{min, max}
 }
 
+func (p *IntBetweenPredicate) String() string {
+	return fmt.Sprintf("IntBetweenPredicate{%d,%d}", p.min, p.max)
+}
+
 func (p *IntBetweenPredicate) KeepColumnChunk(c pq.ColumnChunk) bool {
 
 	if ci := c.ColumnIndex(); ci != nil {
@@ -239,6 +266,10 @@ func NewGenericPredicate[T any](fn func(T) bool, rangeFn func(T, T) bool, extrac
 	return &GenericPredicate[T]{Fn: fn, RangeFn: rangeFn, Extract: extract}
 }
 
+func (p *GenericPredicate[T]) String() string {
+	return "GenericPredicate{}"
+}
+
 func (p *GenericPredicate[T]) KeepColumnChunk(c pq.ColumnChunk) bool {
 	p.helper.setNewRowGroup()
 
@@ -307,6 +338,10 @@ func NewFloatBetweenPredicate(min, max float64) *FloatBetweenPredicate {
 	return &FloatBetweenPredicate{min, max}
 }
 
+func (p *FloatBetweenPredicate) String() string {
+	return fmt.Sprintf("FloatBetweenPredicate{%f,%f}", p.min, p.max)
+}
+
 func (p *FloatBetweenPredicate) KeepColumnChunk(c pq.ColumnChunk) bool {
 
 	if ci := c.ColumnIndex(); ci != nil {
@@ -347,6 +382,14 @@ func NewOrPredicate(preds ...Predicate) *OrPredicate {
 	}
 }
 
+func (p *OrPredicate) String() string {
+	var preds string
+	for _, pred := range p.preds {
+		preds += pred.String() + ","
+	}
+	return fmt.Sprintf("OrPredicate{%s}", p.preds)
+}
+
 func (p *OrPredicate) KeepColumnChunk(c pq.ColumnChunk) bool {
 	ret := false
 	for _, p := range p.preds {
@@ -403,6 +446,13 @@ type InstrumentedPredicate struct {
 
 var _ Predicate = (*InstrumentedPredicate)(nil)
 
+func (p *InstrumentedPredicate) String() string {
+	if p.pred == nil {
+		return fmt.Sprintf("InstrumentedPredicate{%d, nil}", p.InspectedValues)
+	}
+	return fmt.Sprintf("InstrumentedPredicate{%d, %s}", p.InspectedValues, p.pred)
+}
+
 func (p *InstrumentedPredicate) KeepColumnChunk(c pq.ColumnChunk) bool {
 	p.InspectedColumnChunks++
 
@@ -489,6 +539,10 @@ func NewSkipNilsPredicate() *SkipNilsPredicate {
 	return &SkipNilsPredicate{}
 }
 
+func (p *SkipNilsPredicate) String() string {
+	return "SkipNilsPredicate{}"
+}
+
 func (p *SkipNilsPredicate) KeepColumnChunk(c pq.ColumnChunk) bool {
 	return true
 }

pkg/traceql/ast.go

@@ -14,7 +14,7 @@ type Element interface {
 type pipelineElement interface {
 	Element
 	extractConditions(request *FetchSpansRequest)
-	evaluate([]Spanset) ([]Spanset, error)
+	evaluate([]*Spanset) ([]*Spanset, error)
 }
 
 type typedExpression interface {
@@ -85,7 +85,7 @@ func (p Pipeline) extractConditions(req *FetchSpansRequest) {
 	}
 }
 
-func (p Pipeline) evaluate(input []Spanset) (result []Spanset, err error) {
+func (p Pipeline) evaluate(input []*Spanset) (result []*Spanset, err error) {
 	result = input
 
 	for _, element := range p.Elements {
@@ -95,7 +95,7 @@ func (p Pipeline) evaluate(input []Spanset) (result []Spanset, err error) {
 		}
 
 		if len(result) == 0 {
-			return []Spanset{}, nil
+			return []*Spanset{}, nil
 		}
 	}
 
@@ -116,7 +116,7 @@ func (o GroupOperation) extractConditions(request *FetchSpansRequest) {
 	o.Expression.extractConditions(request)
 }
 
-func (GroupOperation) evaluate(ss []Spanset) ([]Spanset, error) {
+func (GroupOperation) evaluate(ss []*Spanset) ([]*Spanset, error) {
 	return ss, nil
 }
 
@@ -130,7 +130,7 @@ func newCoalesceOperation() CoalesceOperation {
 func (o CoalesceOperation) extractConditions(request *FetchSpansRequest) {
 }
 
-func (CoalesceOperation) evaluate(ss []Spanset) ([]Spanset, error) {
+func (CoalesceOperation) evaluate(ss []*Spanset) ([]*Spanset, error) {
 	return ss, nil
 }
 
@@ -257,8 +257,8 @@ func newSpansetFilter(e FieldExpression) SpansetFilter {
 // nolint: revive
 func (SpansetFilter) __spansetExpression() {}
 
-func (f SpansetFilter) evaluate(input []Spanset) ([]Spanset, error) {
-	var output []Spanset
+func (f SpansetFilter) evaluate(input []*Spanset) ([]*Spanset, error) {
+	var output []*Spanset
 
 	for _, ss := range input {
 		if len(ss.Spans) == 0 {
@@ -287,9 +287,9 @@ func (f SpansetFilter) evaluate(input []Spanset) ([]Spanset, error) {
 			continue
 		}
 
-		matchingSpanset := ss
+		matchingSpanset := *ss
 		matchingSpanset.Spans = matchingSpans
-		output = append(output, matchingSpanset)
+		output = append(output, &matchingSpanset)
 	}
 
 	return output, nil

pkg/traceql/ast_execute.go

@@ -10,14 +10,14 @@ import (
 	"github.com/grafana/tempo/pkg/util/log"
 )
 
-func appendSpans(buffer []Span, input []Spanset) []Span {
+func appendSpans(buffer []Span, input []*Spanset) []Span {
 	for _, i := range input {
 		buffer = append(buffer, i.Spans...)
 	}
 	return buffer
 }
 
-func (o SpansetOperation) evaluate(input []Spanset) (output []Spanset, err error) {
+func (o SpansetOperation) evaluate(input []*Spanset) (output []*Spanset, err error) {
 
 	for i := range input {
 		curr := input[i : i+1]
@@ -57,7 +57,7 @@ func (o SpansetOperation) evaluate(input []Spanset) (output []Spanset, err error
 	return output, nil
 }
 
-func (f ScalarFilter) evaluate(input []Spanset) (output []Spanset, err error) {
+func (f ScalarFilter) evaluate(input []*Spanset) (output []*Spanset, err error) {
 
 	// TODO we solve this gap where pipeline elements and scalar binary
 	// operations meet in a generic way. For now we only support well-defined
@@ -105,7 +105,7 @@ func (f SpansetFilter) matches(span Span) (bool, error) {
 	return static.B, nil
 }
 
-func (a Aggregate) evaluate(input []Spanset) (output []Spanset, err error) {
+func (a Aggregate) evaluate(input []*Spanset) (output []*Spanset, err error) {
 
 	for _, ss := range input {
 		switch a.op {
@@ -265,18 +265,19 @@ func (s Static) execute(span Span) (Static, error) {
 }
 
 func (a Attribute) execute(span Span) (Static, error) {
-	static, ok := span.Attributes[a]
+	atts := span.Attributes()
+	static, ok := atts[a]
 	if ok {
 		return static, nil
 	}
 
 	if a.Scope == AttributeScopeNone {
-		for attribute, static := range span.Attributes {
+		for attribute, static := range atts {
 			if a.Name == attribute.Name && attribute.Scope == AttributeScopeSpan {
 				return static, nil
 			}
 		}
-		for attribute, static := range span.Attributes {
+		for attribute, static := range atts {
 			if a.Name == attribute.Name {
 				return static, nil
 			}