SD-9457: Update AggregateRel.CopyWithExpressionRewrite to rewrite mea…

…sure functions in addition to filters
substrait-io · Jan 29, 2025 · 7a051b7 · 7a051b7
1 parent d6e63d9
commit 7a051b7
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 6 deletions.
diff --git a/expr/functions.go b/expr/functions.go
@@ -811,6 +811,11 @@ func (a *AggregateFunction) IntermediateType() (types.FuncDefArgType, error) {
 	return a.declaration.Intermediate()
 }
 
+// SetArg sets the specified argument to the provided value.  The index is not checked for validity.
+func (a *AggregateFunction) SetArg(i int, arg types.FuncArg) {
+	a.args[i] = arg
+}
+
 func (a *AggregateFunction) String() string {
 	var b strings.Builder
 

diff --git a/expr/utils.go b/expr/utils.go
@@ -32,6 +32,7 @@ func (e *ExtensionRegistry) LookupScalarFunction(anchor uint32) (*extensions.Sca
 	return e.Set.LookupScalarFunction(anchor, e.c)
 }
 
+// LookupAggregateFunction returns an AggregateFunctionVariant associated with a previously used function's anchor.
 func (e *ExtensionRegistry) LookupAggregateFunction(anchor uint32) (*extensions.AggregateFunctionVariant, bool) {
 	return e.Set.LookupAggregateFunction(anchor, e.c)
 }

diff --git a/plan/relations.go b/plan/relations.go
@@ -1168,6 +1168,30 @@ func (ar *AggregateRel) Copy(newInputs ...Rel) (Rel, error) {
 	return &aggregate, nil
 }
 
+func (ar *AggregateRel) rewriteAggregateFunc(rewriteFunc RewriteFunc, f *expr.AggregateFunction) (*expr.AggregateFunction, error) {
+	if f == nil {
+		return f, nil
+	}
+	newF := f
+	argsAreEqual := true
+	for i := 0; i < f.NArgs(); i++ {
+		arg := f.Arg(i)
+		if exp, ok := arg.(expr.Expression); ok {
+			var newExp expr.Expression
+			var err error
+			if newExp, err = rewriteFunc(exp); err != nil {
+				return nil, err
+			}
+			newF.SetArg(i, newExp)
+			argsAreEqual = argsAreEqual && exp == newExp
+		}
+	}
+	if argsAreEqual {
+		return f, nil
+	}
+	return newF, nil
+}
+
 func (ar *AggregateRel) CopyWithExpressionRewrite(rewriteFunc RewriteFunc, newInputs ...Rel) (Rel, error) {
 	if len(newInputs) != 1 {
 		return nil, substraitgo.ErrInvalidInputCount
@@ -1187,8 +1211,10 @@ func (ar *AggregateRel) CopyWithExpressionRewrite(rewriteFunc RewriteFunc, newIn
 		if newMeasures[i].filter, err = rewriteFunc(m.filter); err != nil {
 			return nil, err
 		}
-		measuresAreEqual = measuresAreEqual && newMeasures[i].filter == m.filter
-		newMeasures[i].measure = m.measure
+		if newMeasures[i].measure, err = ar.rewriteAggregateFunc(rewriteFunc, m.measure); err != nil {
+			return nil, err
+		}
+		measuresAreEqual = measuresAreEqual && newMeasures[i].filter == m.filter && newMeasures[i].measure == m.measure
 	}
 	if groupsAreEqual && measuresAreEqual && newInputs[0] == ar.input {
 		return ar, nil

diff --git a/plan/relations_test.go b/plan/relations_test.go
@@ -1,6 +1,8 @@
 package plan
 
 import (
+	"github.com/stretchr/testify/require"
+	"github.com/substrait-io/substrait-go/v3/extensions"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -26,8 +28,19 @@ func createPrimitiveBool(value bool) expr.Expression {
 }
 
 func TestRelations_Copy(t *testing.T) {
-	aggregateRel := &AggregateRel{input: createVirtualTableReadRel(1), groupingExpressions: []expr.Expression{createPrimitiveFloat(1.0)}, groupingReferences: [][]uint32{{0}},
-		measures: []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(false, false)}}}
+	extReg := expr.NewExtensionRegistry(extensions.NewSet(), &extensions.DefaultCollection)
+	aggregateFnID := extensions.ID{
+		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		Name: "avg",
+	}
+	aggregateFn, err := expr.NewAggregateFunc(extReg,
+		aggregateFnID, nil, types.AggInvocationAll,
+		types.AggPhaseInitialToResult, nil, createPrimitiveFloat(1.0))
+	require.NoError(t, err)
+
+	aggregateRel := &AggregateRel{input: createVirtualTableReadRel(1),
+		groupingExpressions: []expr.Expression{createPrimitiveFloat(1.0)}, groupingReferences: [][]uint32{{0}},
+		measures: []AggRelMeasure{{measure: aggregateFn, filter: expr.NewPrimitiveLiteral(false, false)}}}
 	crossRel := &CrossRel{left: createVirtualTableReadRel(1), right: createVirtualTableReadRel(2)}
 	extensionLeafRel := &ExtensionLeafRel{}
 	extensionMultiRel := &ExtensionMultiRel{inputs: []Rel{createVirtualTableReadRel(1), createVirtualTableReadRel(2)}}
@@ -91,8 +104,10 @@ func TestRelations_Copy(t *testing.T) {
 				}
 				panic("unexpected expression type")
 			},
-			expectedRel: &AggregateRel{input: createVirtualTableReadRel(8), groupingExpressions: []expr.Expression{createPrimitiveFloat(9.0)}, groupingReferences: [][]uint32{{0}},
-				measures: []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(true, false)}}},
+			expectedRel: &AggregateRel{input: createVirtualTableReadRel(8),
+				groupingExpressions: []expr.Expression{createPrimitiveFloat(9.0)},
+				groupingReferences:  [][]uint32{{0}},
+				measures:            []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(true, false)}}},
 		},
 		{
 			name:            "ExtensionLeafRel Copy with new inputs",