fix(core)!: correct SUM and SUM0 helper results

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.

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<Arguments> 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<Type> 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<Aggregate.Measure> 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());
+  }
+}