add timezone logic to batch readers, tests passing

Author: auden-woolfson

presto-parquet/src/main/java/com/facebook/presto/parquet/batchreader/Int64TimeAndTimestampMicrosFlatBatchReader.java

@@ -16,6 +16,7 @@
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.LongArrayBlock;
 import com.facebook.presto.common.block.RunLengthEncodedBlock;
+import com.facebook.presto.common.type.TimestampWithTimeZoneType;
 import com.facebook.presto.parquet.ColumnReader;
 import com.facebook.presto.parquet.DataPage;
 import com.facebook.presto.parquet.DictionaryPage;
@@ -163,7 +164,12 @@ private ColumnChunk readWithNull()
             totalNonNullCount += nonNullCount;
 
             if (nonNullCount > 0) {
-                valuesDecoder.readNext(values, startOffset, nonNullCount);
+                if (field.getType() instanceof TimestampWithTimeZoneType) {
+                    valuesDecoder.readNextWithTimezone(values, startOffset, nonNullCount);
+                }
+                else {
+                    valuesDecoder.readNext(values, startOffset, nonNullCount);
+                }
 
                 int valueDestinationIndex = startOffset + chunkSize - 1;
                 int valueSourceIndex = startOffset + nonNullCount - 1;
@@ -211,7 +217,12 @@ private ColumnChunk readWithoutNull()
 
             int chunkSize = Math.min(remainingCountInPage, remainingInBatch);
 
-            valuesDecoder.readNext(values, startOffset, chunkSize);
+            if (field.getType() instanceof TimestampWithTimeZoneType) {
+                valuesDecoder.readNextWithTimezone(values, startOffset, chunkSize);
+            }
+            else {
+                valuesDecoder.readNext(values, startOffset, chunkSize);
+            }
             startOffset += chunkSize;
             remainingInBatch -= chunkSize;
             remainingCountInPage -= chunkSize;

presto-parquet/src/main/java/com/facebook/presto/parquet/batchreader/Int64TimeAndTimestampMicrosNestedBatchReader.java

@@ -16,6 +16,7 @@
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.LongArrayBlock;
 import com.facebook.presto.common.block.RunLengthEncodedBlock;
+import com.facebook.presto.common.type.TimestampWithTimeZoneType;
 import com.facebook.presto.parquet.RichColumnDescriptor;
 import com.facebook.presto.parquet.batchreader.decoders.ValuesDecoder.Int64TimeAndTimestampMicrosValuesDecoder;
 import com.facebook.presto.parquet.reader.ColumnChunk;
@@ -64,7 +65,12 @@ protected ColumnChunk readNestedWithNull()
         boolean[] isNull = new boolean[newBatchSize];
         int offset = 0;
         for (ValuesDecoderContext valuesDecoderContext : definitionLevelDecodingContext.getValuesDecoderContexts()) {
-            ((Int64TimeAndTimestampMicrosValuesDecoder) valuesDecoderContext.getValuesDecoder()).readNext(values, offset, valuesDecoderContext.getNonNullCount());
+            if (field.getType() instanceof TimestampWithTimeZoneType) {
+                ((Int64TimeAndTimestampMicrosValuesDecoder) valuesDecoderContext.getValuesDecoder()).readNextWithTimezone(values, offset, valuesDecoderContext.getNonNullCount());
+            }
+            else {
+                ((Int64TimeAndTimestampMicrosValuesDecoder) valuesDecoderContext.getValuesDecoder()).readNext(values, offset, valuesDecoderContext.getNonNullCount());
+            }
 
             int valueDestinationIndex = offset + valuesDecoderContext.getValueCount() - 1;
             int valueSourceIndex = offset + valuesDecoderContext.getNonNullCount() - 1;
@@ -112,7 +118,12 @@ protected ColumnChunk readNestedNoNull()
         long[] values = new long[newBatchSize];
         int offset = 0;
         for (ValuesDecoderContext valuesDecoderContext : definitionLevelDecodingContext.getValuesDecoderContexts()) {
-            ((Int64TimeAndTimestampMicrosValuesDecoder) valuesDecoderContext.getValuesDecoder()).readNext(values, offset, valuesDecoderContext.getNonNullCount());
+            if (field.getType() instanceof TimestampWithTimeZoneType) {
+                ((Int64TimeAndTimestampMicrosValuesDecoder) valuesDecoderContext.getValuesDecoder()).readNextWithTimezone(values, offset, valuesDecoderContext.getNonNullCount());
+            }
+            else {
+                ((Int64TimeAndTimestampMicrosValuesDecoder) valuesDecoderContext.getValuesDecoder()).readNext(values, offset, valuesDecoderContext.getNonNullCount());
+            }
             offset += valuesDecoderContext.getValueCount();
         }