feat(isthmus): support for PrecisionTimestamp conversions (#262)

* support for PrecisionTimestampTZ conversions
substrait-io · Jun 12, 2024 · e726904 · e726904
1 parent 77e79ad
commit e726904
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 29 deletions.
diff --git a/isthmus/src/main/java/io/substrait/isthmus/TypeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/TypeConverter.java
@@ -89,27 +89,9 @@ private Type toSubstrait(RelDataType type, List<String> names) {
       }
       case SYMBOL -> creator.STRING;
       case DATE -> creator.DATE;
-      case TIME -> {
-        if (type.getPrecision() != 6) {
-          throw new UnsupportedOperationException(
-              "unsupported time precision " + type.getPrecision());
-        }
-        yield creator.TIME;
-      }
-      case TIMESTAMP -> {
-        if (type.getPrecision() != 6) {
-          throw new UnsupportedOperationException(
-              "unsupported timestamp precision " + type.getPrecision());
-        }
-        yield creator.TIMESTAMP;
-      }
-      case TIMESTAMP_WITH_LOCAL_TIME_ZONE -> {
-        if (type.getPrecision() != 6) {
-          throw new UnsupportedOperationException(
-              "unsupported timestamptz precision " + type.getPrecision());
-        }
-        yield creator.TIMESTAMP_TZ;
-      }
+      case TIME -> creator.TIME;
+      case TIMESTAMP -> creator.precisionTimestamp(type.getPrecision());
+      case TIMESTAMP_WITH_LOCAL_TIME_ZONE -> creator.precisionTimestampTZ(type.getPrecision());
       case INTERVAL_YEAR, INTERVAL_YEAR_MONTH, INTERVAL_MONTH -> creator.INTERVAL_YEAR;
       case INTERVAL_DAY,
           INTERVAL_DAY_HOUR,
@@ -241,6 +223,31 @@ public RelDataType visit(Type.Timestamp expr) {
       return t(n(expr), SqlTypeName.TIMESTAMP, 6);
     }
 
+    @Override
+    public RelDataType visit(Type.PrecisionTimestamp expr) {
+      int maxPrecision = typeFactory.getTypeSystem().getMaxPrecision(SqlTypeName.TIMESTAMP);
+      if (expr.precision() > maxPrecision) {
+        throw new UnsupportedOperationException(
+            String.format(
+                "unsupported precision_timestamp precision %s, max precision in Calcite type system is set to %s",
+                expr.precision(), maxPrecision));
+      }
+      return t(n(expr), SqlTypeName.TIMESTAMP, expr.precision());
+    }
+
+    @Override
+    public RelDataType visit(Type.PrecisionTimestampTZ expr) throws RuntimeException {
+      int maxPrecision =
+          typeFactory.getTypeSystem().getMaxPrecision(SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE);
+      if (expr.precision() > maxPrecision) {
+        throw new UnsupportedOperationException(
+            String.format(
+                "unsupported precision_timestamp_tz precision %s, max precision in Calcite type system is set to %s",
+                expr.precision(), maxPrecision));
+      }
+      return t(n(expr), SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE, expr.precision());
+    }
+
     @Override
     public RelDataType visit(Type.IntervalYear expr) {
       return typeFactory.createTypeWithNullability(

diff --git a/isthmus/src/test/java/io/substrait/isthmus/CalciteTypeTest.java b/isthmus/src/test/java/io/substrait/isthmus/CalciteTypeTest.java
@@ -83,18 +83,26 @@ void time(boolean nullable) {
 
   @ParameterizedTest
   @ValueSource(booleans = {true, false})
-  void timestamp(boolean nullable) {
-    testType(Type.withNullability(nullable).TIMESTAMP, SqlTypeName.TIMESTAMP, nullable, 6);
+  void precisionTimeStamp(boolean nullable) {
+    for (int precision : new int[] {0, 3, 6}) {
+      testType(
+          Type.withNullability(nullable).precisionTimestamp(precision),
+          SqlTypeName.TIMESTAMP,
+          nullable,
+          precision);
+    }
   }
 
   @ParameterizedTest
   @ValueSource(booleans = {true, false})
-  void timestamptz(boolean nullable) {
-    testType(
-        Type.withNullability(nullable).TIMESTAMP_TZ,
-        SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE,
-        nullable,
-        6);
+  void precisionTimestamptz(boolean nullable) {
+    for (int precision : new int[] {0, 3, 6}) {
+      testType(
+          Type.withNullability(nullable).precisionTimestampTZ(precision),
+          SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE,
+          nullable,
+          precision);
+    }
   }
 
   @ParameterizedTest