added column count func to SelectStatement, refactored query projecti…

…on and order by on join with left side handling

Signed-off-by: Harshit Gangal <harshit@planetscale.com>

go/vt/sqlparser/ast.go

@@ -46,6 +46,7 @@ type (
 		SetLimit(*Limit)
 		SetLock(lock Lock)
 		MakeDistinct()
+		GetColumnCount() int
 	}
 
 	// DDLStatement represents any DDL Statement

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)

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()
+}

go/vt/vtgate/planbuilder/queryprojection.go

@@ -26,20 +26,16 @@ 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
 }
 
 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 {

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
 		}

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
+
 }

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