Merge pull request #749 from dolthub/max/groupby-flatten-bug

GroupBy normalization rule maintains parent Projection dependencies

Author: max-hoffman

enginetest/queries.go

@@ -5779,6 +5779,14 @@ var QueryTests = []QueryTest{
 			{3, 3},
 		},
 	},
+	{
+		Query: "select sum(x.i) + y.i from mytable as x, mytable as y where x.i = y.i GROUP BY x.i",
+		Expected: []sql.Row{
+			{float64(2)},
+			{float64(4)},
+			{float64(6)},
+		},
+	},
 	{
 		Query:    "SELECT 2.0 + CAST(5 AS DECIMAL)",
 		Expected: []sql.Row{{float64(7)}},

sql/analyzer/aggregations.go

@@ -75,13 +75,18 @@ func flattenedGroupBy(ctx *sql.Context, projection, grouping []sql.Expression, c
 func replaceAggregatesWithGetFieldProjections(ctx *sql.Context, projection []sql.Expression) (projections, aggregations []sql.Expression, err error) {
 	var newProjection = make([]sql.Expression, len(projection))
 	var newAggregates []sql.Expression
-
+	allGetFields := make(map[int]sql.Expression)
+	projDeps := make(map[int]struct{})
 	for i, p := range projection {
 		var transformed bool
 		e, err := expression.TransformUp(p, func(e sql.Expression) (sql.Expression, error) {
 			switch e := e.(type) {
 			case sql.Aggregation, sql.WindowAggregation:
-				// continue on
+			// continue on
+			case *expression.GetField:
+				allGetFields[e.Index()] = e
+				projDeps[e.Index()] = struct{}{}
+				return e, nil
 			default:
 				return e, nil
 			}
@@ -107,6 +112,24 @@ func replaceAggregatesWithGetFieldProjections(ctx *sql.Context, projection []sql
 		}
 	}
 
+	// find subset of allGetFields not covered by newAggregates
+	newAggDeps := make(map[int]struct{}, 0)
+	for _, agg := range newAggregates {
+		_, _ = expression.TransformUp(agg, func(e sql.Expression) (sql.Expression, error) {
+			switch e := e.(type) {
+			case *expression.GetField:
+				newAggDeps[e.Index()] = struct{}{}
+			}
+			return e, nil
+		})
+	}
+	for i, _ := range projDeps {
+		if _, ok := newAggDeps[i]; !ok {
+			// add pass-through dependency
+			newAggregates = append(newAggregates, allGetFields[i])
+		}
+	}
+
 	return newProjection, newAggregates, nil
 }
 

sql/analyzer/aggregations_test.go

@@ -168,6 +168,46 @@ func TestFlattenAggregationExprs(t *testing.T) {
 				),
 			),
 		},
+		{
+			name: "aggregate with pass through column dependency",
+			node: plan.NewGroupBy(
+				[]sql.Expression{
+					expression.NewArithmetic(
+						aggregation.NewSum(
+							expression.NewGetFieldWithTable(0, sql.Int64, "foo", "a", false),
+						),
+						expression.NewGetFieldWithTable(1, sql.Int64, "bar", "a", false),
+						"+",
+					),
+				},
+				[]sql.Expression{
+					expression.NewGetFieldWithTable(0, sql.Int64, "foo", "a", false),
+				},
+				plan.NewResolvedTable(table, nil, nil),
+			),
+
+			expected: plan.NewProject(
+				[]sql.Expression{
+					expression.NewArithmetic(
+						expression.NewGetField(0, sql.Float64, "SUM(foo.a)", false),
+						expression.NewGetFieldWithTable(1, sql.Int64, "bar", "a", false),
+						"+",
+					),
+				},
+				plan.NewGroupBy(
+					[]sql.Expression{
+						aggregation.NewSum(
+							expression.NewGetFieldWithTable(0, sql.Int64, "foo", "a", false),
+						),
+						expression.NewGetFieldWithTable(1, sql.Int64, "bar", "a", false),
+					},
+					[]sql.Expression{
+						expression.NewGetFieldWithTable(0, sql.Int64, "foo", "a", false),
+					},
+					plan.NewResolvedTable(table, nil, nil),
+				),
+			),
+		},
 	}
 
 	for _, test := range tests {