From 5c578d074d33344e5314c27742d5c36750e96fe5 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Wed, 1 Nov 2023 16:29:28 +0100
Subject: [PATCH] refactor: prepare code to so it's easy to add avg

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 go/vt/vtgate/planbuilder/operators/phases.go | 49 ++++++++++++--------
 1 file changed, 29 insertions(+), 20 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/phases.go b/go/vt/vtgate/planbuilder/operators/phases.go
index ba13a828d0b..80f5ce71512 100644
--- a/go/vt/vtgate/planbuilder/operators/phases.go
+++ b/go/vt/vtgate/planbuilder/operators/phases.go
@@ -84,7 +84,7 @@ func (p Phase) act(ctx *plancontext.PlanningContext, op ops.Operator) (ops.Opera
 	case delegateAggregation:
 		return enableDelegateAggregation(ctx, op)
 	case addAggrOrdering:
-		return addOrderBysForAggregations(ctx, op)
+		return addOrderingForAllAggregations(ctx, op)
 	case cleanOutPerfDistinct:
 		return removePerformanceDistinctAboveRoute(ctx, op)
 	case subquerySettling:
@@ -120,41 +120,50 @@ func enableDelegateAggregation(ctx *plancontext.PlanningContext, op ops.Operator
 	return addColumnsToInput(ctx, op)
 }
 
-func addOrderBysForAggregations(ctx *plancontext.PlanningContext, root ops.Operator) (ops.Operator, error) {
+// addOrderingForAllAggregations is run we have pushed down Aggregators as far down as possible.
+func addOrderingForAllAggregations(ctx *plancontext.PlanningContext, root ops.Operator) (ops.Operator, error) {
 	visitor := func(in ops.Operator, _ semantics.TableSet, isRoot bool) (ops.Operator, *rewrite.ApplyResult, error) {
 		aggrOp, ok := in.(*Aggregator)
 		if !ok {
 			return in, rewrite.SameTree, nil
 		}
 
+		var res *rewrite.ApplyResult
+
 		requireOrdering, err := needsOrdering(ctx, aggrOp)
 		if err != nil {
 			return nil, nil, err
 		}
-		if !requireOrdering {
-			return in, rewrite.SameTree, nil
-		}
-		orderBys := slice.Map(aggrOp.Grouping, func(from GroupBy) ops.OrderBy {
-			return from.AsOrderBy()
-		})
-		if aggrOp.DistinctExpr != nil {
-			orderBys = append(orderBys, ops.OrderBy{
-				Inner: &sqlparser.Order{
-					Expr: aggrOp.DistinctExpr,
-				},
-				SimplifiedExpr: aggrOp.DistinctExpr,
-			})
-		}
-		aggrOp.Source = &Ordering{
-			Source: aggrOp.Source,
-			Order:  orderBys,
+
+		if requireOrdering {
+			addOrderingFor(aggrOp)
+			res = rewrite.NewTree("added ordering before aggregation", in)
 		}
-		return in, rewrite.NewTree("added ordering before aggregation", in), nil
+
+		return in, res, nil
 	}
 
 	return rewrite.BottomUp(root, TableID, visitor, stopAtRoute)
 }
 
+func addOrderingFor(aggrOp *Aggregator) {
+	orderBys := slice.Map(aggrOp.Grouping, func(from GroupBy) ops.OrderBy {
+		return from.AsOrderBy()
+	})
+	if aggrOp.DistinctExpr != nil {
+		orderBys = append(orderBys, ops.OrderBy{
+			Inner: &sqlparser.Order{
+				Expr: aggrOp.DistinctExpr,
+			},
+			SimplifiedExpr: aggrOp.DistinctExpr,
+		})
+	}
+	aggrOp.Source = &Ordering{
+		Source: aggrOp.Source,
+		Order:  orderBys,
+	}
+}
+
 func needsOrdering(ctx *plancontext.PlanningContext, in *Aggregator) (bool, error) {
 	requiredOrder := slice.Map(in.Grouping, func(from GroupBy) sqlparser.Expr {
 		return from.SimplifiedExpr