add unpacking to expression converter

prestodb · Sep 24, 2024 · d3e9dbe · d3e9dbe
1 parent 0429f72
commit d3e9dbe
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 4 deletions.
diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/ExpressionConverter.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/ExpressionConverter.java
@@ -46,6 +46,7 @@
 import static com.facebook.presto.common.predicate.Marker.Bound.ABOVE;
 import static com.facebook.presto.common.predicate.Marker.Bound.BELOW;
 import static com.facebook.presto.common.predicate.Marker.Bound.EXACTLY;
+import static com.facebook.presto.common.type.DateTimeEncoding.unpackMillisUtc;
 import static com.facebook.presto.iceberg.IcebergColumnHandle.getPushedDownSubfield;
 import static com.facebook.presto.iceberg.IcebergColumnHandle.isPushedDownSubfield;
 import static com.facebook.presto.parquet.ParquetTypeUtils.columnPathFromSubfield;
@@ -198,10 +199,14 @@ private static Object getIcebergLiteralValue(Type type, Marker marker)
             return toIntExact(((Long) marker.getValue()));
         }
 
-        if (type instanceof TimestampType || type instanceof TimeType || type instanceof TimestampWithTimeZoneType) {
+        if (type instanceof TimestampType || type instanceof TimeType) {
             return MILLISECONDS.toMicros((Long) marker.getValue());
         }
 
+        if (type instanceof TimestampWithTimeZoneType) {
+            return MILLISECONDS.toMicros(unpackMillisUtc((Long) marker.getValue()));
+        }
+
         if (type instanceof VarcharType) {
             return ((Slice) marker.getValue()).toStringUtf8();
         }

diff --git a/...uet/src/main/java/com/facebook/presto/parquet/reader/LongTimestampMicrosColumnReader.java b/...uet/src/main/java/com/facebook/presto/parquet/reader/LongTimestampMicrosColumnReader.java
@@ -14,9 +14,12 @@
 package com.facebook.presto.parquet.reader;
 
 import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.type.TimestampWithTimeZoneType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.parquet.RichColumnDescriptor;
 
+import static com.facebook.presto.common.type.DateTimeEncoding.packDateTimeWithZone;
+import static com.facebook.presto.common.type.TimeZoneKey.UTC_KEY;
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
 public class LongTimestampMicrosColumnReader
@@ -33,7 +36,12 @@ protected void readValue(BlockBuilder blockBuilder, Type type)
         if (definitionLevel == columnDescriptor.getMaxDefinitionLevel()) {
             long utcMillis = MICROSECONDS.toMillis(valuesReader.readLong());
             // TODO: specialize the class at creation time
-            type.writeLong(blockBuilder, utcMillis);
+            if (type instanceof TimestampWithTimeZoneType) {
+                type.writeLong(blockBuilder, packDateTimeWithZone(utcMillis, UTC_KEY));
+            }
+            else {
+                type.writeLong(blockBuilder, utcMillis);
+            }
         }
         else if (isValueNull()) {
             blockBuilder.appendNull();

diff --git a/...et/src/main/java/com/facebook/presto/parquet/writer/valuewriter/TimestampValueWriter.java b/...et/src/main/java/com/facebook/presto/parquet/writer/valuewriter/TimestampValueWriter.java
@@ -41,7 +41,6 @@ public void write(Block block)
         for (int i = 0; i < block.getPositionCount(); i++) {
             if (!block.isNull(i)) {
                 long value = type.getLong(block, i);
-                // Convert to proper encoding here
                 long scaledValue = writeMicroseconds ? MILLISECONDS.toMicros(value) : value;
                 getValueWriter().writeLong(scaledValue);
                 getStatistics().updateStats(scaledValue);

diff --git a/...java/com/facebook/presto/parquet/writer/valuewriter/TimestampWithTimezoneValueWriter.java b/...java/com/facebook/presto/parquet/writer/valuewriter/TimestampWithTimezoneValueWriter.java
@@ -19,6 +19,7 @@
 import org.apache.parquet.schema.OriginalType;
 import org.apache.parquet.schema.PrimitiveType;
 
+import static com.facebook.presto.common.type.DateTimeEncoding.unpackMillisUtc;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
@@ -40,7 +41,7 @@ public void write(Block block)
     {
         for (int i = 0; i < block.getPositionCount(); i++) {
             if (!block.isNull(i)) {
-                long value = type.getLong(block, i);
+                long value = unpackMillisUtc(type.getLong(block, i));
                 long scaledValue = writeMicroseconds ? MILLISECONDS.toMicros(value) : value;
                 getValueWriter().writeLong(scaledValue);
                 getStatistics().updateStats(scaledValue);