fix: order by subquery planning

Signed-off-by: Harshit Gangal <harshit@planetscale.com>

go/vt/vtgate/planbuilder/operators/horizon_expanding.go

@@ -115,33 +115,33 @@ func expandSelectHorizon(ctx *plancontext.PlanningContext, horizon *Horizon, sel
 }
 
 func expandOrderBy(ctx *plancontext.PlanningContext, op Operator, qp *QueryProjection) Operator {
-	proj := newAliasedProjection(op)
 	var newOrder []OrderBy
 	sqc := &SubQueryBuilder{}
+	proj, ok := op.(*Projection)
 	for _, expr := range qp.OrderExprs {
 		newExpr, subqs := sqc.pullOutValueSubqueries(ctx, expr.SimplifiedExpr, TableID(op), false)
 		if newExpr == nil {
 			// no subqueries found, let's move on
 			newOrder = append(newOrder, expr)
 			continue
 		}
+		if !ok {
+			panic(vterrors.VT12001("subquery with aggregation in order by"))
+		}
 		proj.addSubqueryExpr(aeWrap(newExpr), newExpr, subqs...)
 		newOrder = append(newOrder, OrderBy{
 			Inner: &sqlparser.Order{
 				Expr:      newExpr,
 				Direction: expr.Inner.Direction,
 			},
-			SimplifiedExpr: newExpr,
+			SimplifiedExpr:     newExpr,
+			SubQueryExpression: subqs,
 		})
 
 	}
-
-	if len(proj.Columns.GetColumns()) > 0 {
-		// if we had to project columns for the ordering,
-		// we need the projection as source
-		op = proj
+	if proj != nil {
+		proj.Source = sqc.getRootOperator(proj.Source, nil)
 	}
-
 	return &Ordering{
 		Source: op,
 		Order:  newOrder,

go/vt/vtgate/planbuilder/operators/operator.go

@@ -84,6 +84,8 @@ type (
 
 		// See GroupBy#SimplifiedExpr for more details about this
 		SimplifiedExpr sqlparser.Expr
+
+		SubQueryExpression []*SubQuery
 	}
 )
 

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -371,6 +371,15 @@ func tryPushOrdering(ctx *plancontext.PlanningContext, in *Ordering) (Operator,
 				return in, NoRewrite
 			}
 		}
+		ap, ok := src.Columns.(AliasedProjections)
+		if !ok {
+			return in, NoRewrite
+		}
+		for _, projExpr := range ap {
+			if projExpr.Info != nil {
+				return in, NoRewrite
+			}
+		}
 		return Swap(in, src, "push ordering under projection")
 	case *Aggregator:
 		if !src.QP.AlignGroupByAndOrderBy(ctx) && !overlaps(ctx, in.Order, src.Grouping) {
@@ -390,11 +399,12 @@ func tryPushOrdering(ctx *plancontext.PlanningContext, in *Ordering) (Operator,
 		return src, Rewrote("push ordering into outer side of subquery")
 	case *SubQuery:
 		outerTableID := TableID(src.Outer)
-		for _, order := range in.Order {
+		for idx, order := range in.Order {
 			deps := ctx.SemTable.RecursiveDeps(order.Inner.Expr)
 			if !deps.IsSolvedBy(outerTableID) {
 				return in, NoRewrite
 			}
+			in.Order[idx].Inner.Expr = rewriteColNameToArgument(ctx, order.Inner.Expr, order.SubQueryExpression, order.SubQueryExpression...)
 		}
 		src.Outer, in.Source = in, src.Outer
 		return src, Rewrote("push ordering into outer side of subquery")

go/vt/vtgate/planbuilder/testdata/select_cases.json

@@ -2282,12 +2282,12 @@
         "ColumnNames": [
           "0:a"
         ],
-        "Columns": "1",
+        "Columns": "0",
         "Inputs": [
           {
             "OperatorType": "Join",
             "Variant": "Join",
-            "JoinColumnIndexes": "L:0",
+            "JoinColumnIndexes": "L:0,L:1",
             "TableName": "`user`_user_extra",
             "Inputs": [
               {
@@ -2309,24 +2309,52 @@
                           "Name": "user",
                           "Sharded": true
                         },
-                        "FieldQuery": "select col from `user` where 1 != 1",
-                        "Query": "select col from `user` limit 1",
+                        "FieldQuery": "select `user`.col from `user` where 1 != 1",
+                        "Query": "select `user`.col from `user` limit 1",
                         "Table": "`user`"
                       }
                     ]
                   },
                   {
                     "InputName": "Outer",
-                    "OperatorType": "Route",
-                    "Variant": "Scatter",
-                    "Keyspace": {
-                      "Name": "user",
-                      "Sharded": true
-                    },
-                    "FieldQuery": "select :__sq1 as __sq1, weight_string(:__sq1) from `user` where 1 != 1",
-                    "OrderBy": "(0|1) ASC",
-                    "Query": "select :__sq1 as __sq1, weight_string(:__sq1) from `user` order by __sq1 asc",
-                    "Table": "`user`"
+                    "OperatorType": "UncorrelatedSubquery",
+                    "Variant": "PulloutValue",
+                    "PulloutVars": [
+                      "__sq1"
+                    ],
+                    "Inputs": [
+                      {
+                        "InputName": "SubQuery",
+                        "OperatorType": "Limit",
+                        "Count": "1",
+                        "Inputs": [
+                          {
+                            "OperatorType": "Route",
+                            "Variant": "Scatter",
+                            "Keyspace": {
+                              "Name": "user",
+                              "Sharded": true
+                            },
+                            "FieldQuery": "select col from `user` where 1 != 1",
+                            "Query": "select col from `user` limit 1",
+                            "Table": "`user`"
+                          }
+                        ]
+                      },
+                      {
+                        "InputName": "Outer",
+                        "OperatorType": "Route",
+                        "Variant": "Scatter",
+                        "Keyspace": {
+                          "Name": "user",
+                          "Sharded": true
+                        },
+                        "FieldQuery": "select :__sq1 as a, :__sq1 as __sq1, weight_string(:__sq1) from `user` where 1 != 1",
+                        "OrderBy": "(1|2) ASC",
+                        "Query": "select :__sq1 as a, :__sq1 as __sq1, weight_string(:__sq1) from `user` order by :__sq1 asc",
+                        "Table": "`user`"
+                      }
+                    ]
                   }
                 ]
               },