feat(isthmus): observe scalar function types

Assisted-by: gpt-5.6-sol (OpenAI Codex)

Author: alexandrefimov

isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java

@@ -11,6 +11,7 @@
 import io.substrait.isthmus.expression.ScalarFunctionConverter;
 import io.substrait.isthmus.expression.SqlArrayValueConstructorCallConverter;
 import io.substrait.isthmus.expression.SqlMapValueConstructorCallConverter;
+import io.substrait.isthmus.expression.TypeObserver;
 import io.substrait.isthmus.expression.WindowFunctionConverter;
 import io.substrait.plan.ImmutableExecutionBehavior;
 import io.substrait.plan.Plan;
@@ -330,11 +331,23 @@ public ExpressionRexConverter getExpressionRexConverter(
             getTypeFactory(),
             getScalarFunctionConverter(),
             getWindowFunctionConverter(),
-            getTypeConverter());
+            getTypeConverter(),
+            getTypeObserver());
     erc.setRelNodeConverter(relNodeConverter);
     return erc;
   }
 
+  /**
+   * Returns the observer for supplied and independently inferred expression types.
+   *
+   * <p>Override to collect type observations during Substrait-to-Calcite conversion.
+   *
+   * @return a no-op observer by default
+   */
+  public TypeObserver getTypeObserver() {
+    return TypeObserver.NOOP;
+  }
+
   /**
    * A {@link RelBuilder} is a Calcite class used for creating {@link
    * org.apache.calcite.rel.RelNode}s.

isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java

@@ -26,8 +26,10 @@
 import java.math.BigDecimal;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -95,6 +97,9 @@ public class ExpressionRexConverter
   /** Converter for Substrait window function invocations to Calcite {@link SqlOperator}s. */
   protected final WindowFunctionConverter windowFunctionConverter;
 
+  /** Observer for supplied and inferred expression types. */
+  protected final TypeObserver typeObserver;
+
   /** Converter for Substrait relational nodes to Calcite {@link RelNode}s, used for subqueries. */
   protected SubstraitRelNodeConverter relNodeConverter;
 
@@ -113,11 +118,35 @@ public ExpressionRexConverter(
       ScalarFunctionConverter scalarFunctionConverter,
       WindowFunctionConverter windowFunctionConverter,
       TypeConverter typeConverter) {
+    this(
+        typeFactory,
+        scalarFunctionConverter,
+        windowFunctionConverter,
+        typeConverter,
+        TypeObserver.NOOP);
+  }
+
+  /**
+   * Creates an {@code ExpressionRexConverter} with type observation enabled.
+   *
+   * @param typeFactory Calcite type factory for type creation
+   * @param scalarFunctionConverter converter for scalar function invocations
+   * @param windowFunctionConverter converter for window function invocations
+   * @param typeConverter converter for Substrait and Calcite type mappings
+   * @param typeObserver observer for supplied and independently inferred expression types
+   */
+  public ExpressionRexConverter(
+      RelDataTypeFactory typeFactory,
+      ScalarFunctionConverter scalarFunctionConverter,
+      WindowFunctionConverter windowFunctionConverter,
+      TypeConverter typeConverter,
+      TypeObserver typeObserver) {
     this.typeFactory = typeFactory;
     this.typeConverter = typeConverter;
     this.rexBuilder = new RexBuilder(typeFactory);
     this.scalarFunctionConverter = scalarFunctionConverter;
     this.windowFunctionConverter = windowFunctionConverter;
+    this.typeObserver = Objects.requireNonNull(typeObserver, "typeObserver");
   }
 
   /**
@@ -519,7 +548,33 @@ public RexNode visit(Expression.ScalarFunctionInvocation expr, Context context)
             .collect(Collectors.toList());
 
     RelDataType returnType = typeConverter.toCalcite(typeFactory, expr.outputType());
-    return rexBuilder.makeCall(returnType, operator, args);
+    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));
+    return rexCall;
+  }
+
+  private void observeScalarType(
+      Expression.ScalarFunctionInvocation expression, Supplier<RexNode> inferredCallSupplier) {
+    TypeObservation observation;
+    RexNode inferredCall;
+    try {
+      inferredCall = inferredCallSupplier.get();
+    } catch (RuntimeException inferenceFailure) {
+      observation =
+          TypeObservation.failure(
+              TypeObservation.Source.SCALAR_FUNCTION, expression, inferenceFailure);
+      typeObserver.observe(observation);
+      return;
+    }
+    observation =
+        TypeObservation.success(
+            TypeObservation.Source.SCALAR_FUNCTION, expression, inferredCall.getType());
+    typeObserver.observe(observation);
   }
 
   private String callConversionFailureMessage(

isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java

@@ -0,0 +1,110 @@
+package io.substrait.isthmus.expression;
+
+import io.substrait.expression.Expression;
+import io.substrait.type.Type;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.calcite.rel.type.RelDataType;
+
+/**
+ * The result of attempting to independently infer a Calcite type during Substrait-to-Calcite
+ * expression conversion. Exactly one of {@link #inferredType()} and {@link #inferenceFailure()} is
+ * present.
+ */
+public final class TypeObservation {
+  /** The expression category that produced an observation. */
+  public enum Source {
+    /** A scalar function invocation. */
+    SCALAR_FUNCTION
+  }
+
+  private final Source source;
+  private final Expression expression;
+  private final RelDataType inferredType;
+  private final RuntimeException inferenceFailure;
+
+  /**
+   * Creates a successful type observation.
+   *
+   * @param source expression category that produced the observation
+   * @param expression expression that produced the observation
+   * @param inferredType type independently inferred by Calcite
+   * @return a successful type observation
+   */
+  static TypeObservation success(Source source, Expression expression, RelDataType inferredType) {
+    return new TypeObservation(source, expression, inferredType, null);
+  }
+
+  /**
+   * Creates a failed type observation.
+   *
+   * @param source expression category that produced the observation
+   * @param expression expression that produced the observation
+   * @param inferenceFailure failure to independently infer a Calcite type
+   * @return a failed type observation
+   */
+  static TypeObservation failure(
+      Source source, Expression expression, RuntimeException inferenceFailure) {
+    return new TypeObservation(source, expression, null, inferenceFailure);
+  }
+
+  private TypeObservation(
+      Source source,
+      Expression expression,
+      RelDataType inferredType,
+      RuntimeException inferenceFailure) {
+    this.source = Objects.requireNonNull(source, "source");
+    this.expression = Objects.requireNonNull(expression, "expression");
+    if ((inferredType == null) == (inferenceFailure == null)) {
+      throw new IllegalArgumentException(
+          "Exactly one of inferredType and inferenceFailure must be present");
+    }
+    this.inferredType = inferredType;
+    this.inferenceFailure = inferenceFailure;
+  }
+
+  /**
+   * Returns the expression category that produced this observation.
+   *
+   * @return the expression category
+   */
+  public Source source() {
+    return source;
+  }
+
+  /**
+   * Returns the expression that produced this observation.
+   *
+   * @return the observed expression
+   */
+  public Expression expression() {
+    return expression;
+  }
+
+  /**
+   * Returns the type supplied by Substrait.
+   *
+   * @return the supplied type
+   */
+  public Type suppliedType() {
+    return expression.getType();
+  }
+
+  /**
+   * Returns the type independently inferred by Calcite.
+   *
+   * @return the inferred type, or empty if inference failed
+   */
+  public Optional<RelDataType> inferredType() {
+    return Optional.ofNullable(inferredType);
+  }
+
+  /**
+   * Returns the Calcite inference failure.
+   *
+   * @return the inference failure, or empty if inference succeeded
+   */
+  public Optional<RuntimeException> inferenceFailure() {
+    return Optional.ofNullable(inferenceFailure);
+  }
+}

isthmus/src/main/java/io/substrait/isthmus/expression/TypeObserver.java

@@ -0,0 +1,17 @@
+package io.substrait.isthmus.expression;
+
+/** Receives type observations while converting Substrait expressions to Calcite. */
+@FunctionalInterface
+public interface TypeObserver {
+  /** Observer that disables type inference and discards all observations. */
+  TypeObserver NOOP = observation -> {};
+
+  /**
+   * Receives the result of attempting to observe an expression's inferred type.
+   *
+   * <p>Exceptions thrown by an observer are propagated to the conversion caller.
+   *
+   * @param observation supplied type and either an inferred type or inference failure
+   */
+  void observe(TypeObservation observation);
+}