Add utc_timestamps option to the chunked Parquet writer

python/cudf/cudf/io/parquet.py

@@ -1858,6 +1858,12 @@ class ParquetWriter:
     store_schema : bool, default False
         If ``True``, enable computing and writing arrow schema to Parquet
         file footer's key-value metadata section for faithful round-tripping.
+    utc_timestamps : bool, default True
+        Value of ``isAdjustedToUTC`` in the Parquet logical type of
+        timestamp columns, and of duration columns written as ``TIME``
+        (with ``store_schema=True``, durations other than days are written
+        as unannotated ``INT64`` and are unaffected). The values themselves
+        are not shifted. Set to ``False`` to match ``DataFrame.to_parquet``.
 
     See Also
     --------
@@ -1877,6 +1883,7 @@ def __init__(
         max_dictionary_size: int = 1048576,
         use_dictionary: bool = True,
         store_schema: bool = False,
+        utc_timestamps: bool = True,
     ):
         filepaths_or_buffers = (
             list(filepath_or_buffer)
@@ -1895,6 +1902,7 @@ def __init__(
         self.max_dictionary_size = max_dictionary_size
         self.use_dictionary = use_dictionary
         self.write_arrow_schema = store_schema
+        self.utc_timestamps = utc_timestamps
 
     def write_table(self, table, partitions_info=None) -> None:
         """Writes a single table to the file"""
@@ -2001,6 +2009,7 @@ def _initialize_chunked_state(
             .max_page_size_rows(self.max_page_size_rows)
             .max_dictionary_size(self.max_dictionary_size)
             .write_arrow_schema(self.write_arrow_schema)
+            .utc_timestamps(self.utc_timestamps)
             .build()
         )
         options.set_dictionary_policy(dict_policy)

python/cudf/cudf/tests/input_output/test_parquet.py

@@ -5,6 +5,7 @@
 import decimal
 import glob
 import hashlib
+import json
 import math
 import os
 import pathlib
@@ -1706,6 +1707,42 @@ def test_parquet_writer_int96_timestamps(tmp_path, pdf, gdf):
     assert_eq(expect, got, check_categorical=False, check_dtype=False)
 
 
+@pytest.mark.parametrize("store_schema", [False, True])
+@pytest.mark.parametrize(
+    "kwargs, expected", [({}, True), ({"utc_timestamps": False}, False)]
+)
+def test_parquet_writer_chunked_utc_timestamps(
+    tmp_path, kwargs, expected, store_schema
+):
+    gdf = cudf.DataFrame(
+        {
+            "ts": cudf.Series([0, 1_000_000, None], dtype="datetime64[us]"),
+            "td": cudf.Series([0, 1_000_000, None], dtype="timedelta64[us]"),
+        }
+    )
+    fname = tmp_path / "gdf.parquet"
+
+    with ParquetWriter(fname, store_schema=store_schema, **kwargs) as writer:
+        writer.write_table(gdf)
+        writer.write_table(gdf)
+
+    schema = pq.ParquetFile(fname).schema
+    ts_type = json.loads(schema.column(0).logical_type.to_json())
+    td_type = json.loads(schema.column(1).logical_type.to_json())
+    assert ts_type["isAdjustedToUTC"] is expected
+    if store_schema:
+        # Durations are written as unannotated INT64 with the arrow schema
+        assert td_type["Type"] == "None"
+        arrow_tz = pq.read_schema(fname).field("ts").type.tz
+        assert arrow_tz == ("UTC" if expected else None)
+    else:
+        assert td_type["Type"] == "Time"
+        assert td_type["isAdjustedToUTC"] is expected
+    assert_eq(
+        cudf.read_parquet(fname), cudf.concat([gdf, gdf], ignore_index=True)
+    )
+
+
 def test_multifile_parquet_folder(tmp_path):
     test_pdf1 = pd.DataFrame(
         {

python/pylibcudf/pylibcudf/io/parquet.pxd

@@ -143,6 +143,8 @@ cdef class ChunkedParquetWriterOptionsBuilder:
 
     cpdef ChunkedParquetWriterOptionsBuilder max_dictionary_size(self, size_t val)
 
+    cpdef ChunkedParquetWriterOptionsBuilder utc_timestamps(self, bool enabled)
+
     cpdef ChunkedParquetWriterOptionsBuilder write_arrow_schema(self, bool enabled)
 
     cpdef ChunkedParquetWriterOptions build(self)

python/pylibcudf/pylibcudf/io/parquet.pyi

@@ -202,6 +202,9 @@ class ChunkedParquetWriterOptionsBuilder:
     def max_dictionary_size(
         self, val: int
     ) -> ChunkedParquetWriterOptionsBuilder: ...
+    def utc_timestamps(
+        self, enabled: bool
+    ) -> ChunkedParquetWriterOptionsBuilder: ...
     def write_arrow_schema(
         self, enabled: bool
     ) -> ChunkedParquetWriterOptionsBuilder: ...

python/pylibcudf/pylibcudf/io/parquet.pyx

@@ -1123,6 +1123,30 @@ cdef class ChunkedParquetWriterOptionsBuilder:
         self.c_obj.max_dictionary_size(val)
         return self
 
+    cpdef ChunkedParquetWriterOptionsBuilder utc_timestamps(self, bool enabled):
+        """
+        Set to true if timestamps are to be annotated as UTC
+        (``isAdjustedToUTC=true``).
+
+        The same flag applies to durations written with the Parquet ``TIME``
+        logical type. When the arrow schema is written, it also sets the
+        timestamp timezone in that schema (``"UTC"`` or none); durations other
+        than days are then written as unannotated ``INT64`` and are unaffected.
+        Only the metadata is affected; values are written as-is and are not
+        converted.
+
+        Parameters
+        ----------
+        enabled : bool
+            Boolean value to enable/disable annotating timestamps as UTC.
+
+        Returns
+        -------
+        ChunkedParquetWriterOptionsBuilder
+        """
+        self.c_obj.utc_timestamps(enabled)
+        return self
+
     cpdef ChunkedParquetWriterOptionsBuilder write_arrow_schema(self, bool enabled):
         """
         Set to true if arrow schema is to be written.

python/pylibcudf/tests/io/test_parquet.py

@@ -1545,6 +1545,34 @@ def test_write_parquet(
     assert isinstance(result, memoryview)
 
 
+@pytest.mark.parametrize("utc_timestamps", [True, False])
+def test_chunked_parquet_writer_utc_timestamps(utc_timestamps):
+    pa_table = pa.table(
+        {"ts": pa.array([0, 1_000_000, None], type=pa.timestamp("us"))}
+    )
+    plc_table = plc.Table.from_arrow(pa_table)
+    table_meta = plc.io.types.TableInputMetadata(plc_table)
+    table_meta.column_metadata[0].set_name("ts")
+    buf = io.BytesIO()
+    options = (
+        plc.io.parquet.ChunkedParquetWriterOptions.builder(
+            plc.io.SinkInfo([buf])
+        )
+        .metadata(table_meta)
+        .utc_timestamps(utc_timestamps)
+        .build()
+    )
+    writer = plc.io.parquet.ChunkedParquetWriter.from_options(options)
+    writer.write(plc_table)
+    writer.write(plc_table)
+    writer.close([])
+
+    logical_type = json.loads(
+        pq.ParquetFile(buf).schema.column(0).logical_type.to_json()
+    )
+    assert logical_type["isAdjustedToUTC"] is utc_timestamps
+
+
 # cupy allocates on its own stream, so this cannot honor the injected default
 # stream used by the stream-validation test pass.
 @pytest.mark.uses_custom_stream