logicalplan: fix correctness issues in distributed optimizer (#702)

Fix isBinaryExpressionWithDistributableMatching to check all partition
labels instead of only the first one. The loop had a return statement
inside the body, causing it to exit on the first iteration. With
multiple partition labels (e.g. region + datacenter), only one arbitrary
label was checked, potentially allowing incorrect distribution.
Move LIMIT_RATIO from the mathematically distributive aggregation list
to the non-distributive group. Applying limit_ratio remotely and then
locally results in double-application (e.g. 0.5 * 0.5 = 0.25 ratio
instead of the intended 0.5). It can still be pushed as-is when
partition labels are preserved.
Add label_join handling alongside label_replace in both isDistributive
and preservesPartitionLabels. Previously, label_join targeting a
partition label (e.g. label_join(metric, "region", ...)) was not
detected, allowing incorrect distribution or partition label loss.

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

Author: MichaHoffmann

logicalplan/distribute.go

@@ -656,7 +656,7 @@ func preservesPartitionLabels(expr Node, partitionLabels map[string]struct{}) bo
 		return preservesPartitionLabels(e.LHS, partitionLabels) &&
 			preservesPartitionLabels(e.RHS, partitionLabels)
 	case *FunctionCall:
-		if e.Func.Name == "label_replace" {
+		if e.Func.Name == "label_replace" || e.Func.Name == "label_join" {
 			if _, ok := partitionLabels[UnsafeUnwrapString(e.Args[1])]; ok {
 				return false
 			}
@@ -705,18 +705,19 @@ func (m DistributedExecutionOptimizer) isDistributive(expr *Node, engineLabels m
 		// 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:
+			parser.TOPK, parser.BOTTOMK, parser.LIMITK:
 		// 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:
+		case parser.AVG, parser.QUANTILE, parser.STDDEV, parser.STDVAR,
+			parser.COUNT_VALUES, parser.LIMIT_RATIO:
 			if !preservesPartitionLabels(e, engineLabels) {
 				return false
 			}
 		default:
 			return false
 		}
 	case *FunctionCall:
-		if e.Func.Name == "label_replace" {
+		if e.Func.Name == "label_replace" || e.Func.Name == "label_join" {
 			targetLabel := UnsafeUnwrapString(e.Args[1])
 			if _, ok := engineLabels[targetLabel]; ok {
 				warns.Add(RewrittenExternalLabelWarning)
@@ -783,11 +784,10 @@ func isBinaryExpressionWithDistributableMatching(expr *Binary, engineLabels map[
 		// changes match cardinality semantics. Each partition only sees one value for
 		// that label, so what's many-to-many globally may become one-to-one per partition,
 		// producing results instead of errors (or vice versa).
-		return !inInclude && inMatching == expr.VectorMatching.On
+		if inInclude || inMatching != expr.VectorMatching.On {
+			return false
+		}
 	}
-	// At this point, partition labels are in the matching set (either via on() or
-	// by not being in ignoring()). This means or/unless can be safely distributed
-	// because the matching ensures series are paired by partition.
 	return true
 }
 

logicalplan/distribute_test.go

@@ -356,6 +356,20 @@ label_replace(max by (location) (dedup(
 )), "region", "$1", "location", "(.*)")`,
 			expectWarn: true,
 		},
+		{
+			name: "label join targeting non-partition label distributes",
+			expr: `label_join(http_requests_total, "zone", ",", "pod")`,
+			expected: `
+dedup(
+  remote(label_join(http_requests_total, "zone", ",", "pod")),
+  remote(label_join(http_requests_total, "zone", ",", "pod")))`,
+		},
+		{
+			name:       "label join targeting partition label does not distribute",
+			expr:       `max by (location) (label_join(http_requests_total, "region", ",", "pod"))`,
+			expected:   `max by (location) (label_join(dedup(remote(http_requests_total), remote(http_requests_total)), "region", ",", "pod"))`,
+			expectWarn: true,
+		},
 		{
 			name: "binary operation in the operand path",
 			expr: `max by (pod) (metric_a / metric_b)`,
@@ -1056,6 +1070,87 @@ sum by (pod) (dedup(
 	}
 }
 
+func TestDistributedExecutionMultiplePartitionLabels(t *testing.T) {
+	// Engines partitioned by both region and datacenter.
+	engines := []api.RemoteEngine{
+		newEngineMock(math.MinInt64, math.MaxInt64, []labels.Labels{labels.FromStrings("region", "east", "datacenter", "dc1")}),
+		newEngineMock(math.MinInt64, math.MaxInt64, []labels.Labels{labels.FromStrings("region", "west", "datacenter", "dc2")}),
+	}
+	optimizers := []Optimizer{
+		DistributedExecutionOptimizer{Endpoints: api.NewStaticEndpoints(engines)},
+	}
+
+	cases := []struct {
+		name     string
+		expr     string
+		expected string
+	}{
+		{
+			name: "on() must include all partition labels to distribute",
+			expr: `metric_a + on (region) metric_b`,
+			expected: `
+dedup(remote(metric_a), remote(metric_a))
++ on (region)
+dedup(remote(metric_b), remote(metric_b))`,
+		},
+		{
+			name: "on() with all partition labels distributes the binary",
+			expr: `metric_a + on (region, datacenter) metric_b`,
+			expected: `
+dedup(
+  remote(metric_a + on (region, datacenter) metric_b),
+  remote(metric_a + on (region, datacenter) metric_b))`,
+		},
+		{
+			name: "ignoring() must not include any partition label to distribute",
+			expr: `metric_a + ignoring (region) metric_b`,
+			expected: `
+dedup(remote(metric_a), remote(metric_a))
++ ignoring (region)
+dedup(remote(metric_b), remote(metric_b))`,
+		},
+		{
+			name: "ignoring() non-partition label distributes the binary",
+			expr: `metric_a + ignoring (pod) metric_b`,
+			expected: `
+dedup(
+  remote(metric_a + ignoring (pod) metric_b),
+  remote(metric_a + ignoring (pod) metric_b))`,
+		},
+		{
+			name: "sum must include all partition labels to preserve",
+			expr: `sum by (region) (metric_a)`,
+			expected: `
+sum by (region) (
+  dedup(
+    remote(sum by (datacenter, region) (metric_a)),
+    remote(sum by (datacenter, region) (metric_a))))`,
+		},
+		{
+			name: "sum by all partition labels preserves",
+			expr: `max(sum by (region, datacenter) (metric_a))`,
+			expected: `
+max(
+  dedup(
+    remote(max by (datacenter, region) (sum by (region, datacenter) (metric_a))),
+    remote(max by (datacenter, region) (sum by (region, datacenter) (metric_a)))))`,
+		},
+	}
+
+	for _, tcase := range cases {
+		t.Run(tcase.name, func(t *testing.T) {
+			expr, err := parser.ParseExpr(tcase.expr)
+			testutil.Ok(t, err)
+
+			plan, err := NewFromAST(expr, &query.Options{Start: time.Unix(0, 0), End: time.Unix(0, 0)}, PlanOptions{})
+			testutil.Ok(t, err)
+			optimizedPlan, _ := plan.Optimize(optimizers)
+			expectedPlan := cleanUp(replacements, tcase.expected)
+			testutil.Equals(t, expectedPlan, renderExprTree(optimizedPlan.Root()))
+		})
+	}
+}
+
 func TestDistributedExecutionClonesNodes(t *testing.T) {
 	var (
 		start    = time.Unix(0, 0)
@@ -1243,6 +1338,16 @@ func TestPreservesPartitionLabels(t *testing.T) {
 			expr:     `label_replace(metric, "zone", "$1", "pod", "(.*)")`,
 			expected: true,
 		},
+		{
+			name:     "label_join targeting partition label does not preserve",
+			expr:     `label_join(metric, "region", ",", "pod")`,
+			expected: false,
+		},
+		{
+			name:     "label_join targeting non-partition label preserves",
+			expr:     `label_join(metric, "zone", ",", "pod")`,
+			expected: true,
+		},
 		{
 			name:     "rate preserves",
 			expr:     `rate(metric[5m])`,