TraceQL: Selection

CHANGELOG.md

@@ -26,6 +26,7 @@ To make use of filtering, configure `autocomplete_filtering_enabled`.
 * [FEATURE] Add support for `by()` and `coalesce()` to TraceQL. [#2490](https:/grafana/tempo/pull/2490) 
 * [FEATURE] Add a GRPC streaming endpoint for traceql search [#2366](https:/grafana/tempo/pull/2366) (@joe-elliott)
 * [FEATURE] Add new API to summarize span metrics from generators [#2481](https:/grafana/tempo/pull/2481) (@zalegrala)
+* [FEATURE] Add `select()` to TraceQL [#2494](https:/grafana/tempo/pull/2494) (@joe-elliott)
 * [ENHANCEMENT] Add `scope` parameter to `/api/search/tags` [#2282](https:/grafana/tempo/pull/2282) (@joe-elliott)
  Create new endpoint `/api/v2/search/tags` that returns all tags organized by scope.
 * [ENHANCEMENT] Ability to toggle off latency or count metrics in metrics-generator [#2070](https:/grafana/tempo/pull/2070) (@AlexDHoffer)

docs/sources/tempo/traceql/_index.md

@@ -237,6 +237,14 @@ to compare the ratios of two span attributes:
 ```
 or anything else that comes to mind.
 
+## Selection
+
+TraceQL can select arbitrary fields from spans. This is particularly performant b/c
+the selected fields are not retrieved until all other criteria is met.
+```
+{ status=error } | select(span.http.status_code, span.http.url)
+```
+
 ## Examples
 
 ### Find traces of a specific operation

pkg/traceql/ast.go

@@ -129,6 +129,16 @@ func newCoalesceOperation() CoalesceOperation {
 func (o CoalesceOperation) extractConditions(request *FetchSpansRequest) {
 }
 
+type SelectOperation struct {
+ exprs []FieldExpression
+}
+
+func newSelectOperation(exprs []FieldExpression) SelectOperation {
+ return SelectOperation{
+ exprs: exprs,
+ }
+}
+
 // **********************
 // Scalars
 // **********************

pkg/traceql/ast_conditions.go

@@ -4,6 +4,16 @@ func (f SpansetFilter) extractConditions(request *FetchSpansRequest) {
  f.Expression.extractConditions(request)
 }
 
+// extractConditions on Select puts its conditions into the SecondPassConditions
+func (o SelectOperation) extractConditions(request *FetchSpansRequest) {
+ selectR := &FetchSpansRequest{}
+ for _, expr := range o.exprs {
+ expr.extractConditions(selectR)
+ }
+ // copy any conditions to the normal request's SecondPassConditions
+ request.SecondPassConditions = append(request.SecondPassConditions, selectR.Conditions...)
+}
+
 func (o BinaryOperation) extractConditions(request *FetchSpansRequest) {
  // TODO we can further optimise this by attempting to execute every FieldExpression, if they only contain statics it should resolve
  switch o.LHS.(type) {

pkg/traceql/ast_conditions_test.go

@@ -104,6 +104,7 @@ func TestSpansetFilter_extractConditions(t *testing.T) {
  spansetFilter.extractConditions(req)
 
  assert.Equal(t, tt.conditions, req.Conditions)
+ assert.Nil(t, req.SecondPassConditions)
  assert.Equal(t, tt.allConditions, req.AllConditions, "FetchSpansRequest.AllConditions")
  })
  }
@@ -143,6 +144,52 @@ func TestScalarFilter_extractConditions(t *testing.T) {
  expr.Pipeline.extractConditions(req)
 
  assert.Equal(t, tt.conditions, req.Conditions)
+ assert.Nil(t, req.SecondPassConditions)
+ assert.Equal(t, tt.allConditions, req.AllConditions, "FetchSpansRequest.AllConditions")
+ })
+ }
+}
+
+func TestSelect_extractConditions(t *testing.T) {
+ tests := []struct {
+ query string
+ conditions []Condition
+ secondPassConditions []Condition
+ allConditions bool
+ }{
+ {
+ query: `{ .foo = "a" } | select(resource.service.name)`,
+ conditions: []Condition{
+ newCondition(NewAttribute("foo"), OpEqual, NewStaticString("a")),
+ },
+ secondPassConditions: []Condition{
+ newCondition(NewScopedAttribute(AttributeScopeResource, false, "service.name"), OpNone),
+ },
+ allConditions: false,
+ },
+ {
+ query: `{ } | select(.name,name)`,
+ conditions: []Condition{},
+ secondPassConditions: []Condition{
+ newCondition(NewAttribute("name"), OpNone),
+ newCondition(NewIntrinsic(IntrinsicName), OpNone),
+ },
+ allConditions: false,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.query, func(t *testing.T) {
+ expr, err := Parse(tt.query)
+ require.NoError(t, err)
+
+ req := &FetchSpansRequest{
+ Conditions: []Condition{},
+ AllConditions: true,
+ }
+ expr.Pipeline.extractConditions(req)
+
+ assert.Equal(t, tt.conditions, req.Conditions)
+ assert.Equal(t, tt.secondPassConditions, req.SecondPassConditions)
  assert.Equal(t, tt.allConditions, req.AllConditions, "FetchSpansRequest.AllConditions")
  })
  }

pkg/traceql/ast_execute.go

@@ -107,6 +107,11 @@ func (o SpansetOperation) evaluate(input []*Spanset) (output []*Spanset, err err
  return output, nil
 }
 
+// SelectOperation evaluate is a no-op b/c the fetch layer has already decorated the spans with the requested attributes
+func (o SelectOperation) evaluate(input []*Spanset) (output []*Spanset, err error) {
+ return input, nil
+}
+
 func (f ScalarFilter) evaluate(input []*Spanset) (output []*Spanset, err error) {
 
  // TODO we solve this gap where pipeline elements and scalar binary

pkg/traceql/ast_stringer.go

@@ -26,6 +26,14 @@ func (o CoalesceOperation) String() string {
  return "coalesce()"
 }
 
+func (o SelectOperation) String() string {
+ s := make([]string, 0, len(o.exprs))
+ for _, e := range o.exprs {
+ s = append(s, e.String())
+ }
+ return "select(" + strings.Join(s, ", ") + ")"
+}
+
 func (o ScalarOperation) String() string {
  return binaryOp(o.Op, o.LHS, o.RHS)
 }

pkg/traceql/ast_validate.go

@@ -42,6 +42,16 @@ func (o CoalesceOperation) validate() error {
  return nil
 }
 
+func (o SelectOperation) validate() error {
+ for _, e := range o.exprs {
+ if err := e.validate(); err != nil {
+ return err
+ }
+ }
+
+ return nil
+}
+
 func (o ScalarOperation) validate() error {
  if err := o.LHS.validate(); err != nil {
  return err

pkg/traceql/engine.go

@@ -70,7 +70,7 @@ func (e *Engine) ExecuteSearch(ctx context.Context, searchReq *tempopb.SearchReq
 
  spansetsEvaluated := 0
  // set up the expression evaluation as a filter to reduce data pulled
- fetchSpansRequest.SecondPassConditions = metaConditions
+ fetchSpansRequest.SecondPassConditions = append(fetchSpansRequest.SecondPassConditions, metaConditions...)
  fetchSpansRequest.SecondPass = func(inSS *Spanset) ([]*Spanset, error) {
  if len(inSS.Spans) == 0 {
  return nil, nil

pkg/traceql/expr.y

@@ -12,6 +12,8 @@ import (
  root RootExpr
  groupOperation GroupOperation
  coalesceOperation CoalesceOperation
+ selectOperation SelectOperation
+ selectArgs []FieldExpression
 
  spansetExpression SpansetExpression
  spansetPipelineExpression SpansetExpression
@@ -43,6 +45,8 @@ import (
 %type <RootExpr> root
 %type <groupOperation> groupOperation
 %type <coalesceOperation> coalesceOperation
+%type <selectOperation> selectOperation
+%type <selectArgs> selectArgs
 
 %type <spansetExpression> spansetExpression
 %type <spansetPipelineExpression> spansetPipelineExpression
@@ -68,13 +72,13 @@ import (
 %token <staticInt> INTEGER
 %token <staticFloat> FLOAT
 %token <staticDuration> DURATION
-%token <val> DOT OPEN_BRACE CLOSE_BRACE OPEN_PARENS CLOSE_PARENS
+%token <val> DOT OPEN_BRACE CLOSE_BRACE OPEN_PARENS CLOSE_PARENS COMMA
  NIL TRUE FALSE STATUS_ERROR STATUS_OK STATUS_UNSET
  KIND_UNSPECIFIED KIND_INTERNAL KIND_SERVER KIND_CLIENT KIND_PRODUCER KIND_CONSUMER
  IDURATION CHILDCOUNT NAME STATUS PARENT KIND
  PARENT_DOT RESOURCE_DOT SPAN_DOT
  COUNT AVG MAX MIN SUM
- BY COALESCE
+ BY COALESCE SELECT
  END_ATTRIBUTE
 
 // Operators are listed with increasing precedence.
@@ -116,10 +120,12 @@ spansetPipeline:
  spansetExpression { $$ = newPipeline($1) }
  | scalarFilter { $$ = newPipeline($1) }
  | groupOperation { $$ = newPipeline($1) }
+ | selectOperation { $$ = newPipeline($1) }
  | spansetPipeline PIPE spansetExpression { $$ = $1.addItem($3) }
  | spansetPipeline PIPE scalarFilter { $$ = $1.addItem($3) }
  | spansetPipeline PIPE groupOperation { $$ = $1.addItem($3) }
  | spansetPipeline PIPE coalesceOperation { $$ = $1.addItem($3) }
+ | spansetPipeline PIPE selectOperation { $$ = $1.addItem($3) }
  ;
 
 groupOperation:
@@ -130,6 +136,15 @@ coalesceOperation:
  COALESCE OPEN_PARENS CLOSE_PARENS { $$ = newCoalesceOperation() }
  ;
 
+selectOperation:
+ SELECT OPEN_PARENS selectArgs CLOSE_PARENS { $$ = newSelectOperation($3) }
+ ;
+
+selectArgs:
+ fieldExpression { $$ = []FieldExpression{$1} }
+ | selectArgs COMMA fieldExpression { $$ = append($1, $3) }
+ ;
+
 spansetExpression: // shares the same operators as scalarPipelineExpression. split out for readability
  OPEN_PARENS spansetExpression CLOSE_PARENS { $$ = $2 }
  | spansetExpression AND spansetExpression { $$ = newSpansetOperation(OpSpansetAnd, $1, $3) }