Fix subquery over avg (#499)

* 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]>

* Fix lint

Signed-off-by: Filip Petkovski <[email protected]>

---------

Signed-off-by: Filip Petkovski <[email protected]>

Author: fpetkovski

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])`},

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) {