Skip to content

Commit

Permalink
Cherry-pick of #14935
Browse files Browse the repository at this point in the history 
Signed-off-by: Andres Taylor <andres@planetscale.com>
  • Loading branch information
@systay
systay committed Jan 25, 2024
1 parent cc346e2 commit 00bef5d
Show file tree
Hide file tree
Showing 18 changed files with 485 additions and 439 deletions.
15 changes: 5 additions & 10 deletions go/vt/vtgate/planbuilder/operators/aggregation_pushing.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -432,17 +432,12 @@ var errAbortAggrPushing = fmt.Errorf("abort aggregation pushing")
func addColumnsFromLHSInJoinPredicates(ctx *plancontext.PlanningContext, rootAggr *Aggregator, join *ApplyJoin, lhs *joinPusher) error {
for _, pred := range join.JoinPredicates {
for _, bve := range pred.LHSExprs {
expr := bve.Expr
wexpr, err := rootAggr.QP.GetSimplifiedExpr(ctx, expr)
if err != nil {
return err
}
idx, found := canReuseColumn(ctx, lhs.pushed.Columns, expr, extractExpr)
idx, found := canReuseColumn(ctx, lhs.pushed.Columns, bve.Expr, extractExpr)
if !found {
idx = len(lhs.pushed.Columns)
lhs.pushed.Columns = append(lhs.pushed.Columns, aeWrap(expr))
lhs.pushed.Columns = append(lhs.pushed.Columns, aeWrap(bve.Expr))
}
_, found = canReuseColumn(ctx, lhs.pushed.Grouping, wexpr, func(by GroupBy) sqlparser.Expr {
_, found = canReuseColumn(ctx, lhs.pushed.Grouping, bve.Expr, func(by GroupBy) sqlparser.Expr {
return by.SimplifiedExpr
})

Expand All @@ -451,8 +446,8 @@ func addColumnsFromLHSInJoinPredicates(ctx *plancontext.PlanningContext, rootAgg
}

lhs.pushed.Grouping = append(lhs.pushed.Grouping, GroupBy{
Inner: expr,
SimplifiedExpr: wexpr,
Inner: bve.Expr,
SimplifiedExpr: bve.Expr,
ColOffset: idx,
WSOffset: -1,
})
Expand Down
11 changes: 3 additions & 8 deletions go/vt/vtgate/planbuilder/operators/distinct.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,16 +52,11 @@ func (d *Distinct) planOffsets(ctx *plancontext.PlanningContext) error {
return err
}
for idx, col := range columns {
e, err := d.QP.GetSimplifiedExpr(ctx, col.Expr)
if err != nil {
// ambiguous columns are not a problem for DISTINCT
e = col.Expr
}
var wsCol *int
typ, coll, _ := ctx.SemTable.TypeForExpr(e)
typ, coll, _ := ctx.SemTable.TypeForExpr(col.Expr)

if ctx.SemTable.NeedsWeightString(e) {
offset, err := d.Source.AddColumn(ctx, true, false, aeWrap(weightStringFor(e)))
if ctx.SemTable.NeedsWeightString(col.Expr) {
offset, err := d.Source.AddColumn(ctx, true, false, aeWrap(weightStringFor(col.Expr)))
if err != nil {
return err
}
Expand Down
147 changes: 12 additions & 135 deletions go/vt/vtgate/planbuilder/operators/queryprojection.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import (
"fmt"
"io"
"slices"
"sort"
"strings"

"vitess.io/vitess/go/mysql/collations"
Expand Down Expand Up @@ -216,9 +215,7 @@ func createQPFromSelect(ctx *plancontext.PlanningContext, sel *sqlparser.Select)
if err := qp.addGroupBy(ctx, sel.GroupBy); err != nil {
return nil, err
}
if err := qp.addOrderBy(ctx, sel.OrderBy); err != nil {
return nil, err
}
qp.addOrderBy(ctx, sel.OrderBy)
if !qp.HasAggr && sel.Having != nil {
qp.HasAggr = containsAggr(sel.Having.Expr)
}
Expand Down Expand Up @@ -340,10 +337,7 @@ func createQPFromUnion(ctx *plancontext.PlanningContext, union *sqlparser.Union)
return nil, err
}

err = qp.addOrderBy(ctx, union.OrderBy)
if err != nil {
return nil, err
}
qp.addOrderBy(ctx, union.OrderBy)

return qp, nil
}
Expand All @@ -365,29 +359,24 @@ func (es *expressionSet) add(ctx *plancontext.PlanningContext, e sqlparser.Expr)
return true
}

func (qp *QueryProjection) addOrderBy(ctx *plancontext.PlanningContext, orderBy sqlparser.OrderBy) error {
func (qp *QueryProjection) addOrderBy(ctx *plancontext.PlanningContext, orderBy sqlparser.OrderBy) {
canPushSorting := true
es := &expressionSet{}
for _, order := range orderBy {
simpleExpr, err := qp.GetSimplifiedExpr(ctx, order.Expr)
if err != nil {
return err
}
if sqlparser.IsNull(simpleExpr) {
if sqlparser.IsNull(order.Expr) {
// ORDER BY null can safely be ignored
continue
}
if !es.add(ctx, simpleExpr) {
if !es.add(ctx, order.Expr) {
continue
}
qp.OrderExprs = append(qp.OrderExprs, ops.OrderBy{
Inner: sqlparser.CloneRefOfOrder(order),
SimplifiedExpr: simpleExpr,
SimplifiedExpr: order.Expr,
})
canPushSorting = canPushSorting && !containsAggr(simpleExpr)
canPushSorting = canPushSorting && !containsAggr(order.Expr)
}
qp.CanPushSorting = canPushSorting
return nil
}

func (qp *QueryProjection) calculateDistinct(ctx *plancontext.PlanningContext) error {
Expand Down Expand Up @@ -436,20 +425,16 @@ func (qp *QueryProjection) addGroupBy(ctx *plancontext.PlanningContext, groupBy
es := &expressionSet{}
for _, group := range groupBy {
selectExprIdx, aliasExpr := qp.FindSelectExprIndexForExpr(ctx, group)
simpleExpr, err := qp.GetSimplifiedExpr(ctx, group)
if err != nil {
return err
}

if err = checkForInvalidGroupingExpressions(simpleExpr); err != nil {
if err := checkForInvalidGroupingExpressions(group); err != nil {
return err
}

if !es.add(ctx, simpleExpr) {
if !es.add(ctx, group) {
continue
}

groupBy := NewGroupBy(group, simpleExpr, aliasExpr)
groupBy := NewGroupBy(group, group, aliasExpr)
groupBy.InnerIndex = selectExprIdx

qp.groupByExprs = append(qp.groupByExprs, groupBy)
Expand All @@ -471,64 +456,6 @@ func (qp *QueryProjection) isExprInGroupByExprs(ctx *plancontext.PlanningContext
return false
}

// GetSimplifiedExpr takes an expression used in ORDER BY or GROUP BY, and returns an expression that is simpler to evaluate
func (qp *QueryProjection) GetSimplifiedExpr(ctx *plancontext.PlanningContext, e sqlparser.Expr) (found sqlparser.Expr, err error) {
if qp == nil {
return e, nil
}
// If the ORDER BY is against a column alias, we need to remember the expression
// behind the alias. The weightstring(.) calls needs to be done against that expression and not the alias.
// Eg - select music.foo as bar, weightstring(music.foo) from music order by bar

in, isColName := e.(*sqlparser.ColName)
if !(isColName && in.Qualifier.IsEmpty()) {
// we are only interested in unqualified column names. if it's not a column name and not unqualified, we're done
return e, nil
}

check := func(e sqlparser.Expr) error {
if found != nil && !ctx.SemTable.EqualsExprWithDeps(found, e) {
return &semantics.AmbiguousColumnError{Column: sqlparser.String(in)}
}
found = e
return nil
}

for _, selectExpr := range qp.SelectExprs {
ae, ok := selectExpr.Col.(*sqlparser.AliasedExpr)
if !ok {
continue
}
aliased := !ae.As.IsEmpty()
if aliased {
if in.Name.Equal(ae.As) {
err = check(ae.Expr)
if err != nil {
return nil, err
}
}
} else {
seCol, ok := ae.Expr.(*sqlparser.ColName)
if !ok {
continue
}
if seCol.Name.Equal(in.Name) {
// If the column name matches, we have a match, even if the table name is not listed
err = check(ae.Expr)
if err != nil {
return nil, err
}
}
}
}

if found == nil {
found = e
}

return found, nil
}

// toString should only be used for tests
func (qp *QueryProjection) toString() string {
type output struct {
Expand Down Expand Up @@ -794,46 +721,6 @@ func (qp *QueryProjection) FindSelectExprIndexForExpr(ctx *plancontext.PlanningC
return nil, nil
}

// OldAlignGroupByAndOrderBy TODO Remove once all of horizon planning is done on the operators
func (qp *QueryProjection) OldAlignGroupByAndOrderBy(ctx *plancontext.PlanningContext) {
// The ORDER BY can be performed before the OA

var newGrouping []GroupBy
if len(qp.OrderExprs) == 0 {
// The query didn't ask for any particular order, so we are free to add arbitrary ordering.
// We'll align the grouping and ordering by the output columns
newGrouping = qp.GetGrouping()
SortGrouping(newGrouping)
for _, groupBy := range newGrouping {
qp.OrderExprs = append(qp.OrderExprs, groupBy.AsOrderBy())
}
} else {
// Here we align the GROUP BY and ORDER BY.
// First step is to make sure that the GROUP BY is in the same order as the ORDER BY
used := make([]bool, len(qp.groupByExprs))
for _, orderExpr := range qp.OrderExprs {
for i, groupingExpr := range qp.groupByExprs {
if !used[i] && ctx.SemTable.EqualsExpr(groupingExpr.SimplifiedExpr, orderExpr.SimplifiedExpr) {
newGrouping = append(newGrouping, groupingExpr)
used[i] = true
}
}
}
if len(newGrouping) != len(qp.groupByExprs) {
// we are missing some groupings. We need to add them both to the new groupings list, but also to the ORDER BY
for i, added := range used {
if !added {
groupBy := qp.groupByExprs[i]
newGrouping = append(newGrouping, groupBy)
qp.OrderExprs = append(qp.OrderExprs, groupBy.AsOrderBy())
}
}
}
}

qp.groupByExprs = newGrouping
}

// AlignGroupByAndOrderBy aligns the group by and order by columns, so they are in the same order
// The GROUP BY clause is a set - the order between the elements does not make any difference,
// so we can simply re-arrange the column order
Expand Down Expand Up @@ -903,18 +790,14 @@ func (qp *QueryProjection) useGroupingOverDistinct(ctx *plancontext.PlanningCont
// not an alias Expr, cannot continue forward.
return false, nil
}
sExpr, err := qp.GetSimplifiedExpr(ctx, ae.Expr)
if err != nil {
return false, err
}
// check if the grouping already exists on that column.
found := slices.IndexFunc(qp.groupByExprs, func(gb GroupBy) bool {
return ctx.SemTable.EqualsExprWithDeps(gb.SimplifiedExpr, sExpr)
return ctx.SemTable.EqualsExprWithDeps(gb.SimplifiedExpr, ae.Expr)
})
if found != -1 {
continue
}
groupBy := NewGroupBy(ae.Expr, sExpr, ae)
groupBy := NewGroupBy(ae.Expr, ae.Expr, ae)
selectExprIdx := idx
groupBy.InnerIndex = &selectExprIdx

Expand All @@ -938,12 +821,6 @@ func checkForInvalidGroupingExpressions(expr sqlparser.Expr) error {
}, expr)
}

func SortGrouping(a []GroupBy) {
sort.Slice(a, func(i, j int) bool {
return CompareRefInt(a[i].InnerIndex, a[j].InnerIndex)
})
}

// CompareRefInt compares two references of integers.
// In case either one is nil, it is considered to be smaller
func CompareRefInt(a *int, b *int) bool {
Expand Down
44 changes: 34 additions & 10 deletions go/vt/vtgate/planbuilder/testdata/aggr_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1054,8 +1054,8 @@
"Name": "user",
"Sharded": true
},
"FieldQuery": "select `user`.col1 as a from `user` where 1 != 1 group by a collate utf8_general_ci",
"Query": "select `user`.col1 as a from `user` where `user`.id = 5 group by a collate utf8_general_ci",
"FieldQuery": "select `user`.col1 as a from `user` where 1 != 1 group by `user`.col1 collate utf8_general_ci",
"Query": "select `user`.col1 as a from `user` where `user`.id = 5 group by `user`.col1 collate utf8_general_ci",
"Table": "`user`",
"Values": [
"INT64(5)"
Expand Down Expand Up @@ -1602,9 +1602,9 @@
"Name": "user",
"Sharded": true
},
"FieldQuery": "select lower(textcol1) as v, count(*), weight_string(lower(textcol1)) from `user` where 1 != 1 group by v, weight_string(lower(textcol1))",
"FieldQuery": "select lower(textcol1) as v, count(*), weight_string(lower(textcol1)) from `user` where 1 != 1 group by lower(textcol1), weight_string(lower(textcol1))",
"OrderBy": "(0|2) ASC",
"Query": "select lower(textcol1) as v, count(*), weight_string(lower(textcol1)) from `user` group by v, weight_string(lower(textcol1)) order by v asc",
"Query": "select lower(textcol1) as v, count(*), weight_string(lower(textcol1)) from `user` group by lower(textcol1), weight_string(lower(textcol1)) order by lower(textcol1) asc",
"Table": "`user`"
}
]
Expand Down Expand Up @@ -1634,9 +1634,9 @@
"Name": "user",
"Sharded": true
},
"FieldQuery": "select char_length(texcol1) as a, count(*), weight_string(char_length(texcol1)) from `user` where 1 != 1 group by a, weight_string(char_length(texcol1))",
"FieldQuery": "select char_length(texcol1) as a, count(*), weight_string(char_length(texcol1)) from `user` where 1 != 1 group by char_length(texcol1), weight_string(char_length(texcol1))",
"OrderBy": "(0|2) ASC",
"Query": "select char_length(texcol1) as a, count(*), weight_string(char_length(texcol1)) from `user` group by a, weight_string(char_length(texcol1)) order by a asc",
"Query": "select char_length(texcol1) as a, count(*), weight_string(char_length(texcol1)) from `user` group by char_length(texcol1), weight_string(char_length(texcol1)) order by char_length(texcol1) asc",
"Table": "`user`"
}
]
Expand Down Expand Up @@ -2638,9 +2638,9 @@
"Name": "user",
"Sharded": true
},
"FieldQuery": "select ascii(val1) as a, count(*), weight_string(ascii(val1)) from `user` where 1 != 1 group by a, weight_string(ascii(val1))",
"FieldQuery": "select ascii(val1) as a, count(*), weight_string(ascii(val1)) from `user` where 1 != 1 group by ascii(val1), weight_string(ascii(val1))",
"OrderBy": "(0|2) ASC",
"Query": "select ascii(val1) as a, count(*), weight_string(ascii(val1)) from `user` group by a, weight_string(ascii(val1)) order by a asc",
"Query": "select ascii(val1) as a, count(*), weight_string(ascii(val1)) from `user` group by ascii(val1), weight_string(ascii(val1)) order by ascii(val1) asc",
"Table": "`user`"
}
]
Expand Down Expand Up @@ -3356,9 +3356,9 @@
"Name": "user",
"Sharded": true
},
"FieldQuery": "select `user`.id, weight_string(`user`.id) from `user` where 1 != 1 group by id, weight_string(`user`.id)",
"FieldQuery": "select `user`.id, weight_string(`user`.id) from `user` where 1 != 1 group by `user`.id, weight_string(`user`.id)",
"OrderBy": "(0|1) ASC",
"Query": "select `user`.id, weight_string(`user`.id) from `user` group by id, weight_string(`user`.id) order by id asc",
"Query": "select `user`.id, weight_string(`user`.id) from `user` group by `user`.id, weight_string(`user`.id) order by `user`.id asc",
"Table": "`user`"
},
{
Expand Down Expand Up @@ -6371,5 +6371,29 @@
"user.user"
]
}
},
{
"comment": "valid but slightly confusing query should work - col in the order by should not get expanded to the column alias col",
"query": "select id, from_unixtime(min(col)) as col from user group by id order by min(col)",
"plan": {
"QueryType": "SELECT",
"Original": "select id, from_unixtime(min(col)) as col from user group by id order by min(col)",
"Instructions": {
"OperatorType": "Route",
"Variant": "Scatter",
"Keyspace": {
"Name": "user",
"Sharded": true
},
"FieldQuery": "select id, from_unixtime(min(col)) as col, min(col) from `user` where 1 != 1 group by id",
"OrderBy": "2 ASC",
"Query": "select id, from_unixtime(min(col)) as col, min(col) from `user` group by id order by min(col) asc",
"ResultColumns": 2,
"Table": "`user`"
},
"TablesUsed": [
"user.user"
]
}
}
]
2 changes: 1 addition & 1 deletion go/vt/vtgate/planbuilder/testdata/ddl_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -260,7 +260,7 @@
"Name": "main",
"Sharded": false
},
"Query": "create view view_a as select a.col1, a.col2 from (select col1, col2 from unsharded where id = 1 union select col1, col2 from unsharded where id = 3) as a"
"Query": "create view view_a as select col1, col2 from (select col1, col2 from unsharded where id = 1 union select col1, col2 from unsharded where id = 3) as a"
},
"TablesUsed": [
"main.view_a"
Expand Down
Loading

0 comments on commit 00bef5d

Please sign in to comment.