Fix subquery planning having an aggregation that is used in order by …

…as long as we can merge it all into a single route (#16402)

Signed-off-by: Manan Gupta <[email protected]>

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

@@ -138,13 +138,18 @@ func expandOrderBy(ctx *plancontext.PlanningContext, op Operator, qp *QueryProje
 			continue
 		}
 
-		// If the operator is not a projection, we cannot handle subqueries with aggregation
+		// If the operator is not a projection, we cannot handle subqueries with aggregation if we are unable to push everything into a single route.
 		if !ok {
-			panic(vterrors.VT12001("subquery with aggregation in order by"))
+			ctx.SemTable.NotSingleRouteErr = vterrors.VT12001("subquery with aggregation in order by")
+			return &Ordering{
+				Source: op,
+				Order:  qp.OrderExprs,
+			}
+		} else {
+			// Add the new subquery expression to the projection
+			proj.addSubqueryExpr(ctx, aeWrap(newExpr), newExpr, subqs...)
 		}
 
-		// Add the new subquery expression to the projection
-		proj.addSubqueryExpr(ctx, aeWrap(newExpr), newExpr, subqs...)
 		// Replace the original order expression with the new expression containing subqueries
 		newOrder = append(newOrder, OrderBy{
 			Inner: &sqlparser.Order{

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

@@ -736,6 +736,33 @@
       ]
     }
   },
+  {
+    "comment": "subquery with an aggregation in order by that can be merged into a single route",
+    "query": "select col, trim((select user_name from user where id = 3)) val from user_extra where user_id = 3 group by col order by val",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select col, trim((select user_name from user where id = 3)) val from user_extra where user_id = 3 group by col order by val",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select col, trim((select user_name from `user` where 1 != 1)) as val from user_extra where 1 != 1 group by col",
+        "Query": "select col, trim((select user_name from `user` where id = 3)) as val from user_extra where user_id = 3 group by col order by trim((select `user`.user_name from `user` where `user`.id = 3)) asc",
+        "Table": "user_extra",
+        "Values": [
+          "3"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user",
+        "user.user_extra"
+      ]
+    }
+  },
   {
     "comment": "Jumbled references",
     "query": "select user.col, user_extra.id, user.col2 from user join user_extra",

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

@@ -49,6 +49,11 @@
     "query": "update user set id = 1 where id = 1",
     "plan": "VT12001: unsupported: you cannot UPDATE primary vindex columns; invalid update on vindex: user_index"
   },
+  {
+    "comment": "subquery with an aggregation in order by that cannot be merged into a single route",
+    "query": "select col, trim((select user_name from user where col = 'a')) val from user_extra where user_id = 3 group by col order by val",
+    "plan": "VT12001: unsupported: subquery with aggregation in order by"
+  },
   {
     "comment": "update change in multicol vindex column",
     "query": "update multicol_tbl set colc = 5, colb = 4 where cola = 1 and colb = 2",