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]> Signed-off-by: Harshit Gangal <[email protected]>
planetscale · Dec 20, 2024 · d506249 · d506249
1 parent a6c381a
commit d506249
Show file tree

Hide file tree

Showing 20 changed files with 206 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
@@ -319,13 +319,21 @@ 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, 19, "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)")
-	mcmp.ExecWithColumnCompare("select * from t1 t left join tbl on t.id1 = 666 and t.id2 = tbl.id")
+
+	// 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")
 }
 
 func TestAlterTableWithView(t *testing.T) {

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
@@ -423,18 +423,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
 }
@@ -469,7 +474,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
 		}
@@ -499,6 +504,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/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go
@@ -674,7 +674,7 @@ func (s *planTestSuite) testFile(filename string, vschema *vschemawrapper.VSchem
 					if tcase.Skip {
 						t.Skip(message)
 					} else {
-						t.Errorf(message)
+						t.Error(message)
 					}
 				} else if tcase.Skip {
 					t.Errorf("query is correct even though it is skipped:\n %s", tcase.Query)

diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
@@ -3644,6 +3644,9 @@
       "Original": "select * from (select id from user having count(*) = 1) s",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "id"
+        ],
         "Columns": [
           0
         ],

diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json
@@ -370,6 +370,9 @@
       "Original": "with s as (select id from user having count(*) = 1) select * from s",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "id"
+        ],
         "Columns": [
           0
         ],
@@ -1300,6 +1303,10 @@
       "Original": "with t as (select user.id from user join user_extra) select id, t.id from  t",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "id",
+          "id"
+        ],
         "Columns": [
           0,
           0

diff --git a/go/vt/vtgate/planbuilder/testdata/from_cases.json b/go/vt/vtgate/planbuilder/testdata/from_cases.json
@@ -2673,6 +2673,10 @@
       "Original": "select id, t.id from (select user.id from user join user_extra) as t",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "id",
+          "id"
+        ],
         "Columns": [
           0,
           0
@@ -3989,8 +3993,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"
           },
           {
@@ -4000,8 +4004,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/memory_sort_cases.json b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json
@@ -289,6 +289,11 @@
       "Original": "select user.col1 as a, user.col2 b, music.col3 c from user, music where user.id = music.id and user.id = 1 order by c",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "a",
+          "b",
+          "c"
+        ],
         "Columns": [
           0,
           1,
@@ -360,6 +365,11 @@
       "Original": "select user.col1 as a, user.col2, music.col3 from user join music on user.id = music.id where user.id = 1 order by 1 asc, 3 desc, 2 asc",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "a",
+          "col2",
+          "col3"
+        ],
         "Columns": [
           0,
           1,

diff --git a/go/vt/vtgate/planbuilder/testdata/postprocess_cases.json b/go/vt/vtgate/planbuilder/testdata/postprocess_cases.json
@@ -1586,6 +1586,10 @@
       "Original": "select name, name from user, music order by name",
       "Instructions": {
         "OperatorType": "SimpleProjection",
+        "ColumnNames": [
+          "name",
+          "name"
+        ],
         "Columns": [
           0,
           0

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"
                   },
                   {