Fix aliasing in routes that have a derived table (#15550)

Signed-off-by: Manan Gupta <[email protected]>

go/test/endtoend/vtgate/queries/reference/reference_test.go

@@ -83,6 +83,25 @@ func TestReferenceRouting(t *testing.T) {
 		`[[INT64(0)]]`,
 	)
 
+	t.Run("Complex reference query", func(t *testing.T) {
+		utils.SkipIfBinaryIsBelowVersion(t, 20, "vtgate")
+		// Verify a complex query using reference tables with a left join having a derived table with an order by clause works as intended.
+		utils.AssertMatches(
+			t,
+			conn,
+			`SELECT t.id FROM (
+                        SELECT zd.id, zd.zip_id
+                        FROM `+shardedKeyspaceName+`.zip_detail AS zd
+                        WHERE zd.id IN (2)
+                        ORDER BY zd.discontinued_at
+                        LIMIT 1
+                ) AS t
+                LEFT JOIN `+shardedKeyspaceName+`.zip_detail AS t0 ON t.zip_id = t0.zip_id
+                ORDER BY t.id`,
+			`[[INT64(2)]]`,
+		)
+	})
+
 	// UPDATE should route an unqualified zip_detail to unsharded keyspace.
 	utils.Exec(t, conn,
 		"UPDATE zip_detail SET discontinued_at = NULL WHERE id = 2")

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -471,7 +471,7 @@ func createProjectionWithTheseColumns(
 	if len(p.columns) == 0 {
 		return src, nil
 	}
-	proj, err := createProjection(ctx, src)
+	proj, err := createProjection(ctx, src, "")
 	if err != nil {
 		return nil, err
 	}

go/vt/vtgate/planbuilder/operators/route.go

@@ -535,14 +535,27 @@ func (r *Route) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Ex
 	return r, err
 }
 
-func createProjection(ctx *plancontext.PlanningContext, src ops.Operator) (*Projection, error) {
+func createProjection(ctx *plancontext.PlanningContext, src ops.Operator, derivedName string) (*Projection, error) {
 	proj := newAliasedProjection(src)
 	cols, err := src.GetColumns(ctx)
 	if err != nil {
 		return nil, err
 	}
 	for _, col := range cols {
-		_, err := proj.addUnexploredExpr(col, col.Expr)
+		if derivedName == "" {
+			_, err = proj.addUnexploredExpr(col, col.Expr)
+			if err != nil {
+				return nil, err
+			}
+			continue
+		}
+
+		// for derived tables, we want to use the exposed colname
+		tableName := sqlparser.NewTableName(derivedName)
+		columnName := col.ColumnName()
+		colName := sqlparser.NewColNameWithQualifier(columnName, tableName)
+		ctx.SemTable.CopySemanticInfo(col.Expr, colName)
+		_, err = proj.addUnexploredExpr(aeWrap(colName), colName)
 		if err != nil {
 			return nil, err
 		}
@@ -565,14 +578,14 @@ func (r *Route) AddColumn(ctx *plancontext.PlanningContext, reuse bool, gb bool,
 
 	// if at least one column is not already present, we check if we can easily find a projection
 	// or aggregation in our source that we can add to
-	op, ok, offsets := addMultipleColumnsToInput(ctx, r.Source, reuse, []bool{gb}, []*sqlparser.AliasedExpr{expr})
+	derived, op, ok, offsets := addMultipleColumnsToInput(ctx, r.Source, reuse, []bool{gb}, []*sqlparser.AliasedExpr{expr})
 	r.Source = op
 	if ok {
 		return offsets[0], nil
 	}
 
 	// If no-one could be found, we probably don't have one yet, so we add one here
-	src, err := createProjection(ctx, r.Source)
+	src, err := createProjection(ctx, r.Source, derived)
 	if err != nil {
 		return 0, err
 	}
@@ -592,58 +605,73 @@ type selectExpressions interface {
 // addColumnToInput adds a column to an operator without pushing it down.
 // It will return a bool indicating whether the addition was successful or not,
 // and an offset to where the column can be found
-func addMultipleColumnsToInput(ctx *plancontext.PlanningContext, operator ops.Operator, reuse bool, addToGroupBy []bool, exprs []*sqlparser.AliasedExpr) (ops.Operator, bool, []int) {
+func addMultipleColumnsToInput(
+	ctx *plancontext.PlanningContext,
+	operator ops.Operator,
+	reuse bool,
+	addToGroupBy []bool,
+	exprs []*sqlparser.AliasedExpr,
+) (derivedName string, // if we found a derived table, this will contain its name
+	projection ops.Operator, // if an operator needed to be built, it will be returned here
+	found bool, // whether a matching op was found or not
+	offsets []int, // the offsets the expressions received
+) {
 	switch op := operator.(type) {
 	case *SubQuery:
-		src, added, offset := addMultipleColumnsToInput(ctx, op.Outer, reuse, addToGroupBy, exprs)
+		derivedName, src, added, offset := addMultipleColumnsToInput(ctx, op.Outer, reuse, addToGroupBy, exprs)
 		if added {
 			op.Outer = src
 		}
-		return op, added, offset
+		return derivedName, op, added, offset
 
 	case *Distinct:
-		src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
+		derivedName, src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
 		if added {
 			op.Source = src
 		}
-		return op, added, offset
+		return derivedName, op, added, offset
 
 	case *Limit:
-		src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
+		derivedName, src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
 		if added {
 			op.Source = src
 		}
-		return op, added, offset
+		return derivedName, op, added, offset
 
 	case *Ordering:
-		src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
+		derivedName, src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
 		if added {
 			op.Source = src
 		}
-		return op, added, offset
+		return derivedName, op, added, offset
 
 	case *LockAndComment:
-		src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
+		derivedName, src, added, offset := addMultipleColumnsToInput(ctx, op.Source, reuse, addToGroupBy, exprs)
 		if added {
 			op.Source = src
 		}
-		return op, added, offset
+		return derivedName, op, added, offset
+
+	case *Horizon:
+		// if the horizon has an alias, then it is a derived table,
+		// we have to add a new projection and can't build on this one
+		return op.Alias, op, false, nil
 
 	case selectExpressions:
 		if op.isDerived() {
 			// if the only thing we can push to is a derived table,
 			// we have to add a new projection and can't build on this one
-			return op, false, nil
+			return "", op, false, nil
 		}
 		offset, _ := op.addColumnsWithoutPushing(ctx, reuse, addToGroupBy, exprs)
-		return op, true, offset
+		return "", op, true, offset
 
 	case *Union:
 		tableID := semantics.SingleTableSet(len(ctx.SemTable.Tables))
 		ctx.SemTable.Tables = append(ctx.SemTable.Tables, nil)
 		unionColumns, err := op.GetColumns(ctx)
 		if err != nil {
-			return op, false, nil
+			return "", op, false, nil
 		}
 		proj := &Projection{
 			Source:  op,
@@ -655,7 +683,7 @@ func addMultipleColumnsToInput(ctx *plancontext.PlanningContext, operator ops.Op
 		}
 		return addMultipleColumnsToInput(ctx, proj, reuse, addToGroupBy, exprs)
 	default:
-		return op, false, nil
+		return "", op, false, nil
 	}
 }
 

go/vt/vtgate/planbuilder/testdata/aggr_cases.json

@@ -5961,8 +5961,8 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select table_name from (select table_name from `user` where 1 != 1 group by table_name) as `tables` where 1 != 1",
-            "Query": "select table_name from (select table_name from `user` where id = 143 group by table_name) as `tables`",
+            "FieldQuery": "select `tables`.table_name from (select table_name from `user` where 1 != 1 group by table_name) as `tables` where 1 != 1",
+            "Query": "select `tables`.table_name from (select table_name from `user` where id = 143 group by table_name) as `tables`",
             "Table": "`user`",
             "Values": [
               "INT64(143)"

go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json

@@ -1162,8 +1162,8 @@
               "Name": "main",
               "Sharded": false
             },
-            "FieldQuery": "select table_name from (select table_name from information_schema.`tables` where 1 != 1) as `tables` where 1 != 1",
-            "Query": "select table_name from (select table_name from information_schema.`tables` where table_schema != 'information_schema' limit 1) as `tables`",
+            "FieldQuery": "select `tables`.table_name from (select table_name from information_schema.`tables` where 1 != 1) as `tables` where 1 != 1",
+            "Query": "select `tables`.table_name from (select table_name from information_schema.`tables` where table_schema != 'information_schema' limit 1) as `tables`",
             "Table": "information_schema.`tables`"
           },
           {

go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json

@@ -1284,8 +1284,8 @@
               "Name": "main",
               "Sharded": false
             },
-            "FieldQuery": "select table_name from (select table_name from information_schema.`tables` where 1 != 1) as `tables` where 1 != 1",
-            "Query": "select table_name from (select table_name from information_schema.`tables` where table_schema != 'information_schema' limit 1) as `tables`",
+            "FieldQuery": "select `tables`.table_name from (select table_name from information_schema.`tables` where 1 != 1) as `tables` where 1 != 1",
+            "Query": "select `tables`.table_name from (select table_name from information_schema.`tables` where table_schema != 'information_schema' limit 1) as `tables`",
             "Table": "information_schema.`tables`"
           },
           {

go/vt/vtgate/planbuilder/testdata/reference_cases.json

@@ -161,6 +161,50 @@
       ]
     }
   },
+  {
+    "comment": "Reference tables using left join with a derived table having a limit clause",
+    "query": "SELECT u.id FROM ( SELECT a.id, a.u_id FROM user.ref_with_source AS a WHERE a.id IN (3) ORDER BY a.d_at LIMIT 1) as u LEFT JOIN user.ref_with_source AS u0 ON u.u_id = u0.u_uid ORDER BY u.id",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "SELECT u.id FROM ( SELECT a.id, a.u_id FROM user.ref_with_source AS a WHERE a.id IN (3) ORDER BY a.d_at LIMIT 1) as u LEFT JOIN user.ref_with_source AS u0 ON u.u_id = u0.u_uid ORDER BY u.id",
+      "Instructions": {
+        "OperatorType": "Join",
+        "Variant": "LeftJoin",
+        "JoinColumnIndexes": "L:0",
+        "JoinVars": {
+          "u_u_id": 1
+        },
+        "TableName": "ref_with_source_ref_with_source",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "Reference",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select u.id, u.u_id from (select a.id, a.u_id from ref_with_source as a where 1 != 1) as u where 1 != 1",
+            "Query": "select u.id, u.u_id from (select a.id, a.u_id from ref_with_source as a where a.id in (3) order by a.d_at asc limit 1) as u order by u.id asc",
+            "Table": "ref_with_source"
+          },
+          {
+            "OperatorType": "Route",
+            "Variant": "Reference",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select 1 from ref_with_source as u0 where 1 != 1",
+            "Query": "select 1 from ref_with_source as u0 where u0.u_uid = :u_u_id",
+            "Table": "ref_with_source"
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.ref_with_source"
+      ]
+    }
+  },
   {
     "comment": "insert into qualified ambiguous reference table routes to source",
     "query": "insert into user.ambiguous_ref_with_source(col) values(1)",

go/vt/vtgate/planbuilder/testdata/select_cases.json

@@ -4782,8 +4782,8 @@
               "Name": "main",
               "Sharded": false
             },
-            "FieldQuery": "select table_name from (select table_name from unsharded where 1 != 1) as `tables` where 1 != 1",
-            "Query": "select table_name from (select table_name from unsharded limit 1) as `tables`",
+            "FieldQuery": "select `tables`.table_name from (select table_name from unsharded where 1 != 1) as `tables` where 1 != 1",
+            "Query": "select `tables`.table_name from (select table_name from unsharded limit 1) as `tables`",
             "Table": "unsharded"
           },
           {