From b89e2450bdac4e7f9dbe39a48f7f24b6c4c5a458 Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Fri, 28 Mar 2025 15:53:18 +0100 Subject: [PATCH] fix(isthmus): support more datetime extract variants Signed-off-by: Niels Pardon --- .../io/substrait/dsl/SubstraitBuilder.java | 3 +- .../isthmus/expression/EnumConverter.java | 77 +++++-- .../expression/ExpressionRexConverter.java | 74 +++++-- .../isthmus/expression/FunctionConverter.java | 42 ++-- .../isthmus/ExpressionConvertabilityTest.java | 103 +++++++++ .../isthmus/FunctionConversionTest.java | 197 ++++++++++++++++++ 6 files changed, 442 insertions(+), 54 deletions(-) diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 220f03b07..204fceb1c 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -8,6 +8,7 @@ import io.substrait.expression.Expression.IfThen; import io.substrait.expression.Expression.SwitchClause; import io.substrait.expression.FieldReference; +import io.substrait.expression.FunctionArg; import io.substrait.expression.ImmutableExpression.Cast; import io.substrait.expression.ImmutableExpression.SingleOrList; import io.substrait.expression.ImmutableExpression.Switch; @@ -640,7 +641,7 @@ public Expression.ScalarFunctionInvocation or(Expression... args) { } public Expression.ScalarFunctionInvocation scalarFn( - String namespace, String key, Type outputType, Expression... args) { + String namespace, String key, Type outputType, FunctionArg... args) { var declaration = extensions.getScalarFunction(SimpleExtension.FunctionAnchor.of(namespace, key)); return Expression.ScalarFunctionInvocation.builder() diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/EnumConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/EnumConverter.java index 178ddb140..baedc9621 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/EnumConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/EnumConverter.java @@ -1,10 +1,13 @@ package io.substrait.isthmus.expression; -import com.google.common.collect.BiMap; -import com.google.common.collect.HashBiMap; import io.substrait.expression.EnumArg; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.SimpleExtension; +import io.substrait.extension.SimpleExtension.Argument; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.function.Supplier; import org.apache.calcite.avatica.util.TimeUnitRange; @@ -25,16 +28,34 @@ */ public class EnumConverter { - private static final BiMap, ArgAnchor> calciteEnumMap = HashBiMap.create(); + private static final Map>> calciteEnumMap = new HashMap<>(); static { + // deprecated {@link io.substrait.type.Type.Timestamp} calciteEnumMap.put( - TimeUnitRange.class, - argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_ts", 0)); + argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_ts", 0), + TimeUnitRange.class); + // deprecated {@link io.substrait.type.Type.TimestampTZ} + calciteEnumMap.put( + argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_tstz_str", 0), + TimeUnitRange.class); + + calciteEnumMap.put( + argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_pts", 0), + TimeUnitRange.class); + calciteEnumMap.put( + argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_ptstz_str", 0), + TimeUnitRange.class); + calciteEnumMap.put( + argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_date", 0), + TimeUnitRange.class); + calciteEnumMap.put( + argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_time", 0), + TimeUnitRange.class); } - private static Optional constructValue( - Class cls, Supplier> option) { + private static Optional> constructValue( + Class> cls, Supplier> option) { if (cls.isAssignableFrom(TimeUnitRange.class)) { return option.get().map(TimeUnitRange::valueOf); } else { @@ -44,8 +65,9 @@ private static Optional constructValue( static Optional toRex( RexBuilder rexBuilder, SimpleExtension.Function fnDef, int argIdx, EnumArg e) { - var aAnch = argAnchor(fnDef, argIdx); - var v = Optional.ofNullable(calciteEnumMap.inverse().getOrDefault(aAnch, null)); + ArgAnchor aAnch = argAnchor(fnDef, argIdx); + Optional>> v = + Optional.ofNullable(calciteEnumMap.getOrDefault(aAnch, null)); Supplier> sOptionVal = () -> { @@ -66,11 +88,11 @@ private static Optional findEnumArg( return Optional.empty(); } else { - var args = function.args(); + List args = function.args(); if (args.size() <= enumAnchor.argIdx) { return Optional.empty(); } - var arg = args.get(enumAnchor.argIdx); + Argument arg = args.get(enumAnchor.argIdx); if (arg instanceof SimpleExtension.EnumArgument ea) { return Optional.of(ea); } else { @@ -79,17 +101,15 @@ private static Optional findEnumArg( } } - static Optional fromRex(SimpleExtension.Function function, RexLiteral literal) { + static Optional fromRex( + SimpleExtension.Function function, RexLiteral literal, int argIdx) { return switch (literal.getType().getSqlTypeName()) { case SYMBOL -> { Object v = literal.getValue(); if (!literal.isNull() && (v instanceof Enum)) { - Enum value = (Enum) v; - Optional enumAnchor = - Optional.ofNullable(calciteEnumMap.getOrDefault(value.getClass(), null)); - yield enumAnchor - .flatMap(en -> findEnumArg(function, en)) - .map(ea -> EnumArg.of(ea, value.name())); + Enum value = (Enum) v; + ArgAnchor enumAnchor = argAnchor(function, argIdx); + yield findEnumArg(function, enumAnchor).map(ea -> EnumArg.of(ea, value.name())); } else { yield Optional.empty(); } @@ -98,8 +118,8 @@ static Optional fromRex(SimpleExtension.Function function, RexLiteral l }; } - static boolean canConvert(Enum value) { - return value != null && calciteEnumMap.containsKey(value.getClass()); + static boolean canConvert(Enum value) { + return value != null && calciteEnumMap.containsValue(value.getClass()); } static boolean isEnumValue(RexNode value) { @@ -116,6 +136,23 @@ public ArgAnchor(final SimpleExtension.FunctionAnchor fn, final int argIdx) { this.fn = fn; this.argIdx = argIdx; } + + @Override + public int hashCode() { + return Objects.hash(fn, argIdx); + } + + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof ArgAnchor)) { + return false; + } + ArgAnchor other = (ArgAnchor) obj; + return Objects.equals(fn, other.fn) && argIdx == other.argIdx; + } } private static ArgAnchor argAnchor(String fnNS, String fnSig, int argIdx) { 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 90e244cce..f7dd76f6e 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -5,10 +5,13 @@ import io.substrait.expression.EnumArg; import io.substrait.expression.Expression; import io.substrait.expression.Expression.FailureBehavior; +import io.substrait.expression.Expression.PrecisionTimestampLiteral; +import io.substrait.expression.Expression.PrecisionTimestampTZLiteral; import io.substrait.expression.Expression.ScalarSubquery; import io.substrait.expression.Expression.SetPredicate; import io.substrait.expression.Expression.SingleOrList; import io.substrait.expression.Expression.Switch; +import io.substrait.expression.Expression.TimestampTZLiteral; import io.substrait.expression.FieldReference; import io.substrait.expression.FunctionArg; import io.substrait.expression.WindowBound; @@ -209,20 +212,65 @@ public RexNode visit(Expression.DateLiteral expr) throws RuntimeException { @Override public RexNode visit(Expression.TimestampLiteral expr) throws RuntimeException { - // Expression.TimestampLiteral is microseconds - // Construct a TimeStampString : - // 1. Truncate microseconds to seconds - // 2. Get the fraction seconds in precision of nanoseconds. - // 3. Construct TimeStampString : seconds + fraction_seconds part. - long microSec = expr.value(); - long seconds = TimeUnit.MICROSECONDS.toSeconds(microSec); - int fracSecondsInNano = - (int) (TimeUnit.MICROSECONDS.toNanos(microSec) - TimeUnit.SECONDS.toNanos(seconds)); + return rexBuilder.makeLiteral( + getTimestampString(expr.value()), typeConverter.toCalcite(typeFactory, expr.getType())); + } - TimestampString tsString = - TimestampString.fromMillisSinceEpoch(TimeUnit.SECONDS.toMillis(seconds)) - .withNanos(fracSecondsInNano); - return rexBuilder.makeLiteral(tsString, typeConverter.toCalcite(typeFactory, expr.getType())); + @Override + public RexNode visit(TimestampTZLiteral expr) throws RuntimeException { + return rexBuilder.makeLiteral( + getTimestampString(expr.value()), typeConverter.toCalcite(typeFactory, expr.getType())); + } + + @Override + public RexNode visit(PrecisionTimestampLiteral expr) throws RuntimeException { + return rexBuilder.makeLiteral( + getTimestampString(expr.value(), expr.precision()), + typeConverter.toCalcite(typeFactory, expr.getType())); + } + + @Override + public RexNode visit(PrecisionTimestampTZLiteral expr) throws RuntimeException { + return rexBuilder.makeLiteral( + getTimestampString(expr.value(), expr.precision()), + typeConverter.toCalcite(typeFactory, expr.getType())); + } + + private TimestampString getTimestampString(long microSec) { + return getTimestampString(microSec, 6); + } + + private TimestampString getTimestampString(long value, int precision) { + switch (precision) { + case 0: + return TimestampString.fromMillisSinceEpoch(TimeUnit.SECONDS.toMillis(value)); + case 3: + { + long seconds = TimeUnit.MILLISECONDS.toSeconds(value); + int fracSecondsInNano = + (int) (TimeUnit.MILLISECONDS.toNanos(value) - TimeUnit.SECONDS.toNanos(seconds)); + return TimestampString.fromMillisSinceEpoch(TimeUnit.SECONDS.toMillis(seconds)) + .withNanos(fracSecondsInNano); + } + case 6: + { + long seconds = TimeUnit.MICROSECONDS.toSeconds(value); + int fracSecondsInNano = + (int) (TimeUnit.MICROSECONDS.toNanos(value) - TimeUnit.SECONDS.toNanos(seconds)); + return TimestampString.fromMillisSinceEpoch(TimeUnit.SECONDS.toMillis(seconds)) + .withNanos(fracSecondsInNano); + } + case 9: + { + long seconds = TimeUnit.NANOSECONDS.toSeconds(value); + int fracSecondsInNano = (int) (value - TimeUnit.SECONDS.toNanos(seconds)); + return TimestampString.fromMillisSinceEpoch(TimeUnit.SECONDS.toMillis(seconds)) + .withNanos(fracSecondsInNano); + } + default: + throw new UnsupportedOperationException( + String.format("Cannot handle PrecisionTimestamp with precision %d.", precision)); + } } @Override diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionConverter.java index 8d6ae849c..9e0e496b0 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionConverter.java @@ -29,6 +29,7 @@ import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.IntStream; import java.util.stream.Stream; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; @@ -124,7 +125,7 @@ public Optional getSqlOperatorFromSubstraitFunc(String key, Type ou operator -> resolver.containsKey(operator) && resolver.get(operator).types().contains(outputTypeStr)) - .collect(java.util.stream.Collectors.toList()); + .collect(Collectors.toList()); // only one SqlOperator is possible if (resolvedOperators.size() == 1) { return Optional.of(resolvedOperators.get(0)); @@ -331,7 +332,7 @@ private Stream matchKeys(List rexOperands, List opTypes } return isOption ? List.of("req", "opt") : List.of(opType); }) - .collect(java.util.stream.Collectors.toList()); + .collect(Collectors.toList()); return Utils.crossProduct(argTypeLists) .map(typList -> typList.stream().collect(Collectors.joining("_"))); @@ -346,42 +347,43 @@ public Optional attemptMatch(C call, Function topLevelCo * Once a FunctionVariant is resolved we can map the String Literal * to a EnumArg. */ - var operands = - call.getOperands().map(topLevelConverter).collect(java.util.stream.Collectors.toList()); - var opTypes = - operands.stream().map(Expression::getType).collect(java.util.stream.Collectors.toList()); + List operandsList = call.getOperands().collect(Collectors.toList()); + List operands = + call.getOperands().map(topLevelConverter).collect(Collectors.toList()); + List opTypes = operands.stream().map(Expression::getType).collect(Collectors.toList()); - var outputType = typeConverter.toSubstrait(call.getType()); + Type outputType = typeConverter.toSubstrait(call.getType()); // try to do a direct match - var typeStrings = + List typeStrings = opTypes.stream().map(t -> t.accept(ToTypeString.INSTANCE)).collect(Collectors.toList()); - var possibleKeys = - matchKeys(call.getOperands().collect(java.util.stream.Collectors.toList()), typeStrings); + Stream possibleKeys = + matchKeys(call.getOperands().collect(Collectors.toList()), typeStrings); - var directMatchKey = + Optional directMatchKey = possibleKeys .map(argList -> name + ":" + argList) .filter(k -> directMap.containsKey(k)) .findFirst(); if (directMatchKey.isPresent()) { - var variant = directMap.get(directMatchKey.get()); + F variant = directMap.get(directMatchKey.get()); variant.validateOutputType(operands, outputType); - List funcArgs = - Streams.zip( - call.getOperands(), - operands.stream(), - (r, o) -> { + IntStream.range(0, operandsList.size()) + .mapToObj( + i -> { + RexNode r = operandsList.get(i); + Expression o = operands.get(i); if (EnumConverter.isEnumValue(r)) { - return EnumConverter.fromRex(variant, (RexLiteral) r).orElseGet(() -> null); + return EnumConverter.fromRex(variant, (RexLiteral) r, i) + .orElseGet(() -> null); } else { return o; } }) - .collect(java.util.stream.Collectors.toList()); - var allArgsMapped = funcArgs.stream().filter(e -> e == null).findFirst().isEmpty(); + .collect(Collectors.toList()); + boolean allArgsMapped = funcArgs.stream().filter(e -> e == null).findFirst().isEmpty(); if (allArgsMapped) { return Optional.of(generateBinding(call, variant, funcArgs, outputType)); } else { diff --git a/isthmus/src/test/java/io/substrait/isthmus/ExpressionConvertabilityTest.java b/isthmus/src/test/java/io/substrait/isthmus/ExpressionConvertabilityTest.java index 0e877be3b..1f54c445f 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/ExpressionConvertabilityTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/ExpressionConvertabilityTest.java @@ -3,6 +3,8 @@ import static io.substrait.isthmus.expression.CallConverters.CASE; import static io.substrait.isthmus.expression.CallConverters.CREATE_SEARCH_CONV; 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.dsl.SubstraitBuilder; import io.substrait.expression.Expression; @@ -19,6 +21,7 @@ import java.io.IOException; import java.util.List; import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.parser.SqlParseException; import org.junit.jupiter.api.Test; @@ -130,4 +133,104 @@ void assertExpressionEquality(Expression expected, Expression actual) { assertEquals( expected.accept(expressionProtoConverter), actual.accept(expressionProtoConverter)); } + + @Test + public void supportedPrecisionForPrecisionTimestampLiteral() { + assertPrecisionTimestampLiteral(0); + assertPrecisionTimestampLiteral(3); + assertPrecisionTimestampLiteral(6); + } + + void assertPrecisionTimestampLiteral(int precision) { + RexNode calciteExpr = + Expression.PrecisionTimestampLiteral.builder() + .value(0) + .precision(precision) + .build() + .accept(converter); + assertInstanceOf(RexLiteral.class, calciteExpr); + } + + @Test + public void supportedPrecisionForPrecisionTimestampTZLiteral() { + assertPrecisionTimestampTZLiteral(0); + assertPrecisionTimestampTZLiteral(3); + assertPrecisionTimestampTZLiteral(6); + } + + void assertPrecisionTimestampTZLiteral(int precision) { + RexNode calciteExpr = + Expression.PrecisionTimestampTZLiteral.builder() + .value(0) + .precision(precision) + .build() + .accept(converter); + assertInstanceOf(RexLiteral.class, calciteExpr); + } + + @Test + public void unsupportedPrecisionForPrecisionTimestampLiteral() { + // test different edge case precision values + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(-1); + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(1); + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(2); + + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(4); + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(5); + + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(7); + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(8); + + // this would be nanoseconds which are supported in Substrait but not in Calcite + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(9); + + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(10); + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(11); + + // this would be picoseconds which are supported in Substrait but not in Calcite + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(12); + + // everything above 12 is neither supported in Substrait nor in Calcite + assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(13); + } + + void assertThrowsUnsupportedPrecisionPrecisionTimestampLiteral(int precision) { + assertThrowsExpressionLiteral( + Expression.PrecisionTimestampLiteral.builder().value(0).precision(precision).build()); + } + + @Test + public void unsupportedPrecisionPrecisionTimestampTZLiteral() { + // test different edge case precision values + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(-1); + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(1); + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(2); + + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(4); + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(5); + + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(7); + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(8); + + // this would be nanoseconds which are supported in Substrait but not in Calcite + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(9); + + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(10); + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(11); + + // this would be picoseconds which are supported in Substrait but not in Calcite + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(12); + + // everything above 12 is neither supported in Substrait nor in Calcite + assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(13); + } + + void assertThrowsUnsupportedPrecisionPrecisionTimestampTZLiteral(int precision) { + assertThrowsExpressionLiteral( + Expression.PrecisionTimestampTZLiteral.builder().value(0).precision(precision).build()); + } + + void assertThrowsExpressionLiteral(Expression.Literal expr) { + assertThrows(UnsupportedOperationException.class, () -> expr.accept(converter)); + } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java b/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java index 4673a48ed..4c06d91b1 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/FunctionConversionTest.java @@ -1,11 +1,14 @@ 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.dsl.SubstraitBuilder; import io.substrait.expression.Expression; +import io.substrait.expression.Expression.ScalarFunctionInvocation; import io.substrait.expression.ExpressionCreator; +import io.substrait.expression.ImmutableEnumArg; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.isthmus.expression.CallConverters; import io.substrait.isthmus.expression.ExpressionRexConverter; @@ -13,6 +16,9 @@ import io.substrait.isthmus.expression.ScalarFunctionConverter; import io.substrait.isthmus.expression.WindowFunctionConverter; import io.substrait.type.TypeCreator; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; import org.junit.jupiter.api.Test; /** @@ -61,4 +67,195 @@ public void subtractDateIDay() { // TODO: remove once this can be converted back to Substrait assertThrows(IllegalArgumentException.class, () -> calciteExpr.accept(rexExpressionConverter)); } + + @Test + public void extractTimestampTzScalarFunction() { + ScalarFunctionInvocation reqTstzFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_tstz_str", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + Expression.TimestampTZLiteral.builder().value(0).build(), + Expression.StrLiteral.builder().value("GMT").build()); + + RexNode calciteExpr = reqTstzFn.accept(expressionRexConverter); + assertEquals(SqlKind.EXTRACT, calciteExpr.getKind()); + assertInstanceOf(RexCall.class, calciteExpr); + + RexCall extract = (RexCall) calciteExpr; + assertEquals( + "EXTRACT(FLAG(MONTH), 1970-01-01 00:00:00:TIMESTAMP_WITH_LOCAL_TIME_ZONE(6), 'GMT':VARCHAR)", + extract.toString()); + } + + @Test + public void extractPrecisionTimestampTzScalarFunction() { + ScalarFunctionInvocation reqPtstzFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_ptstz_str", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + Expression.PrecisionTimestampTZLiteral.builder().value(0).precision(6).build(), + Expression.StrLiteral.builder().value("GMT").build()); + + RexNode calciteExpr = reqPtstzFn.accept(expressionRexConverter); + assertEquals(SqlKind.EXTRACT, calciteExpr.getKind()); + assertInstanceOf(RexCall.class, calciteExpr); + + RexCall extract = (RexCall) calciteExpr; + assertEquals( + "EXTRACT(FLAG(MONTH), 1970-01-01 00:00:00:TIMESTAMP_WITH_LOCAL_TIME_ZONE(6), 'GMT':VARCHAR)", + extract.toString()); + } + + @Test + public void extractTimestampScalarFunction() { + ScalarFunctionInvocation reqTsFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_ts", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + Expression.TimestampLiteral.builder().value(0).build()); + + RexNode calciteExpr = reqTsFn.accept(expressionRexConverter); + assertEquals(SqlKind.EXTRACT, calciteExpr.getKind()); + assertInstanceOf(RexCall.class, calciteExpr); + + RexCall extract = (RexCall) calciteExpr; + assertEquals("EXTRACT(FLAG(MONTH), 1970-01-01 00:00:00:TIMESTAMP(6))", extract.toString()); + } + + @Test + public void extractPrecisionTimestampScalarFunction() { + ScalarFunctionInvocation reqPtsFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_pts", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + Expression.PrecisionTimestampLiteral.builder().value(0).precision(6).build()); + + RexNode calciteExpr = reqPtsFn.accept(expressionRexConverter); + assertEquals(SqlKind.EXTRACT, calciteExpr.getKind()); + assertInstanceOf(RexCall.class, calciteExpr); + + RexCall extract = (RexCall) calciteExpr; + assertEquals("EXTRACT(FLAG(MONTH), 1970-01-01 00:00:00:TIMESTAMP(6))", extract.toString()); + } + + @Test + public void extractDateScalarFunction() { + ScalarFunctionInvocation reqDateFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_date", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + Expression.DateLiteral.builder().value(0).build()); + + RexNode calciteExpr = reqDateFn.accept(expressionRexConverter); + assertEquals(SqlKind.EXTRACT, calciteExpr.getKind()); + assertInstanceOf(RexCall.class, calciteExpr); + + RexCall extract = (RexCall) calciteExpr; + assertEquals("EXTRACT(FLAG(MONTH), 1970-01-01)", extract.toString()); + } + + @Test + public void extractTimeScalarFunction() { + ScalarFunctionInvocation reqTimeFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_time", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MINUTE").build(), + Expression.TimeLiteral.builder().value(0).build()); + + RexNode calciteExpr = reqTimeFn.accept(expressionRexConverter); + assertEquals(SqlKind.EXTRACT, calciteExpr.getKind()); + assertInstanceOf(RexCall.class, calciteExpr); + + RexCall extract = (RexCall) calciteExpr; + assertEquals("EXTRACT(FLAG(MINUTE), 00:00:00:TIME(6))", extract.toString()); + } + + @Test + public void unsupportedExtractTimestampTzWithIndexing() { + ScalarFunctionInvocation reqReqTstzFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_req_tstz_str", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + ImmutableEnumArg.builder().value("ONE").build(), + Expression.TimestampTZLiteral.builder().value(0).build(), + Expression.StrLiteral.builder().value("GMT").build()); + + assertThrows( + UnsupportedOperationException.class, () -> reqReqTstzFn.accept(expressionRexConverter)); + } + + @Test + public void unsupportedExtractPrecisionTimestampTzWithIndexing() { + ScalarFunctionInvocation reqReqPtstzFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_req_ptstz_str", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + ImmutableEnumArg.builder().value("ONE").build(), + Expression.PrecisionTimestampTZLiteral.builder().value(0).precision(6).build(), + Expression.StrLiteral.builder().value("GMT").build()); + + assertThrows( + UnsupportedOperationException.class, () -> reqReqPtstzFn.accept(expressionRexConverter)); + } + + @Test + public void unsupportedExtractTimestampWithIndexing() { + ScalarFunctionInvocation reqReqTsFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_req_ts", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + ImmutableEnumArg.builder().value("ONE").build(), + Expression.TimestampLiteral.builder().value(0).build()); + + assertThrows( + UnsupportedOperationException.class, () -> reqReqTsFn.accept(expressionRexConverter)); + } + + @Test + public void unsupportedExtractPrecisionTimestampWithIndexing() { + ScalarFunctionInvocation reqReqPtsFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_req_pts", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + ImmutableEnumArg.builder().value("ONE").build(), + Expression.PrecisionTimestampLiteral.builder().value(0).precision(6).build()); + + assertThrows( + UnsupportedOperationException.class, () -> reqReqPtsFn.accept(expressionRexConverter)); + } + + @Test + public void unsupportedExtractDateWithIndexing() { + ScalarFunctionInvocation reqReqDateFn = + b.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_DATETIME, + "extract:req_req_date", + TypeCreator.REQUIRED.I64, + ImmutableEnumArg.builder().value("MONTH").build(), + ImmutableEnumArg.builder().value("ONE").build(), + Expression.DateLiteral.builder().value(0).build()); + + assertThrows( + UnsupportedOperationException.class, () -> reqReqDateFn.accept(expressionRexConverter)); + } }