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);
+ }
+ };
+ }
}