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