sharding operators: topk / bottomk

CHANGELOG.md

@@ -3,7 +3,7 @@
 ## main / unreleased
 
 ### Grafana Mimir
-
+* [FEATURE] PromQL: support to sharding topk/bottomk. #7582
 * [BUGFIX] Rules: improve error handling when querier is local to the ruler. #7567
 
 ### Mixin

pkg/frontend/querymiddleware/astmapper/parallel.go

@@ -14,12 +14,14 @@ import (
 )
 
 var summableAggregates = map[parser.ItemType]struct{}{
-	parser.GROUP: {},
-	parser.SUM:   {},
-	parser.MIN:   {},
-	parser.MAX:   {},
-	parser.COUNT: {},
-	parser.AVG:   {},
+	parser.GROUP:   {},
+	parser.SUM:     {},
+	parser.MIN:     {},
+	parser.MAX:     {},
+	parser.COUNT:   {},
+	parser.AVG:     {},
+	parser.TOPK:    {},
+	parser.BOTTOMK: {},
 }
 
 // NonParallelFuncs is the list of functions that shouldn't be parallelized.
@@ -66,6 +68,13 @@ func CanParallelize(expr parser.Expr, logger log.Logger) bool {
 			return false
 		}
 
+		if e.Op == parser.TOPK || e.Op == parser.BOTTOMK {
+			// example: topk(scalar(metric), up)
+			if _, isNum := e.Param.(*parser.NumberLiteral); !isNum {
+				return false
+			}
+		}
+
 		// Ensure there are no nested aggregations
 		nestedAggrs, err := anyNode(e.Expr, isAggregateExpr)
 

pkg/frontend/querymiddleware/astmapper/sharding.go

@@ -293,6 +293,12 @@ func (summer *shardSummer) shardAggregate(expr *parser.AggregateExpr) (mapped pa
 			return nil, false, err
 		}
 		return mapped, true, nil
+	case parser.TOPK, parser.BOTTOMK:
+		mapped, err = summer.shardTopkBottomk(expr)
+		if err != nil {
+			return nil, false, err
+		}
+		return mapped, true, nil
 	}
 
 	// If the aggregation operation is not shardable, we have to return the input
@@ -417,6 +423,29 @@ func (summer *shardSummer) shardMinMax(expr *parser.AggregateExpr) (result parse
 	}, nil
 }
 
+// shardTopkBottomk attempts to shard the given TOPK/BOTTOMK aggregation expression.
+func (summer *shardSummer) shardTopkBottomk(expr *parser.AggregateExpr) (result parser.Expr, err error) {
+	// We expect the given aggregation is either a topk or bottomk.
+	if expr.Op != parser.TOPK && expr.Op != parser.BOTTOMK {
+		return nil, errors.Errorf("expected TOPK or BOTTOMK aggregation while got %s", expr.Op.String())
+	}
+
+	// The TOPK/BOTTOMK aggregation can be parallelized as the TOPK/BOTTOMK of per-shard TOPK/BOTTOMK.
+	// Create a TOPK/BOTTOMK sub-query for each shard and squash it into a CONCAT expression.
+	sharded, err := summer.shardAndSquashAggregateExpr(expr, expr.Op)
+	if err != nil {
+		return nil, err
+	}
+
+	return &parser.AggregateExpr{
+		Op:       expr.Op,
+		Expr:     sharded,
+		Param:    expr.Param,
+		Grouping: expr.Grouping,
+		Without:  expr.Without,
+	}, nil
+}
+
 // shardAvg attempts to shard the given AVG aggregation expression.
 func (summer *shardSummer) shardAvg(expr *parser.AggregateExpr) (result parser.Expr, err error) {
 	// The AVG aggregation can be parallelized as per-shard SUM() divided by per-shard COUNT().
@@ -458,6 +487,7 @@ func (summer *shardSummer) shardAndSquashAggregateExpr(expr *parser.AggregateExp
 		children = append(children, &parser.AggregateExpr{
 			Op:       op,
 			Expr:     sharded,
+			Param:    expr.Param,
 			Grouping: expr.Grouping,
 			Without:  expr.Without,
 		})

pkg/frontend/querymiddleware/astmapper/sharding_test.go

@@ -518,6 +518,16 @@ func TestShardSummer(t *testing.T) {
 			out:                    `count(group without() (` + concatShards(3, `group without() ({namespace="foo",__query_shard__="x_of_y"})`) + `))`,
 			expectedShardedQueries: 3,
 		},
+		{
+			in:                     `topk(3, up)`,
+			out:                    `topk(3, ` + concatShards(3, `topk(3, up{__query_shard__="x_of_y"})`) + `)`,
+			expectedShardedQueries: 3,
+		},
+		{
+			in:                     `bottomk(4, up)`,
+			out:                    `bottomk(4, ` + concatShards(3, `bottomk(4, up{__query_shard__="x_of_y"})`) + `)`,
+			expectedShardedQueries: 3,
+		},
 	} {
 		tt := tt
 

pkg/frontend/querymiddleware/querysharding_test.go

@@ -452,11 +452,11 @@ func TestQuerySharding_Correctness(t *testing.T) {
 		},
 		"topk()": {
 			query:                  `topk(2, metric_counter{const="fixed"})`,
-			expectedShardedQueries: 0,
+			expectedShardedQueries: 1,
 		},
 		"bottomk()": {
 			query:                  `bottomk(2, metric_counter{const="fixed"})`,
-			expectedShardedQueries: 0,
+			expectedShardedQueries: 1,
 		},
 		"vector()": {
 			query:                  `vector(1)`,