Fix subquery over avg

The refactor in #477 introduces a bug where subquery on top of avg produces wrong results. This is because the parents are calculated before the query is rewritten, so the new nodes end up with no parents. Signed-off-by: Filip Petkovski <[email protected]>
thanos-io · Dec 3, 2024 · 1bfb5f9 · 1bfb5f9
1 parent 097e6e9
commit 1bfb5f9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/engine/distributed_test.go b/engine/distributed_test.go
@@ -232,6 +232,8 @@ func TestDistributedAggregations(t *testing.T) {
 		{name: "absent for existing metric with aggregation", query: `sum(absent(foo))`},
 		{name: "absent for existing metric", query: `absent(bar{pod="nginx-1"})`},
 		{name: "absent for existing metric with aggregation", query: `sum(absent(bar{pod="nginx-1"}))`},
+		{name: "subquery with sum/count", query: `max_over_time((sum(bar)/count(bar))[30s:15s])`},
+		{name: "subquery with avg", query: `max_over_time(avg(bar)[30s:15s])`},
 		{name: "subquery with window within engine range", query: `max_over_time(sum_over_time(bar[30s])[30s:15s])`},
 		{name: "subquery with window outside of engine range", query: `max_over_time(sum_over_time(bar[1m])[10m:1m])`},
 		{name: "subquery with misaligned ranges", rangeStart: time.Unix(7, 0), query: `max_over_time(sum(bar)[30s:15s])`},

diff --git a/logicalplan/distribute.go b/logicalplan/distribute.go
@@ -178,13 +178,6 @@ func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options)
 	}
 	minEngineOverlap := labelRanges.minOverlap()
 
-	// 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
-	})
-
 	// Preprocess rewrite distributable averages as sum/count
 	var warns = annotations.New()
 	TraverseBottomUp(nil, &plan, func(parent, current *Node) (stop bool) {
@@ -217,6 +210,12 @@ func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options)
 		return !(isDistributive(parent, m.SkipBinaryPushdown, engineLabels, warns) || isAvgAggregation(parent))
 	})
 
+	// 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) {
 		// If the current operation is not distributive, stop the traversal.
 		if !isDistributive(current, m.SkipBinaryPushdown, engineLabels, warns) {