feat: augment simple projection to also alias columns and fix left jo…

…in aliasing

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

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,14 @@ 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 {
+		if sc.ColNames[idx] == "" {
+			fields = append(fields, inner.Fields[col])
+			continue
+		}
+		field := proto.Clone(inner.Fields[col]).(*querypb.Field)
+		field.Name = sc.ColNames[idx]
+		fields = append(fields, field)
 	}
 	return fields
 }
@@ -114,6 +124,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,

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -355,10 +355,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()
@@ -403,7 +403,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
@@ -412,7 +412,8 @@ func useSimpleProjection(ctx *plancontext.PlanningContext, op *operators.Project
 	return &simpleProjection{
 		logicalPlanCommon: newBuilderCommon(src),
 		eSimpleProj: &engine.SimpleProjection{
-			Cols: cols,
+			Cols:     cols,
+			ColNames: colNames,
 		},
 	}, nil
 }

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:

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

@@ -4997,5 +4997,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"
+      ]
+    }
   }
 ]

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

@@ -1490,40 +1490,62 @@
         "TableName": "`user`_music",
         "Inputs": [
           {
-            "OperatorType": "Distinct",
-            "Collations": [
-              "(0:1)"
+            "OperatorType": "SimpleProjection",
+            "ColumnNames": [
+              "foo"
+            ],
+            "Columns": [
+              0
             ],
             "Inputs": [
               {
-                "OperatorType": "Route",
-                "Variant": "Scatter",
-                "Keyspace": {
-                  "Name": "user",
-                  "Sharded": true
-                },
-                "FieldQuery": "select dt.foo, weight_string(dt.foo) from (select foo from `user` where 1 != 1 union select foo from `user` where 1 != 1) as dt where 1 != 1",
-                "Query": "select dt.foo, weight_string(dt.foo) from (select foo from `user` where bar = 12 union select foo from `user` where bar = 134) as dt",
-                "Table": "`user`"
+                "OperatorType": "Distinct",
+                "Collations": [
+                  "(0:1)"
+                ],
+                "Inputs": [
+                  {
+                    "OperatorType": "Route",
+                    "Variant": "Scatter",
+                    "Keyspace": {
+                      "Name": "user",
+                      "Sharded": true
+                    },
+                    "FieldQuery": "select dt.foo, weight_string(dt.foo) from (select foo from `user` where 1 != 1 union select foo from `user` where 1 != 1) as dt where 1 != 1",
+                    "Query": "select dt.foo, weight_string(dt.foo) from (select foo from `user` where bar = 12 union select foo from `user` where bar = 134) as dt",
+                    "Table": "`user`"
+                  }
+                ]
               }
             ]
           },
           {
-            "OperatorType": "Distinct",
-            "Collations": [
-              "(0:1)"
+            "OperatorType": "SimpleProjection",
+            "ColumnNames": [
+              "bar"
+            ],
+            "Columns": [
+              0
             ],
             "Inputs": [
               {
-                "OperatorType": "Route",
-                "Variant": "Scatter",
-                "Keyspace": {
-                  "Name": "user",
-                  "Sharded": true
-                },
-                "FieldQuery": "select dt.bar, weight_string(dt.bar) from (select bar from music where 1 != 1 union select bar from music where 1 != 1) as dt where 1 != 1",
-                "Query": "select dt.bar, weight_string(dt.bar) from (select bar from music where foo = 12 and bar = :t1_foo union select bar from music where foo = 1234 and bar = :t1_foo) as dt",
-                "Table": "music"
+                "OperatorType": "Distinct",
+                "Collations": [
+                  "(0:1)"
+                ],
+                "Inputs": [
+                  {
+                    "OperatorType": "Route",
+                    "Variant": "Scatter",
+                    "Keyspace": {
+                      "Name": "user",
+                      "Sharded": true
+                    },
+                    "FieldQuery": "select dt.bar, weight_string(dt.bar) from (select bar from music where 1 != 1 union select bar from music where 1 != 1) as dt where 1 != 1",
+                    "Query": "select dt.bar, weight_string(dt.bar) from (select bar from music where foo = 12 and bar = :t1_foo union select bar from music where foo = 1234 and bar = :t1_foo) as dt",
+                    "Table": "music"
+                  }
+                ]
               }
             ]
           }