wip

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
vitessio · Jun 5, 2024 · 5036a2f · 5036a2f
1 parent 4d23f76
commit 5036a2f
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 13 deletions.
diff --git a/go/vt/vtgate/planbuilder/operators/aggregation_pushing.go b/go/vt/vtgate/planbuilder/operators/aggregation_pushing.go
@@ -107,14 +107,13 @@ func pushAggregationThroughSubquery(
 			}
 			pushedAggr.addColumnWithoutPushing(ctx, aeWrap(colName), true)
 		}
+		for _, aggregation := range pushedAggr.Aggregations {
+			aggregation.Original.Expr = subQuery.rewriteColNameToArgument(aggregation.Original.Expr)
+		}
 	}
 
 	src.Outer = pushedAggr
 
-	for _, aggregation := range pushedAggr.Aggregations {
-		aggregation.Original.Expr = rewriteColNameToArgument(ctx, aggregation.Original.Expr, aggregation.SubQueryExpression, src.Inner...)
-	}
-
 	if !rootAggr.Original {
 		return src, Rewrote("push Aggregation under subquery - keep original")
 	}

diff --git a/go/vt/vtgate/planbuilder/operators/subquery.go b/go/vt/vtgate/planbuilder/operators/subquery.go
@@ -31,7 +31,7 @@ import (
 
 // SubQuery represents a subquery used for filtering rows in an
 // outer query through a join.
-type /**/ SubQuery struct {
+type SubQuery struct {
 	// Fields filled in at the time of construction:
 	Outer             Operator             // Outer query operator.
 	Subquery          Operator             // Subquery operator.
@@ -53,7 +53,7 @@ type /**/ SubQuery struct {
 	// We use this information to fail the planning if we are unable to merge the subquery with a route.
 	correlated bool
 
-	IsProjection bool
+	IsArgument bool
 }
 
 func (sq *SubQuery) planOffsets(ctx *plancontext.PlanningContext) Operator {
@@ -156,8 +156,8 @@ func (sq *SubQuery) SetInputs(inputs []Operator) {
 
 func (sq *SubQuery) ShortDescription() string {
 	var typ string
-	if sq.IsProjection {
-		typ = "PROJ"
+	if sq.IsArgument {
+		typ = "ARG"
 	} else {
 		typ = "FILTER"
 	}
@@ -210,7 +210,7 @@ func (sq *SubQuery) settle(ctx *plancontext.PlanningContext, outer Operator) Ope
 	if sq.correlated && sq.FilterType != opcode.PulloutExists {
 		panic(correlatedSubqueryErr)
 	}
-	if sq.IsProjection {
+	if sq.IsArgument {
 		if len(sq.GetMergePredicates()) > 0 {
 			// this means that we have a correlated subquery on our hands
 			panic(correlatedSubqueryErr)

diff --git a/go/vt/vtgate/planbuilder/operators/subquery_builder.go b/go/vt/vtgate/planbuilder/operators/subquery_builder.go
@@ -181,7 +181,7 @@ func createSubquery(
 		Original:         original,
 		ArgName:          argName,
 		originalSubquery: originalSq,
-		IsProjection:     isProjection,
+		IsArgument:       isProjection,
 		TopLevel:         topLevel,
 		JoinColumns:      joinCols,
 		correlated:       correlated,

diff --git a/go/vt/vtgate/planbuilder/operators/subquery_planning.go b/go/vt/vtgate/planbuilder/operators/subquery_planning.go
@@ -477,7 +477,7 @@ func tryMergeSubqueryWithOuter(ctx *plancontext.PlanningContext, subQuery *SubQu
 	if op == nil {
 		return outer, NoRewrite
 	}
-	if !subQuery.IsProjection {
+	if !subQuery.IsArgument {
 		op.Source = newFilter(outer.Source, subQuery.Original)
 	}
 	ctx.MergedSubqueries = append(ctx.MergedSubqueries, subQuery.originalSubquery)
@@ -582,7 +582,7 @@ func (s *subqueryRouteMerger) merge(ctx *plancontext.PlanningContext, inner, out
 	var src Operator
 	if isSharded {
 		src = s.outer.Source
-		if !s.subq.IsProjection {
+		if !s.subq.IsArgument {
 			src = newFilter(s.outer.Source, s.original)
 		}
 	} else {
@@ -643,7 +643,7 @@ func (s *subqueryRouteMerger) rewriteASTExpression(ctx *plancontext.PlanningCont
 		panic(err)
 	}
 
-	if s.subq.IsProjection {
+	if s.subq.IsArgument {
 		ctx.SemTable.CopySemanticInfo(s.subq.originalSubquery.Select, subqStmt)
 		s.subq.originalSubquery.Select = subqStmt
 	} else {