feat: AggregateRel should use grouping references

plan/builders.go

@@ -75,6 +75,7 @@ type Builder interface {
 	ProjectRemap(input Rel, remap []int32, exprs ...expr.Expression) (*ProjectRel, error)
 	// Deprecated: Use AggregateColumns(...).Remap() instead.
 	AggregateColumnsRemap(input Rel, remap []int32, measures []AggRelMeasure, groupByCols ...int32) (*AggregateRel, error)
+	// Deprecated: Use Aggregate(...) instead.
 	AggregateColumns(input Rel, measures []AggRelMeasure, groupByCols ...int32) (*AggregateRel, error)
 	// Deprecated: Use AggregateExprs(...).Remap() instead.
 	AggregateExprsRemap(input Rel, remap []int32, measures []AggRelMeasure, groups ...[]expr.Expression) (*AggregateRel, error)
@@ -286,11 +287,13 @@ func (b *builder) AggregateColumnsRemap(input Rel, remap []int32, measures []Agg
 		}
 	}
 
+	groupingExpressions, groupingReferences := groupingExprs(exprs)
 	return &AggregateRel{
-		RelCommon: RelCommon{mapping: remap},
-		input:     input,
-		groups:    exprs,
-		measures:  measures,
+		RelCommon:           RelCommon{mapping: remap},
+		input:               input,
+		measures:            measures,
+		groupingReferences:  groupingReferences,
+		groupingExpressions: groupingExpressions,
 	}, nil
 }
 
@@ -321,18 +324,53 @@ func (b *builder) AggregateExprsRemap(input Rel, remap []int32, measures []AggRe
 		}
 	}
 
+	groupingExpressions, groupingReferences := groupingExprs(groups)
 	return &AggregateRel{
-		RelCommon: RelCommon{mapping: remap},
-		input:     input,
-		groups:    groups,
-		measures:  measures,
+		RelCommon:           RelCommon{mapping: remap},
+		input:               input,
+		measures:            measures,
+		groupingReferences:  groupingReferences,
+		groupingExpressions: groupingExpressions,
 	}, nil
 }
 
 func (b *builder) AggregateExprs(input Rel, measures []AggRelMeasure, groups ...[]expr.Expression) (*AggregateRel, error) {
 	return b.AggregateExprsRemap(input, nil, measures, groups...)
 }
 
+func (b *builder) Aggregate(input Rel, measures []AggRelMeasure, groupingExpressions []expr.Expression, groupingReferences [][]uint32) (*AggregateRel, error) {
+	if input == nil {
+		return nil, errNilInputRel
+	}
+
+	if (len(measures) + len(groupingExpressions)) == 0 {
+		return nil, fmt.Errorf("%w: must have at least one grouping expression or measure for AggregateRel", substraitgo.ErrInvalidRel)
+	}
+
+	if slices.ContainsFunc(groupingExpressions, func(expr expr.Expression) bool {
+		return expr == nil
+	}) {
+		return nil, fmt.Errorf("%w: groupings cannot contain empty expression list or nil expression", substraitgo.ErrInvalidRel)
+	}
+
+	// Validate groupingReferences indices.
+	for _, refList := range groupingReferences {
+		for _, ref := range refList {
+			if ref >= uint32(len(groupingExpressions)) {
+				return nil, fmt.Errorf("%w: groupingReferences contains invalid indices", substraitgo.ErrInvalidRel)
+			}
+		}
+	}
+
+	return &AggregateRel{
+		RelCommon:           RelCommon{},
+		input:               input,
+		measures:            measures,
+		groupingReferences:  groupingReferences,
+		groupingExpressions: groupingExpressions,
+	}, nil
+}
+
 func (b *builder) CreateTableAsSelectRemap(input Rel, remap []int32, tableName []string, schema types.NamedStruct) (*NamedTableWriteRel, error) {
 	if input == nil {
 		return nil, errNilInputRel

plan/plan.go

@@ -50,6 +50,33 @@ func init() {
 	}
 }
 
+// groupingExprs takes 2-dimensional slice of expressions and returns
+// a single slice of unique expressions and a slice of references to
+// the unique expressions for each group.
+func groupingExprs(groups [][]expr.Expression) ([]expr.Expression, [][]uint32) {
+	groupingExpressions := make([]expr.Expression, 0)
+	groupingReferences := make([][]uint32, 0)
+	for _, group := range groups {
+		refs := make([]uint32, 0)
+		for _, expr := range group {
+			existingExpr := false
+			for eIndex, existing := range groupingExpressions {
+				if existing.Equals(expr) {
+					existingExpr = true
+					refs = append(refs, uint32(eIndex))
+					break
+				}
+			}
+			if !existingExpr {
+				groupingExpressions = append(groupingExpressions, expr)
+				refs = append(refs, uint32(len(groupingExpressions)-1))
+			}
+		}
+		groupingReferences = append(groupingReferences, refs)
+	}
+	return groupingExpressions, groupingReferences
+}
+
 // Relation is either a Root relation (a relation + list of column names)
 // or another relation (such as a CTE or other reference).
 type Relation struct {
@@ -401,16 +428,32 @@ func RelFromProto(rel *proto.Rel, reg expr.ExtensionRegistry) (Rel, error) {
 		}
 
 		base := input.RecordType()
-		groups := make([][]expr.Expression, len(rel.Aggregate.Groupings))
-		for i, g := range rel.Aggregate.Groupings {
-			groups[i] = make([]expr.Expression, len(g.GroupingExpressions))
-			for j, e := range g.GroupingExpressions {
-				groups[i][j], err = expr.ExprFromProto(e, &base, reg)
+		var groupingExpressions []expr.Expression
+		var groupingReferences [][]uint32
+		if len(rel.Aggregate.GroupingExpressions) > 0 {
+			for _, e := range rel.Aggregate.GroupingExpressions {
+				expr, err := expr.ExprFromProto(e, &base, reg)
 				if err != nil {
-					return nil, fmt.Errorf("error getting grouping expr [%d][%d] for AggregateRel: %w",
-						i, j, err)
+					return nil, fmt.Errorf("error getting grouping expr for AggregateRel: %w", err)
 				}
+				groupingExpressions = append(groupingExpressions, expr)
+			}
+			for _, g := range rel.Aggregate.Groupings {
+				groupingReferences = append(groupingReferences, g.ExpressionReferences)
 			}
+		} else { // support old style grouping for backward compatibility
+			groups := make([][]expr.Expression, len(rel.Aggregate.Groupings))
+			for i, g := range rel.Aggregate.Groupings {
+				groups[i] = make([]expr.Expression, len(g.GroupingExpressions))
+				for j, e := range g.GroupingExpressions {
+					groups[i][j], err = expr.ExprFromProto(e, &base, reg)
+					if err != nil {
+						return nil, fmt.Errorf("error getting grouping expr [%d][%d] for AggregateRel: %w",
+							i, j, err)
+					}
+				}
+			}
+			groupingExpressions, groupingReferences = groupingExprs(groups)
 		}
 
 		measures := make([]AggRelMeasure, len(rel.Aggregate.Measures))
@@ -429,10 +472,11 @@ func RelFromProto(rel *proto.Rel, reg expr.ExtensionRegistry) (Rel, error) {
 		}
 
 		out := &AggregateRel{
-			input:        input,
-			groups:       groups,
-			measures:     measures,
-			advExtension: rel.Aggregate.AdvancedExtension,
+			input:               input,
+			measures:            measures,
+			groupingReferences:  groupingReferences,
+			groupingExpressions: groupingExpressions,
+			advExtension:        rel.Aggregate.AdvancedExtension,
 		}
 		out.fromProtoCommon(rel.Aggregate.Common)
 		return out, nil

plan/plan_builder_test.go

@@ -218,15 +218,18 @@ func TestAggregateRelPlan(t *testing.T) {
 									"namedTable": { "names": [ "test" ]}
 								}
 							},
+							"groupingExpressions": [
+								{
+									"selection": {
+										"rootReference": {},
+										"directReference": { "structField": { "field": 0 }}
+									}
+								}
+							],
 							"groupings": [
 								{
-									"groupingExpressions": [
-										{
-											"selection": {
-												"rootReference": {},
-												"directReference": { "structField": { "field": 0 }}
-											}
-										}
+									"expressionReferences": [
+										0
 									]
 								}
 							],
@@ -265,6 +268,20 @@ func TestAggregateRelPlan(t *testing.T) {
 	assert.Equal(t, "NSTRUCT<val: string, cnt: i64>", p.GetRoots()[0].RecordType().String())
 
 	checkRoundTrip(t, expectedJSON, p)
+
+	// Test with grouping expressions and references
+	ref, err := b.RootFieldRef(scan, 0)
+	require.NoError(t, err)
+	exprs := make([]expr.Expression, 0)
+	exprs = append(exprs, ref)
+	root, err = b.AggregateExprs(scan, []plan.AggRelMeasure{b.Measure(aggCount, nil)}, [][]expr.Expression{exprs}...)
+	require.NoError(t, err)
+
+	p, err = b.Plan(root, []string{"val", "cnt"})
+	require.NoError(t, err)
+	assert.Equal(t, "NSTRUCT<val: string, cnt: i64>", p.GetRoots()[0].RecordType().String())
+
+	checkRoundTrip(t, expectedJSON, p)
 }
 
 func TestAggregateNoGrouping(t *testing.T) {

plan/relations.go

@@ -969,17 +969,19 @@ func (am *AggRelMeasure) ToProto() *proto.AggregateRel_Measure {
 type AggregateRel struct {
 	RelCommon
 
-	input        Rel
-	groups       [][]expr.Expression
-	measures     []AggRelMeasure
-	advExtension *extensions.AdvancedExtension
+	input               Rel
+	measures            []AggRelMeasure
+	groupingExpressions []expr.Expression
+	groupingReferences  [][]uint32
+	advExtension        *extensions.AdvancedExtension
 }
 
 func (ar *AggregateRel) directOutputSchema() types.RecordType {
-	groupTypes := make([]types.Type, 0, len(ar.groups)+len(ar.measures))
-	for _, g := range ar.groups {
+	groupTypes := make([]types.Type, 0, len(ar.groupingReferences)+len(ar.measures))
+	for _, g := range ar.groupingReferences {
 		for _, e := range g {
-			groupTypes = append(groupTypes, e.GetType())
+			expression := ar.groupingExpressions[e]
+			groupTypes = append(groupTypes, expression.GetType())
 		}
 	}
 
@@ -996,23 +998,23 @@ func (ar *AggregateRel) RecordType() types.RecordType {
 
 func (ar *AggregateRel) Input() Rel { return ar.input }
 
-// Groupings is a list of expression groupings that the aggregation measures should
-// be calculated for.
-func (ar *AggregateRel) Groupings() [][]expr.Expression { return ar.groups }
-func (ar *AggregateRel) Measures() []AggRelMeasure      { return ar.measures }
+func (ar *AggregateRel) GroupingExpressions() []expr.Expression { return ar.groupingExpressions }
+func (ar *AggregateRel) GroupingReferences() [][]uint32         { return ar.groupingReferences }
+func (ar *AggregateRel) Measures() []AggRelMeasure              { return ar.measures }
 func (ar *AggregateRel) GetAdvancedExtension() *extensions.AdvancedExtension {
 	return ar.advExtension
 }
 
 func (ar *AggregateRel) ToProto() *proto.Rel {
-	groupings := make([]*proto.AggregateRel_Grouping, len(ar.groups))
-	for i, ex := range ar.groups {
-		groupExprs := make([]*proto.Expression, len(ex))
-		for j, e := range ex {
-			groupExprs[j] = e.ToProto()
-		}
+	groupingExpressionsProto := make([]*proto.Expression, len(ar.groupingExpressions))
+	for i, e := range ar.groupingExpressions {
+		groupingExpressionsProto[i] = e.ToProto()
+	}
+
+	groupings := make([]*proto.AggregateRel_Grouping, len(ar.groupingReferences))
+	for i, _ := range ar.groupingReferences {
 		groupings[i] = &proto.AggregateRel_Grouping{
-			GroupingExpressions: groupExprs,
+			ExpressionReferences: ar.groupingReferences[i],
 		}
 	}
 
@@ -1024,11 +1026,12 @@ func (ar *AggregateRel) ToProto() *proto.Rel {
 	return &proto.Rel{
 		RelType: &proto.Rel_Aggregate{
 			Aggregate: &proto.AggregateRel{
-				Common:            ar.toProto(),
-				Input:             ar.input.ToProto(),
-				Groupings:         groupings,
-				Measures:          measures,
-				AdvancedExtension: ar.advExtension,
+				Common:              ar.toProto(),
+				Input:               ar.input.ToProto(),
+				GroupingExpressions: groupingExpressionsProto,
+				Groupings:           groupings,
+				Measures:            measures,
+				AdvancedExtension:   ar.advExtension,
 			},
 		},
 	}
@@ -1061,17 +1064,13 @@ func (ar *AggregateRel) CopyWithExpressionRewrite(rewriteFunc RewriteFunc, newIn
 	}
 	var err error
 	groupsAreEqual := true
-	newGroups := make([][]expr.Expression, len(ar.groups))
-	for i, g := range ar.groups {
-		newGroups[i] = make([]expr.Expression, len(g))
-		for j, e := range g {
-			if newGroups[i][j], err = rewriteFunc(e); err != nil {
-				return nil, err
-			}
-			groupsAreEqual = groupsAreEqual && newGroups[i][j] == e
+	newGroupingExpressions := make([]expr.Expression, len(ar.groupingExpressions))
+	for i, e := range ar.groupingExpressions {
+		if newGroupingExpressions[i], err = rewriteFunc(e); err != nil {
+			return nil, err
 		}
+		groupsAreEqual = groupsAreEqual && newGroupingExpressions[i] == e
 	}
-
 	measuresAreEqual := true
 	newMeasures := make([]AggRelMeasure, len(ar.measures))
 	for i, m := range ar.measures {
@@ -1086,7 +1085,8 @@ func (ar *AggregateRel) CopyWithExpressionRewrite(rewriteFunc RewriteFunc, newIn
 	}
 	aggregate := *ar
 	aggregate.input = newInputs[0]
-	aggregate.groups = newGroups
+	aggregate.groupingExpressions = newGroupingExpressions
+	aggregate.measures = newMeasures
 	return &aggregate, nil
 }
 

plan/relations_test.go

@@ -26,8 +26,8 @@ func createPrimitiveBool(value bool) expr.Expression {
 }
 
 func TestRelations_Copy(t *testing.T) {
-	aggregateRel := &AggregateRel{input: createVirtualTableReadRel(1), groups: [][]expr.Expression{{createPrimitiveFloat(1.0)}},
-		measures: []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(true, false)}}}
+	aggregateRel := &AggregateRel{input: createVirtualTableReadRel(1), groupingExpressions: []expr.Expression{createPrimitiveFloat(1.0)}, groupingReferences: [][]uint32{{0}},
+		measures: []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(false, false)}}}
 	crossRel := &CrossRel{left: createVirtualTableReadRel(1), right: createVirtualTableReadRel(2)}
 	extensionLeafRel := &ExtensionLeafRel{}
 	extensionMultiRel := &ExtensionMultiRel{inputs: []Rel{createVirtualTableReadRel(1), createVirtualTableReadRel(2)}}
@@ -57,7 +57,7 @@ func TestRelations_Copy(t *testing.T) {
 			name:        "AggregateRel Copy with new inputs",
 			relation:    aggregateRel,
 			newInputs:   []Rel{createVirtualTableReadRel(6)},
-			expectedRel: &AggregateRel{input: createVirtualTableReadRel(6), groups: aggregateRel.groups, measures: aggregateRel.measures},
+			expectedRel: &AggregateRel{input: createVirtualTableReadRel(6), groupingReferences: aggregateRel.groupingReferences, groupingExpressions: aggregateRel.groupingExpressions, measures: aggregateRel.measures},
 		},
 		{
 			name:            "AggregateRel Copy with same inputs and noOpRewrite",
@@ -70,7 +70,7 @@ func TestRelations_Copy(t *testing.T) {
 			name:        "AggregateRel Copy with new Inputs and noOpReWrite",
 			relation:    aggregateRel,
 			newInputs:   []Rel{createVirtualTableReadRel(7)},
-			expectedRel: &AggregateRel{input: createVirtualTableReadRel(7), groups: aggregateRel.groups, measures: aggregateRel.measures},
+			expectedRel: &AggregateRel{input: createVirtualTableReadRel(7), groupingExpressions: aggregateRel.groupingExpressions, groupingReferences: aggregateRel.groupingReferences, measures: aggregateRel.measures},
 		},
 		{
 			name:      "AggregateRel Copy with new Inputs and reWriteFunc",
@@ -85,7 +85,7 @@ func TestRelations_Copy(t *testing.T) {
 				}
 				panic("unexpected expression type")
 			},
-			expectedRel: &AggregateRel{input: createVirtualTableReadRel(8), groups: [][]expr.Expression{{createPrimitiveFloat(9.0)}},
+			expectedRel: &AggregateRel{input: createVirtualTableReadRel(8), groupingExpressions: []expr.Expression{createPrimitiveFloat(9.0)}, groupingReferences: [][]uint32{{0}},
 				measures: []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(true, false)}}},
 		},
 		{