Merge pull request #353 from dolthub/zachmu/subquery-col-aliases

Column aliases for subqueries

enginetest/queries.go

@@ -270,12 +270,81 @@ var QueryTests = []QueryTest{
 			{"third row", int64(3)},
 		},
 	},
+	{
+		Query: "SELECT a,b FROM (select i,s FROM mytable) mt (a,b) order by 1;",
+		Expected: []sql.Row{
+			{1, "first row"},
+			{2, "second row"},
+			{3, "third row"},
+		},
+	},
+	{
+		Query: "SELECT a,b FROM (select i,s FROM mytable) mt (a,b) order by a desc;",
+		Expected: []sql.Row{
+			{3, "third row"},
+			{2, "second row"},
+			{1, "first row"},
+		},
+	},
+	{
+		Query: "SELECT a,b FROM (select i,s FROM mytable order by i desc) mt (a,b);",
+		Expected: []sql.Row{
+			{3, "third row"},
+			{2, "second row"},
+			{1, "first row"},
+		},
+		ExpectedColumns: sql.Schema{
+			{
+				Name: "a",
+				Type: sql.Int64,
+			},
+			{
+				Name: "b",
+				Type: sql.MustCreateStringWithDefaults(sqltypes.VarChar, 20),
+			},
+		},
+	},
+	{
+		Query: "SELECT a FROM (select i,s FROM mytable) mt (a,b) order by a desc;",
+		Expected: []sql.Row{
+			{3},
+			{2},
+			{1},
+		},
+	},
 	{
 		Query: `SELECT * FROM (values row(1+1,2+2), row(floor(1.5),concat("a","b"))) a order by 1`,
 		Expected: []sql.Row{
 			{1.0, "ab"},
 			{2, 4},
 		},
+		ExpectedColumns: sql.Schema{
+			{
+				Name: "column_0",
+				Type: sql.Int64,
+			},
+			{
+				Name: "column_1",
+				Type: sql.Int64,
+			},
+		},
+	},
+	{
+		Query: `SELECT * FROM (values row(1+1,2+2), row(floor(1.5),concat("a","b"))) a (c,d) order by 1`,
+		Expected: []sql.Row{
+			{1.0, "ab"},
+			{2, 4},
+		},
+		ExpectedColumns: sql.Schema{
+			{
+				Name: "c",
+				Type: sql.Int64,
+			},
+			{
+				Name: "d",
+				Type: sql.Int64,
+			},
+		},
 	},
 	{
 		Query: `SELECT column_0 FROM (values row(1+1,2+2), row(floor(1.5),concat("a","b"))) a order by 1`,
@@ -294,6 +363,16 @@ var QueryTests = []QueryTest{
 			{3, 9.0},
 		},
 	},
+	{
+		Query: `SELECT B, sum(C) FROM 
+			(values row(1,1), row(1,3), row(2,2), row(2,5), row(3,9)) a (b,c) 
+			group by 1 order by 1`,
+		Expected: []sql.Row{
+			{1, 4.0},
+			{2, 7.0},
+			{3, 9.0},
+		},
+	},
 	{
 		Query: `SELECT i, sum(i) FROM mytable group by 1 having avg(i) > 1 order by 1`,
 		Expected: []sql.Row{
@@ -5464,6 +5543,26 @@ var errorQueries = []QueryErrorTest{
 						ORDER BY 3`,
 		ExpectedErr: sql.ErrAmbiguousColumnInOrderBy,
 	},
+	{
+		Query:       "SELECT C FROM (select i,s FROM mytable) mt (a,b) order by a desc;",
+		ExpectedErr: sql.ErrColumnNotFound,
+	},
+	{
+		Query:       "SELECT i FROM (select i,s FROM mytable) mt (a,b) order by a desc;",
+		ExpectedErr: sql.ErrColumnNotFound,
+	},
+	{
+		Query:       "SELECT mt.i FROM (select i,s FROM mytable) mt (a,b) order by a desc;",
+		ExpectedErr: sql.ErrTableColumnNotFound,
+	},
+	{
+		Query:       "SELECT a FROM (select i,s FROM mytable) mt (a) order by a desc;",
+		ExpectedErr: sql.ErrColumnCountMismatch,
+	},
+	{
+		Query:       "SELECT a FROM (select i,s FROM mytable) mt (a,b,c) order by a desc;",
+		ExpectedErr: sql.ErrColumnCountMismatch,
+	},
 }
 
 // WriteQueryTest is a query test for INSERT, UPDATE, etc. statements. It has a query to run and a select query to

enginetest/script_queries.go

@@ -523,14 +523,6 @@ var ScriptTests = []ScriptTest{
 				Query:    "select row_count()",
 				Expected: []sql.Row{{-1}},
 			},
-			{
-				Query:    "replace into b values (1)",
-				Expected: []sql.Row{{sql.NewOkResult(1)}},
-			},
-			{
-				Query:    "select row_count()",
-				Expected: []sql.Row{{1}},
-			},
 			{
 				Query:    "alter table b add column y int null",
 				Expected: []sql.Row{},

go.mod

@@ -4,7 +4,7 @@ require (
 	github.com/VividCortex/gohistogram v1.0.0 // indirect
 	github.com/cespare/xxhash v1.1.0
 	github.com/dolthub/sqllogictest/go v0.0.0-20201105013724-5123fc66e12c
-	github.com/dolthub/vitess v0.0.0-20210330232114-7db3f027a698
+	github.com/dolthub/vitess v0.0.0-20210401223343-5adfdbfa58b0
 	github.com/fastly/go-utils v0.0.0-20180712184237-d95a45783239 // indirect
 	github.com/go-kit/kit v0.9.0
 	github.com/go-sql-driver/mysql v1.4.1

go.sum

@@ -14,8 +14,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dolthub/sqllogictest/go v0.0.0-20201105013724-5123fc66e12c h1:ZIo6IOXU3/rJK4lp83QRq1zGhQrjQQtlmE2b7H1Vv/k=
 github.com/dolthub/sqllogictest/go v0.0.0-20201105013724-5123fc66e12c/go.mod h1:siLfyv2c92W1eN/R4QqG/+RjjX5W2+gCTRjZxBjI3TY=
-github.com/dolthub/vitess v0.0.0-20210330232114-7db3f027a698 h1:/EPY8hZbFGikWDunqOW9IaQKggEhp2cYIDUC1JUwu/w=
-github.com/dolthub/vitess v0.0.0-20210330232114-7db3f027a698/go.mod h1:hUE8oSk2H5JZnvtlLBhJPYC8WZCA5AoSntdLTcBvdBM=
+github.com/dolthub/vitess v0.0.0-20210401223343-5adfdbfa58b0 h1:OGMXiIpVEVwWf/3z+mYPCMN6CqIRZE/UkqLEMPSdXGM=
+github.com/dolthub/vitess v0.0.0-20210401223343-5adfdbfa58b0/go.mod h1:hUE8oSk2H5JZnvtlLBhJPYC8WZCA5AoSntdLTcBvdBM=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/fastly/go-utils v0.0.0-20180712184237-d95a45783239 h1:Ghm4eQYC0nEPnSJdVkTrXpu9KtoVCSo1hg7mtI7G9KU=

sql/analyzer/prune_columns.go

@@ -116,21 +116,22 @@ func pruneSubqueryColumns(
 
 	columns := make(usedColumns)
 
-	// The columns coming from the parent have the subquery alias name as the
-	// source. We need to find the real table in order to prune the subquery
-	// correctly.
+	// The columns coming from the parent have the subquery alias name as the source. We need to find the real table in
+	// order to prune the subquery correctly. The columns might also have been renamed.
 	tableByCol := make(map[string]string)
-	for _, col := range n.Child.Schema() {
-		tableByCol[col.Name] = col.Source
+	for i, col := range n.Child.Schema() {
+		name := col.Name
+		if len(n.Columns) > 0 {
+			name = n.Columns[i]
+		}
+		tableByCol[name] = col.Source
 	}
 
 	for col := range parentColumns[n.Name()] {
 		table, ok := tableByCol[col]
 		if !ok {
-			// This should never happen, but better be safe than sorry.
 			return nil, fmt.Errorf("this is likely a bug: missing projected column %q on subquery %q", col, n.Name())
 		}
-
 		columns.add(table, col)
 	}
 

sql/analyzer/resolve_ctes.go

@@ -18,7 +18,6 @@ import (
 	"strings"
 
 	"github.com/dolthub/go-mysql-server/sql"
-	"github.com/dolthub/go-mysql-server/sql/expression"
 	"github.com/dolthub/go-mysql-server/sql/plan"
 )
 
@@ -101,14 +100,6 @@ func resolveCtesInNode(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope,
 	return cur, nil
 }
 
-func unalias(p sql.Expression) sql.Expression {
-	a, ok := p.(*expression.Alias)
-	if !ok {
-		return p
-	}
-	return a.Child
-}
-
 func stripWith(ctx *sql.Context, a *Analyzer, n sql.Node, ctes map[string]sql.Node) (sql.Node, error) {
 	with, ok := n.(*plan.With)
 	if !ok {
@@ -125,44 +116,7 @@ func stripWith(ctx *sql.Context, a *Analyzer, n sql.Node, ctes map[string]sql.No
 				return nil, sql.ErrColumnCountMismatch.New()
 			}
 
-			selector := func(parent sql.Node, child sql.Node, childNum int) bool {
-				switch parent.(type) {
-				case *plan.Project, *plan.GroupBy, *plan.Window:
-					return false
-				}
-				return true
-			}
-
-			child, err := plan.TransformUpWithSelector(subquery.Child, selector, func(n sql.Node) (sql.Node, error) {
-				switch n := n.(type) {
-				case *plan.Project:
-					projections := make([]sql.Expression, len(cte.Columns))
-					for i, p := range n.Projections {
-						projections[i] = expression.NewAlias(cte.Columns[i], unalias(p))
-					}
-					return n.WithExpressions(projections...)
-				case *plan.GroupBy:
-					projections := make([]sql.Expression, len(cte.Columns))
-					for i, p := range n.SelectedExprs {
-						projections[i] = expression.NewAlias(cte.Columns[i], unalias(p))
-					}
-					return plan.NewGroupBy(projections, n.GroupByExprs, n.Child), nil
-				case *plan.Window:
-					projections := make([]sql.Expression, len(cte.Columns))
-					for i, p := range n.SelectExprs {
-						projections[i] = expression.NewAlias(cte.Columns[i], unalias(p))
-					}
-					return n.WithExpressions(projections...)
-				default:
-					return n, nil
-				}
-			})
-
-			if err != nil {
-				return nil, err
-			}
-
-			subquery.Child = child
+			subquery = subquery.WithColumns(cte.Columns)
 		}
 
 		ctes[strings.ToLower(cteName)] = subquery

sql/analyzer/resolve_subqueries.go

@@ -35,6 +35,13 @@ func resolveSubqueries(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope)
 				return nil, err
 			}
 
+			if len(n.Columns) > 0 {
+				schemaLen := schemaLength(n.Child)
+				if schemaLen != len(n.Columns) {
+					return nil, sql.ErrColumnCountMismatch.New()
+				}
+			}
+
 			return n.WithChildren(stripQueryProcess(child))
 		default:
 			return n, nil

sql/expression/binary.go

@@ -58,4 +58,3 @@ func (b *Binary) WithChildren(children ...sql.Expression) (sql.Expression, error
 
 	return NewBinary(children[0]), nil
 }
-

sql/parse/parse.go

@@ -1725,10 +1725,18 @@ func tableExprToTable(
 			}
 
 			if t.As.IsEmpty() {
+				// This should be caught by the parser, but here just in case
 				return nil, ErrUnsupportedFeature.New("subquery without alias")
 			}
 
-			return plan.NewSubqueryAlias(t.As.String(), sqlparser.String(e.Select), node), nil
+			sq := plan.NewSubqueryAlias(t.As.String(), sqlparser.String(e.Select), node)
+
+			if len(e.Columns) > 0 {
+				columns := columnsToStrings(e.Columns)
+				sq = sq.WithColumns(columns)
+			}
+
+			return sq, nil
 		case *sqlparser.ValuesStatement:
 			if t.As.IsEmpty() {
 				// Parser should enforce this, but just to be safe
@@ -1739,7 +1747,14 @@ func tableExprToTable(
 				return nil, err
 			}
 
-			return plan.NewValueDerivedTable(values, t.As.String()), nil
+			vdt := plan.NewValueDerivedTable(values, t.As.String())
+
+			if len(e.Columns) > 0 {
+				columns := columnsToStrings(e.Columns)
+				vdt = vdt.WithColumns(columns)
+			}
+
+			return vdt, nil
 		default:
 			return nil, ErrUnsupportedSyntax.New(sqlparser.String(te))
 		}