feat(isthmus): observe window function return types

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.
    *
-   * <p>Override to collect type observations during Substrait-to-Calcite conversion.
+   * <p>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> 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.

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<RexNode> inferredCallSupplier) {
+  private void observeType(
+      Expression expression,
+      TypeObservation.Source source,
+      Supplier<RelDataType> 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<SqlKind> asSqlKind(Expression.SortDirection direction) {

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;

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

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<TypeObservation> 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<TypeObservation> 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<TypeObservation> 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<TypeObservation> 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<TypeObservation> 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<SqlOperator> 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<SqlOperator> getSqlOperatorFromSubstraitFunc(String key, Type outputType) {
+        return Optional.of(operator);
+      }
+    };
+  }
 }