diff --git a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java index 90d943366..40b354678 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java +++ b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java @@ -88,6 +88,9 @@ public class ConverterProvider { /** The Calcite SQL parser configuration, controlling parsing behaviour like identifier casing. */ protected final SqlParser.Config sqlParserConfig; + /** Observer for supplied and independently inferred expression types. */ + protected final TypeObserver typeObserver; + /** Converter for Substrait scalar functions. */ protected ScalarFunctionConverter scalarFunctionConverter; @@ -221,6 +224,7 @@ protected ConverterProvider(Builder builder) { .unquotedCasing .map(builder.sqlParserConfig::withUnquotedCasing) .orElse(builder.sqlParserConfig); + this.typeObserver = builder.typeObserver; this.scalarFunctionConverter = builder.scalarFunctionConverter.orElseGet( @@ -429,12 +433,13 @@ public ExpressionRexConverter getExpressionRexConverter( /** * Returns the observer for supplied and independently inferred expression types. * - *

Override to collect type observations during Substrait-to-Calcite conversion. + *

Configure via {@link Builder#typeObserver(TypeObserver)} or override this method to collect + * type observations during Substrait-to-Calcite conversion. * * @return a no-op observer by default */ public TypeObserver getTypeObserver() { - return TypeObserver.NOOP; + return typeObserver; } /** @@ -565,6 +570,7 @@ public static class Builder { private TypeConverter typeConverter = TypeConverter.DEFAULT; private Plan.ExecutionBehavior executionBehavior = createDefaultExecutionBehavior(); private SqlParser.Config sqlParserConfig = DEFAULT_SQL_PARSER_CONFIG; + private TypeObserver typeObserver = TypeObserver.NOOP; private Optional unquotedCasing = Optional.empty(); // Derived from the extensions and type factory at build time when left unset. @@ -646,6 +652,17 @@ public Builder unquotedCasing(Casing unquotedCasing) { return this; } + /** + * Sets the observer for supplied and independently inferred expression types. + * + * @param typeObserver the type observer + * @return this builder + */ + public Builder typeObserver(TypeObserver typeObserver) { + this.typeObserver = typeObserver; + return this; + } + /** * Sets the scalar function converter. When left unset, it is derived from the configured * extensions and type factory. diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java index e3c42d58e..f64ff9184 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -39,6 +39,7 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCallBinding; import org.apache.calcite.rex.RexFieldCollation; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLambdaRef; @@ -51,6 +52,7 @@ import org.apache.calcite.sql.SqlIntervalQualifier; import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.SqlWindow; import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.parser.SqlParserPos; import org.apache.calcite.sql.type.SqlTypeName; @@ -549,31 +551,32 @@ public RexNode visit(Expression.ScalarFunctionInvocation expr, Context context) RelDataType returnType = typeConverter.toCalcite(typeFactory, expr.outputType()); RexNode rexCall = rexBuilder.makeCall(returnType, operator, args); - // If type observations are not needed, avoid recomputing the RexCall with Calcite's - // independently inferred return type. - if (typeObserver == TypeObserver.NOOP) { - return rexCall; - } - observeScalarType(expr, () -> rexBuilder.makeCall(operator, args)); + observeType( + expr, + TypeObservation.Source.SCALAR_FUNCTION, + () -> rexBuilder.deriveReturnType(operator, args)); return rexCall; } - private void observeScalarType( - Expression.ScalarFunctionInvocation expression, Supplier inferredCallSupplier) { + private void observeType( + Expression expression, + TypeObservation.Source source, + Supplier inferredTypeSupplier) { + // The conversion result never depends on the independently inferred type, so when no observer + // is installed skip Calcite's inference entirely. + if (typeObserver == TypeObserver.NOOP) { + return; + } TypeObservation observation; - RexNode inferredCall; + RelDataType inferredType; try { - inferredCall = inferredCallSupplier.get(); + inferredType = inferredTypeSupplier.get(); } catch (RuntimeException inferenceFailure) { - observation = - TypeObservation.failure( - TypeObservation.Source.SCALAR_FUNCTION, expression, inferenceFailure); + observation = TypeObservation.failure(source, expression, inferenceFailure); typeObserver.observe(observation); return; } - observation = - TypeObservation.success( - TypeObservation.Source.SCALAR_FUNCTION, expression, inferredCall.getType()); + observation = TypeObservation.success(source, expression, inferredType); typeObserver.observe(observation); } @@ -631,19 +634,52 @@ public RexNode visit(Expression.WindowFunctionInvocation expr, Context context) boolean nullWhenCountZero = false; boolean allowPartial = true; - return rexBuilder.makeOver( - outputType, - (SqlAggFunction) operator, - args, - partitionKeys, - orderKeys, - lowerBound, - upperBound, - rowMode, - allowPartial, - nullWhenCountZero, - distinct, - ignoreNulls); + RexNode rexOver = + rexBuilder.makeOver( + outputType, + (SqlAggFunction) operator, + args, + partitionKeys, + orderKeys, + lowerBound, + upperBound, + rowMode, + allowPartial, + nullWhenCountZero, + distinct, + ignoreNulls); + observeType( + expr, + TypeObservation.Source.WINDOW_FUNCTION, + () -> operator.inferReturnType(windowBinding(operator, args, lowerBound, upperBound))); + return rexOver; + } + + /** + * Builds the operand binding Calcite itself uses when inferring the return type of a windowed + * aggregate. + * + *

{@link RexBuilder#deriveReturnType(SqlOperator, java.util.List)} binds the operands with a + * plain {@link RexCallBinding}, whose {@link + * org.apache.calcite.sql.SqlOperatorBinding#hasEmptyGroup()} is always {@code false}. Calcite's + * validator instead derives that flag from the window bounds, so any return type strategy that + * consults it (for example {@code ARG0_NULLABLE_IF_EMPTY} or {@code AGG_SUM}) widens its result + * to nullable over a frame that may be empty. Binding without the flag would report a + * non-nullable inferred type for such an operator, so a genuine nullability deviation would be + * observed as a match. + */ + private RexCallBinding windowBinding( + SqlOperator operator, + List operands, + RexWindowBound lowerBound, + RexWindowBound upperBound) { + boolean emptyGroup = !SqlWindow.isAlwaysNonEmpty(lowerBound, upperBound); + return new RexCallBinding(typeFactory, operator, operands, ImmutableList.of()) { + @Override + public boolean hasEmptyGroup() { + return emptyGroup; + } + }; } private Set asSqlKind(Expression.SortDirection direction) { diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java index 372da8065..764c99032 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java @@ -15,7 +15,10 @@ public final class TypeObservation { /** The expression category that produced an observation. */ public enum Source { /** A scalar function invocation. */ - SCALAR_FUNCTION + SCALAR_FUNCTION, + + /** A window function invocation. */ + WINDOW_FUNCTION } private final Source source; diff --git a/isthmus/src/test/java/io/substrait/isthmus/ConverterProviderBuilderTest.java b/isthmus/src/test/java/io/substrait/isthmus/ConverterProviderBuilderTest.java index eaed78393..e2942d7a7 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ConverterProviderBuilderTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ConverterProviderBuilderTest.java @@ -11,6 +11,7 @@ import io.substrait.extension.SimpleExtension; import io.substrait.isthmus.expression.AggregateFunctionConverter; import io.substrait.isthmus.expression.ScalarFunctionConverter; +import io.substrait.isthmus.expression.TypeObserver; import io.substrait.isthmus.expression.WindowFunctionConverter; import org.apache.calcite.avatica.util.Casing; import org.apache.calcite.rel.type.RelDataTypeFactory; @@ -33,6 +34,7 @@ void derivesFunctionConvertersWhenUnset() { assertNotNull(provider.getAggregateFunctionConverter()); assertNotNull(provider.getWindowFunctionConverter()); assertEquals(ConverterProvider.DEFAULT_SQL_PARSER_CONFIG, provider.getSqlParserConfig()); + assertSame(TypeObserver.NOOP, provider.getTypeObserver()); } @Test diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java index 24e2fa474..0a6a36d45 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java @@ -1,6 +1,7 @@ package io.substrait.isthmus; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertSame; @@ -30,6 +31,7 @@ import org.apache.calcite.rel.logical.LogicalProject; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlAggFunction; import org.apache.calcite.sql.SqlFunction; import org.apache.calcite.sql.SqlFunctionCategory; import org.apache.calcite.sql.SqlKind; @@ -277,6 +279,28 @@ void observeVariadicConcatOnce() { @Test void injectTypeObserverThroughConverterProvider() { + AtomicReference observed = new AtomicReference<>(); + ConverterProvider observingProvider = + ConverterProvider.builder().typeObserver(observed::set).build(); + Expression.ScalarFunctionInvocation expr = + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "add:i32_i32", + R.I32, + sb.i32(7), + sb.i32(42)); + Project query = sb.project(input -> List.of(expr), sb.emptyVirtualTableScan()); + + new SubstraitToCalcite(observingProvider).convert(query); + + assertEquals(R.I32, observed.get().suppliedType()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), + observed.get().inferredType().orElseThrow()); + } + + @Test + void injectTypeObserverByOverridingConverterProvider() { AtomicReference observed = new AtomicReference<>(); ConverterProvider observingProvider = new ConverterProvider() { @@ -297,9 +321,6 @@ public TypeObserver getTypeObserver() { new SubstraitToCalcite(observingProvider).convert(query); assertEquals(R.I32, observed.get().suppliedType()); - assertEquals( - TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), - observed.get().inferredType().orElseThrow()); } @Test @@ -357,12 +378,57 @@ void propagateObserverExceptionAfterInferenceFailure() { @Test void useSubstraitReturnTypeDuringWindowFunctionConversion() { + // THIS IS (INTENTIONALLY) THE WRONG OUTPUT TYPE + // SHOULD BE R.I64 + Expression.WindowFunctionInvocation expr = rowNumberWithReturnType(R.STRING); + + RexNode calciteExpr = expr.accept(expressionRexConverter, Context.newContext()); + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.STRING), calciteExpr.getType()); + } + + @Test + void observeSuppliedAndInferredWindowFunctionTypes() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + Expression.WindowFunctionInvocation expr = rowNumberWithReturnType(R.STRING); + + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + TypeObservation observation = observed.get(); + assertEquals(TypeObservation.Source.WINDOW_FUNCTION, observation.source()); + assertSame(expr, observation.expression()); + assertEquals(R.STRING, observation.suppliedType()); + assertTrue(observation.inferenceFailure().isEmpty()); + assertNotEquals(calciteExpr.getType(), observation.inferredType().orElseThrow()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I64), + observation.inferredType().orElseThrow()); + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.STRING), calciteExpr.getType()); + } + + @Test + void observeMatchingWindowFunctionTypes() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + Expression.WindowFunctionInvocation expr = rowNumberWithReturnType(R.I64); + + expr.accept(observingConverter, Context.newContext()); + + assertSame(expr, observed.get().expression()); + assertEquals(R.I64, observed.get().suppliedType()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I64), + observed.get().inferredType().orElseThrow()); + } + + @Test + void observeArgumentDependentWindowFunctionType() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); Expression.WindowFunctionInvocation expr = sb.windowFn( DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, - "row_number:", - // THIS IS (INTENTIONALLY) THE WRONG OUTPUT TYPE - // SHOULD BE R.I64 + "lag:any", R.STRING, Expression.AggregationPhase.INITIAL_TO_RESULT, Expression.AggregationInvocation.ALL, @@ -371,10 +437,97 @@ void useSubstraitReturnTypeDuringWindowFunctionConversion() { WindowBound.UNBOUNDED, sb.i32(42)); - RexNode calciteExpr = expr.accept(expressionRexConverter, Context.newContext()); + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + // LAG with no default operand forces a nullable return type. + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, N.I32), + observed.get().inferredType().orElseThrow()); assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.STRING), calciteExpr.getType()); } + @Test + void observeNullableWindowTypeOverPossiblyEmptyWindow() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + // first_value infers ARG0 widened to nullable whenever the frame may be empty, and a frame + // bounded by a row offset carries no compile-time guarantee that it is not. + Expression.WindowFunctionInvocation expr = + firstValueOver(WindowBound.Preceding.of(1), WindowBound.Preceding.of(1)); + + expr.accept(observingConverter, Context.newContext()); + + assertTrue(observed.get().inferredType().orElseThrow().isNullable()); + } + + @Test + void observeNonNullableWindowTypeOverAlwaysNonEmptyWindow() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + // A frame running from the start of the partition to the current row always holds that row, + // so first_value keeps the non-nullable argument type. + Expression.WindowFunctionInvocation expr = + firstValueOver(WindowBound.UNBOUNDED, WindowBound.CURRENT_ROW); + + expr.accept(observingConverter, Context.newContext()); + + assertFalse(observed.get().inferredType().orElseThrow().isNullable()); + } + + @Test + void skipWindowTypeInferenceForNoopObserver() { + AtomicInteger inferenceCalls = new AtomicInteger(); + ExpressionRexConverter nonObservingConverter = + new ExpressionRexConverter( + typeFactory, + new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory), + countingInferenceWindowFunctionConverter(inferenceCalls), + TypeConverter.DEFAULT); + Expression.WindowFunctionInvocation expr = rowNumberWithReturnType(R.STRING); + + RexNode calciteExpr = expr.accept(nonObservingConverter, Context.newContext()); + + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.STRING), calciteExpr.getType()); + assertEquals(0, inferenceCalls.get()); + } + + @Test + void reportWindowInferenceFailureWithoutFailingConversion() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = + observingConverter(failingInferenceWindowFunctionConverter(), observed::set); + Expression.WindowFunctionInvocation expr = rowNumberWithReturnType(R.STRING); + + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.STRING), calciteExpr.getType()); + assertSame(expr, observed.get().expression()); + assertEquals(R.STRING, observed.get().suppliedType()); + assertTrue(observed.get().inferredType().isEmpty()); + IllegalStateException failure = + assertInstanceOf( + IllegalStateException.class, observed.get().inferenceFailure().orElseThrow()); + assertEquals("controlled window inference failure", failure.getMessage()); + } + + @Test + void propagateWindowObserverException() { + ExpressionRexConverter observingConverter = + observingConverter( + new WindowFunctionConverter(extensions.windowFunctions(), typeFactory), + observation -> { + throw new IllegalStateException("window observer failure"); + }); + Expression.WindowFunctionInvocation expr = rowNumberWithReturnType(R.I64); + + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> expr.accept(observingConverter, Context.newContext())); + + assertEquals("window observer failure", failure.getMessage()); + } + void assertTypeMatch(RelDataType actual, Type expected) { Type type = TypeConverter.DEFAULT.toSubstrait(actual); assertEquals(expected, type); @@ -389,6 +542,32 @@ private Expression.ScalarFunctionInvocation integerAddWithReturnType(Type output sb.i32(42)); } + private Expression.WindowFunctionInvocation rowNumberWithReturnType(Type outputType) { + return sb.windowFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "row_number:", + outputType, + Expression.AggregationPhase.INITIAL_TO_RESULT, + Expression.AggregationInvocation.ALL, + Expression.WindowBoundsType.RANGE, + WindowBound.UNBOUNDED, + WindowBound.UNBOUNDED); + } + + private Expression.WindowFunctionInvocation firstValueOver( + WindowBound lowerBound, WindowBound upperBound) { + return sb.windowFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "first_value:any", + R.I32, + Expression.AggregationPhase.INITIAL_TO_RESULT, + Expression.AggregationInvocation.ALL, + Expression.WindowBoundsType.ROWS, + lowerBound, + upperBound, + sb.i32(42)); + } + private ExpressionRexConverter observingConverter(TypeObserver observer) { return observingConverter( new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory), observer); @@ -404,6 +583,16 @@ private ExpressionRexConverter observingConverter( observer); } + private ExpressionRexConverter observingConverter( + WindowFunctionConverter windowFunctionConverter, TypeObserver observer) { + return new ExpressionRexConverter( + typeFactory, + new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory), + windowFunctionConverter, + TypeConverter.DEFAULT, + observer); + } + private ScalarFunctionConverter failingInferenceScalarFunctionConverter() { SqlFunction failingOperator = new SqlFunction( @@ -443,4 +632,43 @@ public Optional getSqlOperatorFromSubstraitFunc(String key, Type ou } }; } + + private WindowFunctionConverter failingInferenceWindowFunctionConverter() { + SqlAggFunction failingOperator = + new SqlAggFunction( + "controlled_window_inference_failure", + SqlKind.OTHER_FUNCTION, + binding -> { + throw new IllegalStateException("controlled window inference failure"); + }, + null, + null, + SqlFunctionCategory.USER_DEFINED_FUNCTION) {}; + return windowFunctionConverter(failingOperator); + } + + private WindowFunctionConverter countingInferenceWindowFunctionConverter( + AtomicInteger inferenceCalls) { + SqlAggFunction countingOperator = + new SqlAggFunction( + "counting_window_inference", + SqlKind.OTHER_FUNCTION, + binding -> { + inferenceCalls.incrementAndGet(); + return TypeConverter.DEFAULT.toCalcite(typeFactory, R.I64); + }, + null, + null, + SqlFunctionCategory.USER_DEFINED_FUNCTION) {}; + return windowFunctionConverter(countingOperator); + } + + private WindowFunctionConverter windowFunctionConverter(SqlAggFunction operator) { + return new WindowFunctionConverter(extensions.windowFunctions(), typeFactory) { + @Override + public Optional getSqlOperatorFromSubstraitFunc(String key, Type outputType) { + return Optional.of(operator); + } + }; + } }