feat(isthmus): support month()-style datetime operators (#643)

The Substrait Function for Month style mappings doesn't work as the
calcite only has 2 arguments, but substrait has three.

Isthmus (in either direction) assumes that there is a 1:1 mapping betwen a
substrait function and a calicte function. With the same number of arguments
in each. (although types can be coerced if needeed)

In this case, the 3 arg Substrait function needs to be mapped to a 2 arg
calcite one. In effect this is working by swapping out the function to a different
one and then letting the rest of the logic work as before.

A concern is the increased complexity - and the code as a whole wasn't built with this in mind.

Signed-off-by: MBWhite <whitemat@uk.ibm.com>

Author: mbwhite

core/src/main/java/io/substrait/expression/EnumArg.java

@@ -26,7 +26,10 @@ default <R, C extends VisitationContext, E extends Throwable> R accept(
   }
 
   static EnumArg of(SimpleExtension.EnumArgument enumArg, String option) {
-    assert (enumArg.options().contains(option));
+    if (!enumArg.options().contains(option)) {
+      throw new IllegalArgumentException(
+          String.format("EnumArg value %s not valid for options: %s", option, enumArg.options()));
+    }
     return builder().value(Optional.of(option)).build();
   }
 

examples/isthmus-api/.gitignore

@@ -2,3 +2,4 @@ _apps
 _data
 **/*/bin
 build
+substrait.plan

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

@@ -72,6 +72,12 @@ public class EnumConverter {
     calciteEnumMap.put(
         argAnchor(DefaultExtensionCatalog.FUNCTIONS_STRING, "rtrim:str_str", 0),
         SqlTrimFunction.Flag.class);
+    calciteEnumMap.put(
+        argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_req_date", 0),
+        TimeUnitRange.class);
+    calciteEnumMap.put(
+        argAnchor(DefaultExtensionCatalog.FUNCTIONS_DATETIME, "extract:req_req_date", 1),
+        ExtractIndexing.class);
   }
 
   private static Optional<Enum<?>> constructValue(
@@ -84,6 +90,10 @@ private static Optional<Enum<?>> constructValue(
       return option.get().map(SqlTrimFunction.Flag::valueOf);
     }
 
+    // ExtractIndexing does not need to be converted here. Calcite
+    // doesn't have the concept of the indexing. It's date
+    // functions are all indexed from 1
+
     return Optional.empty();
   }
 

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

@@ -0,0 +1,96 @@
+package io.substrait.isthmus.expression;
+
+import io.substrait.expression.Expression;
+import io.substrait.expression.FunctionArg;
+import io.substrait.extension.SimpleExtension.ScalarFunctionVariant;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import org.apache.calcite.avatica.util.TimeUnitRange;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+
+/**
+ * Custom mapping for the Calcite MONTH/DAY/QUARTER functions.
+ *
+ * <p>These come from Calcite as 2 argument functions (like YEAR) but in Substrait these functions
+ * are 3 arguments; the additional being if this is a 0 or 1 based value. Calcite is a 1 based value
+ * in this case.
+ *
+ * <p>We need to therefore map the MONTH etc functions to a different Substrait function.
+ */
+final class ExtractDateFunctionMapper implements ScalarFunctionMapper {
+
+  private final Map<String, ScalarFunctionVariant> extractFunctions;
+
+  public ExtractDateFunctionMapper(List<ScalarFunctionVariant> functions) {
+
+    Map<String, ScalarFunctionVariant> fns =
+        functions.stream()
+            .filter(f -> "extract".equals(f.name()))
+            .collect(Collectors.toMap(Object::toString, f -> f));
+
+    this.extractFunctions = Collections.unmodifiableMap(fns);
+  }
+
+  @Override
+  public Optional<SubstraitFunctionMapping> toSubstrait(final RexCall call) {
+    if (!SqlStdOperatorTable.EXTRACT.equals(call.op)) {
+      return Optional.empty();
+    }
+
+    if (call.operandCount() < 2) {
+      return Optional.empty();
+    }
+
+    RexNode extractType = call.operands.get(0);
+    if (extractType.getType().getSqlTypeName() != SqlTypeName.SYMBOL) {
+      return Optional.empty();
+    }
+
+    final RexNode dataType = call.operands.get(1);
+    if (!dataType.getType().getSqlTypeName().equals(SqlTypeName.DATE)) {
+      return Optional.empty();
+    }
+
+    TimeUnitRange value = ((RexLiteral) extractType).getValueAs(TimeUnitRange.class);
+
+    switch (value) {
+      case QUARTER:
+      case MONTH:
+      case DAY:
+        {
+          final List<RexNode> newOperands = new LinkedList<>(call.operands);
+          newOperands.add(1, RexBuilder.DEFAULT.makeFlag(ExtractIndexing.ONE));
+
+          final ScalarFunctionVariant substraitFn =
+              this.extractFunctions.get("extract:req_req_date");
+          return Optional.of(
+              new SubstraitFunctionMapping("extract", newOperands, List.of(substraitFn)));
+        }
+      default:
+        return Optional.empty();
+    }
+  }
+
+  @Override
+  public Optional<List<FunctionArg>> getExpressionArguments(
+      final Expression.ScalarFunctionInvocation expression) {
+    String name = expression.declaration().toString();
+
+    if ("extract:req_req_date".equals(name)) {
+      final List<FunctionArg> newArgs = new LinkedList<>(expression.arguments());
+      newArgs.remove(1);
+
+      return Optional.of(newArgs);
+    }
+    return Optional.empty();
+  }
+}

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

@@ -0,0 +1,11 @@
+package io.substrait.isthmus.expression;
+
+/**
+ * Enum to define the INDEXING property on the date functions.
+ *
+ * <p>Controls if the number used for example in months is 0 or 1 based.
+ */
+public enum ExtractIndexing {
+  ONE,
+  ZERO
+}

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

@@ -88,7 +88,7 @@ public Expression.Literal convert(RexLiteral literal) {
         return ExpressionCreator.bool(n, literal.getValueAs(Boolean.class));
       case CHAR:
         {
-          Comparable val = literal.getValue();
+          Comparable<?> val = literal.getValue();
           if (val instanceof NlsString) {
             NlsString nls = (NlsString) val;
             return ExpressionCreator.fixedChar(n, nls.getValue());
@@ -127,11 +127,11 @@ public Expression.Literal convert(RexLiteral literal) {
       case SYMBOL:
         {
           Object value = literal.getValue();
-          // case TimeUnitRange tur -> string(n, tur.name());
           if (value instanceof NlsString) {
             return ExpressionCreator.string(n, ((NlsString) value).getValue());
           } else if (value instanceof Enum) {
             Enum<?> v = (Enum<?>) value;
+
             Optional<Expression.Literal> r =
                 EnumConverter.canConvert(v)
                     ? Optional.of(ExpressionCreator.string(n, v.name()))

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

@@ -24,7 +24,6 @@ public class ScalarFunctionConverter
         Expression,
         ScalarFunctionConverter.WrappedScalarCall>
     implements CallConverter {
-
   /**
    * Function mappers provide a hook point for any custom mapping to Substrait functions and
    * arguments.
@@ -43,7 +42,11 @@ public ScalarFunctionConverter(
       TypeConverter typeConverter) {
     super(functions, additionalSignatures, typeFactory, typeConverter);
 
-    mappers = List.of(new TrimFunctionMapper(functions), new SqrtFunctionMapper(functions));
+    mappers =
+        List.of(
+            new TrimFunctionMapper(functions),
+            new SqrtFunctionMapper(functions),
+            new ExtractDateFunctionMapper(functions));
   }
 
   @Override
@@ -68,6 +71,7 @@ private Optional<SubstraitFunctionMapping> getMappingForCall(final RexCall call)
         .orElse(Optional.empty());
   }
 
+  /** Application of the more complex mappings. */
   private Optional<Expression> mappedConvert(
       SubstraitFunctionMapping mapping,
       RexCall call,
@@ -85,6 +89,7 @@ public Stream<RexNode> getOperands() {
     return attemptMatch(finder, wrapped, topLevelConverter);
   }
 
+  /** Default conversion for functions that have simple 1:1 mappings. */
   private Optional<Expression> defaultConvert(
       RexCall call, Function<RexNode, Expression> topLevelConverter) {
     FunctionFinder finder = signatures.get(call.op);

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

@@ -48,9 +48,8 @@ class FunctionConversionTest extends PlanTestBase {
   @Test
   void subtractDateIDay() {
     // When this function is converted to Calcite, if the Calcite type derivation is used an
-    // java.lang.ArrayIndexOutOfBoundsException is thrown. It is quite likely that this is being
-    // mapped to the wrong
-    // Calcite function.
+    // java.lang.ArrayIndexOutOfBoundsException is thrown. It is quite likely that
+    // this is being mapped to the wrong Calcite function.
     // TODO: https://github.com/substrait-io/substrait-java/issues/377
     Expression.ScalarFunctionInvocation expr =
         b.scalarFn(
@@ -183,6 +182,25 @@ void extractTimeScalarFunction() {
     assertEquals("EXTRACT(FLAG(MINUTE), 00:00:00:TIME(6))", extract.toString());
   }
 
+  @Test
+  void extractDateWithIndexing() {
+    ScalarFunctionInvocation reqReqDateFn =
+        b.scalarFn(
+            DefaultExtensionCatalog.FUNCTIONS_DATETIME,
+            "extract:req_req_date",
+            TypeCreator.REQUIRED.I64,
+            EnumArg.builder().value("MONTH").build(),
+            EnumArg.builder().value("ONE").build(),
+            Expression.DateLiteral.builder().value(0).build());
+
+    RexNode calciteExpr = reqReqDateFn.accept(expressionRexConverter, Context.newContext());
+    assertEquals(SqlKind.EXTRACT, calciteExpr.getKind());
+    assertInstanceOf(RexCall.class, calciteExpr);
+
+    RexCall extract = (RexCall) calciteExpr;
+    assertEquals("EXTRACT(FLAG(MONTH), 1970-01-01)", extract.toString());
+  }
+
   @Test
   void unsupportedExtractTimestampTzWithIndexing() {
     ScalarFunctionInvocation reqReqTstzFn =
@@ -249,22 +267,6 @@ void unsupportedExtractPrecisionTimestampWithIndexing() {
         () -> reqReqPtsFn.accept(expressionRexConverter, Context.newContext()));
   }
 
-  @Test
-  void unsupportedExtractDateWithIndexing() {
-    ScalarFunctionInvocation reqReqDateFn =
-        b.scalarFn(
-            DefaultExtensionCatalog.FUNCTIONS_DATETIME,
-            "extract:req_req_date",
-            TypeCreator.REQUIRED.I64,
-            EnumArg.builder().value("MONTH").build(),
-            EnumArg.builder().value("ONE").build(),
-            Expression.DateLiteral.builder().value(0).build());
-
-    assertThrows(
-        UnsupportedOperationException.class,
-        () -> reqReqDateFn.accept(expressionRexConverter, Context.newContext()));
-  }
-
   @Test
   void concatStringLiteralAndVarchar() throws Exception {
     assertProtoPlanRoundrip("select 'part_'||P_NAME from PART");

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

@@ -212,4 +212,11 @@ void caseWhenTest() throws Exception {
     assertFullRoundTrip(
         "select case when p_size > 100 then 'large' when p_size > 50 then 'medium' else 'small' end from part");
   }
+
+  @Test
+  void dateFunctions() throws Exception {
+    assertSqlSubstraitRelRoundTrip("select month(o_orderdate),year(o_orderdate) from orders");
+    assertSqlSubstraitRelRoundTrip(
+        "select extract(month from o_orderdate),extract(year from o_orderdate),extract(day from o_orderdate) from orders");
+  }
 }