logiclaplan: distribute aggregations that preserve partition labels (#699)

Signed-off-by: Michael Hoffmann <mhoffmann@cloudflare.com>

Author: MichaHoffmann

logicalplan/distribute.go

@@ -143,20 +143,6 @@ func (r Noop) ReturnType() parser.ValueType { return parser.ValueTypeVector }
 
 func (r Noop) Type() NodeType { return NoopNode }
 
-// distributiveAggregations are all PromQL aggregations which support
-// distributed execution.
-var distributiveAggregations = map[parser.ItemType]struct{}{
-	parser.SUM:         {},
-	parser.MIN:         {},
-	parser.MAX:         {},
-	parser.GROUP:       {},
-	parser.COUNT:       {},
-	parser.BOTTOMK:     {},
-	parser.TOPK:        {},
-	parser.LIMITK:      {},
-	parser.LIMIT_RATIO: {},
-}
-
 // DistributedExecutionOptimizer produces a logical plan suitable for
 // distributed Query execution.
 type DistributedExecutionOptimizer struct {
@@ -185,81 +171,138 @@ func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options)
 		}
 	}
 
-	var warns = annotations.New()
+	warns := annotations.New()
+
+	parents := computeParents(&plan)
+	distributionPoints := m.computeDistributionPoints(&plan, parents, engineLabels, warns)
 
-	// TODO(fpetkovski): Consider changing TraverseBottomUp to pass in a list of parents in the transform function.
-	parents := make(map[*Node]*Node)
-	TraverseBottomUp(nil, &plan, func(parent, current *Node) (stop bool) {
-		parents[current] = parent
-		return false
-	})
 	TraverseBottomUp(nil, &plan, func(parent, current *Node) (stop bool) {
-		// Handle avg() specially - it's not distributive but can be distributed as sum/count.
-		if isAvgAggregation(current) {
-			*current = m.distributeAvg(*current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
-			return true
+		if _, distributeNow := distributionPoints[current]; !distributeNow {
+			return false
 		}
 
-		// If the current operation is not distributive, stop the traversal.
-		if !isDistributive(current, m.SkipBinaryPushdown, engineLabels, warns) {
+		if isAvgAggregation(current) && !preservesPartitionLabels(*current, engineLabels) {
+			// avg without partition labels: rewrite as sum/count.
+			*current = m.distributeAvg(*current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
 			return true
 		}
 
-		// 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 and there's a
-			// distributive aggregation ancestor, continue up to let it handle distribution.
-			// This enables patterns like:
-			//   - topk(10, sum by (P, instance) (X))
-			//   - sum(metric_a * group by (P) (metric_b))
-			//   - max(sum by (P, instance) (X))
-			// where P is a partition label - we can push the entire expression
-			// to remote engines.
-			//
-			// We need to check ancestors (not just immediate parent) because the
-			// aggregation might be nested inside a binary expression that is itself
-			// inside another aggregation: sum(A * group by (P) (B))
+		if isAggregation(current) {
 			if preservesPartitionLabels(*current, engineLabels) {
-				if hasDistributiveAncestor(parents, current, m.SkipBinaryPushdown, engineLabels, warns) {
-					return false
-				}
-			}
-			localAggregation := aggr.Op
-			if aggr.Op == parser.COUNT {
-				localAggregation = parser.SUM
-			}
-
-			remoteAggregation := newRemoteAggregation(aggr, engines)
-			subQueries := m.distributeQuery(&remoteAggregation, engines, m.subqueryOpts(parents, current, opts), labelRanges)
-			*current = &Aggregation{
-				Op:       localAggregation,
-				Expr:     subQueries,
-				Param:    aggr.Param,
-				Grouping: aggr.Grouping,
-				Without:  aggr.Without,
+				// Partition-preserving aggregation: push as-is since each engine
+				// computes over disjoint partition values.
+				*current = m.distributeQuery(current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
+			} else {
+				// Distributive aggregation that drops partition labels: use a
+				// two-level split with local_agg(remote_agg(X)).
+				*current = m.distributeAggregation((*current).(*Aggregation), engines, m.subqueryOpts(parents, current, opts), labelRanges)
 			}
 			return true
 		}
 
-		// If the parent operation is distributive or is an avg (which we handle specially),
-		// continue the traversal.
-		if isDistributive(parent, m.SkipBinaryPushdown, engineLabels, warns) || isAvgAggregation(parent) {
-			return false
-		}
-
 		*current = m.distributeQuery(current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
 		return true
 	})
 	return plan, *warns
 }
 
+func (m DistributedExecutionOptimizer) distributeAggregation(aggr *Aggregation, engines []api.RemoteEngine, opts *query.Options, labelRanges labelSetRanges) Node {
+	localAggregation := aggr.Op
+	if aggr.Op == parser.COUNT {
+		localAggregation = parser.SUM
+	}
+	remoteAggregation := newRemoteAggregation(aggr, engines)
+	subQueries := m.distributeQuery(&remoteAggregation, engines, opts, labelRanges)
+	return &Aggregation{
+		Op:       localAggregation,
+		Expr:     subQueries,
+		Param:    aggr.Param,
+		Grouping: aggr.Grouping,
+		Without:  aggr.Without,
+	}
+}
+
+func computeParents(plan *Node) map[*Node]*Node {
+	parents := make(map[*Node]*Node)
+	TraverseBottomUp(nil, plan, func(parent, current *Node) (stop bool) {
+		parents[current] = parent
+		return false
+	})
+	return parents
+}
+
+func (m DistributedExecutionOptimizer) computeDistributionPoints(plan *Node, parents map[*Node]*Node, engineLabels map[string]struct{}, warns *annotations.Annotations) map[*Node]struct{} {
+	marks := make(map[*Node]struct{})
+
+	// First pass: mark distribution points (aggregations, absent functions).
+	Traverse(plan, func(current *Node) {
+		if isAbsent(current) {
+			if m.isDistributive(current, engineLabels, warns) {
+				marks[current] = struct{}{}
+			}
+			return
+		}
+		if isAggregation(current) {
+			// Non-distributive aggregations that don't preserve partition labels
+			// cannot be distributed, except for avg which gets rewritten as sum/count.
+			if !m.isDistributive(current, engineLabels, warns) {
+				if isAvgAggregation(current) {
+					marks[current] = struct{}{}
+				}
+				return
+			}
+			// Distributive aggregations (standard or partition-preserving):
+			// defer to ancestor if possible.
+			if preservesPartitionLabels(*current, engineLabels) {
+				if m.hasDistributiveAncestor(parents, current, engineLabels, warns) {
+					return
+				}
+			}
+			marks[current] = struct{}{}
+		}
+	})
+
+	// Second pass: for nodes whose siblings have marks, mark them too so both
+	// sides of a binary expression get distributed.
+	Traverse(plan, func(current *Node) {
+		if _, ok := marks[current]; ok {
+			return
+		}
+		if subtreeHasMark(current, marks) {
+			return
+		}
+		if !m.isDistributive(current, engineLabels, warns) {
+			return
+		}
+		parent := parents[current]
+		if parent != nil && (m.isDistributive(parent, engineLabels, warns) || isAvgAggregation(parent)) {
+			if !subtreeHasMark(parent, marks) {
+				return
+			}
+		}
+		marks[current] = struct{}{}
+	})
+
+	return marks
+}
+
+func subtreeHasMark(node *Node, marks map[*Node]struct{}) bool {
+	for _, child := range (*node).Children() {
+		if _, ok := marks[child]; ok {
+			return true
+		}
+		if subtreeHasMark(child, marks) {
+			return true
+		}
+	}
+	return false
+}
+
 func (m DistributedExecutionOptimizer) subqueryOpts(parents map[*Node]*Node, current *Node, opts *query.Options) *query.Options {
 	subqueryParents := make([]*Subquery, 0, len(parents))
 	for p := parents[current]; p != nil; p = parents[p] {
@@ -641,7 +684,7 @@ func preservesPartitionLabels(expr Node, partitionLabels map[string]struct{}) bo
 	}
 }
 
-func isDistributive(expr *Node, skipBinaryPushdown bool, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
+func (m DistributedExecutionOptimizer) isDistributive(expr *Node, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
 	if expr == nil {
 		return false
 	}
@@ -653,13 +696,23 @@ func isDistributive(expr *Node, skipBinaryPushdown bool, engineLabels map[string
 		if isBinaryExpressionWithOneScalarSide(e) {
 			return true
 		}
-		return !skipBinaryPushdown &&
+		return !m.SkipBinaryPushdown &&
 			isBinaryExpressionWithDistributableMatching(e, engineLabels) &&
-			isDistributive(&e.LHS, skipBinaryPushdown, engineLabels, warns) &&
-			isDistributive(&e.RHS, skipBinaryPushdown, engineLabels, warns)
+			m.isDistributive(&e.LHS, engineLabels, warns) &&
+			m.isDistributive(&e.RHS, engineLabels, warns)
 	case *Aggregation:
-		// Certain aggregations are currently not supported.
-		if _, ok := distributiveAggregations[e.Op]; !ok {
+		switch e.Op {
+		// Mathematically distributive: can be split into local_agg(remote_agg(X))
+		// regardless of partition labels.
+		case parser.SUM, parser.MIN, parser.MAX, parser.GROUP, parser.COUNT,
+			parser.TOPK, parser.BOTTOMK, parser.LIMITK, parser.LIMIT_RATIO:
+		// Non-distributive: can only be pushed as-is when they preserve
+		// partition labels (each engine computes over disjoint data).
+		case parser.AVG, parser.QUANTILE, parser.STDDEV, parser.STDVAR, parser.COUNT_VALUES:
+			if !preservesPartitionLabels(e, engineLabels) {
+				return false
+			}
+		default:
 			return false
 		}
 	case *FunctionCall:
@@ -840,9 +893,9 @@ func matchesExternalLabels(ms []*labels.Matcher, externalLabels labels.Labels) b
 // parent chain from the current node that can handle distribution.
 // We must have an unbroken chain of distributive nodes to the ancestor for it to
 // be able to handle distribution on our behalf.
-func hasDistributiveAncestor(parents map[*Node]*Node, current *Node, skipBinaryPushdown bool, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
+func (m DistributedExecutionOptimizer) hasDistributiveAncestor(parents map[*Node]*Node, current *Node, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
 	for p := parents[current]; p != nil; p = parents[p] {
-		if !isDistributive(p, skipBinaryPushdown, engineLabels, warns) {
+		if !m.isDistributive(p, engineLabels, warns) {
 			// We hit a non-distributive node, so we can't push through it.
 			// No ancestor can help us distribute.
 			return false

logicalplan/distribute_test.go

@@ -83,19 +83,13 @@ sum(
     remote(count by (pod, region) (metric_a)))))`,
 		},
 		{
-			name: "avg in binary expression with outer sum distributes avg independently",
+			name: "avg in binary expression with outer sum pushes entire expression",
 			expr: `sum(metric_a * avg by (region) (metric_b))`,
 			expected: `
 sum(
-  dedup(remote(metric_a), remote(metric_a))
-  *
-  sum by (region) (dedup(
-    remote(sum by (region) (metric_b)),
-    remote(sum by (region) (metric_b))))
-  / on (region)
-  sum by (region) (dedup(
-    remote(count by (region) (metric_b)),
-    remote(count by (region) (metric_b)))))`,
+  dedup(
+    remote(sum by (region) (metric_a * avg by (region) (metric_b))),
+    remote(sum by (region) (metric_a * avg by (region) (metric_b)))))`,
 		},
 		{
 			name:     "selector",
@@ -171,21 +165,12 @@ sum by (pod) (
 )`,
 		},
 		{
-			name: "avg with without-grouping",
+			name: "avg with without-grouping preserving partition labels",
 			expr: `avg without (pod) (http_requests_total)`,
 			expected: `
-sum without (pod) (
-  dedup(
-    remote(sum without (pod) (http_requests_total)),
-    remote(sum without (pod) (http_requests_total))
-  )
-) / ignoring (pod)
-sum without (pod) (
-  dedup(
-    remote(count without (pod) (http_requests_total)),
-    remote(count without (pod) (http_requests_total))
-  )
-)`,
+dedup(
+  remote(avg without (pod) (http_requests_total)),
+  remote(avg without (pod) (http_requests_total)))`,
 		},
 		{
 			name: "avg with prior aggregation",
@@ -218,6 +203,57 @@ sum by (pod) (
   )
 )`,
 		},
+		{
+			name: "avg by partition label pushes as-is",
+			expr: `avg by (region) (http_requests_total)`,
+			expected: `
+dedup(
+  remote(avg by (region) (http_requests_total)),
+  remote(avg by (region) (http_requests_total)))`,
+		},
+		{
+			name: "avg by partition label defers to distributive ancestor",
+			expr: `max(avg by (region) (http_requests_total))`,
+			expected: `
+max(
+  dedup(
+    remote(max by (region) (avg by (region) (http_requests_total))),
+    remote(max by (region) (avg by (region) (http_requests_total)))))`,
+		},
+		{
+			name: "avg over subquery with inner aggregations pushes entire expression",
+			expr: `avg by (region) (quantile_over_time(0.9, (sum by (region) (rate(metric_a[2m])) / sum by (region) (metric_b))[1h:1m]))`,
+			expected: `
+dedup(
+  remote(avg by (region) (quantile_over_time(0.9, (sum by (region) (rate(metric_a[2m])) / sum by (region) (metric_b))[1h:1m]))),
+  remote(avg by (region) (quantile_over_time(0.9, (sum by (region) (rate(metric_a[2m])) / sum by (region) (metric_b))[1h:1m]))))`,
+		},
+		{
+			name: "quantile by partition label pushes as-is",
+			expr: `quantile by (region) (0.9, http_requests_total)`,
+			expected: `
+dedup(
+  remote(quantile by (region) (0.9, http_requests_total)),
+  remote(quantile by (region) (0.9, http_requests_total)))`,
+		},
+		{
+			name:     "quantile by non-partition label is not distributed",
+			expr:     `quantile by (pod) (0.9, http_requests_total)`,
+			expected: `quantile by (pod) (0.9, dedup(remote(http_requests_total), remote(http_requests_total)))`,
+		},
+		{
+			name: "stddev by partition label pushes as-is",
+			expr: `stddev by (region) (http_requests_total)`,
+			expected: `
+dedup(
+  remote(stddev by (region) (http_requests_total)),
+  remote(stddev by (region) (http_requests_total)))`,
+		},
+		{
+			name:     "stddev by non-partition label is not distributed",
+			expr:     `stddev by (pod) (http_requests_total)`,
+			expected: `stddev by (pod) (dedup(remote(http_requests_total), remote(http_requests_total)))`,
+		},
 		{
 			name: "two-level aggregation",
 			expr: `max by (pod) (sum by (pod) (http_requests_total))`,
@@ -518,15 +554,22 @@ count by (cluster) (
 		{
 			name:              "skip binary pushdown with nested aggregation",
 			expr:              `sum(metric_a * group by (region) (metric_b))`,
-			expected:          `sum(dedup(remote(metric_a), remote(metric_a)) * group by (region) (dedup(remote(group by (region) (metric_b)), remote(group by (region) (metric_b)))))`,
+			expected:          `sum(dedup(remote(metric_a), remote(metric_a)) * dedup(remote(group by (region) (metric_b)), remote(group by (region) (metric_b))))`,
 			skipBinopPushdown: true,
 		},
 		{
 			name:              "skip binary pushdown with outer aggregation",
 			expr:              `max(metric_a + sum by (region, pod) (metric_b))`,
-			expected:          `max(dedup(remote(metric_a), remote(metric_a)) + sum by (region, pod) (dedup(remote(sum by (pod, region) (metric_b)), remote(sum by (pod, region) (metric_b)))))`,
+			expected:          `max(dedup(remote(metric_a), remote(metric_a)) + dedup(remote(sum by (region, pod) (metric_b)), remote(sum by (region, pod) (metric_b))))`,
 			skipBinopPushdown: true,
 		},
+		{
+			// When the RHS of unless has an aggregation that drops the partition label,
+			// both sides should still be distributed independently.
+			name:     "unless with aggregation that drops partition label distributes both sides",
+			expr:     `group by (region, instance) (metric_a unless on (region, instance) max by (instance) (metric_b))`,
+			expected: `group by (region, instance) (dedup(remote(metric_a), remote(metric_a)) unless on (region, instance) max by (instance) (dedup(remote(max by (instance, region) (metric_b)), remote(max by (instance, region) (metric_b)))))`,
+		},
 		{
 			// group_left/group_right with partition label cannot be distributed because
 			// match cardinality changes when each partition only sees one value for that label.