From db81ad0c3aa660d7918f00e44d35838b2f0c95da Mon Sep 17 00:00:00 2001 From: Roman Kondakov Date: Fri, 17 Jul 2026 13:29:31 +0700 Subject: [PATCH] fix(isthmus): preserve declared aggregate output types --- .../substrait/isthmus/AggregateFunctions.java | 133 ++++++++++++++++++ .../isthmus/SubstraitRelNodeConverter.java | 25 +++- .../AggregateFunctionConverter.java | 3 +- .../SubstraitRelNodeConverterTest.java | 30 ++++ 4 files changed, 188 insertions(+), 3 deletions(-) diff --git a/isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java b/isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java index 8ea324948..98b919229 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java +++ b/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. + * + *

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. + * + *

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. + * + *

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 @Nullable T unwrap(Class 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); + } + } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 019ffcd8c..bad94096f 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -339,10 +339,11 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti List groupExprs = groupExprLists.stream().flatMap(Collection::stream).collect(Collectors.toList()); RelBuilder.GroupKey groupKey = relBuilder.groupKey(groupExprs, groupExprLists); + boolean hasEmptyGroup = groupExprLists.stream().anyMatch(List::isEmpty); List aggregateCalls = aggregate.getMeasures().stream() - .map(measure -> fromMeasure(measure, context)) + .map(measure -> fromMeasure(measure, context, child, hasEmptyGroup)) .collect(java.util.stream.Collectors.toList()); Optional 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 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, diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/AggregateFunctionConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/AggregateFunctionConverter.java index 0b9b7e63d..bc8fb0de0 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/AggregateFunctionConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/AggregateFunctionConverter.java @@ -195,7 +195,8 @@ private List 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; } diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java index 7ac6f9f66..deda4fd47 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java +++ b/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