Distribute avg aggregation

engine/distributed_test.go

@@ -206,7 +206,8 @@ func TestDistributedAggregations(t *testing.T) {
 		expectFallback bool
 	}{
 		{name: "sum", query: `sum by (pod) (bar)`},
-		{name: "avg", query: `avg by (pod) (bar)`},
+		{name: "avg", query: `avg(bar)`},
+		{name: "avg with grouping", query: `avg by (pod) (bar)`},
 		{name: "count", query: `count by (pod) (bar)`},
 		{name: "count by __name__", query: `count by (__name__) ({__name__=~".+"})`},
 		{name: "group", query: `group by (pod) (bar)`},

engine/engine.go

@@ -142,6 +142,7 @@ type distributedEngine struct {
 func NewDistributedEngine(opts Opts, endpoints api.RemoteEndpoints) v1.QueryEngine {
 	opts.LogicalOptimizers = []logicalplan.Optimizer{
 		logicalplan.PassthroughOptimizer{Endpoints: endpoints},
+		logicalplan.DistributeAvgOptimizer{},
 		logicalplan.DistributedExecutionOptimizer{Endpoints: endpoints},
 	}
 

logicalplan/distribute.go

@@ -129,6 +129,7 @@ var distributiveAggregations = map[parser.ItemType]struct{}{
 	parser.SUM:     {},
 	parser.MIN:     {},
 	parser.MAX:     {},
+	parser.AVG:     {},
 	parser.GROUP:   {},
 	parser.COUNT:   {},
 	parser.BOTTOMK: {},

logicalplan/distribute_avg.go

@@ -0,0 +1,44 @@
+package logicalplan
+
+import (
+	"github.com/prometheus/prometheus/promql/parser"
+
+	"github.com/thanos-io/promql-engine/query"
+)
+
+// DistributeAvgOptimizer rewrites an AVG aggregation into a SUM/COUNT aggregation so that
+// it can be executed in a distributed manner.
+type DistributeAvgOptimizer struct{}
+
+func (r DistributeAvgOptimizer) Optimize(plan parser.Expr, _ *query.Options) parser.Expr {
+	TraverseBottomUp(nil, &plan, func(parent, current *parser.Expr) (stop bool) {
+		// If the current operation is not distributive, stop the traversal.
+		if !isDistributive(current) {
+			return true
+		}
+
+		// If the current node is avg(), distribute the operation and
+		// stop the traversal.
+		if aggr, ok := (*current).(*parser.AggregateExpr); ok {
+			if aggr.Op != parser.AVG {
+				return true
+			}
+
+			sum := *(*current).(*parser.AggregateExpr)
+			sum.Op = parser.SUM
+			count := *(*current).(*parser.AggregateExpr)
+			count.Op = parser.COUNT
+			*current = &parser.BinaryExpr{
+				Op:             parser.DIV,
+				LHS:            &sum,
+				RHS:            &count,
+				VectorMatching: &parser.VectorMatching{},
+			}
+			return true
+		}
+
+		// If the parent operation is distributive, continue the traversal.
+		return !isDistributive(parent)
+	})
+	return plan
+}

logicalplan/distribute_test.go

@@ -55,12 +55,48 @@ sum without (pod, region) (
 		},
 		{
 			name: "avg",
+			expr: `avg(http_requests_total)`,
+			expected: `
+sum(
+  dedup(
+    remote(sum by (region) (http_requests_total)),
+    remote(sum by (region) (http_requests_total))
+  )
+) /
+sum(
+  dedup(
+    remote(count by (region) (http_requests_total)),
+    remote(count by (region) (http_requests_total))
+  )
+)`,
+		},
+		{
+			name: "avg with grouping",
 			expr: `avg by (pod) (http_requests_total)`,
 			expected: `
-avg by (pod) (
+sum by (pod) (
   dedup(
-    remote(http_requests_total),
-    remote(http_requests_total)
+    remote(sum by (pod, region) (http_requests_total)),
+    remote(sum by (pod, region) (http_requests_total))
+  )
+) /
+sum by (pod) (
+  dedup(
+    remote(count by (pod, region) (http_requests_total)),
+    remote(count by (pod, region) (http_requests_total))
+  )
+)`,
+		},
+		{
+			name: "avg with prior aggregation",
+			expr: `avg by (pod) (sum by (pod) (http_requests_total))`,
+			expected: `
+avg by (pod) (
+  sum by (pod) (
+	dedup(
+      remote(sum by (pod, region) (http_requests_total)),
+	  remote(sum by (pod, region) (http_requests_total))
+	)
   )
 )`,
 		},
@@ -90,9 +126,9 @@ max by (pod) (
 		},
 		{
 			name: "unsupported aggregation in the operand path",
-			expr: `max by (pod) (sort(avg(http_requests_total)))`,
+			expr: `max by (pod) (sort(quantile(0.9, http_requests_total)))`,
 			expected: `
-max by (pod) (sort(avg(
+max by (pod) (sort(quantile(0.9,
   dedup(
     remote(http_requests_total),
     remote(http_requests_total)
@@ -193,7 +229,10 @@ remote(sum by (pod, region) (rate(http_requests_total[2m]) * 60))))`,
 		newEngineMock(math.MinInt64, math.MinInt64, []labels.Labels{labels.FromStrings("region", "east"), labels.FromStrings("region", "south")}),
 		newEngineMock(math.MinInt64, math.MinInt64, []labels.Labels{labels.FromStrings("region", "west")}),
 	}
-	optimizers := []Optimizer{DistributedExecutionOptimizer{Endpoints: api.NewStaticEndpoints(engines)}}
+	optimizers := []Optimizer{
+		DistributeAvgOptimizer{},
+		DistributedExecutionOptimizer{Endpoints: api.NewStaticEndpoints(engines)},
+	}
 	replacements := map[string]*regexp.Regexp{
 		" ": spaces,
 		"(": openParenthesis,
@@ -303,7 +342,10 @@ dedup(
 				newEngineMock(tcase.firstEngineOpts.mint(), tcase.firstEngineOpts.maxt(), []labels.Labels{labels.FromStrings("region", "east")}),
 				newEngineMock(tcase.secondEngineOpts.mint(), tcase.secondEngineOpts.maxt(), []labels.Labels{labels.FromStrings("region", "east")}),
 			}
-			optimizers := []Optimizer{DistributedExecutionOptimizer{Endpoints: api.NewStaticEndpoints(engines)}}
+			optimizers := []Optimizer{
+				DistributeAvgOptimizer{},
+				DistributedExecutionOptimizer{Endpoints: api.NewStaticEndpoints(engines)},
+			}
 
 			expr, err := parser.ParseExpr(tcase.expr)
 			testutil.Ok(t, err)