feat: implement fractional second intervals (#167)

* refactor: factorize SqlIntervalQualifier objects
substrait-io · Aug 18, 2023 · 68aa7c4 · 68aa7c4
1 parent 1a4b7b1
commit 68aa7c4
Show file tree

Hide file tree

Showing 10 changed files with 91 additions and 42 deletions.
diff --git a/core/src/main/java/io/substrait/expression/Expression.java b/core/src/main/java/io/substrait/expression/Expression.java
@@ -297,6 +297,8 @@ abstract static class IntervalDayLiteral implements Literal {
 
     public abstract int seconds();
 
+    public abstract int microseconds();
+
     public Type getType() {
       return Type.withNullability(nullable()).INTERVAL_DAY;
     }

diff --git a/core/src/main/java/io/substrait/expression/ExpressionCreator.java b/core/src/main/java/io/substrait/expression/ExpressionCreator.java
@@ -120,10 +120,16 @@ public static Expression.IntervalYearLiteral intervalYear(
   }
 
   public static Expression.IntervalDayLiteral intervalDay(boolean nullable, int days, int seconds) {
+    return intervalDay(nullable, days, seconds, 0);
+  }
+
+  public static Expression.IntervalDayLiteral intervalDay(
+      boolean nullable, int days, int seconds, int microseconds) {
     return Expression.IntervalDayLiteral.builder()
         .nullable(nullable)
         .days(days)
         .seconds(seconds)
+        .microseconds(microseconds)
         .build();
   }
 

diff --git a/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java b/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java
@@ -127,7 +127,8 @@ public Expression visit(io.substrait.expression.Expression.IntervalDayLiteral ex
                 .setIntervalDayToSecond(
                     Expression.Literal.IntervalDayToSecond.newBuilder()
                         .setDays(expr.days())
-                        .setSeconds(expr.seconds())));
+                        .setSeconds(expr.seconds())
+                        .setMicroseconds(expr.microseconds())));
   }
 
   @Override

diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java
@@ -335,7 +335,8 @@ public Expression.Literal from(io.substrait.proto.Expression.Literal literal) {
       case INTERVAL_DAY_TO_SECOND -> ExpressionCreator.intervalDay(
           literal.getNullable(),
           literal.getIntervalDayToSecond().getDays(),
-          literal.getIntervalDayToSecond().getSeconds());
+          literal.getIntervalDayToSecond().getSeconds(),
+          literal.getIntervalDayToSecond().getMicroseconds());
       case FIXED_CHAR -> ExpressionCreator.fixedChar(literal.getNullable(), literal.getFixedChar());
       case VAR_CHAR -> ExpressionCreator.varChar(
           literal.getNullable(), literal.getVarChar().getValue(), literal.getVarChar().getLength());

diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitTypeSystem.java
@@ -1,9 +1,12 @@
 package io.substrait.isthmus;
 
+import org.apache.calcite.avatica.util.TimeUnit;
 import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
 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.SqlIntervalQualifier;
+import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.sql.type.SqlTypeName;
 
 public class SubstraitTypeSystem extends RelDataTypeSystemImpl {
@@ -42,4 +45,12 @@ public int getMaxNumericPrecision() {
   public static RelDataTypeFactory createTypeFactory() {
     return new JavaTypeFactoryImpl(TYPE_SYSTEM);
   }
+
+  // Interval qualifier from year to month
+  public static final SqlIntervalQualifier YEAR_MONTH_INTERVAL =
+      new SqlIntervalQualifier(TimeUnit.YEAR, TimeUnit.MONTH, SqlParserPos.ZERO);
+
+  // Interval qualifier from day to fractional second at microsecond precision
+  public static final SqlIntervalQualifier DAY_SECOND_INTERVAL =
+      new SqlIntervalQualifier(TimeUnit.DAY, -1, TimeUnit.SECOND, 6, SqlParserPos.ZERO);
 }
diff --git a/isthmus/src/main/java/io/substrait/isthmus/TypeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/TypeConverter.java
@@ -1,5 +1,8 @@
 package io.substrait.isthmus;
 
+import static io.substrait.isthmus.SubstraitTypeSystem.DAY_SECOND_INTERVAL;
+import static io.substrait.isthmus.SubstraitTypeSystem.YEAR_MONTH_INTERVAL;
+
 import io.substrait.function.NullableType;
 import io.substrait.function.TypeExpression;
 import io.substrait.type.NamedStruct;
@@ -9,11 +12,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import javax.annotation.Nullable;
-import org.apache.calcite.avatica.util.TimeUnit;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
-import org.apache.calcite.sql.SqlIntervalQualifier;
-import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.sql.type.MapSqlType;
 import org.apache.calcite.sql.type.SqlTypeName;
 
@@ -42,11 +42,6 @@ public TypeConverter(UserTypeMapper userTypeMapper) {
     this.userTypeMapper = userTypeMapper;
   }
 
-  static final SqlIntervalQualifier INTERVAL_YEAR =
-      new SqlIntervalQualifier(TimeUnit.YEAR, TimeUnit.MONTH, SqlParserPos.ZERO);
-  static final SqlIntervalQualifier INTERVAL_DAY =
-      new SqlIntervalQualifier(TimeUnit.DAY, TimeUnit.SECOND, SqlParserPos.ZERO);
-
   public Type toSubstrait(RelDataType type) {
     return toSubstrait(type, new ArrayList<>());
   }
@@ -248,13 +243,13 @@ public RelDataType visit(Type.Timestamp expr) {
     @Override
     public RelDataType visit(Type.IntervalYear expr) {
       return typeFactory.createTypeWithNullability(
-          typeFactory.createSqlIntervalType(INTERVAL_YEAR), n(expr));
+          typeFactory.createSqlIntervalType(YEAR_MONTH_INTERVAL), n(expr));
     }
 
     @Override
     public RelDataType visit(Type.IntervalDay expr) {
       return typeFactory.createTypeWithNullability(
-          typeFactory.createSqlIntervalType(INTERVAL_DAY), n(expr));
+          typeFactory.createSqlIntervalType(DAY_SECOND_INTERVAL), n(expr));
     }
 
     @Override

diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java
@@ -47,6 +47,14 @@ public class ExpressionRexConverter extends AbstractExpressionVisitor<RexNode, R
           -1,
           SqlParserPos.QUOTED_ZERO);
 
+  private static final SqlIntervalQualifier DAY_SECOND_INTERVAL =
+      new SqlIntervalQualifier(
+          org.apache.calcite.avatica.util.TimeUnit.DAY,
+          -1,
+          org.apache.calcite.avatica.util.TimeUnit.SECOND,
+          3, // Calcite only supports millisecond at the moment
+          SqlParserPos.QUOTED_ZERO);
+
   public ExpressionRexConverter(
       RelDataTypeFactory typeFactory,
       ScalarFunctionConverter scalarFunctionConverter,
@@ -178,6 +186,20 @@ public RexNode visit(Expression.IntervalYearLiteral expr) throws RuntimeExceptio
         new BigDecimal(expr.years() * 12 + expr.months()), YEAR_MONTH_INTERVAL);
   }
 
+  private static final long MICROS_IN_DAY = TimeUnit.DAYS.toMicros(1);
+
+  @Override
+  public RexNode visit(Expression.IntervalDayLiteral expr) throws RuntimeException {
+    return rexBuilder.makeIntervalLiteral(
+        // Current Calcite behavior is to store milliseconds since Epoch
+        // microseconds version: new BigDecimal(expr.days() * MICROS_IN_DAY + expr.seconds() *
+        // 100000L + expr.microseconds()), DAY_SECOND_INTERVAL);
+        new BigDecimal(
+            (expr.days() * MICROS_IN_DAY + expr.seconds() * 1_000_000L + expr.microseconds())
+                / 1000L),
+        DAY_SECOND_INTERVAL);
+  }
+
   @Override
   public RexNode visit(Expression.DecimalLiteral expr) throws RuntimeException {
     byte[] value = expr.value().toByteArray();

diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/LiteralConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/LiteralConverter.java
@@ -160,20 +160,8 @@ public Expression.Literal convert(RexLiteral literal) {
         var months = intervalLength - years * 12;
         yield intervalYear(n, (int) years, (int) months);
       }
-      case INTERVAL_DAY -> {
-        // we need to convert to microseconds.
-        int precision = literal.getType().getPrecision();
-        var intervalLength = literal.getValueAs(BigDecimal.class).longValue();
-        var adjustedLength =
-            precision > 6
-                ? intervalLength / ((int) Math.pow(10, precision - 6))
-                : intervalLength * ((int) Math.pow(10, 6 - precision));
-        var days = adjustedLength / MICROS_IN_DAY;
-        var microseconds = adjustedLength - days * MICROS_IN_DAY;
-        yield intervalDay(n, (int) days, (int) microseconds);
-      }
-
-      case INTERVAL_DAY_HOUR,
+      case INTERVAL_DAY,
+          INTERVAL_DAY_HOUR,
           INTERVAL_DAY_MINUTE,
           INTERVAL_DAY_SECOND,
           INTERVAL_HOUR,
@@ -182,7 +170,18 @@ public Expression.Literal convert(RexLiteral literal) {
           INTERVAL_MINUTE,
           INTERVAL_MINUTE_SECOND,
           INTERVAL_SECOND -> {
-        throw new UnsupportedOperationException("Need to implement IntervalDay");
+        // we need to convert to microseconds.
+        int scale = literal.getType().getScale();
+        var intervalLength = literal.getValueAs(BigDecimal.class).longValue();
+        var adjustedLength =
+            scale > 6
+                ? intervalLength / ((int) Math.pow(10, scale - 6))
+                : intervalLength * ((int) Math.pow(10, 6 - scale));
+        var days = adjustedLength / MICROS_IN_DAY;
+        var totalMicroseconds = adjustedLength - days * MICROS_IN_DAY;
+        var seconds = totalMicroseconds / 1_000_000;
+        var microseconds = totalMicroseconds - 1_000_000 * seconds;
+        yield intervalDay(n, (int) days, (int) seconds, (int) microseconds);
       }
 
       case ROW -> {

diff --git a/isthmus/src/test/java/io/substrait/isthmus/CalciteLiteralTest.java b/isthmus/src/test/java/io/substrait/isthmus/CalciteLiteralTest.java
@@ -2,6 +2,7 @@
 
 import static io.substrait.expression.ExpressionCreator.*;
 import static io.substrait.isthmus.SqlToSubstrait.EXTENSION_COLLECTION;
+import static io.substrait.isthmus.SubstraitTypeSystem.YEAR_MONTH_INTERVAL;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import com.google.common.collect.ImmutableMap;
@@ -154,15 +155,7 @@ void tTimestampTZ() {
   @Test
   void tIntervalYearMonth() {
     BigDecimal bd = new BigDecimal(3 * 12 + 5); // '3-5' year to month
-    RexLiteral intervalYearMonth =
-        rex.makeIntervalLiteral(
-            bd,
-            new SqlIntervalQualifier(
-                org.apache.calcite.avatica.util.TimeUnit.YEAR,
-                -1,
-                org.apache.calcite.avatica.util.TimeUnit.MONTH,
-                -1,
-                SqlParserPos.QUOTED_ZERO));
+    RexLiteral intervalYearMonth = rex.makeIntervalLiteral(bd, YEAR_MONTH_INTERVAL);
     var intervalYearMonthExpr = intervalYear(false, 3, 5);
     bitest(intervalYearMonthExpr, intervalYearMonth);
   }
@@ -194,6 +187,29 @@ void tIntervalYearMonthWithPrecision() {
         convertedRex.getValueAs(BigDecimal.class).longValue());
   }
 
+  @Test
+  void tIntervalMillisecond() {
+    // Calcite stores milliseconds since Epoch, so test only millisecond precision
+    BigDecimal bd =
+        new BigDecimal(
+            TimeUnit.DAYS.toMillis(3)
+                + TimeUnit.HOURS.toMillis(5)
+                + TimeUnit.MINUTES.toMillis(7)
+                + TimeUnit.SECONDS.toMillis(9)
+                + 500); // '3-5:7:9.500' day to second (6)
+    RexLiteral intervalDaySecond =
+        rex.makeIntervalLiteral(
+            bd,
+            new SqlIntervalQualifier(
+                org.apache.calcite.avatica.util.TimeUnit.DAY,
+                -1,
+                org.apache.calcite.avatica.util.TimeUnit.SECOND,
+                3,
+                SqlParserPos.ZERO));
+    var intervalDaySecondExpr = intervalDay(false, 3, 5 * 3600 + 7 * 60 + 9, 500_000);
+    bitest(intervalDaySecondExpr, intervalDaySecond);
+  }
+
   @Test
   void tIntervalYear() {
     BigDecimal bd = new BigDecimal(123 * 12); // '123' year(3)
@@ -244,10 +260,6 @@ void tIntervalMonth() {
         convertedRex.getValueAs(BigDecimal.class).longValue());
   }
 
-  @Disabled("NYI")
-  @Test
-  void tIntervalDay() {}
-
   @Test
   void tFixedChar() {
     test(fixedChar(false, "hello "), c("hello ", SqlTypeName.CHAR));

diff --git a/isthmus/src/test/java/io/substrait/isthmus/CalciteTypeTest.java b/isthmus/src/test/java/io/substrait/isthmus/CalciteTypeTest.java
@@ -92,7 +92,7 @@ void timestamptz(boolean nullable) {
   void intervalYear(boolean nullable) {
     testType(
         Type.withNullability(nullable).INTERVAL_YEAR,
-        type.createSqlIntervalType(TypeConverter.INTERVAL_YEAR),
+        type.createSqlIntervalType(SubstraitTypeSystem.YEAR_MONTH_INTERVAL),
         nullable);
   }
 
@@ -101,7 +101,7 @@ void intervalYear(boolean nullable) {
   void intervalDay(boolean nullable) {
     testType(
         Type.withNullability(nullable).INTERVAL_DAY,
-        type.createSqlIntervalType(TypeConverter.INTERVAL_DAY),
+        type.createSqlIntervalType(SubstraitTypeSystem.DAY_SECOND_INTERVAL),
         nullable);
   }