diff --git a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java index 7fd0ab6e6..56f7ea58b 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; @@ -215,6 +218,7 @@ protected ConverterProvider(Builder builder) { this.typeConverter = builder.typeConverter; this.executionBehavior = builder.executionBehavior; this.sqlParserConfig = builder.sqlParserConfig; + this.typeObserver = builder.typeObserver; this.scalarFunctionConverter = builder.scalarFunctionConverter.orElseGet( @@ -421,12 +425,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; } /** @@ -558,6 +563,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; // Derived from the extensions and type factory at build time when left unset. private Optional scalarFunctionConverter = Optional.empty(); @@ -622,6 +628,17 @@ public Builder sqlParserConfig(SqlParser.Config sqlParserConfig) { 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..43336be67 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -554,26 +554,27 @@ public RexNode visit(Expression.ScalarFunctionInvocation expr, Context context) if (typeObserver == TypeObserver.NOOP) { return rexCall; } - observeScalarType(expr, () -> rexBuilder.makeCall(operator, args)); + observeType( + expr, + TypeObservation.Source.SCALAR_FUNCTION, + () -> rexBuilder.makeCall(operator, args).getType()); return rexCall; } - private void observeScalarType( - Expression.ScalarFunctionInvocation expression, Supplier inferredCallSupplier) { + private void observeType( + Expression expression, + TypeObservation.Source source, + Supplier inferredTypeSupplier) { 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 +632,28 @@ 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); + if (typeObserver == TypeObserver.NOOP) { + return rexOver; + } + observeType( + expr, + TypeObservation.Source.WINDOW_FUNCTION, + () -> rexBuilder.deriveReturnType(operator, args)); + return rexOver; } 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 c3d0f7ba2..52f83a667 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.rel.type.RelDataTypeFactory; import org.junit.jupiter.api.Test; @@ -29,6 +30,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..8f316226e 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java @@ -30,6 +30,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; @@ -279,12 +280,7 @@ void observeVariadicConcatOnce() { void injectTypeObserverThroughConverterProvider() { AtomicReference observed = new AtomicReference<>(); ConverterProvider observingProvider = - new ConverterProvider() { - @Override - public TypeObserver getTypeObserver() { - return observed::set; - } - }; + ConverterProvider.builder().typeObserver(observed::set).build(); Expression.ScalarFunctionInvocation expr = sb.scalarFn( DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, @@ -357,12 +353,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,8 +412,64 @@ void useSubstraitReturnTypeDuringWindowFunctionConversion() { WindowBound.UNBOUNDED, sb.i32(42)); - RexNode calciteExpr = expr.accept(expressionRexConverter, Context.newContext()); + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + assertEquals(SqlTypeName.INTEGER, observed.get().inferredType().orElseThrow().getSqlTypeName()); + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.STRING), calciteExpr.getType()); + } + + @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) { @@ -389,6 +486,18 @@ 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 ExpressionRexConverter observingConverter(TypeObserver observer) { return observingConverter( new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory), observer); @@ -404,6 +513,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 +562,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); + } + }; + } }