From 301b6e257ad7eded5a42e4be9b6e177f50ec44c7 Mon Sep 17 00:00:00 2001 From: Harshit Gangal Date: Tue, 4 Jun 2024 18:22:39 +0530 Subject: [PATCH] fix: order by subquery planning Signed-off-by: Harshit Gangal --- .../operators/horizon_expanding.go | 16 +++--- .../vtgate/planbuilder/operators/operator.go | 2 + .../planbuilder/operators/query_planning.go | 12 +++- .../planbuilder/testdata/select_cases.json | 56 ++++++++++++++----- 4 files changed, 63 insertions(+), 23 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operators/horizon_expanding.go b/go/vt/vtgate/planbuilder/operators/horizon_expanding.go index fc980038f7f..7127b7752ec 100644 --- a/go/vt/vtgate/planbuilder/operators/horizon_expanding.go +++ b/go/vt/vtgate/planbuilder/operators/horizon_expanding.go @@ -115,9 +115,9 @@ 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 { @@ -125,23 +125,23 @@ func expandOrderBy(ctx *plancontext.PlanningContext, op Operator, qp *QueryProje 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, diff --git a/go/vt/vtgate/planbuilder/operators/operator.go b/go/vt/vtgate/planbuilder/operators/operator.go index f1a38974c93..52a4672e28d 100644 --- a/go/vt/vtgate/planbuilder/operators/operator.go +++ b/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 } ) diff --git a/go/vt/vtgate/planbuilder/operators/query_planning.go b/go/vt/vtgate/planbuilder/operators/query_planning.go index f2625bcb90b..289460f154a 100644 --- a/go/vt/vtgate/planbuilder/operators/query_planning.go +++ b/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") diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json index c6a91350d89..8c2092bee87 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.json +++ b/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`" + } + ] } ] },