diff --git a/isthmus/src/main/java/io/substrait/isthmus/SimpleExtensionToSqlOperator.java b/isthmus/src/main/java/io/substrait/isthmus/SimpleExtensionToSqlOperator.java index 652c04a34..468d7dc40 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SimpleExtensionToSqlOperator.java +++ b/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. diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ScalarFunctionConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ScalarFunctionConverter.java index 4f29af82d..fa9dc7b2e 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ScalarFunctionConverter.java +++ b/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. * *

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,13 +163,18 @@ 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( @@ -168,11 +182,185 @@ protected Expression generateBinding( SimpleExtension.ScalarFunctionVariant function, List 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

) is a precision_timestamp

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 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 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 dateArithmeticArguments( + WrappedScalarCall call, + SimpleExtension.ScalarFunctionVariant function, + List 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 declaredType( + SimpleExtension.ScalarFunctionVariant function, List arguments) { + List 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 widenedToOnePrecision(List 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 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 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); } /** diff --git a/isthmus/src/main/java/io/substrait/isthmus/sql/SubstraitSqlDialect.java b/isthmus/src/main/java/io/substrait/isthmus/sql/SubstraitSqlDialect.java index 340021739..846ed57b2 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/sql/SubstraitSqlDialect.java +++ b/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. * diff --git a/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java b/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java index d03239408..93356000d 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java +++ b/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 diff --git a/isthmus/src/test/java/io/substrait/isthmus/DateIntervalArithmeticTest.java b/isthmus/src/test/java/io/substrait/isthmus/DateIntervalArithmeticTest.java new file mode 100644 index 000000000..d38cb9435 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/DateIntervalArithmeticTest.java @@ -0,0 +1,167 @@ +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.assertThrows; + +import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; +import io.substrait.expression.FieldReference; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.isthmus.expression.RexExpressionConverter; +import io.substrait.isthmus.expression.ScalarFunctionConverter; +import io.substrait.isthmus.sql.SubstraitCreateStatementParser; +import io.substrait.relation.NamedScan; +import io.substrait.relation.Project; +import io.substrait.type.TypeCreator; +import java.util.List; +import org.apache.calcite.avatica.util.TimeUnitRange; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlIntervalQualifier; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.calcite.sql.type.SqlTypeName; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; + +class DateIntervalArithmeticTest extends PlanTestBase { + + @ParameterizedTest + @CsvSource({ + "+, 36, HOUR, 1", + "+, -36, HOUR, -1", + "-, 36, HOUR, 1", + "-, -36, HOUR, -1", + "+, 23, HOUR, 0", + "+, -23, HOUR, 0", + "-, 23, HOUR, 0", + "-, -23, HOUR, 0", + "+, 2161, MINUTE, 1", + "-, -2161, MINUTE, -1", + "+, 86400.001, SECOND, 1", + "-, -86400.001, SECOND, -1", + "+, 1 12:30:45.123, DAY TO SECOND, 1", + "-, -1 12:30:45.123, DAY TO SECOND, -1" + }) + void subDayLiteralTruncatesTowardsZero( + String operator, String value, String qualifier, int expectedDays) throws Exception { + String query = + "SELECT d " + operator + " INTERVAL '" + value + "' " + qualifier + " FROM events"; + Project project = + assertInstanceOf( + Project.class, + new SqlToSubstrait() + .convert( + query, + SubstraitCreateStatementParser.processCreateStatementsToCatalog( + "CREATE TABLE events (d DATE)")) + .getRoots() + .get(0) + .getInput()); + Expression.Cast cast = assertInstanceOf(Expression.Cast.class, project.getExpressions().get(0)); + assertEquals(N.DATE, cast.getType()); + Expression.ScalarFunctionInvocation arithmetic = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, cast.input()); + assertEquals( + (operator.equals("+") ? "add" : "subtract") + ":date_iday", arithmetic.declaration().key()); + assertEquals(N.precisionTimestamp(6), arithmetic.outputType()); + assertEquals( + ExpressionCreator.intervalDay(false, expectedDays, 0, 0, 6), arithmetic.arguments().get(1)); + assertFullRoundTrip(project); + } + + @ParameterizedTest + @EnumSource( + value = TimeUnitRange.class, + names = { + "DAY", + "DAY_TO_HOUR", + "DAY_TO_MINUTE", + "DAY_TO_SECOND", + "HOUR", + "HOUR_TO_MINUTE", + "HOUR_TO_SECOND", + "MINUTE", + "MINUTE_TO_SECOND", + "SECOND" + }) + void dateResultAcceptsOnlyDayQualifiedIntervalColumns(TimeUnitRange qualifier) { + RexBuilder rex = new RexBuilder(typeFactory); + RelDataType intervalType = + typeFactory.createSqlIntervalType( + new SqlIntervalQualifier(qualifier.startUnit, qualifier.endUnit, SqlParserPos.ZERO)); + RexExpressionConverter converter = + new RexExpressionConverter( + new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory)); + for (SqlOperator operator : + List.of(SqlStdOperatorTable.DATETIME_PLUS, SqlStdOperatorTable.MINUS_DATE)) { + RexNode call = + rex.makeCall( + typeFactory.createSqlType(SqlTypeName.DATE), + operator, + List.of( + rex.makeInputRef(typeFactory.createSqlType(SqlTypeName.DATE), 0), + rex.makeInputRef(intervalType, 1))); + if (qualifier == TimeUnitRange.DAY) { + Expression.Cast cast = assertInstanceOf(Expression.Cast.class, call.accept(converter)); + assertEquals(R.DATE, cast.getType()); + Expression.ScalarFunctionInvocation arithmetic = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, cast.input()); + assertEquals(R.precisionTimestamp(6), arithmetic.outputType()); + assertInstanceOf(FieldReference.class, arithmetic.arguments().get(1)); + } else { + UnsupportedOperationException error = + assertThrows(UnsupportedOperationException.class, () -> call.accept(converter)); + assertEquals( + "DATE arithmetic with a non-literal sub-day interval is not supported", + error.getMessage()); + } + } + } + + @ParameterizedTest + @CsvSource({"add,false", "add,true", "subtract,false", "subtract,true"}) + void timestampResultWithIntervalColumnRoundTrips(String function, boolean nullable) { + TypeCreator types = TypeCreator.of(nullable); + NamedScan table = + sb.namedScan( + List.of("events"), List.of("d", "duration"), List.of(types.DATE, types.intervalDay(6))); + Project project = + sb.project( + input -> + List.of( + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + function + ":date_iday", + types.precisionTimestamp(6), + sb.fieldReference(input, 0), + sb.fieldReference(input, 1))), + sb.remap(2), + table); + assertFullRoundTrip(project); + } + + @ParameterizedTest + @ValueSource(strings = {"add", "subtract"}) + void timestampResultKeepsSubDayLiteral(String function) { + NamedScan table = sb.namedScan(List.of("events"), List.of("d"), List.of(R.DATE)); + Project project = + sb.project( + input -> + List.of( + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + function + ":date_iday", + R.precisionTimestamp(6), + sb.fieldReference(input, 0), + ExpressionCreator.intervalDay(false, -1, -43200, 0, 6))), + sb.remap(1), + table); + assertFullRoundTrip(project); + } +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/DatetimeBindingRegressionTest.java b/isthmus/src/test/java/io/substrait/isthmus/DatetimeBindingRegressionTest.java new file mode 100644 index 000000000..3c0a217c4 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/DatetimeBindingRegressionTest.java @@ -0,0 +1,253 @@ +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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.isthmus.expression.RexExpressionConverter; +import io.substrait.isthmus.sql.SubstraitCreateStatementParser; +import io.substrait.isthmus.sql.SubstraitSqlToCalcite; +import io.substrait.plan.Plan; +import io.substrait.relation.Project; +import io.substrait.type.Type; +import io.substrait.type.TypeCreator; +import java.util.List; +import org.apache.calcite.avatica.util.TimeUnit; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.rel.type.RelDataTypeSystem; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlIntervalQualifier; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.TimestampString; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; + +class DatetimeBindingRegressionTest extends PlanTestBase { + + private static final String CREATES = + "CREATE TABLE events (ts9 TIMESTAMP(9), ts3 TIMESTAMP(3), i INTEGER)"; + + @ParameterizedTest + @ValueSource(strings = {"add", "subtract"}) + void nullableDatetimeCallOverRequiredLiteralsRoundTrips(String function) { + Expression.ScalarFunctionInvocation call = + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + function + ":pts_iday", + N.precisionTimestamp(6), + ExpressionCreator.precisionTimestamp(false, 1_704_067_200_000_000L, 6), + ExpressionCreator.intervalDay(false, 5, 0, 0, 6)); + assertFullRoundTrip(sb.project(input -> List.of(call), sb.emptyVirtualTableScan())); + } + + @ParameterizedTest + @CsvSource({ + "+, TIMESTAMP, false", + "-, TIMESTAMP, false", + "+, TIMESTAMP, true", + "-, TIMESTAMP, true", + "+, TIMESTAMP_WITH_LOCAL_TIME_ZONE, false", + "-, TIMESTAMP_WITH_LOCAL_TIME_ZONE, false", + "+, TIMESTAMP_WITH_LOCAL_TIME_ZONE, true", + "-, TIMESTAMP_WITH_LOCAL_TIME_ZONE, true" + }) + void timestampZeroDayArithmeticWidensBeforeResolving( + String operator, SqlTypeName typeName, boolean nullable) throws Exception { + String inputType = + typeName == SqlTypeName.TIMESTAMP ? "TIMESTAMP(0)" : "TIMESTAMP(0) WITH LOCAL TIME ZONE"; + String creates = "CREATE TABLE events (ts0 " + inputType + (nullable ? "" : " NOT NULL") + ")"; + String query = "SELECT ts0 " + operator + " INTERVAL '5' DAY FROM events"; + Expression.Cast result = + assertInstanceOf(Expression.Cast.class, firstExpression(query, creates)); + Expression.ScalarFunctionInvocation call = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, result.input()); + TypeCreator types = TypeCreator.of(nullable); + Type resultType = + typeName == SqlTypeName.TIMESTAMP + ? types.precisionTimestamp(0) + : types.precisionTimestampTZ(0); + Type widenedType = + typeName == SqlTypeName.TIMESTAMP + ? types.precisionTimestamp(6) + : types.precisionTimestampTZ(6); + assertEquals(resultType, result.getType()); + assertEquals(widenedType, call.outputType()); + assertEquals(widenedType, ((Expression) call.arguments().get(0)).getType()); + assertFullRoundTrip(query, creates); + } + + @ParameterizedTest + @ValueSource(strings = {"+", "-"}) + void widenedIntervalLiteralRoundTrips(String operator) throws Exception { + assertFullRoundTrip("SELECT ts9 " + operator + " INTERVAL '5' DAY FROM events", CREATES); + } + + @ParameterizedTest + @ValueSource(strings = {"0.123", "-0.123"}) + void widenedIntervalLiteralPreservesSubseconds(String value) throws Exception { + String query = "SELECT ts9 + INTERVAL '" + value + "' SECOND FROM events"; + Plan plan = + new SqlToSubstrait() + .convert( + query, SubstraitCreateStatementParser.processCreateStatementsToCatalog(CREATES)); + Project project = assertInstanceOf(Project.class, plan.getRoots().get(0).getInput()); + Expression.ScalarFunctionInvocation call = + assertInstanceOf( + Expression.ScalarFunctionInvocation.class, project.getExpressions().get(0)); + assertEquals( + ExpressionCreator.intervalDay( + false, 0, 0, value.startsWith("-") ? -123_000_000 : 123_000_000, 9), + call.arguments().get(1)); + assertFullRoundTrip(query, CREATES); + } + + @ParameterizedTest + @CsvSource({ + "TIMESTAMP, 2024-01-01 00:00:00.123, 1704067200123000000", + "TIMESTAMP, 1969-12-31 23:59:59.123, -877000000", + "TIMESTAMP_WITH_LOCAL_TIME_ZONE, 2024-01-01 00:00:00.123, 1704067200123000000", + "TIMESTAMP_WITH_LOCAL_TIME_ZONE, 1969-12-31 23:59:59.123, -877000000" + }) + void widenedTimestampLiteralPreservesValue( + SqlTypeName typeName, String timestamp, long expected) { + RexBuilder rex = new RexBuilder(converterProvider.getTypeFactory()); + RexNode literal = + rex.makeLiteral( + new TimestampString(timestamp), + converterProvider.getTypeFactory().createSqlType(typeName, 3), + true); + Expression.Cast result = + assertInstanceOf(Expression.Cast.class, timestampPlusInterval(literal)); + Expression.ScalarFunctionInvocation call = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, result.input()); + assertEquals( + typeName == SqlTypeName.TIMESTAMP + ? ExpressionCreator.precisionTimestamp(false, expected, 9) + : ExpressionCreator.precisionTimestampTZ(false, expected, 9), + call.arguments().get(0)); + } + + @ParameterizedTest + @EnumSource( + value = SqlTypeName.class, + names = {"TIMESTAMP", "TIMESTAMP_WITH_LOCAL_TIME_ZONE"}) + void overflowingTimestampLiteralFailsOnConversion(SqlTypeName typeName) { + RexBuilder rex = new RexBuilder(converterProvider.getTypeFactory()); + RexNode literal = + rex.makeLiteral( + new TimestampString("9999-12-31 00:00:00"), + converterProvider.getTypeFactory().createSqlType(typeName, 3), + true); + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> timestampPlusInterval(literal)); + assertTrue(error.getMessage().contains("precision 9"), error.getMessage()); + assertTrue(error.getMessage().contains("64-bit"), error.getMessage()); + } + + @ParameterizedTest + @EnumSource( + value = SqlTypeName.class, + names = {"TIMESTAMP", "TIMESTAMP_WITH_LOCAL_TIME_ZONE"}) + void widenedTimestampNullRemainsNull(SqlTypeName typeName) { + RexBuilder rex = new RexBuilder(converterProvider.getTypeFactory()); + RexNode literal = + rex.makeNullLiteral(converterProvider.getTypeFactory().createSqlType(typeName, 3)); + Expression.Cast result = + assertInstanceOf(Expression.Cast.class, timestampPlusInterval(literal)); + Expression.ScalarFunctionInvocation call = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, result.input()); + assertEquals( + ExpressionCreator.typedNull( + typeName == SqlTypeName.TIMESTAMP + ? N.precisionTimestamp(9) + : N.precisionTimestampTZ(9)), + call.arguments().get(0)); + } + + private Expression timestampPlusInterval(RexNode timestamp) { + RexBuilder rex = new RexBuilder(converterProvider.getTypeFactory()); + RexNode interval = + rex.makeInputRef( + converterProvider + .getTypeFactory() + .createSqlIntervalType( + new SqlIntervalQualifier( + TimeUnit.DAY, 10, TimeUnit.SECOND, 9, SqlParserPos.ZERO)), + 0); + return rex.makeCall(timestamp.getType(), SqlStdOperatorTable.PLUS, List.of(timestamp, interval)) + .accept(new RexExpressionConverter(converterProvider.getScalarFunctionConverter())); + } + + @ParameterizedTest + @ValueSource(strings = {"+", "-"}) + void widenedIntervalExpressionProducesParseableSql(String operator) throws Exception { + String query = "SELECT ts9 " + operator + " i * INTERVAL '1' DAY FROM events"; + Plan plan = + new SqlToSubstrait() + .convert( + query, SubstraitCreateStatementParser.processCreateStatementsToCatalog(CREATES)); + String sql = toSql(plan); + assertTrue(sql.contains("AS INTERVAL DAY"), sql); + SubstraitSqlToCalcite.convertQuery( + sql, + SubstraitCreateStatementParser.processCreateStatementsToCatalog(CREATES), + converterProvider); + assertFullRoundTrip(query, CREATES); + } + + @ParameterizedTest + @ValueSource(strings = {"TIMESTAMP(3)", "TIMESTAMP(3) WITH LOCAL TIME ZONE"}) + void timestampPrecisionAboveConfiguredLimitFailsOnConversion(String inputType) throws Exception { + ConverterProvider provider = + ConverterProvider.builder() + .typeFactory(new JavaTypeFactoryImpl(RelDataTypeSystem.DEFAULT)) + .build(); + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> + new SqlToSubstrait(provider) + .convert( + "SELECT ts3 - INTERVAL '5' DAY FROM events", + SubstraitCreateStatementParser.processCreateStatementsToCatalog( + "CREATE TABLE events (ts3 " + inputType + ")"))); + assertTrue(error.getMessage().contains("precision 6"), error.getMessage()); + assertTrue( + error.getMessage().contains("max precision in Calcite type system is set to 3"), + error.getMessage()); + } + + @ParameterizedTest + @CsvSource({ + "local_timestamp, TIMESTAMP(6) WITH LOCAL TIME ZONE, pts", + "assume_timezone, TIMESTAMP(6), ptstz" + }) + void dynamicTimezoneFunctionUsesDeclaredType(String function, String inputType, String result) + throws Exception { + ConverterProvider provider = + new AutomaticDynamicFunctionMappingConverterProvider(ConverterProvider.builder()); + Plan plan = + new SqlToSubstrait(provider) + .convert( + "SELECT " + function + "(ts, CAST('Asia/Tokyo' AS VARCHAR)) FROM events", + SubstraitCreateStatementParser.processCreateStatementsToCatalog( + "CREATE TABLE events (ts " + inputType + ")")); + Project project = assertInstanceOf(Project.class, plan.getRoots().get(0).getInput()); + Expression.ScalarFunctionInvocation call = + assertInstanceOf( + Expression.ScalarFunctionInvocation.class, project.getExpressions().get(0)); + assertEquals( + result.equals("pts") ? N.precisionTimestamp(6) : N.precisionTimestampTZ(6), + call.outputType()); + new PlanTestBase(provider) {}.assertFullRoundTrip(project); + } +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java b/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java index 50ad9c711..f31ee54b8 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java @@ -95,8 +95,20 @@ void subtractDateIDay(Type outputType, int intervalPrecision, Type expectedInfer RelDataType inferredType = observation.inferredType().orElseThrow(); assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, expectedInferredType), inferredType); + // The way back writes the type the declaration derives, which is the inferred one here, and + // casts to the supplied type where the two differ. Expression reverse = calciteExpr.accept(rexExpressionConverter); - assertEquals(expr, reverse); + Expression.ScalarFunctionInvocation declared = + Expression.ScalarFunctionInvocation.builder() + .from(expr) + .outputType(expectedInferredType) + .build(); + assertEquals( + outputType.equals(expectedInferredType) + ? declared + : ExpressionCreator.cast( + outputType, declared, Expression.FailureBehavior.THROW_EXCEPTION), + reverse); } static Stream subtractDateIDayTypes() { diff --git a/isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java b/isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java index 1535f770e..c4be9290a 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java +++ b/isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java @@ -17,6 +17,7 @@ import io.substrait.plan.Plan.Root; import io.substrait.plan.PlanProtoConverter; import io.substrait.plan.ProtoPlanConverter; +import io.substrait.relation.Project; import io.substrait.relation.ProtoRelConverter; import io.substrait.relation.Rel; import io.substrait.relation.RelProtoConverter; @@ -361,6 +362,15 @@ protected Plan toSubstraitPlan(String sql, CalciteCatalogReader catalog) return new SqlToSubstrait().convert(sql, catalog); } + protected Expression firstExpression(String query, String createStatements) + throws SqlParseException { + Plan plan = + toSubstraitPlan( + query, + SubstraitCreateStatementParser.processCreateStatementsToCatalog(createStatements)); + return ((Project) plan.getRoots().get(0).getInput()).getExpressions().get(0); + } + protected String toSql(io.substrait.proto.Plan protoPlan) { Plan plan = new ProtoPlanConverter(extensions).from(protoPlan); return toSql(plan); diff --git a/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeAdditionTest.java b/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeAdditionTest.java index ba1527093..4eb6d6d07 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeAdditionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeAdditionTest.java @@ -1,5 +1,9 @@ package io.substrait.isthmus; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; + +import io.substrait.expression.Expression; import org.junit.jupiter.api.Test; /** @@ -35,6 +39,23 @@ void dateAddIntervalYearToMonth() throws Exception { assertFullRoundTrip(query, CREATES); } + /** + * The spec declares add(date, interval_day<P>) -> precision_timestamp<P>, where + * Calcite types the sum a DATE. The call carries the declared type and is cast back. + */ + @Test + void dateAddIntervalDayCarriesTheDeclaredType() throws Exception { + Expression.Cast cast = + assertInstanceOf( + Expression.Cast.class, + firstExpression("SELECT event_date + INTERVAL '5' DAY FROM events", CREATES)); + assertEquals(N.DATE, cast.getType()); + Expression.ScalarFunctionInvocation addition = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, cast.input()); + assertEquals("add:date_iday", addition.declaration().key()); + assertEquals(N.precisionTimestamp(6), addition.outputType()); + } + @Test void dateAddIntervalDay() throws Exception { String query = "SELECT event_date + INTERVAL '5' DAY FROM events"; diff --git a/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeSubtractionTest.java b/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeSubtractionTest.java index 613625e8f..9fe61b6e8 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeSubtractionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/PrecisionTimestampDatetimeSubtractionTest.java @@ -1,6 +1,8 @@ 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.assertTrue; import io.substrait.expression.Expression; @@ -51,6 +53,44 @@ void dateSubtractIntervalDay() throws Exception { assertFullRoundTrip(query, CREATES); } + /** + * The spec declares subtract(date, interval_day<P>) -> precision_timestamp<P>, + * where Calcite types the difference a DATE. The call carries the declared type and is cast back. + */ + @Test + void dateSubtractIntervalDayCarriesTheDeclaredType() throws Exception { + Expression.Cast cast = + assertInstanceOf( + Expression.Cast.class, + firstExpression("SELECT event_date - INTERVAL '5' DAY FROM events", CREATES)); + assertEquals(N.DATE, cast.getType()); + Expression.ScalarFunctionInvocation subtraction = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, cast.input()); + assertEquals("subtract:date_iday", subtraction.declaration().key()); + assertEquals(N.precisionTimestamp(6), subtraction.outputType()); + } + + /** + * A TIMESTAMP(3) and an interval_day<6> would bind P to two values. Widening either operand + * is lossless, so the timestamp takes the interval's precision, and the result is cast back to + * Calcite's TIMESTAMP(3). + */ + @Test + void precisionTimestampSubtractIntervalDayBindsOnePrecision() throws Exception { + Expression.Cast cast = + assertInstanceOf( + Expression.Cast.class, + firstExpression("SELECT event_timestamp - INTERVAL '5' DAY FROM events", CREATES)); + assertEquals(N.precisionTimestamp(3), cast.getType()); + Expression.ScalarFunctionInvocation subtraction = + assertInstanceOf(Expression.ScalarFunctionInvocation.class, cast.input()); + assertEquals("subtract:pts_iday", subtraction.declaration().key()); + assertEquals(N.precisionTimestamp(6), subtraction.outputType()); + Expression.Cast widened = + assertInstanceOf(Expression.Cast.class, subtraction.arguments().get(0)); + assertEquals(N.precisionTimestamp(6), widened.getType()); + } + @Test void dateSubtractIntervalUnparsesAsSqlArithmetic() throws Exception { String query = "SELECT event_date - INTERVAL '5' DAY FROM events";