[release-18.0] planner: Handle ORDER BY inside derived tables (#16353) (

#16358)

Signed-off-by: Andres Taylor <andres@planetscale.com>
Co-authored-by: Andres Taylor <andres@planetscale.com>

go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go

@@ -402,7 +402,12 @@ func TestOrderByCount(t *testing.T) {
 
 	mcmp.Exec("insert into t9(id1, id2, id3) values(1, '1', '1'), (2, '2', '2'), (3, '2', '2'), (4, '3', '3'), (5, '3', '3'), (6, '3', '3')")
 
-	mcmp.AssertMatches("SELECT t9.id2 FROM t9 GROUP BY t9.id2 ORDER BY COUNT(t9.id2) DESC", `[[VARCHAR("3")] [VARCHAR("2")] [VARCHAR("1")]]`)
+	mcmp.Exec("SELECT t9.id2 FROM t9 GROUP BY t9.id2 ORDER BY COUNT(t9.id2) DESC")
+	version, err := cluster.GetMajorVersion("vtgate")
+	require.NoError(t, err)
+	if version == 18 {
+		mcmp.Exec("select COUNT(*) from (select 1 as one FROM t9 WHERE id3 = 3 ORDER BY id1 DESC LIMIT 3 OFFSET 0) subquery_for_count")
+	}
 }
 
 func TestAggregateAnyValue(t *testing.T) {

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

@@ -79,16 +79,23 @@ func expandUnionHorizon(ctx *plancontext.PlanningContext, horizon *Horizon, unio
 }
 
 func expandSelectHorizon(ctx *plancontext.PlanningContext, horizon *Horizon, sel *sqlparser.Select) (ops.Operator, *rewrite.ApplyResult, error) {
-	op, err := createProjectionFromSelect(ctx, horizon)
+	qp, err := horizon.getQP(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	qp, err := horizon.getQP(ctx)
+	if horizon.IsDerived() {
+		// if we are dealing with a derived table, we need to make sure that the ordering columns
+		// are available outside the derived table
+		for _, order := range horizon.Query.GetOrderBy() {
+			qp.addColumn(ctx, order.Expr)
+		}
+	}
+
+	op, err := createProjectionFromSelect(ctx, horizon)
 	if err != nil {
 		return nil, nil, err
 	}
-
 	var extracted []string
 	if qp.HasAggr {
 		extracted = append(extracted, "Aggregation")

go/vt/vtgate/planbuilder/operators/queryprojection.go

@@ -807,6 +807,23 @@ func (qp *QueryProjection) useGroupingOverDistinct(ctx *plancontext.PlanningCont
 	return true, nil
 }
 
+// addColumn adds a column to the QueryProjection if it is not already present
+func (qp *QueryProjection) addColumn(ctx *plancontext.PlanningContext, expr sqlparser.Expr) {
+	for _, selectExpr := range qp.SelectExprs {
+		getExpr, err := selectExpr.GetExpr()
+		if err != nil {
+			continue
+		}
+		if ctx.SemTable.EqualsExprWithDeps(getExpr, expr) {
+			return
+		}
+	}
+	qp.SelectExprs = append(qp.SelectExprs, SelectExpr{
+		Col:  aeWrap(expr),
+		Aggr: containsAggr(expr),
+	})
+}
+
 func checkForInvalidGroupingExpressions(expr sqlparser.Expr) error {
 	return sqlparser.Walk(func(node sqlparser.SQLNode) (bool, error) {
 		if _, isAggregate := node.(sqlparser.AggrFunc); isAggregate {

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

@@ -6413,6 +6413,44 @@
       ]
     }
   },
+  {
+    "comment": "Aggregation over a ORDER BY/LIMIT inside a derived table",
+    "query": "SELECT COUNT(*) FROM (SELECT 1 AS one FROM `user` WHERE `user`.`is_not_deleted` = true ORDER BY id DESC LIMIT 25 OFFSET 0) subquery_for_count",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "SELECT COUNT(*) FROM (SELECT 1 AS one FROM `user` WHERE `user`.`is_not_deleted` = true ORDER BY id DESC LIMIT 25 OFFSET 0) subquery_for_count",
+      "Instructions": {
+        "OperatorType": "Aggregate",
+        "Variant": "Scalar",
+        "Aggregates": "count_star(0) AS count(*)",
+        "ResultColumns": 1,
+        "Inputs": [
+          {
+            "OperatorType": "Limit",
+            "Count": "INT64(25)",
+            "Offset": "INT64(0)",
+            "Inputs": [
+              {
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select 1, id, weight_string(id) from (select 1 as one, id from `user` where 1 != 1) as subquery_for_count where 1 != 1",
+                "OrderBy": "(1|2) DESC",
+                "Query": "select 1, id, weight_string(id) from (select 1 as one, id from `user` where `user`.is_not_deleted = true) as subquery_for_count order by id desc limit :__upper_limit",
+                "Table": "`user`"
+              }
+            ]
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
   {
     "comment": "baz in the HAVING clause can't be accessed because of the GROUP BY",
     "query": "select foo, count(bar) as x from user group by foo having baz > avg(baz) order by x",