Gen4: Support for order by column not present in projection

go/test/endtoend/vtgate/gen4/gen4_test.go

@@ -0,0 +1,71 @@
+/*
+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 vtgate
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/sqltypes"
+)
+
+func TestOrderBy(t *testing.T) {
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	// insert some data.
+	checkedExec(t, conn, `insert into t1(id, col) values (100, 123),(10, 12),(1, 13),(1000, 1234)`)
+
+	// Gen4 only supported query.
+	assertMatches(t, conn, `select col from t1 order by id`, `[[INT64(13)] [INT64(12)] [INT64(123)] [INT64(1234)]]`)
+
+	// Gen4 unsupported query. v3 supported.
+	assertMatches(t, conn, `select col from t1 order by 1`, `[[INT64(12)] [INT64(13)] [INT64(123)] [INT64(1234)]]`)
+
+	// unsupported in v3 and Gen4.
+	_, err = exec(t, conn, `select t1.* from t1 order by id`)
+	require.Error(t, err)
+}
+
+func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
+	t.Helper()
+	qr := checkedExec(t, conn, query)
+	got := fmt.Sprintf("%v", qr.Rows)
+	diff := cmp.Diff(expected, got)
+	if diff != "" {
+		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
+	}
+}
+
+func checkedExec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
+	t.Helper()
+	qr, err := exec(t, conn, query)
+	require.NoError(t, err, "for query: "+query)
+	return qr
+}
+
+func exec(t *testing.T, conn *mysql.Conn, query string) (*sqltypes.Result, error) {
+	t.Helper()
+	return conn.ExecuteFetch(query, 1000, true)
+}

go/test/endtoend/vtgate/gen4/main_test.go

@@ -0,0 +1,99 @@
+/*
+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 vtgate
+
+import (
+	"flag"
+	"os"
+	"testing"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	KeyspaceName    = "ks"
+	Cell            = "test"
+	SchemaSQL       = `create table t1(
+	id bigint,
+	col bigint,
+	primary key(id)
+) Engine=InnoDB;
+`
+
+	VSchema = `
+{
+  "sharded": true,
+  "vindexes": {
+    "xxhash": {
+      "type": "xxhash"
+    }
+  },
+  "tables": {
+    "t1": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ]
+    }
+  }
+}`
+)
+
+func TestMain(m *testing.M) {
+	defer cluster.PanicHandler(nil)
+	flag.Parse()
+
+	exitCode := func() int {
+		clusterInstance = cluster.NewCluster(Cell, "localhost")
+		defer clusterInstance.Teardown()
+
+		// Start topo server
+		err := clusterInstance.StartTopo()
+		if err != nil {
+			return 1
+		}
+
+		// Start keyspace
+		keyspace := &cluster.Keyspace{
+			Name:      KeyspaceName,
+			SchemaSQL: SchemaSQL,
+			VSchema:   VSchema,
+		}
+		err = clusterInstance.StartKeyspace(*keyspace, []string{"-80", "80-"}, 1, true)
+		if err != nil {
+			return 1
+		}
+
+		// Start vtgate
+		clusterInstance.VtGateExtraArgs = []string{"-planner_version", "Gen4Fallback"} // enable Gen4 planner.
+		err = clusterInstance.StartVtgate()
+		if err != nil {
+			return 1
+		}
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+		return m.Run()
+	}()
+	os.Exit(exitCode)
+}

go/vt/vtgate/planbuilder/select.go

@@ -159,15 +159,24 @@ func planAggregations(qp *queryProjection, plan logicalPlan, semTable *semantics
 func planOrderBy(qp *queryProjection, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) {
 	switch plan := plan.(type) {
 	case *route:
+		additionalColAdded := false
 		for _, order := range qp.orderExprs {
 			offset, exists := qp.orderExprColMap[order]
-			if !exists {
-				return nil, semantics.Gen4NotSupportedF("order by column not exists in select list")
-			}
 			colName, ok := order.Expr.(*sqlparser.ColName)
 			if !ok {
 				return nil, semantics.Gen4NotSupportedF("order by non-column expression")
 			}
+			if !exists {
+				expr := &sqlparser.AliasedExpr{
+					Expr: order.Expr,
+				}
+				var err error
+				offset, err = pushProjection(expr, plan, semTable)
+				if err != nil {
+					return nil, err
+				}
+				additionalColAdded = true
+			}
 
 			table := semTable.Dependencies(colName)
 			tableInfo, err := semTable.TableInfoFor(table)
@@ -200,9 +209,7 @@ func planOrderBy(qp *queryProjection, plan logicalPlan, semTable *semantics.SemT
 				if err != nil {
 					return nil, err
 				}
-			}
-			if weightStringOffset != -1 {
-				plan.eroute.TruncateColumnCount = len(qp.selectExprs) + len(qp.aggrExprs)
+				additionalColAdded = true
 			}
 
 			plan.eroute.OrderBy = append(plan.eroute.OrderBy, engine.OrderbyParams{
@@ -212,6 +219,10 @@ func planOrderBy(qp *queryProjection, plan logicalPlan, semTable *semantics.SemT
 			})
 			plan.Select.AddOrder(order)
 		}
+		if additionalColAdded {
+			plan.eroute.TruncateColumnCount = len(qp.selectExprs) + len(qp.aggrExprs)
+		}
+
 		return plan, nil
 	default:
 		return nil, semantics.Gen4NotSupportedF("ordering on complex query")

go/vt/vtgate/planbuilder/testdata/memory_sort_cases.txt

@@ -570,3 +570,24 @@
   }
 }
 Gen4 plan same as above
+
+# scatter order by with order by column not present
+"select col from user order by id"
+"unsupported: in scatter query: order by must reference a column in the select list: id asc"
+{
+  "QueryType": "SELECT",
+  "Original": "select col from user order by id",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "SelectScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select col, id, weight_string(id) from `user` where 1 != 1",
+    "OrderBy": "1 ASC",
+    "Query": "select col, id, weight_string(id) from `user` order by id asc",
+    "ResultColumns": 1,
+    "Table": "`user`"
+  }
+}