fix(isthmus)!: write the output type the datetime extension declares

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

@@ -52,6 +52,19 @@ public final class SimpleExtensionToSqlOperator {
 
   private SimpleExtensionToSqlOperator() {}
 
+  /**
+   * Identifies a dynamically mapped scalar operator whose parameterized result uses the first
+   * operand's type as a placeholder until the Substrait declaration is resolved.
+   *
+   * @param operator the Calcite operator to inspect
+   * @return whether its return type is a placeholder
+   */
+  public static boolean hasPlaceholderReturnType(SqlOperator operator) {
+    SqlReturnTypeInference inference = operator.getReturnTypeInference();
+    return inference instanceof ScalarReturnTypeInference
+        && !(((ScalarReturnTypeInference) inference).function.returnType() instanceof Type);
+  }
+
   /**
    * Converts all functions in a Substrait {@link SimpleExtension.ExtensionCollection} (scalar,
    * aggregate, and window) into Calcite {@link SqlOperator}s using the default type factory.

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

@@ -1,29 +1,38 @@
 package io.substrait.isthmus.expression;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.math.LongMath;
 import io.substrait.expression.Expression;
+import io.substrait.expression.ExpressionCreator;
 import io.substrait.expression.FunctionArg;
+import io.substrait.extension.DefaultExtensionCatalog;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.isthmus.CallConverter;
+import io.substrait.isthmus.SimpleExtensionToSqlOperator;
+import io.substrait.isthmus.SubstraitTypeSystem;
 import io.substrait.isthmus.TypeConverter;
 import io.substrait.type.Type;
+import io.substrait.type.TypeCreator;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
 
 /**
  * Converts Calcite {@link RexCall} scalar functions to Substrait {@link Expression} using known
  * Substrait {@link SimpleExtension.ScalarFunctionVariant} declarations.
  *
  * <p>Supports custom function mappers for special cases (e.g., TRIM, SQRT), and falls back to
- * default signature-based matching. Produces {@link Expression.ScalarFunctionInvocation}.
+ * default signature-based matching. Produces {@link Expression.ScalarFunctionInvocation},
+ * optionally wrapped in a cast to preserve Calcite's inferred result type for datetime calls.
  */
 public class ScalarFunctionConverter
     extends FunctionConverter<
@@ -154,25 +163,204 @@ private boolean isPotentialFunctionMatch(FunctionFinder finder, WrappedScalarCal
   }
 
   /**
-   * Builds an {@link Expression.ScalarFunctionInvocation} for a matched function.
+   * Binds a matched function. Datetime calls carry the declaration's resolved type, with a cast
+   * back when Calcite inferred a different result type. Dynamically mapped placeholder types do not
+   * require a cast.
    *
    * @param call the wrapped Calcite call providing operands and type
    * @param function the Substrait scalar function declaration to invoke
    * @param arguments converted argument list for the invocation
-   * @param outputType the Substrait output type for the invocation
-   * @return a scalar function invocation expression
+   * @param outputType the Calcite call's result type converted to Substrait
+   * @return a scalar invocation, optionally wrapped in a cast to the inferred result type
+   * @throws UnsupportedOperationException for DATE arithmetic with a non-literal sub-day interval
+   * @throws IllegalArgumentException if the resolved timestamp precision is unsupported or a
+   *     widened datetime literal overflows
    */
   @Override
   protected Expression generateBinding(
       WrappedScalarCall call,
       SimpleExtension.ScalarFunctionVariant function,
       List<? extends FunctionArg> arguments,
       Type outputType) {
-    return Expression.ScalarFunctionInvocation.builder()
-        .outputType(outputType)
-        .declaration(function)
-        .addAllArguments(arguments)
-        .build();
+    if (!DefaultExtensionCatalog.FUNCTIONS_DATETIME.equals(function.getAnchor().urn())) {
+      return ExpressionCreator.scalarFunction(function, outputType, arguments);
+    }
+    // The datetime extension declares its results by parameter, where Calcite keeps an operand's
+    // own type: add(date, interval_day<P>) is a precision_timestamp<P> there and a DATE here. The
+    // call carries the declared type, and a cast back to Calcite's keeps the column the type the
+    // query gives it.
+    List<? extends FunctionArg> bound =
+        dateArithmeticArguments(call, function, arguments, outputType);
+    if (!(function.returnType() instanceof Type)) {
+      // Precision-carrying arguments share one parameter in datetime declarations. Comparisons
+      // return bool and keep their operand precisions; widening also narrows a timestamp's range.
+      bound = widenedToOnePrecision(bound);
+    }
+    Optional<Type> declared = declaredType(function, bound);
+    if (declared.isEmpty()) {
+      return ExpressionCreator.scalarFunction(function, outputType, arguments);
+    }
+    requireSupportedTimestampPrecision(declared.get());
+    // Keep the call's nullability when its declared type otherwise matches.
+    boolean sameType = declared.get().equalsIgnoringNullability(outputType);
+    Expression invocation =
+        ExpressionCreator.scalarFunction(function, sameType ? outputType : declared.get(), bound);
+    return sameType
+            || SimpleExtensionToSqlOperator.hasPlaceholderReturnType(call.delegate.getOperator())
+        ? invocation
+        : ExpressionCreator.cast(
+            outputType, invocation, Expression.FailureBehavior.THROW_EXCEPTION);
+  }
+
+  private static List<? extends FunctionArg> dateArithmeticArguments(
+      WrappedScalarCall call,
+      SimpleExtension.ScalarFunctionVariant function,
+      List<? extends FunctionArg> arguments,
+      Type outputType) {
+    // A Substrait-origin call carries a timestamp result and must keep its sub-day interval.
+    if (!(outputType instanceof Type.Date)
+        || !("add:date_iday".equals(function.key())
+            || "subtract:date_iday".equals(function.key()))) {
+      return arguments;
+    }
+    RexNode intervalOperand = call.getOperands().skip(1).findFirst().orElseThrow();
+    if (intervalOperand.getType().getSqlTypeName() == SqlTypeName.INTERVAL_DAY) {
+      return arguments;
+    }
+    FunctionArg interval = arguments.get(1);
+    if (interval instanceof Expression.NullLiteral) {
+      return arguments;
+    }
+    if (!(interval instanceof Expression.IntervalDayLiteral)) {
+      throw new UnsupportedOperationException(
+          "DATE arithmetic with a non-literal sub-day interval is not supported");
+    }
+    // LiteralConverter decomposes the total duration towards zero. Keeping only its days makes
+    // the timestamp call followed by a DATE cast agree with Calcite's DATE arithmetic.
+    Expression.IntervalDayLiteral literal = (Expression.IntervalDayLiteral) interval;
+    return List.of(
+        arguments.get(0),
+        ExpressionCreator.intervalDay(
+            literal.nullable(), literal.days(), 0, 0, literal.precision()));
+  }
+
+  private static Optional<Type> declaredType(
+      SimpleExtension.ScalarFunctionVariant function, List<? extends FunctionArg> arguments) {
+    List<Type> types =
+        arguments.stream()
+            .filter(Expression.class::isInstance)
+            .map(argument -> ((Expression) argument).getType())
+            .collect(Collectors.toList());
+    try {
+      return Optional.of(function.resolveType(types));
+    } catch (UnsupportedOperationException e) {
+      return Optional.empty();
+    }
+  }
+
+  private List<FunctionArg> widenedToOnePrecision(List<? extends FunctionArg> arguments) {
+    int precision =
+        arguments.stream()
+            .filter(Expression.class::isInstance)
+            .map(argument -> precisionOf(((Expression) argument).getType()))
+            .flatMap(Optional::stream)
+            .max(Integer::compare)
+            .orElse(0);
+    return arguments.stream()
+        .map(
+            argument -> {
+              if (!(argument instanceof Expression)) {
+                return argument;
+              }
+              Expression expression = (Expression) argument;
+              Type type = expression.getType();
+              Optional<Integer> own = precisionOf(type);
+              if (own.isEmpty() || own.get() == precision) {
+                return argument;
+              }
+              Type widened = withPrecision(type, precision);
+              requireSupportedTimestampPrecision(widened);
+              if (expression instanceof Expression.NullLiteral) {
+                return ExpressionCreator.typedNull(widened);
+              }
+              try {
+                long factor = LongMath.checkedPow(10, precision - own.get());
+                if (expression instanceof Expression.IntervalDayLiteral) {
+                  Expression.IntervalDayLiteral literal =
+                      (Expression.IntervalDayLiteral) expression;
+                  return ExpressionCreator.intervalDay(
+                      literal.nullable(),
+                      literal.days(),
+                      literal.seconds(),
+                      Math.multiplyExact(literal.subseconds(), factor),
+                      precision);
+                }
+                if (expression instanceof Expression.PrecisionTimestampLiteral) {
+                  Expression.PrecisionTimestampLiteral literal =
+                      (Expression.PrecisionTimestampLiteral) expression;
+                  return ExpressionCreator.precisionTimestamp(
+                      literal.nullable(), Math.multiplyExact(literal.value(), factor), precision);
+                }
+                if (expression instanceof Expression.PrecisionTimestampTZLiteral) {
+                  Expression.PrecisionTimestampTZLiteral literal =
+                      (Expression.PrecisionTimestampTZLiteral) expression;
+                  return ExpressionCreator.precisionTimestampTZ(
+                      literal.nullable(), Math.multiplyExact(literal.value(), factor), precision);
+                }
+              } catch (ArithmeticException e) {
+                throw new IllegalArgumentException(
+                    "Datetime literal cannot be represented at precision "
+                        + precision
+                        + " in a signed 64-bit value",
+                    e);
+              }
+              return ExpressionCreator.cast(
+                  widened, expression, Expression.FailureBehavior.THROW_EXCEPTION);
+            })
+        .collect(Collectors.toList());
+  }
+
+  private void requireSupportedTimestampPrecision(Type type) {
+    if (type instanceof Type.PrecisionTimestamp) {
+      SubstraitTypeSystem.requireSupportedPrecision(
+          typeFactory.getTypeSystem(),
+          SqlTypeName.TIMESTAMP,
+          "precision_timestamp",
+          ((Type.PrecisionTimestamp) type).precision());
+    } else if (type instanceof Type.PrecisionTimestampTZ) {
+      SubstraitTypeSystem.requireSupportedPrecision(
+          typeFactory.getTypeSystem(),
+          SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE,
+          "precision_timestamp_tz",
+          ((Type.PrecisionTimestampTZ) type).precision());
+    }
+  }
+
+  private static Optional<Integer> precisionOf(Type type) {
+    if (type instanceof Type.PrecisionTimestamp) {
+      return Optional.of(((Type.PrecisionTimestamp) type).precision());
+    }
+    if (type instanceof Type.PrecisionTimestampTZ) {
+      return Optional.of(((Type.PrecisionTimestampTZ) type).precision());
+    }
+    if (type instanceof Type.IntervalDay) {
+      return Optional.of(((Type.IntervalDay) type).precision());
+    }
+    return Optional.empty();
+  }
+
+  private static Type withPrecision(Type type, int precision) {
+    TypeCreator creator = TypeCreator.of(type.nullable());
+    if (type instanceof Type.PrecisionTimestamp) {
+      return creator.precisionTimestamp(precision);
+    }
+    if (type instanceof Type.PrecisionTimestampTZ) {
+      return creator.precisionTimestampTZ(precision);
+    }
+    if (type instanceof Type.IntervalDay) {
+      return creator.intervalDay(precision);
+    }
+    throw new IllegalArgumentException("Unsupported datetime type: " + type);
   }
 
   /**

isthmus/src/main/java/io/substrait/isthmus/sql/SubstraitSqlDialect.java

@@ -2,8 +2,10 @@
 
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.rel2sql.RelToSqlConverter;
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.SqlDialect;
 import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.util.SqlString;
 
 /**
@@ -50,6 +52,21 @@ public SubstraitSqlDialect(Context context) {
     super(context);
   }
 
+  /**
+   * Renders interval casts with their SQL qualifier rather than Calcite's internal type name.
+   *
+   * @param type the target type of the cast
+   * @return the interval qualifier, or the default cast specification for other types
+   */
+  @Override
+  public SqlNode getCastSpec(RelDataType type) {
+    if (type.getSqlTypeName().getFamily() == SqlTypeFamily.INTERVAL_DAY_TIME
+        || type.getSqlTypeName().getFamily() == SqlTypeFamily.INTERVAL_YEAR_MONTH) {
+      return type.getIntervalQualifier();
+    }
+    return super.getCastSpec(type);
+  }
+
   /**
    * Indicates whether this dialect supports approximate COUNT(DISTINCT) operations.
    *

isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java

@@ -1,6 +1,7 @@
 package io.substrait.isthmus;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.google.common.collect.ImmutableList;
@@ -16,9 +17,13 @@
 import java.util.function.Consumer;
 import org.apache.calcite.avatica.util.TimeUnitRange;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.util.TimestampString;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 
 class CalciteCallTest extends CalciteObjs {
 
@@ -34,14 +39,56 @@ class CalciteCallTest extends CalciteObjs {
 
   @Test
   void extract() {
-    test(
-        "extract:req_pts",
+    // extract declares i64. A call Calcite types INTEGER carries the declared type and is cast back
+    // to the one Calcite gave it.
+    Expression expression =
         rex.makeCall(
-            t(SqlTypeName.INTEGER),
-            SqlStdOperatorTable.EXTRACT,
-            ImmutableList.of(rex.makeFlag(TimeUnitRange.HOUR), c(10L, SqlTypeName.TIMESTAMP, 10))),
-        func -> {},
-        false);
+                t(SqlTypeName.INTEGER),
+                SqlStdOperatorTable.EXTRACT,
+                ImmutableList.of(
+                    rex.makeFlag(TimeUnitRange.HOUR), c(10L, SqlTypeName.TIMESTAMP, 10)))
+            .accept(rexExpressionConverter);
+    Expression.Cast cast = assertInstanceOf(Expression.Cast.class, expression);
+    assertEquals(TypeCreator.REQUIRED.I32, cast.getType());
+    Expression.ScalarFunctionInvocation func =
+        assertInstanceOf(Expression.ScalarFunctionInvocation.class, cast.input());
+    assertEquals("extract:req_pts", func.declaration().key());
+    assertEquals(TypeCreator.REQUIRED.I64, func.outputType());
+  }
+
+  @ParameterizedTest
+  @CsvSource({"lt,8", "lt,9", "lte,8", "lte,9", "gt,8", "gt,9", "gte,8", "gte,9"})
+  void datetimeComparisonKeepsOperandPrecisions(String function, int precision) {
+    SqlOperator operator;
+    switch (function) {
+      case "lt":
+        operator = SqlStdOperatorTable.LESS_THAN;
+        break;
+      case "lte":
+        operator = SqlStdOperatorTable.LESS_THAN_OR_EQUAL;
+        break;
+      case "gt":
+        operator = SqlStdOperatorTable.GREATER_THAN;
+        break;
+      default:
+        operator = SqlStdOperatorTable.GREATER_THAN_OR_EQUAL;
+    }
+    RexNode timestamp = rex.makeInputRef(t(SqlTypeName.TIMESTAMP, precision), 0);
+    // This bound is valid at millisecond precision but overflows an i64 at precision 8 or 9.
+    RexNode bound = rex.makeTimestampLiteral(new TimestampString("9999-12-31 00:00:00"), 3);
+    Expression.ScalarFunctionInvocation comparison =
+        assertInstanceOf(
+            Expression.ScalarFunctionInvocation.class,
+            rex.makeCall(operator, timestamp, bound).accept(rexExpressionConverter));
+    assertEquals(function + ":pts_pts", comparison.declaration().key());
+    assertEquals(TypeCreator.REQUIRED.BOOLEAN, comparison.outputType());
+    assertEquals(
+        TypeCreator.REQUIRED.precisionTimestamp(precision),
+        ((Expression) comparison.arguments().get(0)).getType());
+    assertInstanceOf(Expression.PrecisionTimestampLiteral.class, comparison.arguments().get(1));
+    assertEquals(
+        TypeCreator.REQUIRED.precisionTimestamp(3),
+        ((Expression) comparison.arguments().get(1)).getType());
   }
 
   @Test