[TraceQL] Improve performance of select() queries (#4438)

* Tighten up allconditions disable so that select expressions don't disable it * add benchmark * changelog * More test cases for scalar expressions
grafana · Dec 13, 2024 · 90b2ac5 · 90b2ac5
1 parent 4e8a96f
commit 90b2ac5
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -57,6 +57,7 @@ querier:
 * [ENHANCEMENT] Reduce allocs related to marshalling dedicated columns repeatedly in the query frontend. [#4007](https://github.com/grafana/tempo/pull/4007) (@joe-elliott)
 * [ENHANCEMENT] Improve performance of TraceQL queries [#4114](https://github.com/grafana/tempo/pull/4114) (@mdisibio)
 * [ENHANCEMENT] Improve performance of TraceQL queries [#4163](https://github.com/grafana/tempo/pull/4163) (@mdisibio) 
+* [ENHANCEMENT] Improve performance of some TraceQL queries using select() operation [#4438](https://github.com/grafana/tempo/pull/4438) (@mdisibio)
 * [ENHANCEMENT] Reduce memory usage of classic histograms in the span-metrics and service-graphs processors [#4232](https://github.com/grafana/tempo/pull/4232) (@mdisibio)
 * [ENHANCEMENT] Implement simple Fetch by key for cache items  [#4032](https://github.com/grafana/tempo/pull/4032) (@javiermolinar)
 * [ENHANCEMENT] Replace Grafana Agent example by Grafana Alloy[#4030](https://github.com/grafana/tempo/pull/4030) (@javiermolinar)

diff --git a/pkg/traceql/ast.go b/pkg/traceql/ast.go
@@ -159,10 +159,6 @@ func (p Pipeline) extractConditions(req *FetchSpansRequest) {
 	for _, element := range p.Elements {
 		element.extractConditions(req)
 	}
-	// TODO this needs to be fine-tuned a bit, e.g. { .foo = "bar" } | by(.namespace), AllConditions can still be true
-	if len(p.Elements) > 1 {
-		req.AllConditions = false
-	}
 }
 
 func (p Pipeline) evaluate(input []*Spanset) (result []*Spanset, err error) {
@@ -197,6 +193,7 @@ func newGroupOperation(e FieldExpression) GroupOperation {
 
 func (o GroupOperation) extractConditions(request *FetchSpansRequest) {
 	o.Expression.extractConditions(request)
+	request.AllConditions = false
 }
 
 type CoalesceOperation struct{}
@@ -265,7 +262,6 @@ func (o ScalarOperation) impliedType() StaticType {
 func (o ScalarOperation) extractConditions(request *FetchSpansRequest) {
 	o.LHS.extractConditions(request)
 	o.RHS.extractConditions(request)
-	request.AllConditions = false
 }
 
 type Aggregate struct {
@@ -292,6 +288,7 @@ func (a Aggregate) impliedType() StaticType {
 }
 
 func (a Aggregate) extractConditions(request *FetchSpansRequest) {
+	request.AllConditions = false
 	if a.e != nil {
 		a.e.extractConditions(request)
 	}
@@ -421,7 +418,6 @@ func (ScalarFilter) __spansetExpression() {}
 func (f ScalarFilter) extractConditions(request *FetchSpansRequest) {
 	f.lhs.extractConditions(request)
 	f.rhs.extractConditions(request)
-	request.AllConditions = false
 }
 
 // **********************

diff --git a/pkg/traceql/ast_conditions_test.go b/pkg/traceql/ast_conditions_test.go
@@ -147,6 +147,21 @@ func TestScalarFilter_extractConditions(t *testing.T) {
 			},
 			allConditions: false,
 		},
+		{
+			query: `({ span.http.status_code = 200 } | count()) > ({ span.http.status_code = 500 } | count())`,
+			conditions: []Condition{
+				newCondition(NewScopedAttribute(AttributeScopeSpan, false, "http.status_code"), OpEqual, NewStaticInt(200)),
+				newCondition(NewScopedAttribute(AttributeScopeSpan, false, "http.status_code"), OpEqual, NewStaticInt(500)),
+			},
+			allConditions: false,
+		},
+		{
+			query: `{ .foo = "a" } | 3 > 2`,
+			conditions: []Condition{
+				newCondition(NewAttribute("foo"), OpEqual, NewStaticString("a")),
+			},
+			allConditions: true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.query, func(t *testing.T) {
@@ -230,7 +245,7 @@ func TestSelect_extractConditions(t *testing.T) {
 			secondPassConditions: []Condition{
 				newCondition(NewScopedAttribute(AttributeScopeResource, false, "service.name"), OpNone),
 			},
-			allConditions: false,
+			allConditions: true,
 		},
 		{
 			query: `{ } | select(.name,name)`,
@@ -241,7 +256,7 @@ func TestSelect_extractConditions(t *testing.T) {
 				newCondition(NewAttribute("name"), OpNone),
 				newCondition(NewIntrinsic(IntrinsicName), OpNone),
 			},
-			allConditions: false,
+			allConditions: true,
 		},
 	}
 	for _, tt := range tests {

diff --git a/pkg/traceql/ast_test.go b/pkg/traceql/ast_test.go
@@ -610,7 +610,7 @@ func TestPipelineExtractConditions(t *testing.T) {
 					newCondition(NewAttribute("foo1"), OpEqual, NewStaticString("a")),
 					newCondition(NewAttribute("foo2"), OpEqual, NewStaticString("b")),
 				},
-				AllConditions: false,
+				AllConditions: true,
 			},
 		},
 		{

diff --git a/tempodb/encoding/vparquet3/block_traceql_test.go b/tempodb/encoding/vparquet3/block_traceql_test.go
@@ -592,6 +592,7 @@ func BenchmarkBackendBlockTraceQL(b *testing.B) {
 		{"||", "{ resource.service.name = `loki-querier` } || { resource.service.name = `loki-gateway` }"},
 		{"mixed", `{resource.namespace!="" && resource.service.name="cortex-gateway" && duration>50ms && resource.cluster=~"prod.*"}`},
 		{"complex", `{resource.cluster=~"prod.*" && resource.namespace = "tempo-prod" && resource.container="query-frontend" && name = "HTTP GET - tempo_api_v2_search_tags" && span.http.status_code = 200 && duration > 1s}`},
+		{"select", `{resource.cluster=~"prod.*" && resource.namespace = "tempo-prod"} | select(resource.container)`},
 	}
 
 	ctx := context.TODO()

diff --git a/tempodb/encoding/vparquet4/block_traceql_test.go b/tempodb/encoding/vparquet4/block_traceql_test.go
@@ -943,6 +943,7 @@ func BenchmarkBackendBlockTraceQL(b *testing.B) {
 		{"||", "{ resource.service.name = `loki-querier` } || { resource.service.name = `loki-gateway` }"},
 		{"mixed", `{resource.namespace!="" && resource.service.name="cortex-gateway" && duration>50ms && resource.cluster=~"prod.*"}`},
 		{"complex", `{resource.cluster=~"prod.*" && resource.namespace = "tempo-prod" && resource.container="query-frontend" && name = "HTTP GET - tempo_api_v2_search_tags" && span.http.status_code = 200 && duration > 1s}`},
+		{"select", `{resource.cluster=~"prod.*" && resource.namespace = "tempo-prod"} | select(resource.container)`},
 	}
 
 	ctx := context.TODO()