Fix a few testing bugs

pkg/logql/downstream.go

@@ -652,7 +652,7 @@ func (ev *DownstreamEvaluator) NewStepEvaluator(
 		}
 
 		if len(results) != 1 {
-			return nil, fmt.Errorf("unexpected results length for sharded quantile: got (%d), want (1)", len(results))
+			return nil, fmt.Errorf("unexpected results length for sharded count min sketch: got (%d), want (1)", len(results))
 		}
 
 		vector, ok := results[0].Data.(*CountMinSketchVector)

pkg/logql/downstream_test.go

@@ -191,7 +191,7 @@ func TestMappingEquivalenceSketches(t *testing.T) {
 	}{
 		//{`quantile_over_time(0.70, {a=~".+"} | logfmt | unwrap value [1s]) by (a)`, 0.05, []QueryRangeType{InstantType, RangeType}},
 		//{`quantile_over_time(0.99, {a=~".+"} | logfmt | unwrap value [1s]) by (a)`, 0.02, []QueryRangeType{InstantType, RangeType}},
-		{`approx_topk(3, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [100s])))`, 0.02, []QueryRangeType{InstantType}},
+		{`topk(3, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [1s])))`, 0.002, []QueryRangeType{InstantType}},
 	} {
 		q := NewMockQuerier(
 			shards,
@@ -251,8 +251,8 @@ func TestMappingEquivalenceSketches(t *testing.T) {
 			// plus set step and interval to 0
 			params, err := NewLiteralParams(
 				tc.query,
-				time.Unix(0, int64(rounds+1)),
-				time.Unix(0, int64(rounds+1)),
+				time.Unix(1, 0),
+				time.Unix(1, 0),
 				0,
 				0,
 				logproto.FORWARD,
@@ -261,11 +261,18 @@ func TestMappingEquivalenceSketches(t *testing.T) {
 				nil,
 			)
 			require.NoError(t, err)
-			qry := regular.Query(params)
+			qry := regular.Query(params.Copy())
 			ctx := user.InjectOrgID(context.Background(), "fake")
 
 			strategy := NewPowerOfTwoStrategy(ConstantShards(shards))
 			mapper := NewShardMapper(strategy, nilShardMetrics, []string{ShardQuantileOverTime})
+
+			// TODO: use different test
+			// Use approximated topk
+			params.queryString = "approx_" + params.queryString
+			params.queryExpr, err = syntax.ParseExpr(params.queryString)
+			require.NoError(t, err)
+
 			_, _, mapped, err := mapper.Parse(params.GetExpression())
 			require.NoError(t, err)
 
@@ -276,6 +283,7 @@ func TestMappingEquivalenceSketches(t *testing.T) {
 
 			res, err := qry.Exec(ctx)
 			require.NoError(t, err)
+			require.NotEmpty(t, res.Data.(promql.Vector))
 
 			shardedRes, err := shardedQry.Exec(ctx)
 			require.NoError(t, err)
@@ -680,7 +688,7 @@ func relativeErrorVector(t *testing.T, expected, actual promql.Vector, alpha flo
 		require.Equal(t, expected[i].Metric, actual[i].Metric)
 
 		e[i] = expected[i].F
-		a[i] = expected[i].F
+		a[i] = actual[i].F
 	}
 	require.InEpsilonSlice(t, e, a, alpha)
 }

pkg/logql/shardmapper.go

@@ -335,6 +335,7 @@ func (m ShardMapper) mapVectorAggregationExpr(expr *syntax.VectorAggregationExpr
 				Left:      sharded,
 				Grouping:  expr.Grouping,
 				Operation: syntax.OpTypeTopK,
+				Params:    expr.Params,
 			}, bytesPerShard, nil
 		default:
 			// this should not be reachable. If an operation is shardable it should

pkg/logql/syntax/ast.go

@@ -2285,6 +2285,8 @@ var shardableOps = map[string]bool{
 	OpTypeMax:   true,
 	OpTypeMin:   true,
 
+	OpTypeApproxTopK: true,
+
 	// range vector ops
 	OpRangeTypeAvg:       true,
 	OpRangeTypeCount:     true,

pkg/logql/test_utils.go

@@ -227,26 +227,19 @@ type MockDownstreamer struct {
 
 func (m MockDownstreamer) Downstreamer(_ context.Context) Downstreamer { return m }
 
-func (m MockDownstreamer) Downstream(ctx context.Context, queries []DownstreamQuery, _ Accumulator) ([]logqlmodel.Result, error) {
-	results := make([]logqlmodel.Result, 0, len(queries))
-	for _, query := range queries {
+func (m MockDownstreamer) Downstream(ctx context.Context, queries []DownstreamQuery, acc Accumulator) ([]logqlmodel.Result, error) {
+	for i, query := range queries {
 		res, err := m.Query(query.Params).Exec(ctx)
 		if err != nil {
 			return nil, err
 		}
-		results = append(results, res)
-	}
-
-	if matrix, ok := results[0].Data.(ProbabilisticQuantileMatrix); ok {
-		if len(results) == 1 {
-			return results, nil
-		}
-		for _, m := range results[1:] {
-			matrix, _ = matrix.Merge(m.Data.(ProbabilisticQuantileMatrix))
+		err = acc.Accumulate(ctx, res, i)
+		if err != nil {
+			return nil, err
 		}
-		return []logqlmodel.Result{{Data: matrix}}, nil
 	}
-	return results, nil
+
+	return acc.Result(), nil
 }
 
 // create nStreams of nEntries with labelNames each where each label value