From 3ae8bece2c1b398a54a0e0017fc61b2c44c05f0b Mon Sep 17 00:00:00 2001
From: "vitess-bot[bot]" <108069721+vitess-bot[bot]@users.noreply.github.com>
Date: Tue, 16 Jul 2024 17:00:09 +0530
Subject: [PATCH] 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 <manan@planetscale.com>
---
 .../operators/horizon_expanding.go            | 13 ++++++---
 .../planbuilder/testdata/select_cases.json    | 27 +++++++++++++++++++
 .../testdata/unsupported_cases.json           |  5 ++++
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/horizon_expanding.go b/go/vt/vtgate/planbuilder/operators/horizon_expanding.go
index dd40dbeb918..fd3d992d0ca 100644
--- a/go/vt/vtgate/planbuilder/operators/horizon_expanding.go
+++ b/go/vt/vtgate/planbuilder/operators/horizon_expanding.go
@@ -141,13 +141,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{
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json
index 97806e67a9d..b22d86a7c1b 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.json
+++ b/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",
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
index 251af436d27..9bd5c1e6d99 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
@@ -24,6 +24,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",