diff --git a/isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java b/isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java index c2c048fef..7f9f791cc 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java +++ b/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); - } } /** diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java index c313b5a12..0e6dcf7f4 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java +++ b/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. + * + *

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()); + } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java b/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java index 8fef49c1a..e05b3d450 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/OutputNamesTest.java +++ b/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())); diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java index e27ba35cd..70c5807dd 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java +++ b/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( diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java index 56d00ace8..98649e283 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitTypeSystemTest.java @@ -6,8 +6,12 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.SimpleExtension; +import io.substrait.function.ToTypeString; import io.substrait.isthmus.sql.SubstraitCreateStatementParser; import io.substrait.plan.Plan; +import io.substrait.type.Type; import io.substrait.type.TypeCreator; import java.util.List; import org.apache.calcite.prepare.CalciteCatalogReader; @@ -21,8 +25,150 @@ class SubstraitTypeSystemTest { + private static final TypeCreator R = TypeCreator.REQUIRED; + private final RelDataTypeSystem typeSystem = SubstraitTypeSystem.TYPE_SYSTEM; + /** + * Decimal addition, subtraction, multiplication, division and modulus derive what the decimal + * extension declares, checked against the catalog's own derivation over precisions and scales on + * both sides of the 38 cap, where the result gives up scale. + */ + @Test + void decimalArithmeticDerivesWhatTheExtensionDeclares() { + List precisions = List.of(1, 5, 7, 12, 19, 28, 37, 38); + for (int p1 : precisions) { + for (int s1 : scales(p1)) { + for (int p2 : precisions) { + for (int s2 : scales(p2)) { + RelDataType left = TYPE_FACTORY.createSqlType(SqlTypeName.DECIMAL, p1, s1); + RelDataType right = TYPE_FACTORY.createSqlType(SqlTypeName.DECIMAL, p2, s2); + List operands = List.of(R.decimal(p1, s1), R.decimal(p2, s2)); + assertEquals( + declared("add:dec_dec", operands), + substrait(typeSystem.deriveDecimalPlusType(TYPE_FACTORY, left, right))); + assertEquals( + declared("subtract:dec_dec", operands), + substrait(typeSystem.deriveDecimalPlusType(TYPE_FACTORY, left, right))); + assertEquals( + declared("multiply:dec_dec", operands), + substrait(typeSystem.deriveDecimalMultiplyType(TYPE_FACTORY, left, right))); + assertEquals( + declared("divide:dec_dec", operands), + substrait(typeSystem.deriveDecimalDivideType(TYPE_FACTORY, left, right))); + assertEquals( + declared("modulus:dec_dec", operands), + substrait(typeSystem.deriveDecimalModType(TYPE_FACTORY, left, right))); + } + } + } + } + } + + /** + * For type derivation, an integer operand counts as the decimal that holds its type: {@code + * decimal(10,0)} for an INTEGER and {@code decimal(19,0)} for a BIGINT, whatever the type + * system's own decimal precision. + */ + @Test + void decimalArithmeticWithAnIntegerOperandDerivesWhatTheExtensionDeclares() { + RelDataType decimal = TYPE_FACTORY.createSqlType(SqlTypeName.DECIMAL, 7, 2); + for (SqlTypeName integer : List.of(SqlTypeName.INTEGER, SqlTypeName.BIGINT)) { + RelDataType other = TYPE_FACTORY.createSqlType(integer); + List operands = + List.of(R.decimal(7, 2), R.decimal(integer == SqlTypeName.INTEGER ? 10 : 19, 0)); + assertEquals( + declared("add:dec_dec", operands), + substrait(typeSystem.deriveDecimalPlusType(TYPE_FACTORY, decimal, other))); + assertEquals( + declared("multiply:dec_dec", operands), + substrait(typeSystem.deriveDecimalMultiplyType(TYPE_FACTORY, decimal, other))); + assertEquals( + declared("divide:dec_dec", operands), + substrait(typeSystem.deriveDecimalDivideType(TYPE_FACTORY, decimal, other))); + } + } + + /** + * SUM and AVG derive what the arithmetic extensions declare. The aggregate function sets the + * result's nullability, so only the type is compared. + */ + @Test + void sumAndAvgDeriveWhatTheExtensionsDeclare() { + List numerics = + List.of( + SqlTypeName.TINYINT, + SqlTypeName.SMALLINT, + SqlTypeName.INTEGER, + SqlTypeName.BIGINT, + SqlTypeName.REAL, + SqlTypeName.DOUBLE); + for (SqlTypeName name : numerics) { + RelDataType argument = TYPE_FACTORY.createSqlType(name); + assertAggregate("sum", argument, typeSystem.deriveSumType(TYPE_FACTORY, argument)); + assertAggregate("avg", argument, typeSystem.deriveAvgAggType(TYPE_FACTORY, argument)); + } + for (int precision : List.of(1, 7, 38)) { + for (int scale : scales(precision)) { + RelDataType argument = TYPE_FACTORY.createSqlType(SqlTypeName.DECIMAL, precision, scale); + assertAggregate("sum", argument, typeSystem.deriveSumType(TYPE_FACTORY, argument)); + assertAggregate("avg", argument, typeSystem.deriveAvgAggType(TYPE_FACTORY, argument)); + } + } + } + + /** + * The same in a query: isthmus's own SUM and AVG take their types from the type system, where + * they used to repeat the argument's. The average stays nullable, the right answer for an empty + * group, although avg:dec declares a required result. + */ + @Test + void sumAndAvgInAQueryTakeTheTypeSystemsTypes() throws Exception { + CalciteCatalogReader catalog = + SubstraitCreateStatementParser.processCreateStatementsToCatalog( + "CREATE TABLE t (d DECIMAL(7, 2), i INT)"); + Plan plan = new SqlToSubstrait().convert("SELECT SUM(d), AVG(d), SUM(i) FROM t", catalog); + + assertEquals( + List.of( + TypeCreator.NULLABLE.decimal(38, 2), + TypeCreator.NULLABLE.decimal(38, 2), + TypeCreator.NULLABLE.I64), + plan.getRoots().get(0).getInput().getRecordType().fields()); + } + + private static List scales(int precision) { + return List.of(0, precision / 2, precision).stream().distinct().toList(); + } + + private static Type declared(String key, List operands) { + return DefaultExtensionCatalog.DEFAULT_COLLECTION + .getScalarFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL, key)) + .resolveType(operands); + } + + private static Type substrait(RelDataType type) { + return TypeConverter.DEFAULT.toSubstrait(type); + } + + private static void assertAggregate(String name, RelDataType argument, RelDataType derived) { + Type operand = substrait(argument); + String urn = + argument.getSqlTypeName() == SqlTypeName.DECIMAL + ? DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL + : DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC; + String key = name + ":" + operand.accept(ToTypeString.INSTANCE); + Type declared = + DefaultExtensionCatalog.DEFAULT_COLLECTION + .getAggregateFunction(SimpleExtension.FunctionAnchor.of(urn, key)) + .resolveType(List.of(operand)); + assertTrue( + declared.equalsIgnoringNullability(substrait(derived)), + () -> key + " declares " + declared + " but derives " + derived.getFullTypeString()); + } + @Test void decimalMaxPrecision() { assertEquals(38, typeSystem.getMaxPrecision(SqlTypeName.DECIMAL));