Gen4: derived table with aliased column names

go/test/endtoend/vtgate/misc_test.go

@@ -544,3 +544,17 @@ func TestSelectEqualUniqueOuterJoinRightPredicate(t *testing.T) {
 	utils.Exec(t, conn, "insert into t2(id3, id4) values (0,20),(1,19),(2,18),(3,17),(4,16),(5,15)")
 	utils.AssertMatches(t, conn, `SELECT id3 FROM t1 LEFT JOIN t2 ON t1.id1 = t2.id3 WHERE t2.id3 = 10`, `[]`)
 }
+
+func TestDerivedTableColumns(t *testing.T) {
+	defer cluster.PanicHandler(t)
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	utils.Exec(t, conn, `delete from t1`)
+	defer utils.Exec(t, conn, `delete from t1`)
+
+	utils.Exec(t, conn, "insert into t1(id1, id2) values (0,10),(1,9),(2,8),(3,7),(4,6),(5,5)")
+	utils.AssertMatches(t, conn, `SELECT /* GEN4_COMPARE_ONLY_GEN4 */ t.id FROM (SELECT id2 FROM t1) AS t(id) ORDER BY t.id DESC`, `[[INT64(10)] [INT64(9)] [INT64(8)] [INT64(7)] [INT64(6)] [INT64(5)]]`)
+}

go/test/endtoend/vtgate/queries/informationschema/informationschema_test.go

@@ -86,8 +86,8 @@ func TestInformationSchemaQueryGetsRoutedToTheRightTableAndKeyspace(t *testing.T
 
 	utils.Exec(t, conn, "insert into t1(id1, id2) values (1, 1), (2, 2), (3,3), (4,4)")
 
-	_ = utils.Exec(t, conn, "SELECT * FROM t1000") // test that the routed table is available to us
-	result := utils.Exec(t, conn, "SELECT * FROM information_schema.tables WHERE table_schema = database() and table_name='t1000'")
+	_ = utils.Exec(t, conn, "SELECT /* GEN4_COMPARE_ONLY_GEN4 */ * FROM t1000") // test that the routed table is available to us
+	result := utils.Exec(t, conn, "SELECT /* GEN4_COMPARE_ONLY_GEN4 */ * FROM information_schema.tables WHERE table_schema = database() and table_name='t1000'")
 	assert.NotEmpty(t, result.Rows)
 }
 

go/vt/vtgate/engine/gen4_compare_v3.go

@@ -71,8 +71,14 @@ func (gc *Gen4CompareV3) TryExecute(
 	bindVars map[string]*querypb.BindVariable,
 	wantfields bool,
 ) (*sqltypes.Result, error) {
-	gen4Result, gen4Err := gc.Gen4.TryExecute(vcursor, bindVars, wantfields)
-	v3Result, v3Err := gc.V3.TryExecute(vcursor, bindVars, wantfields)
+	var v3Err, gen4Err error
+	v3Result, gen4Result := &sqltypes.Result{}, &sqltypes.Result{}
+	if gc.Gen4 != nil {
+		gen4Result, gen4Err = gc.Gen4.TryExecute(vcursor, bindVars, wantfields)
+	}
+	if gc.V3 != nil {
+		v3Result, v3Err = gc.V3.TryExecute(vcursor, bindVars, wantfields)
+	}
 
 	if err := CompareV3AndGen4Errors(v3Err, gen4Err); err != nil {
 		return nil, err
@@ -91,16 +97,21 @@ func (gc *Gen4CompareV3) TryStreamExecute(
 	wantfields bool,
 	callback func(*sqltypes.Result) error,
 ) error {
+	var v3Err, gen4Err error
 	v3Result, gen4Result := &sqltypes.Result{}, &sqltypes.Result{}
 
-	gen4Err := gc.Gen4.TryStreamExecute(vcursor, bindVars, wantfields, func(result *sqltypes.Result) error {
-		gen4Result.AppendResult(result)
-		return nil
-	})
-	v3Err := gc.V3.TryStreamExecute(vcursor, bindVars, wantfields, func(result *sqltypes.Result) error {
-		v3Result.AppendResult(result)
-		return nil
-	})
+	if gc.Gen4 != nil {
+		gen4Err = gc.Gen4.TryStreamExecute(vcursor, bindVars, wantfields, func(result *sqltypes.Result) error {
+			gen4Result.AppendResult(result)
+			return nil
+		})
+	}
+	if gc.V3 != nil {
+		v3Err = gc.V3.TryStreamExecute(vcursor, bindVars, wantfields, func(result *sqltypes.Result) error {
+			v3Result.AppendResult(result)
+			return nil
+		})
+	}
 
 	if err := CompareV3AndGen4Errors(v3Err, gen4Err); err != nil {
 		return err

go/vt/vtgate/planbuilder/abstract/derived.go

@@ -25,9 +25,10 @@ import (
 
 // Derived represents a derived table in the query
 type Derived struct {
-	Sel   sqlparser.SelectStatement
-	Inner Operator
-	Alias string
+	Sel           sqlparser.SelectStatement
+	Inner         Operator
+	Alias         string
+	ColumnAliases sqlparser.Columns
 }
 
 var _ Operator = (*Derived)(nil)

go/vt/vtgate/planbuilder/abstract/operator.go

@@ -82,7 +82,7 @@ func getOperatorFromTableExpr(tableExpr sqlparser.TableExpr, semTable *semantics
 			if err != nil {
 				return nil, err
 			}
-			return &Derived{Alias: tableExpr.As.String(), Inner: inner, Sel: tbl.Select}, nil
+			return &Derived{Alias: tableExpr.As.String(), Inner: inner, Sel: tbl.Select, ColumnAliases: tableExpr.Columns}, nil
 		default:
 			return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unable to use: %T", tbl)
 		}

go/vt/vtgate/planbuilder/derivedtree.go

@@ -24,9 +24,10 @@ import (
 )
 
 type derivedTree struct {
-	query sqlparser.SelectStatement
-	inner queryTree
-	alias string
+	query         sqlparser.SelectStatement
+	inner         queryTree
+	alias         string
+	columnAliases sqlparser.Columns
 
 	// columns needed to feed other plans
 	columns       []*sqlparser.ColName

go/vt/vtgate/planbuilder/gen4_compare_v3_planner.go

@@ -17,10 +17,16 @@ limitations under the License.
 package planbuilder
 
 import (
+	"strings"
+
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vtgate/engine"
 )
 
+type commentDirective struct {
+	onlyV3, onlyGen4 bool
+}
+
 func gen4CompareV3Planner(query string) func(sqlparser.Statement, *sqlparser.ReservedVars, ContextVSchema) (engine.Primitive, error) {
 	return func(statement sqlparser.Statement, vars *sqlparser.ReservedVars, ctxVSchema ContextVSchema) (engine.Primitive, error) {
 		// we will be switching the planner version to Gen4 and V3 in order to
@@ -29,32 +35,37 @@ func gen4CompareV3Planner(query string) func(sqlparser.Statement, *sqlparser.Res
 		defer ctxVSchema.SetPlannerVersion(Gen4CompareV3)
 
 		// preliminary checks on the given statement
-		onlyGen4, hasOrderBy, err := preliminaryChecks(statement)
+		onlyGen4, hasOrderBy, comments, err := preliminaryChecks(statement)
 		if err != nil {
 			return nil, err
 		}
+		cd := parseComments(comments)
 
 		// plan statement using Gen4
 		gen4Primitive, gen4Err := planWithPlannerVersion(statement, vars, ctxVSchema, query, Gen4)
 
+		// if onlyGen4 is set to true or Gen4's instruction contain a lock primitive,
+		// we use only Gen4's primitive and exit early without using V3's.
+		// since lock primitives can imply the creation or deletion of locks,
+		// we want to execute them once using Gen4 to avoid the duplicated locks
+		// or double lock-releases.
+		if !cd.onlyV3 && (onlyGen4 || cd.onlyGen4) || (gen4Primitive != nil && hasLockPrimitive(gen4Primitive)) {
+			return gen4Primitive, gen4Err
+		}
+
 		// get V3's plan
 		v3Primitive, v3Err := planWithPlannerVersion(statement, vars, ctxVSchema, query, V3)
 
+		if cd.onlyV3 && !cd.onlyGen4 && !onlyGen4 {
+			return v3Primitive, v3Err
+		}
+
 		// check potential errors from Gen4 and V3
 		err = engine.CompareV3AndGen4Errors(v3Err, gen4Err)
 		if err != nil {
 			return nil, err
 		}
 
-		// if onlyGen4 is set to true or Gen4's instruction contain a lock primitive,
-		// we use only Gen4's primitive and exit early without using V3's.
-		// since lock primitives can imply the creation or deletion of locks,
-		// we want to execute them once using Gen4 to avoid the duplicated locks
-		// or double lock-releases.
-		if onlyGen4 || hasLockPrimitive(gen4Primitive) {
-			return gen4Primitive, gen4Err
-		}
-
 		return &engine.Gen4CompareV3{
 			V3:         v3Primitive,
 			Gen4:       gen4Primitive,
@@ -63,11 +74,26 @@ func gen4CompareV3Planner(query string) func(sqlparser.Statement, *sqlparser.Res
 	}
 }
 
-func preliminaryChecks(statement sqlparser.Statement) (bool, bool, error) {
+func parseComments(comments []string) commentDirective {
+	cd := commentDirective{}
+	for _, comment := range comments {
+		if strings.Contains(comment, "GEN4_COMPARE_ONLY_V3") {
+			cd.onlyV3 = true
+		}
+		if strings.Contains(comment, "GEN4_COMPARE_ONLY_GEN4") {
+			cd.onlyGen4 = true
+		}
+	}
+	return cd
+}
+
+func preliminaryChecks(statement sqlparser.Statement) (bool, bool, []string, error) {
 	var onlyGen4, hasOrderBy bool
+	var comments []string
 	switch s := statement.(type) {
 	case *sqlparser.Union:
 		hasOrderBy = len(s.OrderBy) > 0
+		comments = s.GetComments()
 
 		// walk through the union and search for select statements that have
 		// a next val select expression, in which case we need to only use
@@ -81,10 +107,12 @@ func preliminaryChecks(statement sqlparser.Statement) (bool, bool, error) {
 			return true, nil
 		}, s)
 		if err != nil {
-			return false, false, err
+			return false, false, nil, err
 		}
 	case *sqlparser.Select:
 		hasOrderBy = len(s.OrderBy) > 0
+		comments = s.GetComments()
+
 		for _, expr := range s.SelectExprs {
 			// we are not executing the plan a second time if the query is a select next val,
 			// since the first execution might increment the `next` value, results will almost
@@ -95,7 +123,7 @@ func preliminaryChecks(statement sqlparser.Statement) (bool, bool, error) {
 			}
 		}
 	}
-	return onlyGen4, hasOrderBy, nil
+	return onlyGen4, hasOrderBy, comments, nil
 }
 
 func planWithPlannerVersion(statement sqlparser.Statement, vars *sqlparser.ReservedVars, ctxVSchema ContextVSchema, query string, version PlannerVersion) (engine.Primitive, error) {

go/vt/vtgate/planbuilder/horizon_planning.go

@@ -180,6 +180,14 @@ func pushProjection(expr *sqlparser.AliasedExpr, plan logicalPlan, semTable *sem
 			return 0, false, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.BadFieldError, "Unknown column '%s' in 'order clause'", sqlparser.String(expr))
 		}
 
+		// if we are trying to push a projection that belongs to a DerivedTable
+		// we rewrite that expression, so it matches the column name used inside
+		// that derived table.
+		err = rewriteProjectionOfDerivedTable(expr, semTable)
+		if err != nil {
+			return 0, false, err
+		}
+
 		offset := len(sel.SelectExprs)
 		sel.SelectExprs = append(sel.SelectExprs, expr)
 		return offset, true, nil
@@ -265,6 +273,22 @@ func pushProjection(expr *sqlparser.AliasedExpr, plan logicalPlan, semTable *sem
 	}
 }
 
+func rewriteProjectionOfDerivedTable(expr *sqlparser.AliasedExpr, semTable *semantics.SemTable) error {
+	var err error
+	ti, _ := semTable.TableInfoForExpr(expr.Expr)
+	if ti == nil {
+		return nil
+	}
+	_, isDerivedTable := ti.(*semantics.DerivedTable)
+	if isDerivedTable {
+		expr.Expr, err = semantics.RewriteDerivedExpression(expr.Expr, ti)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func removeKeyspaceFromColName(expr *sqlparser.AliasedExpr) *sqlparser.AliasedExpr {
 	if _, ok := expr.Expr.(*sqlparser.ColName); ok {
 		expr = sqlparser.CloneRefOfAliasedExpr(expr)

go/vt/vtgate/planbuilder/querytree_transformers.go

@@ -135,8 +135,9 @@ func transformDerivedPlan(ctx *planningContext, n *derivedTree) (logicalPlan, er
 	innerSelect := rb.Select
 	derivedTable := &sqlparser.DerivedTable{Select: innerSelect}
 	tblExpr := &sqlparser.AliasedTableExpr{
-		Expr: derivedTable,
-		As:   sqlparser.NewTableIdent(n.alias),
+		Expr:    derivedTable,
+		As:      sqlparser.NewTableIdent(n.alias),
+		Columns: n.columnAliases,
 	}
 	selectExprs := sqlparser.SelectExprs{}
 	for _, colName := range n.columns {

go/vt/vtgate/planbuilder/route_planning.go

@@ -74,9 +74,10 @@ func optimizeQuery(ctx *planningContext, opTree abstract.Operator) (queryTree, e
 			return nil, err
 		}
 		return &derivedTree{
-			query: op.Sel,
-			inner: treeInner,
-			alias: op.Alias,
+			query:         op.Sel,
+			inner:         treeInner,
+			alias:         op.Alias,
+			columnAliases: op.ColumnAliases,
 		}, nil
 	case *abstract.SubQuery:
 		return optimizeSubQuery(ctx, op)

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

@@ -4161,8 +4161,100 @@ Gen4 plan same as above
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select `user`.id, `user`.col, id + 1 from `user` where 1 != 1",
-            "Query": "select `user`.id, `user`.col, id + 1 from `user`",
+            "FieldQuery": "select `user`.id, `user`.col, `user`.id + 1 from `user` where 1 != 1",
+            "Query": "select `user`.id, `user`.col, `user`.id + 1 from `user`",
+            "Table": "`user`"
+          },
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select 1 from user_extra where 1 != 1",
+            "Query": "select 1 from user_extra",
+            "Table": "user_extra"
+          }
+        ]
+      }
+    ]
+  }
+}
+
+# derived table with aliased columns and outer predicate pushed in derived table
+"select u.a from (select id as b, name from user) u(a, n) where u.n = 1"
+"unsupported: column aliases in derived table"
+{
+  "QueryType": "SELECT",
+  "Original": "select u.a from (select id as b, name from user) u(a, n) where u.n = 1",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "SelectEqual",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select u.a from (select id as b, `name` from `user` where 1 != 1) as u(a, n) where 1 != 1",
+    "Query": "select u.a from (select id as b, `name` from `user` where `name` = 1) as u(a, n)",
+    "Table": "`user`",
+    "Values": [
+      1
+    ],
+    "Vindex": "name_user_map"
+  }
+}
+
+# derived table with aliased columns predicate in both the outer and inner
+"select u.a from (select id as b, name from user where b = 1) u(a, n) where u.n = 1"
+"unsupported: column aliases in derived table"
+{
+  "QueryType": "SELECT",
+  "Original": "select u.a from (select id as b, name from user where b = 1) u(a, n) where u.n = 1",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "SelectEqual",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select u.a from (select id as b, `name` from `user` where 1 != 1) as u(a, n) where 1 != 1",
+    "Query": "select u.a from (select id as b, `name` from `user` where b = 1 and `name` = 1) as u(a, n)",
+    "Table": "`user`",
+    "Values": [
+      1
+    ],
+    "Vindex": "name_user_map"
+  }
+}
+
+# derived table with aliased columns and a join that requires pushProjection
+"select i+1 from (select user.id from user join user_extra) t(i)"
+"unsupported: column aliases in derived table"
+{
+  "QueryType": "SELECT",
+  "Original": "select i+1 from (select user.id from user join user_extra) t(i)",
+  "Instructions": {
+    "OperatorType": "SimpleProjection",
+    "Columns": [
+      1
+    ],
+    "Inputs": [
+      {
+        "OperatorType": "Join",
+        "Variant": "Join",
+        "JoinColumnIndexes": "-1,-2",
+        "TableName": "`user`_user_extra",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select `user`.id, `user`.id + 1 from `user` where 1 != 1",
+            "Query": "select `user`.id, `user`.id + 1 from `user`",
             "Table": "`user`"
           },
           {