feat: add support for current_date and current_timestamp

current_timestamp return type is of unsupported precision 0
substrait-io · Jul 17, 2023 · 079b9ad · 079b9ad
1 parent 9df39ed
commit 079b9ad
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java
@@ -1,11 +1,15 @@
 package io.substrait.isthmus;
 
 import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rel.type.RelDataTypeSystem;
 import org.apache.calcite.rel.type.RelDataTypeSystemImpl;
 import org.apache.calcite.sql.type.SqlTypeName;
 
+import java.lang.reflect.Type;
+import java.util.function.Function;
+
 public class SubstraitTypeSystem extends RelDataTypeSystemImpl {
   static final org.slf4j.Logger logger =
       org.slf4j.LoggerFactory.getLogger(SubstraitTypeSystem.class);
@@ -14,8 +18,7 @@ public class SubstraitTypeSystem extends RelDataTypeSystemImpl {
 
   private SubstraitTypeSystem() {}
 
-  @Override
-  public int getMaxPrecision(final SqlTypeName typeName) {
+  static private int fixedTimePrecision(final SqlTypeName typeName, final int otherwise) {
     switch (typeName) {
       case INTERVAL_DAY:
       case INTERVAL_YEAR:
@@ -25,8 +28,14 @@ public int getMaxPrecision(final SqlTypeName typeName) {
       case TIMESTAMP:
       case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
         return 6;
+      default:
+        return otherwise;
     }
-    return super.getMaxPrecision(typeName);
+  }
+
+  @Override
+  public int getMaxPrecision(final SqlTypeName typeName) {
+    return fixedTimePrecision(typeName, super.getMaxPrecision(typeName));
   }
 
   @Override
@@ -40,6 +49,14 @@ public int getMaxNumericPrecision() {
   }
 
   public static RelDataTypeFactory createTypeFactory() {
-    return new JavaTypeFactoryImpl(TYPE_SYSTEM);
+    return new JavaTypeFactoryImpl(TYPE_SYSTEM) {
+      @Override public RelDataType createType(Type type) {
+        return super.createType(type);
+      }
+
+      @Override public RelDataType createSqlType(SqlTypeName typeName, int precision) {
+        return super.createSqlType(typeName, fixedTimePrecision(typeName, precision));
+      }
+    };
   }
-}
+}
diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionMappings.java b/isthmus/src/main/java/io/substrait/isthmus/expression/FunctionMappings.java
@@ -63,7 +63,9 @@ public class FunctionMappings {
                 s(SqlStdOperatorTable.UPPER, "upper"),
                 s(SqlStdOperatorTable.BETWEEN),
                 s(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM, "is_not_distinct_from"),
-                s(SqlStdOperatorTable.COALESCE, "coalesce"))
+                s(SqlStdOperatorTable.COALESCE, "coalesce"),
+                s(SqlStdOperatorTable.CURRENT_DATE, "current_date"),
+                s(SqlStdOperatorTable.CURRENT_TIMESTAMP, "current_timestamp"))
             .build();
 
     AGGREGATE_SIGS =

diff --git a/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java b/isthmus/src/test/java/io/substrait/isthmus/CalciteCallTest.java
@@ -100,6 +100,16 @@ public void not() {
     test("not:bool", rex.makeCall(NOT, c(false, SqlTypeName.BOOLEAN)));
   }
 
+  @Test
+  public void currentDate() throws IOException {
+    test("current_date:", rex.makeCall(CURRENT_DATE));
+  }
+
+  @Test
+  public void currentTimestamp() throws IOException {
+    test("current_timestamp:", rex.makeCall(CURRENT_TIMESTAMP) );
+  }
+
   private void test(String expectedName, RexNode call) {
     test(expectedName, call, c -> {}, true);
   }