From 0923bb0e2110d9d8d88c5b21e6134599ad0fba98 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Thu, 6 Jun 2024 11:29:50 +0200
Subject: [PATCH] feat: rewrite subqueries that have been merged and are used
 in ordering

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 .../planbuilder/operators/queryprojection.go  | 14 ++++++++--
 .../operators/subquery_planning.go            | 27 ++++++++++++++++++-
 .../planbuilder/testdata/select_cases.json    |  9 +++----
 3 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/queryprojection.go b/go/vt/vtgate/planbuilder/operators/queryprojection.go
index cbacdd25e90..8f54a56fd90 100644
--- a/go/vt/vtgate/planbuilder/operators/queryprojection.go
+++ b/go/vt/vtgate/planbuilder/operators/queryprojection.go
@@ -269,8 +269,7 @@ func (qp *QueryProjection) addOrderBy(ctx *plancontext.PlanningContext, orderBy
 	canPushSorting := true
 	es := &expressionSet{}
 	for _, order := range orderBy {
-		if sqlparser.IsNull(order.Expr) {
-			// ORDER BY null can safely be ignored
+		if notInterestingToOrderBy(ctx, order.Expr) {
 			continue
 		}
 		if !es.add(ctx, order.Expr) {
@@ -284,6 +283,17 @@ func (qp *QueryProjection) addOrderBy(ctx *plancontext.PlanningContext, orderBy
 	}
 }
 
+func notInterestingToOrderBy(ctx *plancontext.PlanningContext, expr sqlparser.Expr) bool {
+	switch expr.(type) {
+	case *sqlparser.NullVal, *sqlparser.Literal, *sqlparser.Argument:
+		return true
+	case *sqlparser.Subquery:
+		return ctx.SemTable.RecursiveDeps(expr).IsEmpty()
+	default:
+		return false
+	}
+}
+
 func (qp *QueryProjection) calculateDistinct(ctx *plancontext.PlanningContext) {
 	if qp.Distinct && !qp.HasAggr {
 		distinct := qp.useGroupingOverDistinct(ctx)
diff --git a/go/vt/vtgate/planbuilder/operators/subquery_planning.go b/go/vt/vtgate/planbuilder/operators/subquery_planning.go
index 3ff3fa519ba..caf3a51536c 100644
--- a/go/vt/vtgate/planbuilder/operators/subquery_planning.go
+++ b/go/vt/vtgate/planbuilder/operators/subquery_planning.go
@@ -102,6 +102,8 @@ func settleSubqueries(ctx *plancontext.PlanningContext, op Operator) Operator {
 					aggr.Original.Expr = newExpr
 				}
 			}
+		case *Ordering:
+			op.settleOrderingExpressions(ctx)
 		}
 		return op, NoRewrite
 	}
@@ -109,6 +111,29 @@ func settleSubqueries(ctx *plancontext.PlanningContext, op Operator) Operator {
 	return BottomUp(op, TableID, visit, nil)
 }
 
+func (o *Ordering) settleOrderingExpressions(ctx *plancontext.PlanningContext) {
+	for idx, order := range o.Order {
+		for _, sq := range ctx.MergedSubqueries {
+			arg := ctx.GetReservedArgumentFor(sq)
+			expr := sqlparser.Rewrite(order.SimplifiedExpr, nil, func(cursor *sqlparser.Cursor) bool {
+				switch expr := cursor.Node().(type) {
+				case *sqlparser.ColName:
+					if expr.Name.String() == arg {
+						cursor.Replace(sq)
+					}
+				case *sqlparser.Argument:
+					if expr.Name == arg {
+						cursor.Replace(sq)
+					}
+				}
+
+				return true
+			})
+			o.Order[idx].SimplifiedExpr = expr.(sqlparser.Expr)
+		}
+	}
+}
+
 func mergeSubqueryExpr(ctx *plancontext.PlanningContext, pe *ProjExpr) {
 	se, ok := pe.Info.(SubQueryExpression)
 	if !ok {
@@ -132,7 +157,7 @@ func rewriteMergedSubqueryExpr(ctx *plancontext.PlanningContext, se SubQueryExpr
 							return true
 						}
 					case *sqlparser.Argument:
-						if expr.Name != sq.ArgName {
+						if expr.Name != sq. /**/ ArgName {
 							return true
 						}
 					default:
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json
index 7c35390e055..199aa3775d2 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json
@@ -2302,8 +2302,8 @@
                       "Name": "user",
                       "Sharded": true
                     },
-                    "FieldQuery": "select `user`.col from `user` where 1 != 1",
-                    "Query": "select `user`.col from `user` limit 1",
+                    "FieldQuery": "select col from `user` where 1 != 1",
+                    "Query": "select col from `user` limit 1",
                     "Table": "`user`"
                   }
                 ]
@@ -2316,9 +2316,8 @@
                   "Name": "user",
                   "Sharded": true
                 },
-                "FieldQuery": "select :__sq1 as a, :__sq1, weight_string(:__sq1) from `user` where 1 != 1",
-                "OrderBy": "(1|2) ASC",
-                "Query": "select :__sq1 as a, :__sq1, weight_string(:__sq1) from `user` order by :__sq1 asc",
+                "FieldQuery": "select :__sq1 as a from `user` where 1 != 1",
+                "Query": "select :__sq1 as a from `user`",
                 "Table": "`user`"
               }
             ]