fix(isthmus)!: derive SUM, AVG and decimal arithmetic types as the extensions declare

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

@@ -8,6 +8,7 @@
 import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlKind;
@@ -70,7 +71,7 @@ public class AggregateFunctions {
   /** Substrait-specific SUM aggregate function (nullable return type). */
   public static final SqlAggFunction SUM = new SubstraitSumAggFunction();
 
-  /** Substrait-specific SUM0 aggregate function (non-null BIGINT return type). */
+  /** Substrait-specific SUM0 aggregate function (non-null sum type, named sum0). */
   public static final SqlAggFunction SUM0 = new SubstraitSumEmptyIsZeroAggFunction();
 
   /**
@@ -247,7 +248,10 @@ public SubstraitSumAggFunction() {
 
     @Override
     public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
-      return ReturnTypes.ARG0_FORCE_NULLABLE.inferReturnType(opBinding);
+      RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
+      return typeFactory.createTypeWithNullability(
+          typeFactory.getTypeSystem().deriveSumType(typeFactory, opBinding.getOperandType(0)),
+          true);
     }
   }
 
@@ -259,14 +263,14 @@ public SubstraitAvgAggFunction(SqlKind kind) {
 
     @Override
     public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
-      return ReturnTypes.ARG0_FORCE_NULLABLE.inferReturnType(opBinding);
+      RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
+      return typeFactory.createTypeWithNullability(
+          typeFactory.getTypeSystem().deriveAvgAggType(typeFactory, opBinding.getOperandType(0)),
+          true);
     }
   }
 
-  /**
-   * Substrait variant of {@link SqlSumEmptyIsZeroAggFunction} that forces BIGINT return type and
-   * uses a user-friendly name.
-   */
+  /** Substrait variant of {@link SqlSumEmptyIsZeroAggFunction} that uses a user-friendly name. */
   private static class SubstraitSumEmptyIsZeroAggFunction
       extends org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction {
     public SubstraitSumEmptyIsZeroAggFunction() {
@@ -278,11 +282,6 @@ public String getName() {
       // Override default `$sum0` with `sum0` for readability
       return "sum0";
     }
-
-    @Override
-    public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
-      return ReturnTypes.BIGINT.inferReturnType(opBinding);
-    }
   }
 
   /**

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

@@ -4,11 +4,13 @@
 import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rel.type.RelDataTypeFactoryImpl;
 import org.apache.calcite.rel.type.RelDataTypeSystem;
 import org.apache.calcite.rel.type.RelDataTypeSystemImpl;
 import org.apache.calcite.sql.SqlIntervalQualifier;
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.sql.type.SqlTypeUtil;
 
 /**
  * Custom {@link RelDataTypeSystem} implementation for Substrait.
@@ -21,6 +23,8 @@ public class SubstraitTypeSystem extends RelDataTypeSystemImpl {
   /** Singleton instance of Substrait type system. */
   public static final RelDataTypeSystem TYPE_SYSTEM = new SubstraitTypeSystem();
 
+  private static final int MAX_DECIMAL_PRECISION = 38;
+
   /** Default type factory using the Substrait type system. */
   public static final RelDataTypeFactory TYPE_FACTORY = new JavaTypeFactoryImpl(TYPE_SYSTEM);
 
@@ -170,4 +174,156 @@ public int getMaxScale(final SqlTypeName typeName) {
   public boolean shouldConvertRaggedUnionTypesToVarying() {
     return true;
   }
+
+  /**
+   * Returns the type of a {@code SUM} (and {@code $SUM0}) as the Substrait extensions declare it:
+   * an integer sum is an {@code i64} and a floating-point one an {@code fp64} ({@code
+   * functions_arithmetic.yaml}), where Calcite's default keeps the argument's own type, which the
+   * sum overflows on real data. A decimal sum is Calcite's default, {@code DECIMAL(38, s)}, which
+   * is what {@code sum:dec} declares.
+   *
+   * @param typeFactory the type factory
+   * @param argumentType the type of the summed values
+   * @return the type of the sum
+   */
+  @Override
+  public RelDataType deriveSumType(RelDataTypeFactory typeFactory, RelDataType argumentType) {
+    RelDataType sum;
+    switch (argumentType.getSqlTypeName()) {
+      case TINYINT:
+      case SMALLINT:
+      case INTEGER:
+      case BIGINT:
+        sum = typeFactory.createSqlType(SqlTypeName.BIGINT);
+        break;
+      case REAL:
+      case FLOAT:
+      case DOUBLE:
+        sum = typeFactory.createSqlType(SqlTypeName.DOUBLE);
+        break;
+      default:
+        return super.deriveSumType(typeFactory, argumentType);
+    }
+    return typeFactory.createTypeWithNullability(sum, argumentType.isNullable());
+  }
+
+  /**
+   * Returns the type of an {@code AVG} as the Substrait extensions declare it: a decimal average
+   * keeps its scale at precision 38, and every other type averages to itself.
+   *
+   * <p>Calcite routes {@code STDDEV_POP}, {@code STDDEV_SAMP}, {@code VAR_POP} and {@code VAR_SAMP}
+   * through this hook as well, and it cannot tell which function called it, so over a decimal they
+   * become {@code DECIMAL(38, s)} too. No extension declares a decimal {@code std_dev} or {@code
+   * variance}: the argument is cast to {@code fp64} and the result cast back.
+   *
+   * @param typeFactory the type factory
+   * @param argumentType the type of the averaged values
+   * @return the type of the average
+   */
+  @Override
+  public RelDataType deriveAvgAggType(RelDataTypeFactory typeFactory, RelDataType argumentType) {
+    if (argumentType.getSqlTypeName() != SqlTypeName.DECIMAL) {
+      return super.deriveAvgAggType(typeFactory, argumentType);
+    }
+    return typeFactory.createTypeWithNullability(
+        typeFactory.createSqlType(
+            SqlTypeName.DECIMAL, MAX_DECIMAL_PRECISION, argumentType.getScale()),
+        argumentType.isNullable());
+  }
+
+  /**
+   * Returns the type of a decimal addition or subtraction as {@code add:dec_dec} and {@code
+   * subtract:dec_dec} declare it.
+   */
+  @Override
+  public RelDataType deriveDecimalPlusType(
+      RelDataTypeFactory typeFactory, RelDataType type1, RelDataType type2) {
+    return decimalResult(
+        typeFactory,
+        type1,
+        type2,
+        (p1, s1, p2, s2) -> {
+          int scale = Math.max(s1, s2);
+          return new int[] {scale + Math.max(p1 - s1, p2 - s2) + 1, scale};
+        });
+  }
+
+  /** Returns the type of a decimal multiplication as {@code multiply:dec_dec} declares it. */
+  @Override
+  public RelDataType deriveDecimalMultiplyType(
+      RelDataTypeFactory typeFactory, RelDataType type1, RelDataType type2) {
+    return decimalResult(
+        typeFactory, type1, type2, (p1, s1, p2, s2) -> new int[] {p1 + p2 + 1, s1 + s2});
+  }
+
+  /** Returns the type of a decimal division as {@code divide:dec_dec} declares it. */
+  @Override
+  public RelDataType deriveDecimalDivideType(
+      RelDataTypeFactory typeFactory, RelDataType type1, RelDataType type2) {
+    return decimalResult(
+        typeFactory,
+        type1,
+        type2,
+        (p1, s1, p2, s2) -> {
+          int scale = Math.max(6, s1 + p2 + 1);
+          return new int[] {p1 - s1 + p2 + scale, scale};
+        });
+  }
+
+  /** Returns the type of a decimal modulus as {@code modulus:dec_dec} declares it. */
+  @Override
+  public RelDataType deriveDecimalModType(
+      RelDataTypeFactory typeFactory, RelDataType type1, RelDataType type2) {
+    return decimalResult(
+        typeFactory,
+        type1,
+        type2,
+        (p1, s1, p2, s2) -> {
+          int scale = Math.max(s1, s2);
+          return new int[] {Math.min(p1 - s1, p2 - s2) + scale, scale};
+        });
+  }
+
+  /** The unbounded precision and scale a decimal operation's declaration starts from. */
+  @FunctionalInterface
+  private interface DecimalRule {
+    int[] apply(int p1, int s1, int p2, int s2);
+  }
+
+  /**
+   * Applies a decimal operation's rule the way the decimal extension does: a result above precision
+   * 38 is capped there and gives up scale for it, down to a scale of 6 or its own scale if that is
+   * smaller. Returns {@code null}, as Calcite's default does, unless both operands are exact
+   * numerics and one of them is a decimal.
+   */
+  private static RelDataType decimalResult(
+      RelDataTypeFactory typeFactory, RelDataType type1, RelDataType type2, DecimalRule rule) {
+    if (!SqlTypeUtil.isExactNumeric(type1)
+        || !SqlTypeUtil.isExactNumeric(type2)
+        || !(SqlTypeUtil.isDecimal(type1) || SqlTypeUtil.isDecimal(type2))) {
+      return null;
+    }
+    RelDataType decimal1 =
+        RelDataTypeFactoryImpl.isJavaType(type1) ? typeFactory.decimalOf(type1) : type1;
+    RelDataType decimal2 =
+        RelDataTypeFactoryImpl.isJavaType(type2) ? typeFactory.decimalOf(type2) : type2;
+    int[] initial =
+        rule.apply(
+            decimal1.getPrecision(),
+            decimal1.getScale(),
+            decimal2.getPrecision(),
+            decimal2.getScale());
+    int initialPrecision = initial[0];
+    int initialScale = initial[1];
+    int precision = Math.min(initialPrecision, MAX_DECIMAL_PRECISION);
+    int scale =
+        initialPrecision > MAX_DECIMAL_PRECISION
+            ? Math.max(
+                initialScale - (initialPrecision - MAX_DECIMAL_PRECISION),
+                Math.min(initialScale, 6))
+            : initialScale;
+    return typeFactory.createTypeWithNullability(
+        typeFactory.createSqlType(SqlTypeName.DECIMAL, precision, scale),
+        type1.isNullable() || type2.isNullable());
+  }
 }

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

@@ -337,13 +337,13 @@ void namesAnAggregateThatEmitsItsGroupingSetIndexThroughAMapping() {
   @Test
   void dropsNamesWhereTheColumnsAreNotTheRelationsColumns() {
     // Under CALCITE_INFERENCE a measure takes the type Calcite infers rather than the one the plan
-    // declares: SUM(i32) declared i64 becomes an INTEGER column, so the names would land on a
+    // declares: a plan typing SUM(i32) as fp64 gets a BIGINT column, so the names would land on a
     // column the plan does not describe. With the declared types they apply.
     Rel scan32 = sb.namedScan(List.of("t32"), List.of("a", "b"), List.of(R.I32, N.STRING));
     Rel aggregate =
         sb.aggregate(
                 input -> List.of(sb.grouping(input, 1)),
-                input -> List.of(declaring(sb.sum(input, 0), N.I64)),
+                input -> List.of(declaring(sb.sum(input, 0), N.FP64)),
                 Optional.of(Rel.Remap.of(List.of(1, 0))),
                 scan32)
             .withHint(Optional.of(Hint.builder().addOutputNames("total", "label").build()));

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

@@ -129,8 +129,8 @@ void declaredMeasureOutputTypes() {
     void declaredDecimalWidthIsPreserved() {
       Rel input =
           sb.namedScan(List.of("example"), List.of("d", "g"), List.of(R.decimal(10, 2), R.STRING));
-      // The standard extension declarations for decimal sum and avg return DECIMAL<38,S>, while
-      // Calcite's inference keeps the argument's precision. The declared width must survive.
+      // Calcite now infers DECIMAL(38,2) for both, as the extensions declare, so the plan declares
+      // another width here: the declared width must survive the conversion.
       Rel aggregate =
           sb.aggregate(
               i -> sb.grouping(i, 1),
@@ -140,18 +140,18 @@ void declaredDecimalWidthIsPreserved() {
                           sb.aggregateFn(
                               DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL,
                               "sum:dec",
-                              N.decimal(38, 2),
+                              N.decimal(30, 2),
                               sb.fieldReference(i, 0))),
                       sb.measure(
                           sb.aggregateFn(
                               DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL,
                               "avg:dec",
-                              N.decimal(38, 2),
+                              N.decimal(30, 2),
                               sb.fieldReference(i, 0)))),
               input);
 
       RelNode relNode = substraitToCalcite.convert(aggregate);
-      assertRowMatch(relNode.getRowType(), R.STRING, N.decimal(38, 2), N.decimal(38, 2));
+      assertRowMatch(relNode.getRowType(), R.STRING, N.decimal(30, 2), N.decimal(30, 2));
       assertFullRoundTrip(aggregate);
     }
 
@@ -425,7 +425,7 @@ void calciteInferenceIgnoresAnUnrepresentablePlanType() {
                   AggregateConversion.OutputTypeSource.CALCITE_INFERENCE,
                   AggregateConversion.FunctionBindingValidation.NONE));
 
-      assertRowMatch(calciteInference.convert(aggregate).getRowType(), N.STRING, N.I32);
+      assertRowMatch(calciteInference.convert(aggregate).getRowType(), N.STRING, N.I64);
 
       // Preserving it, on the other hand, is impossible — and says so instead of substituting.
       assertThrows(
@@ -454,9 +454,9 @@ void unwrappingRestoresTheDelegatesAndReinfersTypes() {
       assertEquals(AggregateFunctions.SUM, unwrappedCall.getAggregation());
       assertEquals(
           Optional.empty(), AggregateFunctions.boundBinding(unwrappedCall.getAggregation()));
-      // sum(i32) re-infers as nullable INTEGER, not the plan's declared i64 — the same loss
+      // sum(i32) re-infers as nullable BIGINT, not the plan's declared required i64 — the same loss
       // CALCITE_INFERENCE chooses up front.
-      assertEquals(SqlTypeName.INTEGER, unwrappedCall.getType().getSqlTypeName());
+      assertEquals(SqlTypeName.BIGINT, unwrappedCall.getType().getSqlTypeName());
       assertTrue(unwrappedCall.getType().isNullable());
       // An aggregate with no bound calls is returned unchanged.
       assertSame(unwrapped, AggregateFunctions.unwrapBound(unwrapped));
@@ -483,14 +483,14 @@ public SubstraitRelNodeConverter getSubstraitRelNodeConverter(RelBuilder relBuil
       Rel aggregate =
           sb.aggregate(
               input -> sb.grouping(input, 2),
-              input -> List.of(withOutputType(sb.sum(input, 0), N.I64)),
+              input -> List.of(withOutputType(sb.sum(input, 0), N.FP64)),
               commonTable);
 
       RelNode relNode = new SubstraitToCalcite(provider).convert(aggregate);
 
       assertEquals(List.of("overridden"), factories);
-      // ...and the configured conversion is applied: Calcite's inference wins over the plan's i64.
-      assertRowMatch(relNode.getRowType(), N.STRING, N.I32);
+      // ...and the configured conversion is applied: Calcite's inference wins over the plan's fp64.
+      assertRowMatch(relNode.getRowType(), N.STRING, N.I64);
     }
 
     private io.substrait.relation.Aggregate.Measure count(Rel input, String overflow) {
@@ -705,13 +705,14 @@ void strictValidationAppliesUnderCalciteInferenceToo() {
                   AggregateConversion.FunctionBindingValidation.EXTENSION_DECLARATION));
 
       // sum(i32) derives i64? from its declaration: a compliant plan converts (to the inferred
-      // type), a non-compliant one is rejected before any type is chosen.
+      // type, which Calcite now derives as the declaration does), a non-compliant one is rejected
+      // before any type is chosen.
       Rel valid =
           sb.aggregate(
               input -> sb.grouping(input, 2),
               input -> List.of(withOutputType(sb.sum(input, 0), N.I64)),
               commonTable);
-      assertRowMatch(strict.convert(valid).getRowType(), N.STRING, N.I32);
+      assertRowMatch(strict.convert(valid).getRowType(), N.STRING, N.I64);
 
       Rel invalid =
           sb.aggregate(
@@ -1059,15 +1060,15 @@ void directConverterConstructorHonorsTheProviderPolicy() {
       Rel aggregate =
           sb.aggregate(
               input -> sb.grouping(input, 2),
-              input -> List.of(withOutputType(sb.sum(input, 0), N.I64)),
+              input -> List.of(withOutputType(sb.sum(input, 0), N.FP64)),
               commonTable);
       RelBuilder relBuilder = provider.getRelBuilder(provider.getSchemaResolver().apply(aggregate));
       SubstraitRelNodeConverter converter = new SubstraitRelNodeConverter(relBuilder, provider);
 
       RelNode relNode = aggregate.accept(converter, SubstraitRelNodeConverter.Context.newContext());
 
-      // Calcite's inference wins over the plan's i64: the provider's policy applied.
-      assertRowMatch(relNode.getRowType(), N.STRING, N.I32);
+      // Calcite's inference wins over the plan's fp64: the provider's policy applied.
+      assertRowMatch(relNode.getRowType(), N.STRING, N.I64);
     }
 
     @Test
@@ -1129,7 +1130,7 @@ void calciteInferenceModeIgnoresPlanType() {
       Rel aggregate =
           sb.aggregate(
               input -> sb.grouping(input, 2),
-              input -> List.of(withOutputType(sb.sum(input, 0), N.I64)),
+              input -> List.of(withOutputType(sb.sum(input, 0), N.FP64)),
               commonTable);
       SubstraitToCalcite calciteInference =
           withAggregateConversion(
@@ -1139,9 +1140,9 @@ void calciteInferenceModeIgnoresPlanType() {
 
       RelNode relNode = calciteInference.convert(aggregate);
 
-      // Calcite's SUM infers nullable i32 for sum(i32), not the plan's i64, and no wrapper is
+      // Calcite's SUM infers nullable i64 for sum(i32), not the plan's fp64, and no wrapper is
       // added.
-      assertRowMatch(relNode.getRowType(), N.STRING, N.I32);
+      assertRowMatch(relNode.getRowType(), N.STRING, N.I64);
       org.apache.calcite.rel.core.Aggregate calciteAgg =
           (org.apache.calcite.rel.core.Aggregate) relNode;
       assertEquals(