From 391028ceba74bf99409e09b4f6e00bddc1832be4 Mon Sep 17 00:00:00 2001 From: Harshit Gangal Date: Wed, 30 Jun 2021 12:23:45 +0530 Subject: [PATCH] added column count func to SelectStatement, refactored query projection and order by on join with left side handling Signed-off-by: Harshit Gangal --- go/vt/sqlparser/ast.go | 1 + go/vt/sqlparser/ast_funcs.go | 15 +++ go/vt/vtgate/planbuilder/memory_sort_gen4.go | 91 +++++++++++++ go/vt/vtgate/planbuilder/queryprojection.go | 12 +- go/vt/vtgate/planbuilder/route_planning.go | 2 +- go/vt/vtgate/planbuilder/selectGen4.go | 123 ++++++++++-------- go/vt/vtgate/planbuilder/testdata/onecase.txt | 43 ++++++ 7 files changed, 223 insertions(+), 64 deletions(-) create mode 100644 go/vt/vtgate/planbuilder/memory_sort_gen4.go diff --git a/go/vt/sqlparser/ast.go b/go/vt/sqlparser/ast.go index ee4682a04e7..90503f46c2c 100644 --- a/go/vt/sqlparser/ast.go +++ b/go/vt/sqlparser/ast.go @@ -46,6 +46,7 @@ type ( SetLimit(*Limit) SetLock(lock Lock) MakeDistinct() + GetColumnCount() int } // DDLStatement represents any DDL Statement diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go index 5029e6d5691..7f8b090c3b9 100644 --- a/go/vt/sqlparser/ast_funcs.go +++ b/go/vt/sqlparser/ast_funcs.go @@ -744,6 +744,11 @@ func (node *Select) MakeDistinct() { node.Distinct = true } +// GetColumnCount return SelectExprs count. +func (node *Select) GetColumnCount() int { + return len(node.SelectExprs) +} + // AddWhere adds the boolean expression to the // WHERE clause as an AND condition. func (node *Select) AddWhere(expr Expr) { @@ -796,6 +801,11 @@ func (node *ParenSelect) MakeDistinct() { node.Select.MakeDistinct() } +// GetColumnCount implements the SelectStatement interface +func (node *ParenSelect) GetColumnCount() int { + return node.Select.GetColumnCount() +} + // AddWhere adds the boolean expression to the // WHERE clause as an AND condition. func (node *Update) AddWhere(expr Expr) { @@ -832,6 +842,11 @@ func (node *Union) MakeDistinct() { node.UnionSelects[len(node.UnionSelects)-1].Distinct = true } +// GetColumnCount implements the SelectStatement interface +func (node *Union) GetColumnCount() int { + return node.FirstStatement.GetColumnCount() +} + //Unionize returns a UNION, either creating one or adding SELECT to an existing one func Unionize(lhs, rhs SelectStatement, distinct bool, by OrderBy, limit *Limit, lock Lock) *Union { union, isUnion := lhs.(*Union) diff --git a/go/vt/vtgate/planbuilder/memory_sort_gen4.go b/go/vt/vtgate/planbuilder/memory_sort_gen4.go new file mode 100644 index 00000000000..9f93df3a49b --- /dev/null +++ b/go/vt/vtgate/planbuilder/memory_sort_gen4.go @@ -0,0 +1,91 @@ +/* +Copyright 2021 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package planbuilder + +import ( + "vitess.io/vitess/go/sqltypes" + vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc" + "vitess.io/vitess/go/vt/sqlparser" + "vitess.io/vitess/go/vt/vterrors" + "vitess.io/vitess/go/vt/vtgate/engine" + "vitess.io/vitess/go/vt/vtgate/semantics" +) + +var _ logicalPlan = (*memorySortGen4)(nil) + +type memorySortGen4 struct { + orderBy []engine.OrderbyParams + input logicalPlan + truncateColumnCount int +} + +func (m *memorySortGen4) Order() int { + panic("implement me") +} + +func (m *memorySortGen4) ResultColumns() []*resultColumn { + panic("implement me") +} + +func (m *memorySortGen4) Reorder(i int) { + panic("implement me") +} + +func (m *memorySortGen4) Wireup(lp logicalPlan, jt *jointab) error { + panic("implement me") +} + +func (m *memorySortGen4) WireupGen4(semTable *semantics.SemTable) error { + return m.input.WireupGen4(semTable) +} + +func (m *memorySortGen4) SupplyVar(from, to int, col *sqlparser.ColName, varname string) { + panic("implement me") +} + +func (m *memorySortGen4) SupplyCol(col *sqlparser.ColName) (rc *resultColumn, colNumber int) { + panic("implement me") +} + +func (m *memorySortGen4) SupplyWeightString(colNumber int) (weightcolNumber int, err error) { + panic("implement me") +} + +func (m *memorySortGen4) Primitive() engine.Primitive { + return &engine.MemorySort{ + UpperLimit: sqltypes.PlanValue{}, + OrderBy: m.orderBy, + Input: m.input.Primitive(), + TruncateColumnCount: m.truncateColumnCount, + } +} + +func (m *memorySortGen4) Inputs() []logicalPlan { + return []logicalPlan{m.input} +} + +func (m *memorySortGen4) Rewrite(inputs ...logicalPlan) error { + if len(inputs) != 1 { + return vterrors.New(vtrpcpb.Code_INTERNAL, "[BUG]: expected only 1 input") + } + m.input = inputs[0] + return nil +} + +func (m *memorySortGen4) ContainsTables() semantics.TableSet { + return m.input.ContainsTables() +} diff --git a/go/vt/vtgate/planbuilder/queryprojection.go b/go/vt/vtgate/planbuilder/queryprojection.go index e72b56fc244..2c8c62625a5 100644 --- a/go/vt/vtgate/planbuilder/queryprojection.go +++ b/go/vt/vtgate/planbuilder/queryprojection.go @@ -26,11 +26,8 @@ import ( ) type queryProjection struct { - selectExprs []*sqlparser.AliasedExpr - aggrExprs []*sqlparser.AliasedExpr - groupOrderingCommonExpr map[sqlparser.Expr]*sqlparser.Order - - orderExprs sqlparser.OrderBy + selectExprs []*sqlparser.AliasedExpr + aggrExprs []*sqlparser.AliasedExpr // orderExprColMap keeps a map between the Order object and the offset into the select expressions list orderExprColMap map[*sqlparser.Order]int @@ -38,8 +35,7 @@ type queryProjection struct { func newQueryProjection() *queryProjection { return &queryProjection{ - groupOrderingCommonExpr: map[sqlparser.Expr]*sqlparser.Order{}, - orderExprColMap: map[*sqlparser.Order]int{}, + orderExprColMap: map[*sqlparser.Order]int{}, } } @@ -65,8 +61,6 @@ func createQPFromSelect(sel *sqlparser.Select) (*queryProjection, error) { qp.selectExprs = append(qp.selectExprs, exp) } - qp.orderExprs = sel.OrderBy - allExpr := append(qp.selectExprs, qp.aggrExprs...) for _, order := range sel.OrderBy { diff --git a/go/vt/vtgate/planbuilder/route_planning.go b/go/vt/vtgate/planbuilder/route_planning.go index 4d2d98d7a12..381bde213cf 100644 --- a/go/vt/vtgate/planbuilder/route_planning.go +++ b/go/vt/vtgate/planbuilder/route_planning.go @@ -196,7 +196,7 @@ func planHorizon(sel *sqlparser.Select, plan logicalPlan, semTable *semantics.Se } } if len(sel.OrderBy) > 0 { - plan, err = planOrderBy(qp, plan, semTable) + plan, err = planOrderBy(qp, sel.OrderBy, plan, semTable) if err != nil { return nil, err } diff --git a/go/vt/vtgate/planbuilder/selectGen4.go b/go/vt/vtgate/planbuilder/selectGen4.go index f012749d8cc..ea389bd4b8c 100644 --- a/go/vt/vtgate/planbuilder/selectGen4.go +++ b/go/vt/vtgate/planbuilder/selectGen4.go @@ -117,37 +117,25 @@ func planAggregations(qp *queryProjection, plan logicalPlan, semTable *semantics return oa, nil } -func planOrderBy(qp *queryProjection, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) { +func planOrderBy(qp *queryProjection, orderExprs sqlparser.OrderBy, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) { switch plan := plan.(type) { case *route: - return planOrderByForRoute(qp, plan, semTable) + return planOrderByForRoute(qp, orderExprs, plan, semTable) case *joinGen4: - return planOrderByForJoin(qp, plan, semTable) + return planOrderByForJoin(qp, orderExprs, plan, semTable) default: return nil, semantics.Gen4NotSupportedF("ordering on complex query") } } -func planOrderByForRoute(qp *queryProjection, plan *route, semTable *semantics.SemTable) (logicalPlan, error) { - additionalColAdded := false - for _, order := range qp.orderExprs { - offset, exists := qp.orderExprColMap[order] - var weightStringExpr sqlparser.Expr - if !exists { - expr := &sqlparser.AliasedExpr{ - Expr: order.Expr, - } - var err error - offset, err = pushProjection(expr, plan, semTable, true) - weightStringExpr = order.Expr - if err != nil { - return nil, err - } - additionalColAdded = true - } else { - weightStringExpr = qp.selectExprs[offset].Expr +func planOrderByForRoute(qp *queryProjection, orderExprs sqlparser.OrderBy, plan *route, semTable *semantics.SemTable) (logicalPlan, error) { + origColCount := plan.Select.GetColumnCount() + for _, order := range orderExprs { + offset, expr, err := getOrProjectExpr(qp, plan, semTable, order) + if err != nil { + return nil, err } - colName, ok := weightStringExpr.(*sqlparser.ColName) + colName, ok := expr.(*sqlparser.ColName) if !ok { return nil, semantics.Gen4NotSupportedF("order by non-column expression") } @@ -157,33 +145,14 @@ func planOrderByForRoute(qp *queryProjection, plan *route, semTable *semantics.S if err != nil { return nil, err } - weightStringNeeded := true - for _, c := range tbl.GetColumns() { - if colName.Name.String() == c.Name { - if sqltypes.IsNumber(c.Type) { - weightStringNeeded = false - } - break - } - } + weightStringNeeded := needsWeightString(tbl, colName) weightStringOffset := -1 if weightStringNeeded { - expr := &sqlparser.AliasedExpr{ - Expr: &sqlparser.FuncExpr{ - Name: sqlparser.NewColIdent("weight_string"), - Exprs: []sqlparser.SelectExpr{ - &sqlparser.AliasedExpr{ - Expr: weightStringExpr, - }, - }, - }, - } - weightStringOffset, err = pushProjection(expr, plan, semTable, true) + weightStringOffset, err = pushExpression(plan, expr, semTable) if err != nil { return nil, err } - additionalColAdded = true } plan.eroute.OrderBy = append(plan.eroute.OrderBy, engine.OrderbyParams{ @@ -193,30 +162,76 @@ func planOrderByForRoute(qp *queryProjection, plan *route, semTable *semantics.S }) plan.Select.AddOrder(order) } - if additionalColAdded { - plan.eroute.TruncateColumnCount = len(qp.selectExprs) + len(qp.aggrExprs) + if origColCount != plan.Select.GetColumnCount() { + plan.eroute.TruncateColumnCount = origColCount } return plan, nil } -func planOrderByForJoin(qp *queryProjection, plan *joinGen4, semTable *semantics.SemTable) (logicalPlan, error) { +func pushExpression(plan *route, expr sqlparser.Expr, semTable *semantics.SemTable) (int, error) { + aliasedExpr := &sqlparser.AliasedExpr{ + Expr: &sqlparser.FuncExpr{ + Name: sqlparser.NewColIdent("weight_string"), + Exprs: []sqlparser.SelectExpr{ + &sqlparser.AliasedExpr{ + Expr: expr, + }, + }, + }, + } + return pushProjection(aliasedExpr, plan, semTable, true) +} + +func needsWeightString(tbl semantics.TableInfo, colName *sqlparser.ColName) bool { + for _, c := range tbl.GetColumns() { + if colName.Name.String() == c.Name { + return !sqltypes.IsNumber(c.Type) + } + } + return true // we didn't find the column. better to add just to be safe1 +} + +// getOrProjectExpr either gets the offset to the expression if it is already projected, +// or pushes the projection if needed +func getOrProjectExpr(qp *queryProjection, plan *route, semTable *semantics.SemTable, order *sqlparser.Order) (offset int, expr sqlparser.Expr, err error) { + offset, exists := qp.orderExprColMap[order] + if exists { + // we are ordering by an expression that is already part of the output + return offset, qp.selectExprs[offset].Expr, nil + } + + aliasedExpr := &sqlparser.AliasedExpr{Expr: order.Expr} + offset, err = pushProjection(aliasedExpr, plan, semTable, true) + if err != nil { + return 0, nil, err + } + + return offset, order.Expr, nil +} + +func planOrderByForJoin(qp *queryProjection, orderExprs sqlparser.OrderBy, plan *joinGen4, semTable *semantics.SemTable) (logicalPlan, error) { isAllLeft := true var err error - for _, expr := range qp.orderExprs { + for _, expr := range orderExprs { exprDependencies := semTable.Dependencies(expr.Expr) - if exprDependencies.IsSolvedBy(plan.Left.ContainsTables()) { - plan.Left, err = planOrderBy(qp, plan.Left, semTable) - if err != nil { - return nil, err - } - } else { + if !exprDependencies.IsSolvedBy(plan.Left.ContainsTables()) { isAllLeft = false break } } if isAllLeft { + plan.Left, err = planOrderBy(qp, orderExprs, plan.Left, semTable) + if err != nil { + return nil, err + } return plan, nil } - return nil, semantics.Gen4NotSupportedF("ordering on vtgate") + + return &memorySortGen4{ + orderBy: nil, + input: plan, + truncateColumnCount: 0, + }, nil + } diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.txt b/go/vt/vtgate/planbuilder/testdata/onecase.txt index e819513f354..79d0039f857 100644 --- a/go/vt/vtgate/planbuilder/testdata/onecase.txt +++ b/go/vt/vtgate/planbuilder/testdata/onecase.txt @@ -1 +1,44 @@ # Add your test case here for debugging and run go test -run=One. +"select user.name, music.id from user, music order by user.name, music.id" +{ + "QueryType": "SELECT", + "Original": "select user.name, music.id from user, music order by user.name, music.id", + "Instructions": { + "OperatorType": "Sort", + "Variant": "Memory", + "OrderBy": "0 ASC, 1 ASC", + "Inputs": [ + { + "OperatorType": "Join", + "Variant": "Join", + "JoinColumnIndexes": "-1,1,-2,2", + "TableName": "`user`_music", + "Inputs": [ + { + "OperatorType": "Route", + "Variant": "SelectScatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select `user`.`name`, weight_string(`user`.`name`) from `user` where 1 != 1", + "Query": "select `user`.`name`, weight_string(`user`.`name`) from `user`", + "Table": "`user`" + }, + { + "OperatorType": "Route", + "Variant": "SelectScatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select music.id, weight_string(music.id) from music where 1 != 1", + "Query": "select music.id, weight_string(music.id) from music", + "Table": "music" + } + ] + } + ] + } +} +Gen4 plan same as above