[GH-15810] Allow the user to adjust parquet import timezone

h2oai · Jun 10, 2024 · b86e8fe · b86e8fe
1 parent 18341e3
commit b86e8fe
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 2 deletions.
diff --git a/h2o-parsers/h2o-parquet-parser/src/main/java/water/parser/parquet/ChunkConverter.java b/h2o-parsers/h2o-parquet-parser/src/main/java/water/parser/parquet/ChunkConverter.java
@@ -13,6 +13,8 @@
 
 import java.time.Instant;
 
+import static water.H2O.OptArgs.SYSTEM_PROP_PREFIX;
+
 /**
  * Implementation of Parquet's GroupConverter for H2O's chunks.
  *
@@ -303,25 +305,42 @@ public void addBinary(Binary value) {
   }
 
   private static class TimestampConverter extends PrimitiveConverter {
+    public final static String TIMESTAMP_ADJUSTMENT = SYSTEM_PROP_PREFIX + "parquet.import.timestamp.adjustment";
+    public static final long MILLIS_IN_AN_HOUR = 60 * 60 * 1000;
 
     private final int _colIdx;
     private final WriterDelegate _writer;
 
+    private int timeStampAdjustmentHours = 0;
+
     TimestampConverter(int _colIdx, WriterDelegate _writer) {
       this._colIdx = _colIdx;
       this._writer = _writer;
+      if (System.getProperty(TIMESTAMP_ADJUSTMENT) != null) {
+        String timestampHours = System.getProperty(TIMESTAMP_ADJUSTMENT); 
+        try {
+          timeStampAdjustmentHours = Integer.parseInt(timestampHours);
+        } catch (NumberFormatException e) {
+          throw new IllegalArgumentException("Error parsing " + TIMESTAMP_ADJUSTMENT + ", it should be hours as integer.");
+        }
+      }
     }
 
     @Override
     public void addLong(long value) {
-      _writer.addNumCol(_colIdx, value, 0);
+      _writer.addNumCol(_colIdx, adjustTimestamp(value), 0);
     }
 
     @Override
     public void addBinary(Binary value) {
       final long timestampMillis = ParquetInt96TimestampConverter.getTimestampMillis(value);
 
-      _writer.addNumCol(_colIdx, timestampMillis);
+      _writer.addNumCol(_colIdx, adjustTimestamp(timestampMillis));
+    }
+
+    private long adjustTimestamp(long value) {
+      long adjustment = (long) timeStampAdjustmentHours * MILLIS_IN_AN_HOUR;
+      return value + adjustment;
     }
   }
 

diff --git a/h2o-py/tests/testdir_parser/pyunit_parquet_adjust_timezone.py b/h2o-py/tests/testdir_parser/pyunit_parquet_adjust_timezone.py
@@ -0,0 +1,24 @@
+import sys
+
+sys.path.insert(1, "../../../")
+from tests import pyunit_utils
+from h2o.frame import H2OFrame
+import tempfile
+import h2o
+
+'''
+Adjust timestamp parquet
+'''
+
+def adjust_timestamp_parquet():
+    with tempfile.TemporaryDirectory() as dir:
+        timestamp_df = H2OFrame({"timestamp": '2024-06-06 06:06:06'})
+        h2o.export_file(timestamp_df, path=dir, format="parquet", write_checksum=False)
+        imported_df = h2o.import_file(dir)
+        expected_df = H2OFrame({"timestamp": '2024-06-06 08:06:06'})
+    pyunit_utils.compare_frames(expected_df, imported_df, numElements=1)
+
+if __name__ == "__main__":
+    pyunit_utils.standalone_test(adjust_timestamp_parquet, init_options={"jvm_custom_args": ["-Dsys.ai.h2o.parquet.import.timestamp.adjustment=2"]})
+else:
+    adjust_timestamp_parquet()