Repository: src-d/go-mysql-server
Pull request #368 (2018-09-19 10:11:05)

sql/analyzer: resolve aggregate aliases in groupby grouping
Closes https://github.com/src-d/gitbase/issues/460


Patch:
diff --git a/sql/analyzer/resolve_columns.go b/sql/analyzer/resolve_columns.go
index e04ac67..0295ea1 100644
--- a/sql/analyzer/resolve_columns.go
+++ b/sql/analyzer/resolve_columns.go
@@ -1,6 +1,8 @@
 package analyzer
 
 import (
+	"fmt"
+	"sort"
 	"strings"
 
 	errors "gopkg.in/src-d/go-errors.v1"
@@ -267,6 +269,167 @@ func resolveColumns(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error)
 	})
 }
 
+// resolveGroupingColumns reorders the aggregation in a groupby so aliases
+// defined in it can be resolved in the grouping of the groupby. To do so,
+// all aliases are pushed down to a projection node under the group by.
+func resolveGroupingColumns(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
+	a.Log("resoving group columns")
+	if n.Resolved() {
+		return n, nil
+	}
+
+	return n.TransformUp(func(n sql.Node) (sql.Node, error) {
+		g, ok := n.(*plan.GroupBy)
+		if n.Resolved() || !ok || len(g.Grouping) == 0 {
+			return n, nil
+		}
+
+		// The reason we have two sets of columns, one for grouping and
+		// one for aggregate is because an alias can redefine a column name
+		// of the child schema. In the grouping, if that column is referenced
+		// it refers to the alias, and not the one in the child. However,
+		// in the aggregate, aliases in that same aggregate cannot be used,
+		// so it refers to the column in the child node.
+		var groupingColumns = make(map[string]struct{})
+		for _, g := range g.Grouping {
+			for _, n := range findAllColumns(g) {
+				groupingColumns[n] = struct{}{}
+			}
+		}
+
+		var aggregateColumns = make(map[string]struct{})
+		for _, agg := range g.Aggregate {
+			// This alias is going to be pushed down, so don't bother gathering
+			// its requirements.
+			if alias, ok := agg.(*expression.Alias); ok {
+				if _, ok := groupingColumns[alias.Name()]; ok {
+					continue
+				}
+			}
+
+			for _, n := range findAllColumns(agg) {
+				aggregateColumns[n] = struct{}{}
+			}
+		}
+
+		var newAggregate []sql.Expression
+		var projection []sql.Expression
+		// Aliases will keep the aliases that have been pushed down and their
+		// index in the new aggregate.
+		var aliases = make(map[string]int)
+
+		var needsReorder bool
+		for _, a := range g.Aggregate {
+			alias, ok := a.(*expression.Alias)
+			// Note that aliases of aggregations cannot be used in the grouping
+			// because the grouping is needed before computing the aggregation.
+			if !ok || containsAggregation(alias) {
+				newAggregate = append(newAggregate, a)
+				continue
+			}
+
+			// Only if the alias is required in the grouping set needsReorder
+			// to true. If it's not required, there's no need for a reorder if
+			// no other alias is required.
+			_, ok = groupingColumns[alias.Name()]
+			if ok {
+				aliases[alias.Name()] = len(newAggregate)
+				needsReorder = true
+				delete(groupingColumns, alias.Name())
+
+				projection = append(projection, a)
+				newAggregate = append(newAggregate, expression.NewUnresolvedColumn(alias.Name()))
+			} else {
+				newAggregate = append(newAggregate, a)
+			}
+		}
+
+		if !needsReorder {
+			return n, nil
+		}
+
+		// Instead of iterating columns directly, we want them sorted so the
+		// executions of the rule are consistent.
+		var missingCols = make([]string, 0, len(aggregateColumns)+len(groupingColumns))
+		for col := range aggregateColumns {
+			missingCols = append(missingCols, col)
+		}
+		for col := range groupingColumns {
+			missingCols = append(missingCols, col)
+		}
+		sort.Strings(missingCols)
+
+		var renames = make(map[string]string)
+		// All columns required by expressions in both grouping and aggregation
+		// must also be projected in the new projection node or they will not
+		// be able to resolve.
+		for _, col := range missingCols {
+			name := col
+			// If an alias has been pushed down with the same name as a missing
+			// column, there will be a conflict of names. We must find an unique name
+			// for the missing column.
+			if _, ok := aliases[col]; ok {
+				for i := 1; ; i++ {
+					name = fmt.Sprintf("%s_%02d", col, i)
+					if !stringContains(missingCols, name) {
+						break
+					}
+				}
+			}
+
+			if name == col {
+				projection = append(projection, expression.NewUnresolvedColumn(col))
+			} else {
+				renames[col] = name
+				projection = append(projection, expression.NewAlias(
+					expression.NewUnresolvedColumn(col),
+					name,
+				))
+			}
+		}
+
+		// If there is any name conflict between columns we need to rename every
+		// usage inside the aggregate.
+		if len(renames) > 0 {
+			for i, expr := range newAggregate {
+				var err error
+				newAggregate[i], err = expr.TransformUp(func(e sql.Expression) (sql.Expression, error) {
+					col, ok := e.(*expression.UnresolvedColumn)
+					if ok {
+						// We need to make sure we don't rename the reference to the
+						// pushed down alias.
+						if to, ok := renames[col.Name()]; ok && aliases[col.Name()] != i {
+							return expression.NewUnresolvedColumn(to), nil
+						}
+					}
+
+					return e, nil
+				})
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		return plan.NewGroupBy(
+			newAggregate, g.Grouping,
+			plan.NewProject(projection, g.Child),
+		), nil
+	})
+}
+
+func findAllColumns(e sql.Expression) []string {
+	var cols []string
+	expression.Inspect(e, func(e sql.Expression) bool {
+		col, ok := e.(*expression.UnresolvedColumn)
+		if ok {
+			cols = append(cols, col.Name())
+		}
+		return true
+	})
+	return cols
+}
+
 func dedupStrings(in []string) []string {
 	var seen = make(map[string]struct{})
 	var result []string
diff --git a/sql/analyzer/rules.go b/sql/analyzer/rules.go
index c2a08a3..47035ac 100644
--- a/sql/analyzer/rules.go
+++ b/sql/analyzer/rules.go
@@ -9,6 +9,7 @@ var DefaultRules = []Rule{
 	{"resolve_natural_joins", resolveNaturalJoins},
 	{"resolve_orderby_literals", resolveOrderByLiterals},
 	{"resolve_orderby", resolveOrderBy},
+	{"resolve_grouping_columns", resolveGroupingColumns},
 	{"qualify_columns", qualifyColumns},
 	{"resolve_columns", resolveColumns},
 	{"resolve_database", resolveDatabase},


Test patch:
diff --git a/engine_test.go b/engine_test.go
index 5e26013..d3372de 100644
--- a/engine_test.go
+++ b/engine_test.go
@@ -127,6 +127,14 @@ var queries = []struct {
 			{int32(1), "third row"},
 		},
 	},
+	{
+		`SELECT COUNT(*) as cnt, s as fi FROM mytable GROUP BY fi`,
+		[]sql.Row{
+			{int32(1), "first row"},
+			{int32(1), "second row"},
+			{int32(1), "third row"},
+		},
+	},
 	{
 		"SELECT CAST(-3 AS UNSIGNED) FROM mytable",
 		[]sql.Row{
diff --git a/sql/analyzer/resolve_columns_test.go b/sql/analyzer/resolve_columns_test.go
index 3d42364..e1165be 100644
--- a/sql/analyzer/resolve_columns_test.go
+++ b/sql/analyzer/resolve_columns_test.go
@@ -226,3 +226,70 @@ func TestResolveColumnsSession(t *testing.T) {
 
 	require.Equal(expected, result)
 }
+
+func TestResolveGroupingColumns(t *testing.T) {
+	require := require.New(t)
+
+	a := NewDefault(nil)
+	node := plan.NewGroupBy(
+		[]sql.Expression{
+			expression.NewAlias(
+				expression.NewUnresolvedFunction("foo", true,
+					expression.NewUnresolvedColumn("c"),
+				),
+				"c",
+			),
+			expression.NewAlias(
+				expression.NewUnresolvedColumn("d"),
+				"b",
+			),
+			expression.NewUnresolvedFunction("bar", false,
+				expression.NewUnresolvedColumn("b"),
+			),
+		},
+		[]sql.Expression{
+			expression.NewUnresolvedColumn("a"),
+			expression.NewUnresolvedColumn("b"),
+		},
+		plan.NewResolvedTable(mem.NewTable("table", nil)),
+	)
+
+	expected := plan.NewGroupBy(
+		[]sql.Expression{
+			expression.NewAlias(
+				expression.NewUnresolvedFunction("foo", true,
+					expression.NewUnresolvedColumn("c"),
+				),
+				"c",
+			),
+			expression.NewUnresolvedColumn("b"),
+			expression.NewUnresolvedFunction("bar", false,
+				expression.NewUnresolvedColumn("b_01"),
+			),
+		},
+		[]sql.Expression{
+			expression.NewUnresolvedColumn("a"),
+			expression.NewUnresolvedColumn("b"),
+		},
+		plan.NewProject(
+			[]sql.Expression{
+				expression.NewAlias(
+					expression.NewUnresolvedColumn("d"),
+					"b",
+				),
+				expression.NewUnresolvedColumn("a"),
+				expression.NewAlias(
+					expression.NewUnresolvedColumn("b"),
+					"b_01",
+				),
+				expression.NewUnresolvedColumn("c"),
+			},
+			plan.NewResolvedTable(mem.NewTable("table", nil)),
+		),
+	)
+
+	result, err := resolveGroupingColumns(sql.NewEmptyContext(), a, node)
+	require.NoError(err)
+
+	require.Equal(expected, result)
+}