logicalplan: distribute throgh nested aggregations

Signed-off-by: Michael Hoffmann <[email protected]>

Author: MichaHoffmann

engine/distributed_test.go

@@ -282,6 +282,28 @@ func TestDistributedAggregations(t *testing.T) {
 		{name: "group_right with partition label in include", query: `bar{pod="nginx-1"} * on (pod) group_right (zone) bar`},
 		{name: "group_left without partition label", query: `bar * on (zone) group_left (pod) bar{zone="east-1"}`},
 		{name: "group_right without partition label", query: `bar{zone="east-1"} * on (zone) group_right (pod) bar`},
+		// Nested aggregation tests - using deterministic aggregations (max/min/sum/count)
+		// instead of topk/bottomk to avoid flaky tests due to tie-breaking.
+		{name: "max over sum by partition", query: `max(sum by (zone, pod) (bar))`},
+		{name: "max over max by partition", query: `max(max by (zone) (bar))`},
+		{name: "min over sum by partition", query: `min(sum by (zone, pod) (bar))`},
+		{name: "min over min by partition", query: `min(min by (zone) (bar))`},
+		{name: "count over sum by partition", query: `count(sum by (zone, pod) (bar))`},
+		{name: "sum over max by partition", query: `sum(max by (zone) (bar))`},
+		{name: "min over count by partition", query: `min(count by (zone) (bar))`},
+		{name: "sum over max by zone", query: `sum(max by (zone) (bar))`},
+		{name: "max over sum over rate by partition", query: `max(sum by (zone) (rate(bar[1m])))`},
+		{name: "min over avg by partition", query: `min(avg by (zone) (bar))`},
+		{name: "max over max over sum by partition", query: `max(max(sum by (zone) (bar)))`},
+		{name: "sum over min over max by partition", query: `sum(min(max by (zone, pod) (bar)))`},
+		{name: "max over sum without partition", query: `max(sum by (pod) (bar))`},
+		{name: "min over max without partition", query: `min(max(bar))`},
+		{name: "max over binary op by partition", query: `max(sum by (zone) (bar) / count by (zone) (bar))`},
+		{name: "count over max by partition", query: `count(max by (zone) (bar))`},
+		{name: "group over sum by partition", query: `group(sum by (zone) (bar))`},
+		{name: "max over binary with on() by partition", query: `max(bar * on (zone, pod) bar)`},
+		{name: "max over sum with without() by partition", query: `max(sum without (pod) (bar))`},
+		{name: "max over sum with without() without partition", query: `max(sum without (zone) (bar))`},
 	}
 
 	lookbackDeltas := []time.Duration{0, 30 * time.Second, 5 * time.Minute}

logicalplan/distribute.go

@@ -221,9 +221,30 @@ func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options)
 			return true
 		}
 
-		// If the current node is an aggregation, distribute the operation and
-		// stop the traversal.
+		// Handle absent functions specially
+		if isAbsent(*current) {
+			*current = m.distributeAbsent(*current, engines, calculateStartOffset(current, opts.LookbackDelta), m.subqueryOpts(parents, current, opts))
+			return true
+		}
+
+		// If the current node is an aggregation, check if we should distribute here
+		// or continue traversing up.
 		if aggr, ok := (*current).(*Aggregation); ok {
+			// If this aggregation preserves partition labels, check if the parent
+			// is also a distributive aggregation that we could push through.
+			// This enables patterns like: topk(10, sum by (P, instance) (X))
+			// where P is a partition label - we can push the entire expression
+			// to remote engines.
+			if preservesPartitionLabels(*current, engineLabels) {
+				if parent != nil {
+					if parentAggr, ok := (*parent).(*Aggregation); ok {
+						if _, ok := distributiveAggregations[parentAggr.Op]; ok {
+							// Parent is a distributive aggregation, continue up
+							return false
+						}
+					}
+				}
+			}
 			localAggregation := aggr.Op
 			if aggr.Op == parser.COUNT {
 				localAggregation = parser.SUM
@@ -240,10 +261,6 @@ func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options)
 			}
 			return true
 		}
-		if isAbsent(*current) {
-			*current = m.distributeAbsent(*current, engines, calculateStartOffset(current, opts.LookbackDelta), m.subqueryOpts(parents, current, opts))
-			return true
-		}
 
 		// If the parent operation is distributive, continue the traversal.
 		if isDistributive(parent, m.SkipBinaryPushdown, engineLabels, warns) {
@@ -523,6 +540,73 @@ func numSteps(start, end time.Time, step time.Duration) int64 {
 	return (end.UnixMilli()-start.UnixMilli())/step.Milliseconds() + 1
 }
 
+// preservesPartitionLabels checks if an expression preserves all partition labels.
+// An expression preserves partition labels if the output series will still have
+// those labels, meaning results from different engines won't overlap and can be
+// coalesced without deduplication.
+//
+// This enables pushing more operations to remote engines. For example:
+//
+//	topk(10, sum by (P, instance) (X))
+//
+// If P is a partition label, the sum preserves P, so topk can also be pushed
+// down since each engine's top 10 won't overlap with other engines' top 10.
+func preservesPartitionLabels(expr Node, partitionLabels map[string]struct{}) bool {
+	if len(partitionLabels) == 0 {
+		return true
+	}
+
+	switch e := expr.(type) {
+	case *VectorSelector, *MatrixSelector, *NumberLiteral, *StringLiteral:
+		return true
+	case *Aggregation:
+		for lbl := range partitionLabels {
+			if slices.Contains(e.Grouping, lbl) == e.Without {
+				return false
+			}
+		}
+		return true
+	case *Binary:
+		if e.VectorMatching != nil {
+			for lbl := range partitionLabels {
+				inMatching := slices.Contains(e.VectorMatching.MatchingLabels, lbl)
+				inInclude := slices.Contains(e.VectorMatching.Include, lbl)
+				if !inInclude && inMatching != e.VectorMatching.On {
+					return false
+				}
+			}
+		}
+		return preservesPartitionLabels(e.LHS, partitionLabels) &&
+			preservesPartitionLabels(e.RHS, partitionLabels)
+	case *FunctionCall:
+		if e.Func.Name == "label_replace" {
+			if _, ok := partitionLabels[UnsafeUnwrapString(e.Args[1])]; ok {
+				return false
+			}
+		}
+		for _, arg := range e.Args {
+			if arg.ReturnType() == parser.ValueTypeVector || arg.ReturnType() == parser.ValueTypeMatrix {
+				if !preservesPartitionLabels(arg, partitionLabels) {
+					return false
+				}
+			}
+		}
+		return true
+	case *Unary:
+		return preservesPartitionLabels(e.Expr, partitionLabels)
+	case *Parens:
+		return preservesPartitionLabels(e.Expr, partitionLabels)
+	case *StepInvariantExpr:
+		return preservesPartitionLabels(e.Expr, partitionLabels)
+	case *CheckDuplicateLabels:
+		return preservesPartitionLabels(e.Expr, partitionLabels)
+	case *Subquery:
+		return preservesPartitionLabels(e.Expr, partitionLabels)
+	default:
+		return false
+	}
+}
+
 func isDistributive(expr *Node, skipBinaryPushdown bool, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
 	if expr == nil {
 		return false

logicalplan/distribute_test.go

@@ -442,8 +442,6 @@ count by (cluster) (
 			skipBinopPushdown: true,
 		},
 		{
-			// group_left/group_right with partition label cannot be distributed because
-			// match cardinality changes when each partition only sees one value for that label.
 			name:     "topk over binary with group_left including partition label does not distribute",
 			expr:     `topk(5, metric_a * on (pod) group_left(region) metric_b)`,
 			expected: `topk(5, dedup(remote(metric_a), remote(metric_a)) * on (pod) group_left (region) dedup(remote(metric_b), remote(metric_b)))`,
@@ -528,6 +526,66 @@ count by (cluster) (
 			expr:     `metric_a unless on () vector(0)`,
 			expected: `dedup(remote(metric_a), remote(metric_a)) unless on () vector(0)`,
 		},
+		{
+			name:     "topk over sum by partition",
+			expr:     `topk(10, sum by (region, instance) (http_requests_total))`,
+			expected: `topk(10, dedup(remote(topk by (region) (10, sum by (region, instance) (http_requests_total))), remote(topk by (region) (10, sum by (region, instance) (http_requests_total)))))`,
+		},
+		{
+			name:     "bottomk over max by partition",
+			expr:     `bottomk(5, max by (region, pod) (cpu_usage))`,
+			expected: `bottomk(5, dedup(remote(bottomk by (region) (5, max by (region, pod) (cpu_usage))), remote(bottomk by (region) (5, max by (region, pod) (cpu_usage)))))`,
+		},
+		{
+			name:     "topk over sum without partition",
+			expr:     `topk(10, sum by (instance) (http_requests_total))`,
+			expected: `topk(10, sum by (instance) (dedup(remote(sum by (instance, region) (http_requests_total)), remote(sum by (instance, region) (http_requests_total)))))`,
+		},
+		{
+			name:     "count over sum by partition",
+			expr:     `count(sum by (region, pod) (http_requests_total))`,
+			expected: `sum(dedup(remote(count by (region) (sum by (region, pod) (http_requests_total))), remote(count by (region) (sum by (region, pod) (http_requests_total)))))`,
+		},
+		{
+			name:     "topk over binary with group_left preserving partition",
+			expr:     `topk(5, metric_a * on (pod) group_left(region) metric_b)`,
+			expected: `topk(5, dedup(remote(topk by (region) (5, metric_a * on (pod) group_left (region) metric_b)), remote(topk by (region) (5, metric_a * on (pod) group_left (region) metric_b))))`,
+		},
+		{
+			name:     "topk over binary with group_right preserving partition",
+			expr:     `topk(5, metric_a * on (pod) group_right(region) metric_b)`,
+			expected: `topk(5, dedup(remote(topk by (region) (5, metric_a * on (pod) group_right (region) metric_b)), remote(topk by (region) (5, metric_a * on (pod) group_right (region) metric_b))))`,
+		},
+		{
+			name:     "topk over binary with on() including partition",
+			expr:     `topk(5, metric_a * on (region, pod) metric_b)`,
+			expected: `topk(5, dedup(remote(topk by (region) (5, metric_a * on (region, pod) metric_b)), remote(topk by (region) (5, metric_a * on (region, pod) metric_b))))`,
+		},
+		{
+			name:     "topk over binary with on() excluding partition",
+			expr:     `topk(5, metric_a * on (pod) metric_b)`,
+			expected: `topk(5, dedup(remote(metric_a), remote(metric_a)) * on (pod) dedup(remote(metric_b), remote(metric_b)))`,
+		},
+		{
+			name:     "topk over binary with ignoring() excluding partition",
+			expr:     `topk(5, metric_a * ignoring (pod) metric_b)`,
+			expected: `topk(5, dedup(remote(topk by (region) (5, metric_a * ignoring (pod) metric_b)), remote(topk by (region) (5, metric_a * ignoring (pod) metric_b))))`,
+		},
+		{
+			name:     "topk over binary with ignoring() including partition",
+			expr:     `topk(5, metric_a * ignoring (region) metric_b)`,
+			expected: `topk(5, dedup(remote(metric_a), remote(metric_a)) * ignoring (region) dedup(remote(metric_b), remote(metric_b)))`,
+		},
+		{
+			name:     "topk over sum with without() excluding partition",
+			expr:     `topk(5, sum without (pod) (metric_a))`,
+			expected: `topk(5, dedup(remote(topk by (region) (5, sum without (pod) (metric_a))), remote(topk by (region) (5, sum without (pod) (metric_a)))))`,
+		},
+		{
+			name:     "topk over sum with without() including partition",
+			expr:     `topk(5, sum without (region) (metric_a))`,
+			expected: `topk(5, sum without (region) (dedup(remote(sum without () (metric_a)), remote(sum without () (metric_a)))))`,
+		},
 	}
 
 	engines := []api.RemoteEngine{