create clone to prevent accidental modification of original argument …

…list
substrait-io · Jan 29, 2025 · 4d751cc · 4d751cc
1 parent fd3495d
commit 4d751cc
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 13 deletions.
diff --git a/expr/functions.go b/expr/functions.go
@@ -811,6 +811,23 @@ func (a *AggregateFunction) IntermediateType() (types.FuncDefArgType, error) {
 	return a.declaration.Intermediate()
 }
 
+func (a *AggregateFunction) Clone() *AggregateFunction {
+	newA := *a
+	newA.args = make([]types.FuncArg, len(a.args), len(a.args))
+	for i := 0; i < len(a.args); i++ {
+		newA.args[i] = a.args[i]
+	}
+	newA.options = make([]*types.FunctionOption, len(a.options), len(a.options))
+	for i := 0; i < len(a.options); i++ {
+		newA.options[i] = a.options[i]
+	}
+	newA.Sorts = make([]SortField, len(a.Sorts), len(a.Sorts))
+	for i := 0; i < len(a.Sorts); i++ {
+		newA.Sorts[i] = a.Sorts[i]
+	}
+	return &newA
+}
+
 // 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

diff --git a/plan/relations.go b/plan/relations.go
@@ -1172,7 +1172,7 @@ func (ar *AggregateRel) rewriteAggregateFunc(rewriteFunc RewriteFunc, f *expr.Ag
 	if f == nil {
 		return f, nil
 	}
-	newF := f
+	newF := f.Clone()
 	argsAreEqual := true
 	for i := 0; i < f.NArgs(); i++ {
 		arg := f.Arg(i)

diff --git a/plan/relations_test.go b/plan/relations_test.go
@@ -37,10 +37,15 @@ func TestRelations_Copy(t *testing.T) {
 		aggregateFnID, nil, types.AggInvocationAll,
 		types.AggPhaseInitialToResult, nil, createPrimitiveFloat(1.0))
 	require.NoError(t, err)
+	aggregateFnRevised, err := expr.NewAggregateFunc(extReg,
+		aggregateFnID, nil, types.AggInvocationAll,
+		types.AggPhaseInitialToResult, nil, createPrimitiveFloat(9.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)}}}
+		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)}}
@@ -73,10 +78,13 @@ func TestRelations_Copy(t *testing.T) {
 	}
 	testCases := []relationTestCase{
 		{
-			name:        "AggregateRel Copy with new inputs",
-			relation:    aggregateRel,
-			newInputs:   []Rel{createVirtualTableReadRel(6)},
-			expectedRel: &AggregateRel{input: createVirtualTableReadRel(6), groupingReferences: aggregateRel.groupingReferences, groupingExpressions: aggregateRel.groupingExpressions, measures: aggregateRel.measures},
+			name:      "AggregateRel Copy with new inputs",
+			relation:  aggregateRel,
+			newInputs: []Rel{createVirtualTableReadRel(6)},
+			expectedRel: &AggregateRel{input: createVirtualTableReadRel(6),
+				groupingReferences:  aggregateRel.groupingReferences,
+				groupingExpressions: aggregateRel.groupingExpressions,
+				measures:            aggregateRel.measures},
 		},
 		{
 			name:            "AggregateRel Copy with same inputs and noOpRewrite",
@@ -86,13 +94,16 @@ func TestRelations_Copy(t *testing.T) {
 			expectedSameRel: true,
 		},
 		{
-			name:        "AggregateRel Copy with new Inputs and noOpReWrite",
-			relation:    aggregateRel,
-			newInputs:   []Rel{createVirtualTableReadRel(7)},
-			expectedRel: &AggregateRel{input: createVirtualTableReadRel(7), groupingExpressions: aggregateRel.groupingExpressions, groupingReferences: aggregateRel.groupingReferences, measures: aggregateRel.measures},
+			name:      "AggregateRel Copy with new Inputs and noOpReWrite",
+			relation:  aggregateRel,
+			newInputs: []Rel{createVirtualTableReadRel(7)},
+			expectedRel: &AggregateRel{input: createVirtualTableReadRel(7),
+				groupingExpressions: aggregateRel.groupingExpressions,
+				groupingReferences:  aggregateRel.groupingReferences,
+				measures:            aggregateRel.measures},
 		},
 		{
-			name:      "AggregateRel Copy with new Inputs and reWriteFunc",
+			name:      "AggregateRel Copy with new Inputs and rewriteFunc",
 			relation:  aggregateRel,
 			newInputs: []Rel{createVirtualTableReadRel(8)},
 			rewriteFunc: func(expression expr.Expression) (expr.Expression, error) {
@@ -107,7 +118,7 @@ func TestRelations_Copy(t *testing.T) {
 			expectedRel: &AggregateRel{input: createVirtualTableReadRel(8),
 				groupingExpressions: []expr.Expression{createPrimitiveFloat(9.0)},
 				groupingReferences:  [][]uint32{{0}},
-				measures:            []AggRelMeasure{{filter: expr.NewPrimitiveLiteral(true, false)}}},
+				measures:            []AggRelMeasure{{measure: aggregateFnRevised, filter: expr.NewPrimitiveLiteral(true, false)}}},
 		},
 		{
 			name:            "ExtensionLeafRel Copy with new inputs",