Enable predicate pushdown in MongoDB temporal types

trinodb · Nov 1, 2022 · 6234975 · 6234975
1 parent 7799b04
commit 6234975
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java
@@ -82,18 +82,23 @@
 import static io.trino.spi.HostAddress.fromParts;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc;
 import static io.trino.spi.type.DateType.DATE;
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.SmallintType.SMALLINT;
 import static io.trino.spi.type.TimeType.TIME_MILLIS;
 import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS;
 import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS;
+import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_SECOND;
+import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_MICROSECOND;
 import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND;
 import static io.trino.spi.type.Timestamps.roundDiv;
 import static io.trino.spi.type.TinyintType.TINYINT;
 import static io.trino.spi.type.VarbinaryType.VARBINARY;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static io.trino.spi.type.VarcharType.createUnboundedVarcharType;
+import static java.lang.Math.floorDiv;
+import static java.lang.Math.floorMod;
 import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
 import static java.time.ZoneOffset.UTC;
@@ -547,13 +552,15 @@ private static Optional<Object> translateValue(Object trinoNativeValue, Type typ
         }
 
         if (type == TIMESTAMP_MILLIS) {
-            long millisUtc = (long) trinoNativeValue;
-            Instant instant = Instant.ofEpochMilli(millisUtc);
+            long epochMicros = (long) trinoNativeValue;
+            long epochSecond = floorDiv(epochMicros, MICROSECONDS_PER_SECOND);
+            int nanoFraction = floorMod(epochMicros, MICROSECONDS_PER_SECOND) * NANOSECONDS_PER_MICROSECOND;
+            Instant instant = Instant.ofEpochSecond(epochSecond, nanoFraction);
             return Optional.of(LocalDateTime.ofInstant(instant, UTC));
         }
 
         if (type == TIMESTAMP_TZ_MILLIS) {
-            long millisUtc = (long) trinoNativeValue;
+            long millisUtc = unpackMillisUtc((long) trinoNativeValue);
             Instant instant = Instant.ofEpochMilli(millisUtc);
             return Optional.of(LocalDateTime.ofInstant(instant, UTC));
         }

diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/TypeUtils.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/TypeUtils.java
@@ -24,9 +24,13 @@
 
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.spi.type.DateType.DATE;
 import static io.trino.spi.type.IntegerType.INTEGER;
 import static io.trino.spi.type.SmallintType.SMALLINT;
 import static io.trino.spi.type.StandardTypes.JSON;
+import static io.trino.spi.type.TimeType.TIME_MILLIS;
+import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS;
+import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS;
 import static io.trino.spi.type.TinyintType.TINYINT;
 
 public final class TypeUtils
@@ -36,7 +40,11 @@ public final class TypeUtils
             TINYINT,
             SMALLINT,
             INTEGER,
-            BIGINT);
+            BIGINT,
+            DATE,
+            TIME_MILLIS,
+            TIMESTAMP_MILLIS,
+            TIMESTAMP_TZ_MILLIS);
 
     private TypeUtils() {}
 

diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/BaseMongoConnectorTest.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/BaseMongoConnectorTest.java
@@ -259,6 +259,10 @@ public Object[][] predicatePushdownProvider()
                 {"bigint '4'"},
                 {"'test'"},
                 {"objectid('6216f0c6c432d45190f25e7c')"},
+                {"date '1970-01-01'"},
+                {"time '00:00:00.000'"},
+                {"timestamp '1970-01-01 00:00:00.000'"},
+                {"timestamp '1970-01-01 00:00:00.000 UTC'"},
         };
     }