prepare the queryproject to make order by easier at plan time

Signed-off-by: Andres Taylor <[email protected]>
vitessio · Jul 1, 2021 · 5c23db4 · 5c23db4
1 parent c4c4696
commit 5c23db4
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 86 deletions.
diff --git a/go/vt/vtgate/planbuilder/queryprojection.go b/go/vt/vtgate/planbuilder/queryprojection.go
@@ -28,15 +28,16 @@ import (
 type queryProjection struct {
 	selectExprs []*sqlparser.AliasedExpr
 	aggrExprs   []*sqlparser.AliasedExpr
+	orderExprs  []orderBy
+}
 
-	// orderExprColMap keeps a map between the Order object and the offset into the select expressions list
-	orderExprColMap map[*sqlparser.Order]int
+type orderBy struct {
+	inner         *sqlparser.Order
+	weightStrExpr sqlparser.Expr
 }
 
 func newQueryProjection() *queryProjection {
-	return &queryProjection{
-		orderExprColMap: map[*sqlparser.Order]int{},
-	}
+	return &queryProjection{}
 }
 
 func createQPFromSelect(sel *sqlparser.Select) (*queryProjection, error) {
@@ -64,7 +65,7 @@ func createQPFromSelect(sel *sqlparser.Select) (*queryProjection, error) {
 	allExpr := append(qp.selectExprs, qp.aggrExprs...)
 
 	for _, order := range sel.OrderBy {
-		findOrderExprInAllExprs(order, allExpr, qp)
+		qp.addOrderBy(order, allExpr)
 	}
 
 	if sel.GroupBy == nil || sel.OrderBy == nil {
@@ -74,33 +75,41 @@ func createQPFromSelect(sel *sqlparser.Select) (*queryProjection, error) {
 	return qp, nil
 }
 
-func findOrderExprInAllExprs(order *sqlparser.Order, allExpr []*sqlparser.AliasedExpr, qp *queryProjection) {
-	colExpr, isColName := order.Expr.(*sqlparser.ColName)
-	// Case 1: Order by expression is an aliased expression in the select clause
-	// Eg - select id as foo from music order by foo
-	if isColName && colExpr.Qualifier.IsEmpty() {
-		for offset, expr := range allExpr {
-			isAliasExpr := !expr.As.IsEmpty()
-			if isAliasExpr && colExpr.Name.Equal(expr.As) {
-				qp.orderExprColMap[order] = offset
-				return
-			}
-		}
-	}
-	// Case 2: Order by is the column offset to be used from the select expressions
+func (qp *queryProjection) addOrderBy(order *sqlparser.Order, allExpr []*sqlparser.AliasedExpr) {
+	// Order by is the column offset to be used from the select expressions
 	// Eg - select id from music order by 1
 	literalExpr, isLiteral := order.Expr.(*sqlparser.Literal)
 	if isLiteral && literalExpr.Type == sqlparser.IntVal {
 		num, _ := strconv.Atoi(literalExpr.Val)
-		qp.orderExprColMap[order] = num - 1
-	} else {
-		// Case 3: Order by is an expression that we already have in the select expression
-		// Eg - select id from music order by id
-		for offset, expr := range allExpr {
-			if sqlparser.EqualsExpr(order.Expr, expr.Expr) {
-				qp.orderExprColMap[order] = offset
+		qp.orderExprs = append(qp.orderExprs, orderBy{
+			inner: &sqlparser.Order{
+				Expr:      allExpr[num-1].Expr,
+				Direction: order.Direction,
+			},
+			weightStrExpr: allExpr[num-1].Expr,
+		})
+		return
+	}
+
+	// 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
+	colExpr, isColName := order.Expr.(*sqlparser.ColName)
+	if isColName && colExpr.Qualifier.IsEmpty() {
+		for _, expr := range allExpr {
+			isAliasExpr := !expr.As.IsEmpty()
+			if isAliasExpr && colExpr.Name.Equal(expr.As) {
+				qp.orderExprs = append(qp.orderExprs, orderBy{
+					inner:         order,
+					weightStrExpr: expr.Expr,
+				})
 				return
 			}
 		}
 	}
+
+	qp.orderExprs = append(qp.orderExprs, orderBy{
+		inner:         order,
+		weightStrExpr: order.Expr,
+	})
 }
diff --git 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, sel.OrderBy, plan, semTable)
+		plan, err = planOrderBy(qp, qp.orderExprs, plan, semTable)
 		if err != nil {
 			return nil, err
 		}

diff --git a/go/vt/vtgate/planbuilder/selectGen4.go b/go/vt/vtgate/planbuilder/selectGen4.go
@@ -131,7 +131,7 @@ func planAggregations(qp *queryProjection, plan logicalPlan, semTable *semantics
 	return oa, nil
 }
 
-func planOrderBy(qp *queryProjection, orderExprs sqlparser.OrderBy, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) {
+func planOrderBy(qp *queryProjection, orderExprs []orderBy, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) {
 	switch plan := plan.(type) {
 	case *route:
 		return planOrderByForRoute(orderExprs, plan, semTable)
@@ -142,15 +142,15 @@ func planOrderBy(qp *queryProjection, orderExprs sqlparser.OrderBy, plan logical
 	}
 }
 
-func planOrderByForRoute(orderExprs sqlparser.OrderBy, plan *route, semTable *semantics.SemTable) (logicalPlan, error) {
+func planOrderByForRoute(orderExprs []orderBy, plan *route, semTable *semantics.SemTable) (logicalPlan, error) {
 	origColCount := plan.Select.GetColumnCount()
 	for _, order := range orderExprs {
-		expr := order.Expr
+		expr := order.weightStrExpr
 		offset, err := wrapExprAndPush(expr, plan, semTable)
 		if err != nil {
 			return nil, err
 		}
-		colName, ok := expr.(*sqlparser.ColName)
+		colName, ok := order.inner.Expr.(*sqlparser.ColName)
 		if !ok {
 			return nil, semantics.Gen4NotSupportedF("order by non-column expression")
 		}
@@ -173,9 +173,9 @@ func planOrderByForRoute(orderExprs sqlparser.OrderBy, plan *route, semTable *se
 		plan.eroute.OrderBy = append(plan.eroute.OrderBy, engine.OrderbyParams{
 			Col:             offset,
 			WeightStringCol: weightStringOffset,
-			Desc:            order.Direction == sqlparser.DescOrder,
+			Desc:            order.inner.Direction == sqlparser.DescOrder,
 		})
-		plan.Select.AddOrder(order)
+		plan.Select.AddOrder(order.inner)
 	}
 	if origColCount != plan.Select.GetColumnCount() {
 		plan.eroute.TruncateColumnCount = origColCount
@@ -211,7 +211,7 @@ func wrapExprAndPush(exp sqlparser.Expr, plan logicalPlan, semTable *semantics.S
 	return offset, err
 }
 
-func planOrderByForJoin(qp *queryProjection, orderExprs sqlparser.OrderBy, plan *joinGen4, semTable *semantics.SemTable) (logicalPlan, error) {
+func planOrderByForJoin(qp *queryProjection, orderExprs []orderBy, plan *joinGen4, semTable *semantics.SemTable) (logicalPlan, error) {
 	if allLeft(orderExprs, semTable, plan.Left.ContainsTables()) {
 		newLeft, err := planOrderBy(qp, orderExprs, plan.Left, semTable)
 		if err != nil {
@@ -231,9 +231,9 @@ func planOrderByForJoin(qp *queryProjection, orderExprs sqlparser.OrderBy, plan
 		eMemorySort: primitive,
 	}
 
-	for _, order := range orderExprs { // order by user.firstName + user.LastName
-		expr := order.Expr
-		offset, err := wrapExprAndPush(order.Expr, plan, semTable)
+	for _, order := range orderExprs {
+		expr := order.inner.Expr
+		offset, err := wrapExprAndPush(order.inner.Expr, plan, semTable)
 		if err != nil {
 			return nil, err
 		}
@@ -250,7 +250,7 @@ func planOrderByForJoin(qp *queryProjection, orderExprs sqlparser.OrderBy, plan
 
 		weightStringOffset := -1
 		if needsWeightString(tbl, col) {
-			weightStringOffset, err = wrapExprAndPush(weightStringFor(expr), plan, semTable)
+			weightStringOffset, err = wrapExprAndPush(weightStringFor(order.weightStrExpr), plan, semTable)
 			if err != nil {
 				return nil, err
 			}
@@ -259,7 +259,7 @@ func planOrderByForJoin(qp *queryProjection, orderExprs sqlparser.OrderBy, plan
 		ms.eMemorySort.OrderBy = append(ms.eMemorySort.OrderBy, engine.OrderbyParams{
 			Col:               offset,
 			WeightStringCol:   weightStringOffset,
-			Desc:              order.Direction == sqlparser.DescOrder,
+			Desc:              order.inner.Direction == sqlparser.DescOrder,
 			StarColFixedIndex: offset,
 		})
 	}
@@ -268,9 +268,9 @@ func planOrderByForJoin(qp *queryProjection, orderExprs sqlparser.OrderBy, plan
 
 }
 
-func allLeft(orderExprs sqlparser.OrderBy, semTable *semantics.SemTable, lhsTables semantics.TableSet) bool {
+func allLeft(orderExprs []orderBy, semTable *semantics.SemTable, lhsTables semantics.TableSet) bool {
 	for _, expr := range orderExprs {
-		exprDependencies := semTable.Dependencies(expr.Expr)
+		exprDependencies := semTable.Dependencies(expr.inner.Expr)
 		if !exprDependencies.IsSolvedBy(lhsTables) {
 			return false
 		}

diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.txt b/go/vt/vtgate/planbuilder/testdata/onecase.txt
@@ -1,44 +1 @@
 # 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
diff --git a/go/vt/vtgate/planbuilder/testdata/postprocess_cases.txt b/go/vt/vtgate/planbuilder/testdata/postprocess_cases.txt
@@ -406,7 +406,23 @@ Gen4 error: Unknown column '2' in 'order clause'
     "Table": "music"
   }
 }
-Gen4 plan same as above
+{
+  "QueryType": "SELECT",
+  "Original": "select id as foo from music order by 1",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "SelectScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select id as foo, weight_string(id) from music where 1 != 1",
+    "OrderBy": "0 ASC",
+    "Query": "select id as foo, weight_string(id) from music order by id asc",
+    "ResultColumns": 1,
+    "Table": "music"
+  }
+}
 
 # ORDER BY NULL
 "select col from user order by null"