feat: literal support for precision timestamp types (#283)

Support for PrecisionTimestamp and PrecisionTimestampTZ literals
substrait-io · Jul 19, 2024 · 94996f9 · 94996f9
1 parent 0f8514a
commit 94996f9
Show file tree

Hide file tree

Showing 7 changed files with 136 additions and 2 deletions.
diff --git a/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java b/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java
@@ -74,6 +74,16 @@ public OUTPUT visit(Expression.TimestampTZLiteral expr) throws EXCEPTION {
     return visitFallback(expr);
   }
 
+  @Override
+  public OUTPUT visit(Expression.PrecisionTimestampLiteral expr) throws EXCEPTION {
+    return visitFallback(expr);
+  }
+
+  @Override
+  public OUTPUT visit(Expression.PrecisionTimestampTZLiteral expr) throws EXCEPTION {
+    return visitFallback(expr);
+  }
+
   @Override
   public OUTPUT visit(Expression.IntervalYearLiteral expr) throws EXCEPTION {
     return visitFallback(expr);

diff --git a/core/src/main/java/io/substrait/expression/Expression.java b/core/src/main/java/io/substrait/expression/Expression.java
@@ -270,6 +270,44 @@ public <R, E extends Throwable> R accept(ExpressionVisitor<R, E> visitor) throws
     }
   }
 
+  @Value.Immutable
+  abstract static class PrecisionTimestampLiteral implements Literal {
+    public abstract long value();
+
+    public abstract int precision();
+
+    public Type getType() {
+      return Type.withNullability(nullable()).precisionTimestamp(precision());
+    }
+
+    public static ImmutableExpression.PrecisionTimestampLiteral.Builder builder() {
+      return ImmutableExpression.PrecisionTimestampLiteral.builder();
+    }
+
+    public <R, E extends Throwable> R accept(ExpressionVisitor<R, E> visitor) throws E {
+      return visitor.visit(this);
+    }
+  }
+
+  @Value.Immutable
+  abstract static class PrecisionTimestampTZLiteral implements Literal {
+    public abstract long value();
+
+    public abstract int precision();
+
+    public Type getType() {
+      return Type.withNullability(nullable()).precisionTimestampTZ(precision());
+    }
+
+    public static ImmutableExpression.PrecisionTimestampTZLiteral.Builder builder() {
+      return ImmutableExpression.PrecisionTimestampTZLiteral.builder();
+    }
+
+    public <R, E extends Throwable> R accept(ExpressionVisitor<R, E> visitor) throws E {
+      return visitor.visit(this);
+    }
+  }
+
   @Value.Immutable
   abstract static class IntervalYearLiteral implements Literal {
     public abstract int years();

diff --git a/core/src/main/java/io/substrait/expression/ExpressionCreator.java b/core/src/main/java/io/substrait/expression/ExpressionCreator.java
@@ -76,17 +76,29 @@ public static Expression.TimeLiteral time(boolean nullable, long value) {
     return Expression.TimeLiteral.builder().nullable(nullable).value(value).build();
   }
 
+  /**
+   * @deprecated Timestamp is deprecated in favor of PrecisionTimestamp
+   */
+  @Deprecated
   public static Expression.TimestampLiteral timestamp(boolean nullable, long value) {
     return Expression.TimestampLiteral.builder().nullable(nullable).value(value).build();
   }
 
+  /**
+   * @deprecated Timestamp is deprecated in favor of PrecisionTimestamp
+   */
+  @Deprecated
   public static Expression.TimestampLiteral timestamp(boolean nullable, LocalDateTime value) {
     var epochMicro =
         TimeUnit.SECONDS.toMicros(value.toEpochSecond(ZoneOffset.UTC))
             + TimeUnit.NANOSECONDS.toMicros(value.toLocalTime().getNano());
     return timestamp(nullable, epochMicro);
   }
 
+  /**
+   * @deprecated Timestamp is deprecated in favor of PrecisionTimestamp
+   */
+  @Deprecated
   public static Expression.TimestampLiteral timestamp(
       boolean nullable,
       int year,
@@ -102,17 +114,43 @@ public static Expression.TimestampLiteral timestamp(
             .withNano((int) TimeUnit.MICROSECONDS.toNanos(micros)));
   }
 
+  /**
+   * @deprecated TimestampTZ is deprecated in favor of PrecisionTimestampTZ
+   */
+  @Deprecated
   public static Expression.TimestampTZLiteral timestampTZ(boolean nullable, long value) {
     return Expression.TimestampTZLiteral.builder().nullable(nullable).value(value).build();
   }
 
+  /**
+   * @deprecated TimestampTZ is deprecated in favor of PrecisionTimestampTZ
+   */
+  @Deprecated
   public static Expression.TimestampTZLiteral timestampTZ(boolean nullable, Instant value) {
     var epochMicro =
         TimeUnit.SECONDS.toMicros(value.getEpochSecond())
             + TimeUnit.NANOSECONDS.toMicros(value.getNano());
     return timestampTZ(nullable, epochMicro);
   }
 
+  public static Expression.PrecisionTimestampLiteral precisionTimestamp(
+      boolean nullable, long value, int precision) {
+    return Expression.PrecisionTimestampLiteral.builder()
+        .nullable(nullable)
+        .value(value)
+        .precision(precision)
+        .build();
+  }
+
+  public static Expression.PrecisionTimestampTZLiteral precisionTimestampTZ(
+      boolean nullable, long value, int precision) {
+    return Expression.PrecisionTimestampTZLiteral.builder()
+        .nullable(nullable)
+        .value(value)
+        .precision(precision)
+        .build();
+  }
+
   public static Expression.IntervalYearLiteral intervalYear(
       boolean nullable, int years, int months) {
     return Expression.IntervalYearLiteral.builder()

diff --git a/core/src/main/java/io/substrait/expression/ExpressionVisitor.java b/core/src/main/java/io/substrait/expression/ExpressionVisitor.java
@@ -31,6 +31,10 @@ public interface ExpressionVisitor<R, E extends Throwable> {
 
   R visit(Expression.TimestampTZLiteral expr) throws E;
 
+  R visit(Expression.PrecisionTimestampLiteral expr) throws E;
+
+  R visit(Expression.PrecisionTimestampTZLiteral expr) throws E;
+
   R visit(Expression.IntervalYearLiteral expr) throws E;
 
   R visit(Expression.IntervalDayLiteral expr) throws E;

diff --git a/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java b/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java
@@ -113,6 +113,32 @@ public Expression visit(io.substrait.expression.Expression.TimestampTZLiteral ex
     return lit(bldr -> bldr.setNullable(expr.nullable()).setTimestampTz(expr.value()));
   }
 
+  @Override
+  public Expression visit(io.substrait.expression.Expression.PrecisionTimestampLiteral expr) {
+    return lit(
+        bldr ->
+            bldr.setNullable(expr.nullable())
+                .setPrecisionTimestamp(
+                    Expression.Literal.PrecisionTimestamp.newBuilder()
+                        .setValue(expr.value())
+                        .setPrecision(expr.precision())
+                        .build())
+                .build());
+  }
+
+  @Override
+  public Expression visit(io.substrait.expression.Expression.PrecisionTimestampTZLiteral expr) {
+    return lit(
+        bldr ->
+            bldr.setNullable(expr.nullable())
+                .setPrecisionTimestampTz(
+                    Expression.Literal.PrecisionTimestamp.newBuilder()
+                        .setValue(expr.value())
+                        .setPrecision(expr.precision())
+                        .build())
+                .build());
+  }
+
   @Override
   public Expression visit(io.substrait.expression.Expression.IntervalYearLiteral expr) {
     return lit(

diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java
@@ -324,6 +324,16 @@ public Expression.Literal from(io.substrait.proto.Expression.Literal literal) {
       case STRING -> ExpressionCreator.string(literal.getNullable(), literal.getString());
       case BINARY -> ExpressionCreator.binary(literal.getNullable(), literal.getBinary());
       case TIMESTAMP -> ExpressionCreator.timestamp(literal.getNullable(), literal.getTimestamp());
+      case TIMESTAMP_TZ -> ExpressionCreator.timestampTZ(
+          literal.getNullable(), literal.getTimestampTz());
+      case PRECISION_TIMESTAMP -> ExpressionCreator.precisionTimestamp(
+          literal.getNullable(),
+          literal.getPrecisionTimestamp().getValue(),
+          literal.getPrecisionTimestamp().getPrecision());
+      case PRECISION_TIMESTAMP_TZ -> ExpressionCreator.precisionTimestampTZ(
+          literal.getNullable(),
+          literal.getPrecisionTimestampTz().getValue(),
+          literal.getPrecisionTimestampTz().getPrecision());
       case DATE -> ExpressionCreator.date(literal.getNullable(), literal.getDate());
       case TIME -> ExpressionCreator.time(literal.getNullable(), literal.getTime());
       case INTERVAL_YEAR_TO_MONTH -> ExpressionCreator.intervalYear(
@@ -354,8 +364,6 @@ public Expression.Literal from(io.substrait.proto.Expression.Literal literal) {
           literal.getNullable(),
           literal.getMap().getKeyValuesList().stream()
               .collect(Collectors.toMap(kv -> from(kv.getKey()), kv -> from(kv.getValue()))));
-      case TIMESTAMP_TZ -> ExpressionCreator.timestampTZ(
-          literal.getNullable(), literal.getTimestampTz());
       case UUID -> ExpressionCreator.uuid(literal.getNullable(), literal.getUuid());
       case NULL -> ExpressionCreator.typedNull(protoTypeConverter.from(literal.getNull()));
       case LIST -> ExpressionCreator.list(

diff --git a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java
@@ -98,6 +98,16 @@ public Optional<Expression> visit(Expression.TimestampTZLiteral expr) throws EXC
     return visitLiteral(expr);
   }
 
+  @Override
+  public Optional<Expression> visit(Expression.PrecisionTimestampLiteral expr) throws EXCEPTION {
+    return visitLiteral(expr);
+  }
+
+  @Override
+  public Optional<Expression> visit(Expression.PrecisionTimestampTZLiteral expr) throws EXCEPTION {
+    return visitLiteral(expr);
+  }
+
   @Override
   public Optional<Expression> visit(Expression.IntervalYearLiteral expr) throws EXCEPTION {
     return visitLiteral(expr);