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