diff --git a/core/src/main/java/io/substrait/relation/AggregateFunctionProtoConverter.java b/core/src/main/java/io/substrait/relation/AggregateFunctionProtoConverter.java index e8ba02c61..c1fd02a36 100644 --- a/core/src/main/java/io/substrait/relation/AggregateFunctionProtoConverter.java +++ b/core/src/main/java/io/substrait/relation/AggregateFunctionProtoConverter.java @@ -6,6 +6,7 @@ import io.substrait.extension.SimpleExtension; import io.substrait.proto.AggregateFunction; import io.substrait.proto.FunctionArgument; +import io.substrait.proto.SortField; import io.substrait.type.proto.TypeProtoConverter; import io.substrait.util.EmptyVisitationContext; import java.util.List; @@ -56,8 +57,21 @@ public AggregateFunction toProto(Aggregate.Measure measure) { args.get(i) .accept(aggFuncDef, i, argVisitor, EmptyVisitationContext.INSTANCE)) .collect(Collectors.toList())) + .addAllSorts( + measure.getFunction().sort().stream() + .map( + sort -> + SortField.newBuilder() + .setExpr(exprProtoConverter.toProto(sort.expr())) + .setDirection(sort.direction().toProto()) + .build()) + .collect(Collectors.toList())) .setFunctionReference( functionCollector.getFunctionReference(measure.getFunction().declaration())) + .addAllOptions( + measure.getFunction().options().stream() + .map(ExpressionProtoConverter::from) + .collect(Collectors.toList())) .build(); } } diff --git a/core/src/test/java/io/substrait/extendedexpression/ExtendedExpressionRoundTripTest.java b/core/src/test/java/io/substrait/extendedexpression/ExtendedExpressionRoundTripTest.java index 9b6ae95a8..18bc43398 100644 --- a/core/src/test/java/io/substrait/extendedexpression/ExtendedExpressionRoundTripTest.java +++ b/core/src/test/java/io/substrait/extendedexpression/ExtendedExpressionRoundTripTest.java @@ -6,8 +6,10 @@ import io.substrait.expression.Expression; import io.substrait.expression.ExpressionCreator; import io.substrait.expression.FieldReference; +import io.substrait.expression.FunctionOption; import io.substrait.expression.ImmutableFieldReference; import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.SimpleExtension; import io.substrait.relation.Aggregate; import io.substrait.type.NamedStruct; import io.substrait.type.Type; @@ -41,6 +43,58 @@ void testRoundTrip(ExtendedExpression.ExpressionReferenceBase expressionReferenc assertExtendedExpressionOperation(expressionReferences, namedStruct); } + @Test + void preservesAggregateOrderingAndFunctionsUsedOnlyInSorts() { + FieldReference value = FieldReference.newRootStructReference(0, R.STRING); + FieldReference sortKey = FieldReference.newRootStructReference(1, R.I64); + AggregateFunctionInvocation function = + AggregateFunctionInvocation.builder() + .declaration( + extensions.getAggregateFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_STRING, "string_agg:str_str"))) + .addArguments(value, ExpressionCreator.string(false, ",")) + .outputType(R.STRING) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .addSort( + Expression.SortField.builder() + .expr(sb.add(sortKey, sb.i64(1))) + .direction(Expression.SortDirection.DESC_NULLS_LAST) + .build(), + Expression.SortField.builder() + .expr(value) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build()) + .build(); + + assertExtendedExpressionOperation( + List.of( + ImmutableAggregateFunctionReference.builder() + .measure(Aggregate.Measure.builder().function(function).build()) + .addOutputNames("concatenated") + .build()), + NamedStruct.of(List.of("value", "sort_key"), R.struct(R.STRING, R.I64))); + } + + @Test + void preservesAggregateOptionPreferences() { + AggregateFunctionInvocation function = + AggregateFunctionInvocation.builder() + .from(sb.sum(FieldReference.newRootStructReference(0, R.I64)).getFunction()) + .addOptions( + FunctionOption.builder().name("overflow").addValues("ERROR", "SATURATE").build()) + .build(); + + assertExtendedExpressionOperation( + List.of( + ImmutableAggregateFunctionReference.builder() + .measure(Aggregate.Measure.builder().function(function).build()) + .addOutputNames("total") + .build()), + NamedStruct.of(List.of("value"), R.struct(R.I64))); + } + @Test void getNoExpressionDefined() { IllegalStateException illegalStateException =