fix(isthmus): preserve declared aggregate output types

isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java

@@ -1,15 +1,22 @@
 package io.substrait.isthmus;
 
+import java.util.Objects;
 import java.util.Optional;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperandCountRange;
 import org.apache.calcite.sql.SqlOperatorBinding;
+import org.apache.calcite.sql.SqlSyntax;
+import org.apache.calcite.sql.SqlWriter;
 import org.apache.calcite.sql.fun.SqlAvgAggFunction;
 import org.apache.calcite.sql.fun.SqlMinMaxAggFunction;
 import org.apache.calcite.sql.fun.SqlSumAggFunction;
 import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction;
 import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.util.Optionality;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Provides Substrait-specific variants of Calcite aggregate functions to ensure type inference
@@ -59,6 +66,38 @@ public class AggregateFunctions {
   /** Substrait-specific SUM0 aggregate function (non-null BIGINT return type). */
   public static final SqlAggFunction SUM0 = new SubstraitSumEmptyIsZeroAggFunction();
 
+  /**
+   * Returns an aggregate function that infers the output type declared by a Substrait invocation.
+   *
+   * <p>Calcite validates an {@code AggregateCall}'s stored type against the function's inference
+   * rule and {@code RelBuilder} re-infers that type when it copies a call. Carrying the declared
+   * type in the function therefore keeps the Substrait contract intact through both operations.
+   *
+   * @param aggFunction aggregate function resolved from the Substrait function declaration
+   * @param outputType output type carried by the Substrait invocation
+   * @return an equivalent aggregate function with the declared output type
+   */
+  public static SqlAggFunction withOutputType(SqlAggFunction aggFunction, RelDataType outputType) {
+    return new DeclaredOutputSqlAggFunction(aggFunction, outputType);
+  }
+
+  /**
+   * Returns the function wrapped by {@link #withOutputType}, or the input unchanged.
+   *
+   * <p>The Calcite-to-Substrait function matcher keys on the configured aggregate function
+   * instances, so it must look through the per-invocation output-type wrapper.
+   *
+   * @param aggFunction aggregate function to inspect
+   * @return the underlying configured aggregate function
+   */
+  public static SqlAggFunction withoutDeclaredOutputType(SqlAggFunction aggFunction) {
+    SqlAggFunction current = aggFunction;
+    while (current instanceof DeclaredOutputSqlAggFunction) {
+      current = ((DeclaredOutputSqlAggFunction) current).delegate;
+    }
+    return current;
+  }
+
   /**
    * Converts default Calcite aggregate functions to Substrait-specific variants when needed.
    *
@@ -154,4 +193,98 @@ public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
       return ReturnTypes.BIGINT.inferReturnType(opBinding);
     }
   }
+
+  /**
+   * Per-invocation aggregate function whose inference returns the output type carried by Substrait.
+   *
+   * <p>All non-type behavior remains that of the configured function. Equality deliberately follows
+   * the configured function only: output-type compatibility is a consumer policy and should not be
+   * smuggled into Calcite operator identity.
+   */
+  private static final class DeclaredOutputSqlAggFunction extends SqlAggFunction {
+    private final SqlAggFunction delegate;
+
+    private DeclaredOutputSqlAggFunction(SqlAggFunction delegate, RelDataType outputType) {
+      super(
+          delegate.getName(),
+          delegate.getSqlIdentifier(),
+          delegate.getKind(),
+          ReturnTypes.explicit(outputType),
+          delegate.getOperandTypeInference(),
+          delegate.getOperandTypeChecker(),
+          delegate.getFunctionType(),
+          delegate.requiresOrder(),
+          delegate.requiresOver(),
+          delegate.requiresGroupOrder());
+      this.delegate = delegate;
+    }
+
+    @Override
+    public SqlOperandCountRange getOperandCountRange() {
+      return delegate.getOperandCountRange();
+    }
+
+    @Override
+    public SqlSyntax getSyntax() {
+      return delegate.getSyntax();
+    }
+
+    @Override
+    public void unparse(SqlWriter writer, SqlCall call, int leftPrec, int rightPrec) {
+      delegate.unparse(writer, call, leftPrec, rightPrec);
+    }
+
+    @Override
+    public Optionality getDistinctOptionality() {
+      return delegate.getDistinctOptionality();
+    }
+
+    @Override
+    public boolean allowsFilter() {
+      return delegate.allowsFilter();
+    }
+
+    @Override
+    public boolean allowsNullTreatment() {
+      return delegate.allowsNullTreatment();
+    }
+
+    @Override
+    public boolean skipsNullInputs() {
+      return delegate.skipsNullInputs();
+    }
+
+    @Override
+    public @Nullable SqlAggFunction getRollup() {
+      return delegate.getRollup();
+    }
+
+    @Override
+    public boolean isPercentile() {
+      return delegate.isPercentile();
+    }
+
+    @Override
+    public boolean allowsFraming() {
+      return delegate.allowsFraming();
+    }
+
+    @Override
+    public <T> @Nullable T unwrap(Class<T> clazz) {
+      T unwrapped = super.unwrap(clazz);
+      return unwrapped != null ? unwrapped : delegate.unwrap(clazz);
+    }
+
+    @Override
+    public boolean equals(@Nullable Object obj) {
+      return obj == this
+          || obj instanceof DeclaredOutputSqlAggFunction
+              && delegate.equals(((DeclaredOutputSqlAggFunction) obj).delegate);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(getClass(), delegate);
+    }
+  }
 }

isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java

@@ -339,10 +339,11 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
     List<RexNode> groupExprs =
         groupExprLists.stream().flatMap(Collection::stream).collect(Collectors.toList());
     RelBuilder.GroupKey groupKey = relBuilder.groupKey(groupExprs, groupExprLists);
+    boolean hasEmptyGroup = groupExprLists.stream().anyMatch(List::isEmpty);
 
     List<AggregateCall> aggregateCalls =
         aggregate.getMeasures().stream()
-            .map(measure -> fromMeasure(measure, context))
+            .map(measure -> fromMeasure(measure, context, child, hasEmptyGroup))
             .collect(java.util.stream.Collectors.toList());
 
     Optional<Remap> remap = aggregate.getRemap();
@@ -383,7 +384,8 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
     return applyRemap(node, remap);
   }
 
-  private AggregateCall fromMeasure(Aggregate.Measure measure, Context context) {
+  private AggregateCall fromMeasure(
+      Aggregate.Measure measure, Context context, RelNode input, boolean hasEmptyGroup) {
     List<FunctionArg> eArgs = measure.getFunction().arguments();
     // Only value (Expression) arguments map to Calcite aggregate operands. Enum arguments such as
     // the std_dev/variance "distribution" are used to disambiguate the operator, not as operands.
@@ -448,6 +450,25 @@ private AggregateCall fromMeasure(Aggregate.Measure measure, Context context) {
                   .collect(Collectors.toList()));
     }
 
+    AggregateCall inferredCall =
+        AggregateCall.create(
+            aggFunction,
+            distinct,
+            false,
+            false,
+            Collections.emptyList(),
+            argIndex,
+            filterArg,
+            null,
+            relCollation,
+            hasEmptyGroup,
+            input,
+            null,
+            null);
+    if (!returnType.equals(inferredCall.getType())) {
+      aggFunction = AggregateFunctions.withOutputType(aggFunction, returnType);
+    }
+
     return AggregateCall.create(
         aggFunction,
         distinct,

isthmus/src/main/java/io/substrait/isthmus/expression/AggregateFunctionConverter.java

@@ -195,7 +195,8 @@ private List<RexNode> leadingEnumArgs(AggregateCall call) {
   protected FunctionFinder getFunctionFinder(AggregateCall call) {
     // replace COUNT() + distinct == true and approximate == true with APPROX_COUNT_DISTINCT
     // before converting into substrait function
-    SqlAggFunction aggFunction = call.getAggregation();
+    SqlAggFunction aggFunction =
+        AggregateFunctions.withoutDeclaredOutputType(call.getAggregation());
     if (aggFunction == SqlStdOperatorTable.COUNT && call.isDistinct() && call.isApproximate()) {
       aggFunction = SqlStdOperatorTable.APPROX_COUNT_DISTINCT;
     }

isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java

@@ -1,5 +1,6 @@
 package io.substrait.isthmus;
 
+import io.substrait.expression.AggregateFunctionInvocation;
 import io.substrait.plan.Plan;
 import io.substrait.relation.Join.JoinType;
 import io.substrait.relation.Rel;
@@ -51,6 +52,35 @@ void emit() {
       RelNode relNode = substraitToCalcite.convert(root.getInput());
       assertRowMatch(relNode.getRowType(), N.STRING, R.I64);
     }
+
+    @Test
+    void declaredMeasureOutputTypes() {
+      Rel aggregate =
+          sb.aggregate(
+              input -> sb.grouping(input, 2),
+              input ->
+                  List.of(
+                      withOutputType(sb.sum(input, 0), R.I64),
+                      withOutputType(sb.avg(input, 1), R.FP32)),
+              commonTable);
+
+      RelNode relNode = substraitToCalcite.convert(aggregate);
+      assertRowMatch(relNode.getRowType(), N.STRING, R.I64, R.FP32);
+      assertFullRoundTrip(aggregate);
+    }
+
+    private io.substrait.relation.Aggregate.Measure withOutputType(
+        io.substrait.relation.Aggregate.Measure measure, Type outputType) {
+      AggregateFunctionInvocation function =
+          AggregateFunctionInvocation.builder()
+              .from(measure.getFunction())
+              .outputType(outputType)
+              .build();
+      return io.substrait.relation.Aggregate.Measure.builder()
+          .from(measure)
+          .function(function)
+          .build();
+    }
   }
 
   @Nested