fix: Implement approx-topk function on querier (#17816)

Author: benclive

pkg/logql/count_min_sketch.go

@@ -2,6 +2,7 @@ package logql
 
 import (
 	"container/heap"
+	"context"
 	"fmt"
 	"slices"
 	"strings"
@@ -322,7 +323,7 @@ type CountMinSketchVectorStepEvaluator struct {
 	buffer []byte
 }
 
-var _ StepEvaluator = NewQuantileSketchVectorStepEvaluator(nil, 0)
+var _ StepEvaluator = NewCountMinSketchVectorStepEvaluator(nil)
 
 func NewCountMinSketchVectorStepEvaluator(vec *CountMinSketchVector) *CountMinSketchVectorStepEvaluator {
 	return &CountMinSketchVectorStepEvaluator{
@@ -357,3 +358,45 @@ func (e *CountMinSketchVectorStepEvaluator) Next() (bool, int64, StepResult) {
 func (*CountMinSketchVectorStepEvaluator) Close() error { return nil }
 
 func (*CountMinSketchVectorStepEvaluator) Error() error { return nil }
+
+var _ StepEvaluator = (*CountMinSketchEvalStepEvaluator)(nil)
+
+// CountMinSketchEvalStepEvaluator transforms a CountMinSketchEvalExpr into a CountMinSketchVector.
+type CountMinSketchEvalStepEvaluator struct {
+	ctx           context.Context
+	nextEvFactory SampleEvaluatorFactory
+	expr          *CountMinSketchEvalExpr
+	params        Params
+}
+
+func NewCountMinSketchEvalStepEvaluator(ctx context.Context, nextEvFactory SampleEvaluatorFactory, expr *CountMinSketchEvalExpr, params Params) (*CountMinSketchEvalStepEvaluator, error) {
+	return &CountMinSketchEvalStepEvaluator{
+		ctx:           ctx,
+		nextEvFactory: nextEvFactory,
+		expr:          expr,
+		params:        params,
+	}, nil
+}
+
+func (e *CountMinSketchEvalStepEvaluator) Next() (bool, int64, StepResult) {
+	nextEv, err := e.nextEvFactory.NewStepEvaluator(e.ctx, e.nextEvFactory, e.expr.SampleExpr, e.params)
+	if err != nil {
+		return false, 0, CountMinSketchVector{}
+	}
+
+	ok, _, results := nextEv.Next()
+	if !ok {
+		return false, 0, CountMinSketchVector{}
+	}
+
+	data := results.CountMinSketchVec()
+	handler := NewCountMinSketchVectorStepEvaluator(&data)
+
+	return handler.Next()
+}
+
+func (*CountMinSketchEvalStepEvaluator) Close() error { return nil }
+
+func (*CountMinSketchEvalStepEvaluator) Error() error { return nil }
+
+func (e *CountMinSketchEvalStepEvaluator) Explain(_ Node) {}

pkg/logql/downstream.go

@@ -414,6 +414,9 @@ func (e CountMinSketchEvalExpr) String() string {
 
 		sb.WriteString(d.String())
 	}
+	if len(e.downstreams) == 0 {
+		sb.WriteString(e.SampleExpr.String())
+	}
 	return fmt.Sprintf("CountMinSketchEval<%s>", sb.String())
 }
 

pkg/logql/evaluator.go

@@ -359,6 +359,8 @@ func (ev *DefaultEvaluator) NewStepEvaluator(
 			})
 		}
 		return newVectorAggEvaluator(ctx, nextEvFactory, e, q, ev.maxCountMinSketchHeapSize)
+	case *CountMinSketchEvalExpr:
+		return NewCountMinSketchEvalStepEvaluator(ctx, nextEvFactory, e, q)
 	case *syntax.RangeAggregationExpr:
 		it, err := ev.querier.SelectSamples(ctx, SelectSampleParams{
 			&logproto.SampleQueryRequest{

pkg/logql/optimize.go

@@ -20,10 +20,36 @@ func optimizeSampleExpr(expr syntax.SampleExpr) (syntax.SampleExpr, error) {
 	if err != nil {
 		return nil, err
 	}
+	replaceApproxTopK(expr)
 	removeLineformat(expr)
 	return expr, nil
 }
 
+// replaceApproxTopKWithTopk replaces all ApproxTopKExpr with TopKExpr.
+// ApproxTopKExpr is not supported by the querier, so we replace it with the implementation if this function reaches the querier.
+func replaceApproxTopK(expr syntax.SampleExpr) {
+	expr.Walk(func(e syntax.Expr) bool {
+		vectorExpr, ok := e.(*syntax.VectorAggregationExpr)
+		if !ok {
+			return true
+		}
+		if vectorExpr.Operation != syntax.OpTypeApproxTopK {
+			return true
+		}
+
+		vectorExpr.Operation = syntax.OpTypeTopK
+		vectorExpr.Left = &CountMinSketchEvalExpr{
+			SampleExpr: &syntax.VectorAggregationExpr{
+				Operation: syntax.OpTypeCountMinSketch,
+				Params:    0,
+				Grouping:  vectorExpr.Grouping,
+				Left:      vectorExpr.Left,
+			},
+		}
+		return true
+	})
+}
+
 // removeLineformat removes unnecessary line_format within a SampleExpr.
 func removeLineformat(expr syntax.SampleExpr) {
 	expr.Walk(func(e syntax.Expr) bool {

pkg/logql/optimize_test.go

@@ -27,6 +27,10 @@ func Test_optimizeSampleExpr(t *testing.T) {
 		{`sum by(name)(rate({region="us-east1"} | json | line_format "something else" | unwrap foo[5m]))`, `sum by (name)(rate({region="us-east1"} | json | unwrap foo[5m]))`},
 		{`quantile_over_time(1,{region="us-east1"} | json | line_format "something else" | unwrap foo[5m])`, `quantile_over_time(1,{region="us-east1"} | json | unwrap foo[5m])`},
 		{`sum by(name)(count_over_time({region="us-east1"} | json | line_format "something else" | label_format foo=bar | line_format "boo"[5m]))`, `sum by (name)(count_over_time({region="us-east1"} | json | label_format foo=bar[5m]))`},
+
+		// Replace approx_topk with __count_min_sketch__ variant
+		{`approx_topk(3, 1+1)`, `topk(3,CountMinSketchEval<__count_min_sketch__(2)>)`},
+		{`approx_topk(3, sum by(name)(rate({region="us-east1"}[5m])))`, `topk(3,CountMinSketchEval<__count_min_sketch__(sum by (name)(rate({region="us-east1"}[5m])))>)`},
 	}
 	for _, tt := range tests {
 		t.Run(tt.in, func(t *testing.T) {