Skip to content

Commit 557faa7

Browse files
a-hirotaclaude
andcommitted
Add utc_timestamps to pylibcudf ChunkedParquetWriterOptionsBuilder
Behavioral change: expose the libcudf `utc_timestamps` builder option on the chunked Parquet writer options builder, matching the existing `ParquetWriterOptionsBuilder.utc_timestamps`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent c87e83d commit 557faa7

4 files changed

Lines changed: 49 additions & 0 deletions

File tree

‎python/pylibcudf/pylibcudf/io/parquet.pxd‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,8 @@ cdef class ChunkedParquetWriterOptionsBuilder:
143143

144144
cpdef ChunkedParquetWriterOptionsBuilder max_dictionary_size(self, size_t val)
145145

146+
cpdef ChunkedParquetWriterOptionsBuilder utc_timestamps(self, bool enabled)
147+
146148
cpdef ChunkedParquetWriterOptionsBuilder write_arrow_schema(self, bool enabled)
147149

148150
cpdef ChunkedParquetWriterOptions build(self)

‎python/pylibcudf/pylibcudf/io/parquet.pyi‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,9 @@ class ChunkedParquetWriterOptionsBuilder:
202202
def max_dictionary_size(
203203
self, val: int
204204
) -> ChunkedParquetWriterOptionsBuilder: ...
205+
def utc_timestamps(
206+
self, enabled: bool
207+
) -> ChunkedParquetWriterOptionsBuilder: ...
205208
def write_arrow_schema(
206209
self, enabled: bool
207210
) -> ChunkedParquetWriterOptionsBuilder: ...

‎python/pylibcudf/pylibcudf/io/parquet.pyx‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1123,6 +1123,22 @@ cdef class ChunkedParquetWriterOptionsBuilder:
11231123
self.c_obj.max_dictionary_size(val)
11241124
return self
11251125

1126+
cpdef ChunkedParquetWriterOptionsBuilder utc_timestamps(self, bool enabled):
1127+
"""
1128+
Set to true if timestamps are to be written as UTC.
1129+
1130+
Parameters
1131+
----------
1132+
enabled : bool
1133+
Boolean value to enable/disable writing of timestamps as UTC.
1134+
1135+
Returns
1136+
-------
1137+
ChunkedParquetWriterOptionsBuilder
1138+
"""
1139+
self.c_obj.utc_timestamps(enabled)
1140+
return self
1141+
11261142
cpdef ChunkedParquetWriterOptionsBuilder write_arrow_schema(self, bool enabled):
11271143
"""
11281144
Set to true if arrow schema is to be written.

‎python/pylibcudf/tests/io/test_parquet.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1545,6 +1545,34 @@ def test_write_parquet(
15451545
assert isinstance(result, memoryview)
15461546

15471547

1548+
@pytest.mark.parametrize("utc_timestamps", [True, False])
1549+
def test_chunked_parquet_writer_utc_timestamps(utc_timestamps):
1550+
pa_table = pa.table(
1551+
{"ts": pa.array([0, 1_000_000, None], type=pa.timestamp("us"))}
1552+
)
1553+
plc_table = plc.Table.from_arrow(pa_table)
1554+
table_meta = plc.io.types.TableInputMetadata(plc_table)
1555+
table_meta.column_metadata[0].set_name("ts")
1556+
buf = io.BytesIO()
1557+
options = (
1558+
plc.io.parquet.ChunkedParquetWriterOptions.builder(
1559+
plc.io.SinkInfo([buf])
1560+
)
1561+
.metadata(table_meta)
1562+
.utc_timestamps(utc_timestamps)
1563+
.build()
1564+
)
1565+
writer = plc.io.parquet.ChunkedParquetWriter.from_options(options)
1566+
writer.write(plc_table)
1567+
writer.write(plc_table)
1568+
writer.close([])
1569+
1570+
logical_type = json.loads(
1571+
pq.ParquetFile(buf).schema.column(0).logical_type.to_json()
1572+
)
1573+
assert logical_type["isAdjustedToUTC"] is utc_timestamps
1574+
1575+
15481576
# cupy allocates on its own stream, so this cannot honor the injected default
15491577
# stream used by the stream-validation test pass.
15501578
@pytest.mark.uses_custom_stream

0 commit comments

Comments
 (0)