diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 141ecf473..fa9c96bb9 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -1471,11 +1471,7 @@ public Aggregate.Measure min(Rel input, int field) { * @return a new {@link Aggregate.Measure} representing MIN */ public Aggregate.Measure min(Expression expr) { - return singleArgumentArithmeticAggregate( - expr, - "min", - // min output is always nullable - TypeCreator.asNullable(expr.getType())); + return singleArgumentArithmeticAggregate(expr, "min"); } /** @@ -1496,11 +1492,7 @@ public Aggregate.Measure max(Rel input, int field) { * @return a new {@link Aggregate.Measure} representing MAX */ public Aggregate.Measure max(Expression expr) { - return singleArgumentArithmeticAggregate( - expr, - "max", - // max output is always nullable - TypeCreator.asNullable(expr.getType())); + return singleArgumentArithmeticAggregate(expr, "max"); } /** @@ -1521,11 +1513,7 @@ public Aggregate.Measure avg(Rel input, int field) { * @return a new {@link Aggregate.Measure} representing AVG */ public Aggregate.Measure avg(Expression expr) { - return singleArgumentArithmeticAggregate( - expr, - "avg", - // avg output is always nullable - TypeCreator.asNullable(expr.getType())); + return singleArgumentArithmeticAggregate(expr, "avg"); } /** @@ -1546,11 +1534,7 @@ public Aggregate.Measure sum(Rel input, int field) { * @return a new {@link Aggregate.Measure} representing SUM */ public Aggregate.Measure sum(Expression expr) { - return singleArgumentArithmeticAggregate( - expr, - "sum", - // sum output is always nullable - TypeCreator.asNullable(expr.getType())); + return singleArgumentArithmeticAggregate(expr, "sum"); } /** @@ -1561,7 +1545,7 @@ public Aggregate.Measure sum(Expression expr) { * @return a new {@link Aggregate.Measure} representing SUM0 */ public Aggregate.Measure sum0(Rel input, int field) { - return sum(fieldReference(input, field)); + return sum0(fieldReference(input, field)); } /** @@ -1571,11 +1555,7 @@ public Aggregate.Measure sum0(Rel input, int field) { * @return a new {@link Aggregate.Measure} representing SUM0 */ public Aggregate.Measure sum0(Expression expr) { - return singleArgumentArithmeticAggregate( - expr, - "sum0", - // sum0 output is always NOT NULL I64 - R.I64); + return singleArgumentArithmeticAggregate(expr, "sum0"); } /** @@ -1766,7 +1746,7 @@ private Aggregate.Measure statisticalAggregate( } private Aggregate.Measure singleArgumentArithmeticAggregate( - Expression expr, String functionName, Type outputType) { + Expression expr, String functionName) { String typeString = ToTypeString.apply(expr.getType()); SimpleExtension.AggregateFunctionVariant declaration = extensions.getAggregateFunction( @@ -1776,7 +1756,7 @@ private Aggregate.Measure singleArgumentArithmeticAggregate( return measure( AggregateFunctionInvocation.builder() .arguments(Arrays.asList(expr)) - .outputType(outputType) + .outputType(declaration.resolveType(Arrays.asList(expr.getType()))) .declaration(declaration) // INITIAL_TO_RESULT is the most restrictive aggregation phase type, // as it does not allow decomposition. Use it as the default for now. diff --git a/core/src/test/java/io/substrait/dsl/SubstraitBuilderAggregateTest.java b/core/src/test/java/io/substrait/dsl/SubstraitBuilderAggregateTest.java new file mode 100644 index 000000000..22d31a003 --- /dev/null +++ b/core/src/test/java/io/substrait/dsl/SubstraitBuilderAggregateTest.java @@ -0,0 +1,79 @@ +package io.substrait.dsl; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import io.substrait.TestBase; +import io.substrait.expression.FieldReference; +import io.substrait.relation.Aggregate; +import io.substrait.relation.NamedScan; +import io.substrait.type.Type; +import java.util.List; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +class SubstraitBuilderAggregateTest extends TestBase { + + static Stream numericTypes() { + return Stream.of(R, N) + .flatMap( + creator -> + Stream.of( + Arguments.of(creator.I8, R.I64), + Arguments.of(creator.I16, R.I64), + Arguments.of(creator.I32, R.I64), + Arguments.of(creator.I64, R.I64), + Arguments.of(creator.FP32, R.FP64), + Arguments.of(creator.FP64, R.FP64))); + } + + static Stream numericInputTypes() { + return numericTypes().map(arguments -> (Type) arguments.get()[0]); + } + + @ParameterizedTest + @MethodSource("numericTypes") + void sumWidensToNullableResult(Type inputType, Type widenedType) { + NamedScan scan = sb.namedScan(List.of("t"), List.of("v"), List.of(inputType)); + FieldReference input = sb.fieldReference(scan, 0); + Aggregate.Measure byExpression = sb.sum(input); + Aggregate.Measure byField = sb.sum(scan, 0); + + assertEquals("sum", byExpression.getFunction().declaration().name()); + assertEquals(widenedType.withNullable(true), byExpression.getFunction().outputType()); + assertEquals(byExpression, byField); + verifyRoundTrip( + Aggregate.builder().input(scan).addGroupings(sb.grouping()).addMeasures(byField).build()); + } + + @ParameterizedTest + @MethodSource("numericTypes") + void sum0WidensToRequiredResult(Type inputType, Type widenedType) { + NamedScan scan = sb.namedScan(List.of("t"), List.of("v"), List.of(inputType)); + FieldReference input = sb.fieldReference(scan, 0); + Aggregate.Measure byExpression = sb.sum0(input); + Aggregate.Measure byField = sb.sum0(scan, 0); + + assertEquals("sum0", byExpression.getFunction().declaration().name()); + assertEquals(widenedType, byExpression.getFunction().outputType()); + assertEquals(byExpression, byField); + verifyRoundTrip( + Aggregate.builder().input(scan).addGroupings(sb.grouping()).addMeasures(byField).build()); + } + + @ParameterizedTest + @MethodSource("numericInputTypes") + void minMaxAndAvgKeepNullableInputType(Type inputType) { + NamedScan scan = sb.namedScan(List.of("t"), List.of("v"), List.of(inputType)); + FieldReference input = sb.fieldReference(scan, 0); + List measures = List.of(sb.min(input), sb.max(input), sb.avg(input)); + + assertEquals(List.of(sb.min(scan, 0), sb.max(scan, 0), sb.avg(scan, 0)), measures); + for (Aggregate.Measure measure : measures) { + assertEquals(inputType.withNullable(true), measure.getFunction().outputType()); + } + verifyRoundTrip( + Aggregate.builder().input(scan).addGroupings(sb.grouping()).measures(measures).build()); + } +}