feat: rewrite subqueries that have been merged and are used in ordering

Signed-off-by: Andres Taylor <andres@planetscale.com>
vitessio · Jun 6, 2024 · 0923bb0 · 0923bb0
1 parent 59bad72
commit 0923bb0
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 8 deletions.
diff --git 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
@@ -102,13 +102,38 @@ func settleSubqueries(ctx *plancontext.PlanningContext, op Operator) Operator {
  aggr.Original.Expr = newExpr
  }
  }
+ case *Ordering:
+ op.settleOrderingExpressions(ctx)
  }
  return op, NoRewrite
  }
 
  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
@@ -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`"
  }
  ]