Fixing Column aliases in outer join queries (vitessio#15384)

Signed-off-by: Manan Gupta <[email protected]> Signed-off-by: Florent Poinsard <[email protected]> Co-authored-by: Florent Poinsard <[email protected]>
GrahamCampbell · Mar 6, 2024 · f3e3f55 · f3e3f55
1 parent e19177b
commit f3e3f55
Show file tree

Hide file tree

Showing 14 changed files with 167 additions and 56 deletions.
diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go
@@ -317,3 +317,22 @@ func TestTransactionModeVar(t *testing.T) {
 		})
 	}
 }
+
+// TestAliasesInOuterJoinQueries tests that aliases work in queries that have outer join clauses.
+func TestAliasesInOuterJoinQueries(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 20, "vtgate")
+
+	mcmp, closer := start(t)
+	defer closer()
+
+	// Insert data into the 2 tables
+	mcmp.Exec("insert into t1(id1, id2) values (1,2), (42,5), (5, 42)")
+	mcmp.Exec("insert into tbl(id, unq_col, nonunq_col) values (1,2,3), (2,5,3), (3, 42, 2)")
+
+	// Check that the select query works as intended and then run it again verifying the column names as well.
+	mcmp.AssertMatches("select t1.id1 as t0, t1.id1 as t1, tbl.unq_col as col from t1 left outer join tbl on t1.id2 = tbl.nonunq_col", `[[INT64(1) INT64(1) INT64(42)] [INT64(5) INT64(5) NULL] [INT64(42) INT64(42) NULL]]`)
+	mcmp.ExecWithColumnCompare("select t1.id1 as t0, t1.id1 as t1, tbl.unq_col as col from t1 left outer join tbl on t1.id2 = tbl.nonunq_col")
+
+	mcmp.AssertMatches("select t1.id1 as t0, t1.id1 as t1, tbl.unq_col as col from t1 left outer join tbl on t1.id2 = tbl.nonunq_col order by t1.id2 limit 2", `[[INT64(1) INT64(1) INT64(42)] [INT64(42) INT64(42) NULL]]`)
+	mcmp.ExecWithColumnCompare("select t1.id1 as t0, t1.id1 as t1, tbl.unq_col as col from t1 left outer join tbl on t1.id2 = tbl.nonunq_col order by t1.id2 limit 2")
+}
diff --git a/go/vt/vtgate/engine/cached_size.go b/go/vt/vtgate/engine/cached_size.go
diff --git a/go/vt/vtgate/engine/simple_projection.go b/go/vt/vtgate/engine/simple_projection.go
@@ -19,6 +19,8 @@ package engine
 import (
 	"context"
 
+	"google.golang.org/protobuf/proto"
+
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 )
@@ -29,8 +31,10 @@ var _ Primitive = (*SimpleProjection)(nil)
 type SimpleProjection struct {
 	// Cols defines the column numbers from the underlying primitive
 	// to be returned.
-	Cols  []int
-	Input Primitive
+	Cols []int
+	// ColNames are the column names to use for the columns.
+	ColNames []string
+	Input    Primitive
 }
 
 // NeedsTransaction implements the Primitive interface
@@ -104,8 +108,13 @@ func (sc *SimpleProjection) buildFields(inner *sqltypes.Result) []*querypb.Field
 		return nil
 	}
 	fields := make([]*querypb.Field, 0, len(sc.Cols))
-	for _, col := range sc.Cols {
-		fields = append(fields, inner.Fields[col])
+	for idx, col := range sc.Cols {
+		field := inner.Fields[col]
+		if sc.ColNames[idx] != "" {
+			field = proto.Clone(field).(*querypb.Field)
+			field.Name = sc.ColNames[idx]
+		}
+		fields = append(fields, field)
 	}
 	return fields
 }
@@ -114,6 +123,16 @@ func (sc *SimpleProjection) description() PrimitiveDescription {
 	other := map[string]any{
 		"Columns": sc.Cols,
 	}
+	emptyColNames := true
+	for _, cName := range sc.ColNames {
+		if cName != "" {
+			emptyColNames = false
+			break
+		}
+	}
+	if !emptyColNames {
+		other["ColumnNames"] = sc.ColNames
+	}
 	return PrimitiveDescription{
 		OperatorType: "SimpleProjection",
 		Other:        other,

diff --git a/go/vt/vtgate/engine/simple_projection_test.go b/go/vt/vtgate/engine/simple_projection_test.go
@@ -44,8 +44,9 @@ func TestSubqueryExecute(t *testing.T) {
 	}
 
 	sq := &SimpleProjection{
-		Cols:  []int{0, 2},
-		Input: prim,
+		Cols:     []int{0, 2},
+		ColNames: []string{"", ""},
+		Input:    prim,
 	}
 
 	bv := map[string]*querypb.BindVariable{
@@ -93,8 +94,9 @@ func TestSubqueryStreamExecute(t *testing.T) {
 	}
 
 	sq := &SimpleProjection{
-		Cols:  []int{0, 2},
-		Input: prim,
+		Cols:     []int{0, 2},
+		ColNames: []string{"", ""},
+		Input:    prim,
 	}
 
 	bv := map[string]*querypb.BindVariable{
@@ -142,8 +144,9 @@ func TestSubqueryGetFields(t *testing.T) {
 	}
 
 	sq := &SimpleProjection{
-		Cols:  []int{0, 2},
-		Input: prim,
+		Cols:     []int{0, 2},
+		ColNames: []string{"", ""},
+		Input:    prim,
 	}
 
 	bv := map[string]*querypb.BindVariable{

diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go
@@ -351,10 +351,10 @@ func transformProjection(ctx *plancontext.PlanningContext, op *operators.Project
 		return nil, err
 	}
 
-	if cols := op.AllOffsets(); cols != nil {
+	if cols, colNames := op.AllOffsets(); cols != nil {
 		// if all this op is doing is passing through columns from the input, we
 		// can use the faster SimpleProjection
-		return useSimpleProjection(ctx, op, cols, src)
+		return useSimpleProjection(ctx, op, cols, colNames, src)
 	}
 
 	ap, err := op.GetAliasedProjections()
@@ -399,7 +399,7 @@ func getEvalEngingeExpr(ctx *plancontext.PlanningContext, pe *operators.ProjExpr
 
 // useSimpleProjection uses nothing at all if the output is already correct,
 // or SimpleProjection when we have to reorder or truncate the columns
-func useSimpleProjection(ctx *plancontext.PlanningContext, op *operators.Projection, cols []int, src logicalPlan) (logicalPlan, error) {
+func useSimpleProjection(ctx *plancontext.PlanningContext, op *operators.Projection, cols []int, colNames []string, src logicalPlan) (logicalPlan, error) {
 	columns := op.Source.GetColumns(ctx)
 	if len(columns) == len(cols) && elementsMatchIndices(cols) {
 		// the columns are already in the right order. we don't need anything at all here
@@ -408,7 +408,8 @@ func useSimpleProjection(ctx *plancontext.PlanningContext, op *operators.Project
 	return &simpleProjection{
 		logicalPlanCommon: newBuilderCommon(src),
 		eSimpleProj: &engine.SimpleProjection{
-			Cols: cols,
+			Cols:     cols,
+			ColNames: colNames,
 		},
 	}, nil
 }

diff --git a/go/vt/vtgate/planbuilder/operators/projection.go b/go/vt/vtgate/planbuilder/operators/projection.go
@@ -399,18 +399,23 @@ func (p *Projection) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
 
 // AllOffsets returns a slice of integer offsets for all columns in the Projection
 // if all columns are of type Offset. If any column is not of type Offset, it returns nil.
-func (p *Projection) AllOffsets() (cols []int) {
+func (p *Projection) AllOffsets() (cols []int, colNames []string) {
 	ap, err := p.GetAliasedProjections()
 	if err != nil {
-		return nil
+		return nil, nil
 	}
 	for _, c := range ap {
 		offset, ok := c.Info.(Offset)
 		if !ok {
-			return nil
+			return nil, nil
+		}
+		colName := ""
+		if c.Original.As.NotEmpty() {
+			colName = c.Original.As.String()
 		}
 
 		cols = append(cols, int(offset))
+		colNames = append(colNames, colName)
 	}
 	return
 }
@@ -445,7 +450,7 @@ func (p *Projection) Compact(ctx *plancontext.PlanningContext) (Operator, *Apply
 	needed := false
 	for i, projection := range ap {
 		e, ok := projection.Info.(Offset)
-		if !ok || int(e) != i {
+		if !ok || int(e) != i || projection.Original.As.NotEmpty() {
 			needed = true
 			break
 		}
@@ -475,6 +480,9 @@ func (p *Projection) compactWithJoin(ctx *plancontext.PlanningContext, join *App
 	for _, col := range ap {
 		switch colInfo := col.Info.(type) {
 		case Offset:
+			if col.Original.As.NotEmpty() {
+				return p, NoRewrite
+			}
 			newColumns = append(newColumns, join.Columns[colInfo])
 			newColumnsAST.add(join.JoinColumns.columns[colInfo])
 		case nil:

diff --git a/go/vt/vtgate/planbuilder/testdata/ddl_cases_no_default_keyspace.json b/go/vt/vtgate/planbuilder/testdata/ddl_cases_no_default_keyspace.json
@@ -144,7 +144,7 @@
           "Name": "user",
           "Sharded": true
         },
-        "Query": "create view view_a as select a.user_id as user_id, a.col1 as col1, a.col2 as col2, b.user_id as user_id, b.col1 as col1, b.col2 as col2 from authoritative as a join authoritative as b on a.user_id = b.user_id"
+        "Query": "create view view_a as select a.user_id, a.col1, a.col2, b.user_id, b.col1, b.col2 from authoritative as a join authoritative as b on a.user_id = b.user_id"
       },
       "TablesUsed": [
         "user.view_a"
@@ -201,7 +201,7 @@
           "Name": "user",
           "Sharded": true
         },
-        "Query": "create view view_a as select `user`.id, a.user_id as user_id, a.col1 as col1, a.col2 as col2, `user`.col1 from authoritative as a join `user` on a.user_id = `user`.id"
+        "Query": "create view view_a as select `user`.id, a.user_id, a.col1, a.col2, `user`.col1 from authoritative as a join `user` on a.user_id = `user`.id"
       },
       "TablesUsed": [
         "user.view_a"

diff --git a/go/vt/vtgate/planbuilder/testdata/from_cases.json b/go/vt/vtgate/planbuilder/testdata/from_cases.json
@@ -3760,8 +3760,8 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select authoritative.col1 as col1, authoritative.user_id as user_id, authoritative.col2 as col2 from authoritative where 1 != 1",
-            "Query": "select authoritative.col1 as col1, authoritative.user_id as user_id, authoritative.col2 as col2 from authoritative",
+            "FieldQuery": "select authoritative.col1, authoritative.user_id, authoritative.col2 from authoritative where 1 != 1",
+            "Query": "select authoritative.col1, authoritative.user_id, authoritative.col2 from authoritative",
             "Table": "authoritative"
           },
           {
@@ -3771,8 +3771,8 @@
               "Name": "main",
               "Sharded": false
             },
-            "FieldQuery": "select unsharded_authoritative.col2 as col2 from unsharded_authoritative where 1 != 1",
-            "Query": "select unsharded_authoritative.col2 as col2 from unsharded_authoritative where unsharded_authoritative.col1 = :authoritative_col1",
+            "FieldQuery": "select unsharded_authoritative.col2 from unsharded_authoritative where 1 != 1",
+            "Query": "select unsharded_authoritative.col2 from unsharded_authoritative where unsharded_authoritative.col1 = :authoritative_col1",
             "Table": "unsharded_authoritative"
           }
         ]

diff --git a/go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json b/go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json
@@ -750,8 +750,8 @@
           "Name": "main",
           "Sharded": false
         },
-        "FieldQuery": "select a.VARIABLE_NAME as VARIABLE_NAME, a.VARIABLE_VALUE as VARIABLE_VALUE, b.CHARACTER_SET_NAME as CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME as DEFAULT_COLLATE_NAME, b.DESCRIPTION as DESCRIPTION, b.MAXLEN as MAXLEN from information_schema.GLOBAL_STATUS as a, information_schema.CHARACTER_SETS as b where 1 != 1",
-        "Query": "select a.VARIABLE_NAME as VARIABLE_NAME, a.VARIABLE_VALUE as VARIABLE_VALUE, b.CHARACTER_SET_NAME as CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME as DEFAULT_COLLATE_NAME, b.DESCRIPTION as DESCRIPTION, b.MAXLEN as MAXLEN from information_schema.GLOBAL_STATUS as a, information_schema.CHARACTER_SETS as b",
+        "FieldQuery": "select a.VARIABLE_NAME, a.VARIABLE_VALUE, b.CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME, b.DESCRIPTION, b.MAXLEN from information_schema.GLOBAL_STATUS as a, information_schema.CHARACTER_SETS as b where 1 != 1",
+        "Query": "select a.VARIABLE_NAME, a.VARIABLE_VALUE, b.CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME, b.DESCRIPTION, b.MAXLEN from information_schema.GLOBAL_STATUS as a, information_schema.CHARACTER_SETS as b",
         "Table": "information_schema.CHARACTER_SETS, information_schema.GLOBAL_STATUS"
       }
     }

diff --git a/go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json b/go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json
@@ -815,8 +815,8 @@
           "Name": "main",
           "Sharded": false
         },
-        "FieldQuery": "select a.CONSTRAINT_CATALOG as CONSTRAINT_CATALOG, a.CONSTRAINT_SCHEMA as CONSTRAINT_SCHEMA, a.CONSTRAINT_NAME as CONSTRAINT_NAME, a.CHECK_CLAUSE as CHECK_CLAUSE, b.CHARACTER_SET_NAME as CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME as DEFAULT_COLLATE_NAME, b.DESCRIPTION as DESCRIPTION, b.MAXLEN as MAXLEN from information_schema.CHECK_CONSTRAINTS as a, information_schema.CHARACTER_SETS as b where 1 != 1",
-        "Query": "select a.CONSTRAINT_CATALOG as CONSTRAINT_CATALOG, a.CONSTRAINT_SCHEMA as CONSTRAINT_SCHEMA, a.CONSTRAINT_NAME as CONSTRAINT_NAME, a.CHECK_CLAUSE as CHECK_CLAUSE, b.CHARACTER_SET_NAME as CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME as DEFAULT_COLLATE_NAME, b.DESCRIPTION as DESCRIPTION, b.MAXLEN as MAXLEN from information_schema.CHECK_CONSTRAINTS as a, information_schema.CHARACTER_SETS as b",
+        "FieldQuery": "select a.CONSTRAINT_CATALOG, a.CONSTRAINT_SCHEMA, a.CONSTRAINT_NAME, a.CHECK_CLAUSE, b.CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME, b.DESCRIPTION, b.MAXLEN from information_schema.CHECK_CONSTRAINTS as a, information_schema.CHARACTER_SETS as b where 1 != 1",
+        "Query": "select a.CONSTRAINT_CATALOG, a.CONSTRAINT_SCHEMA, a.CONSTRAINT_NAME, a.CHECK_CLAUSE, b.CHARACTER_SET_NAME, b.DEFAULT_COLLATE_NAME, b.DESCRIPTION, b.MAXLEN from information_schema.CHECK_CONSTRAINTS as a, information_schema.CHARACTER_SETS as b",
         "Table": "information_schema.CHARACTER_SETS, information_schema.CHECK_CONSTRAINTS"
       }
     }

diff --git a/go/vt/vtgate/planbuilder/testdata/rails_cases.json b/go/vt/vtgate/planbuilder/testdata/rails_cases.json
@@ -50,8 +50,8 @@
                       "Name": "user",
                       "Sharded": true
                     },
-                    "FieldQuery": "select author5s.id as id, author5s.`name` as `name`, author5s.created_at as created_at, author5s.updated_at as updated_at, book6s.supplier5_id, book6s.id from author5s, book6s where 1 != 1",
-                    "Query": "select author5s.id as id, author5s.`name` as `name`, author5s.created_at as created_at, author5s.updated_at as updated_at, book6s.supplier5_id, book6s.id from author5s, book6s where book6s.author5_id = author5s.id",
+                    "FieldQuery": "select author5s.id, author5s.`name`, author5s.created_at, author5s.updated_at, book6s.supplier5_id, book6s.id from author5s, book6s where 1 != 1",
+                    "Query": "select author5s.id, author5s.`name`, author5s.created_at, author5s.updated_at, book6s.supplier5_id, book6s.id from author5s, book6s where book6s.author5_id = author5s.id",
                     "Table": "author5s, book6s"
                   },
                   {

diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json
@@ -361,8 +361,8 @@
           "Name": "user",
           "Sharded": true
         },
-        "FieldQuery": "select a.user_id as user_id, a.col1 as col1, a.col2 as col2, b.user_id as user_id, b.col1 as col1, b.col2 as col2 from authoritative as a, authoritative as b where 1 != 1",
-        "Query": "select a.user_id as user_id, a.col1 as col1, a.col2 as col2, b.user_id as user_id, b.col1 as col1, b.col2 as col2 from authoritative as a, authoritative as b where a.user_id = b.user_id",
+        "FieldQuery": "select a.user_id, a.col1, a.col2, b.user_id, b.col1, b.col2 from authoritative as a, authoritative as b where 1 != 1",
+        "Query": "select a.user_id, a.col1, a.col2, b.user_id, b.col1, b.col2 from authoritative as a, authoritative as b where a.user_id = b.user_id",
         "Table": "authoritative"
       },
       "TablesUsed": [
@@ -433,8 +433,8 @@
           "Name": "user",
           "Sharded": true
         },
-        "FieldQuery": "select `user`.id, a.user_id as user_id, a.col1 as col1, a.col2 as col2, `user`.col1 from authoritative as a, `user` where 1 != 1",
-        "Query": "select `user`.id, a.user_id as user_id, a.col1 as col1, a.col2 as col2, `user`.col1 from authoritative as a, `user` where a.user_id = `user`.id",
+        "FieldQuery": "select `user`.id, a.user_id, a.col1, a.col2, `user`.col1 from authoritative as a, `user` where 1 != 1",
+        "Query": "select `user`.id, a.user_id, a.col1, a.col2, `user`.col1 from authoritative as a, `user` where a.user_id = `user`.id",
         "Table": "`user`, authoritative"
       },
       "TablesUsed": [
@@ -4919,5 +4919,63 @@
         "user.user"
       ]
     }
+  },
+  {
+    "comment": "column name aliases in outer join queries",
+    "query": "select name as t0, name as t1 from user left outer join user_extra on user.cola = user_extra.cola",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select name as t0, name as t1 from user left outer join user_extra on user.cola = user_extra.cola",
+      "Instructions": {
+        "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "t0",
+          "t1"
+        ],
+        "Columns": [
+          0,
+          0
+        ],
+        "Inputs": [
+          {
+            "OperatorType": "Join",
+            "Variant": "LeftJoin",
+            "JoinColumnIndexes": "L:0",
+            "JoinVars": {
+              "user_cola": 1
+            },
+            "TableName": "`user`_user_extra",
+            "Inputs": [
+              {
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select `name`, `user`.cola from `user` where 1 != 1",
+                "Query": "select `name`, `user`.cola from `user`",
+                "Table": "`user`"
+              },
+              {
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select 1 from user_extra where 1 != 1",
+                "Query": "select 1 from user_extra where user_extra.cola = :user_cola",
+                "Table": "user_extra"
+              }
+            ]
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user",
+        "user.user_extra"
+      ]
+    }
   }
 ]
diff --git a/go/vt/vtgate/semantics/early_rewriter.go b/go/vt/vtgate/semantics/early_rewriter.go
@@ -1216,17 +1216,13 @@ type expanderState struct {
 // addColumn adds columns to the expander state. If we have vschema info about the query,
 // we also store which columns were expanded
 func (e *expanderState) addColumn(col ColumnInfo, tbl TableInfo, tblName sqlparser.TableName) {
-	withQualifier := e.needsQualifier
 	var colName *sqlparser.ColName
 	var alias sqlparser.IdentifierCI
-	if withQualifier {
+	if e.needsQualifier {
 		colName = sqlparser.NewColNameWithQualifier(col.Name, tblName)
 	} else {
 		colName = sqlparser.NewColName(col.Name)
 	}
-	if e.needsQualifier {
-		alias = sqlparser.NewIdentifierCI(col.Name)
-	}
 	e.colNames = append(e.colNames, &sqlparser.AliasedExpr{Expr: colName, As: alias})
 	e.storeExpandInfo(tbl, tblName, colName)
 }