From 729315710b60771e5f11a75fff40766fd74d8890 Mon Sep 17 00:00:00 2001 From: Annie Hu <28297557+anniegracehu@users.noreply.github.com> Date: Fri, 14 Aug 2026 16:25:28 -0700 Subject: [PATCH] GH-50866: [Python][Parquet] Add use_threads to ParquetWriter to encode a row group's columns in parallel Every Parquet write from Python is single-threaded: pyarrow only binds FileWriter::WriteTable, which encodes a row group's columns one after another. The C++ writer can encode them in parallel (ArrowWriterProperties::use_threads), but only on the buffered row group path (NewBufferedRowGroup + WriteRecordBatch). ParquetWriter and write_table get a use_threads argument (default False). With it, write_table opens one buffered row group per row_group_size rows and writes that range's batches into it. Row groups, statistics and the stored schema are the same as WriteTable produces. --- python/pyarrow/_parquet.pyx | 71 ++++++++++++++- python/pyarrow/includes/libparquet.pxd | 8 +- python/pyarrow/parquet/core.py | 10 +++ .../tests/parquet/test_parquet_writer.py | 89 +++++++++++++++++++ 4 files changed, 174 insertions(+), 4 deletions(-) diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 932632a50410..42c675ffffb6 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -2283,7 +2283,8 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties( writer_engine_version=None, use_compliant_nested_type=True, store_schema=True, - write_time_adjusted_to_utc=False) except *: + write_time_adjusted_to_utc=False, + use_threads=False) except *: """Arrow writer properties""" cdef: shared_ptr[ArrowWriterProperties] arrow_properties @@ -2334,6 +2335,10 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties( arrow_props.set_time_adjusted_to_utc(write_time_adjusted_to_utc) + # Only honored on the buffered row group path, see + # ParquetWriter._write_table_parallel_columns. + arrow_props.set_use_threads(use_threads) + arrow_properties = arrow_props.build() return arrow_properties @@ -2372,6 +2377,7 @@ cdef class ParquetWriter(_Weakrefable): unique_ptr[FileWriter] writer shared_ptr[COutputStream] sink bint own_sink + bint use_threads def __cinit__(self, where, Schema schema not None, use_dictionary=None, compression=None, version=None, @@ -2398,13 +2404,18 @@ cdef class ParquetWriter(_Weakrefable): store_decimal_as_integer=False, use_content_defined_chunking=False, write_time_adjusted_to_utc=False, - bloom_filter_options=None): + bloom_filter_options=None, + use_threads=False): cdef: shared_ptr[WriterProperties] properties shared_ptr[ArrowWriterProperties] arrow_properties c_string c_where CMemoryPool* pool + # Same fallback as ParquetReader.set_use_threads: a libarrow built + # without threading writes serially whatever the caller asked for. + self.use_threads = bool(use_threads) and is_threading_enabled() + try: where = _stringify_path(where) except TypeError: @@ -2445,6 +2456,7 @@ cdef class ParquetWriter(_Weakrefable): use_compliant_nested_type=use_compliant_nested_type, store_schema=store_schema, write_time_adjusted_to_utc=write_time_adjusted_to_utc, + use_threads=self.use_threads, ) pool = maybe_unbox_memory_pool(memory_pool) @@ -2471,10 +2483,65 @@ cdef class ParquetWriter(_Weakrefable): else: c_row_group_size = row_group_size + if self.use_threads and ctable.num_rows() > 0: + # FileWriter::WriteTable latches chunk_size to + # max_row_group_length; do the same so both paths produce the + # same row groups. + self._write_table_parallel_columns( + table, min(c_row_group_size, _MAX_ROW_GROUP_SIZE)) + return + with nogil: check_status(self.writer.get() .WriteTable(deref(ctable), c_row_group_size)) + cdef _write_table_parallel_columns(self, Table table, + int64_t row_group_size): + """ + Write `table` as row groups of `row_group_size` rows, encoding the + columns of each row group on the CPU thread pool. + + FileWriter::WriteTable encodes the columns of a row group one after + another. Only the buffered row group path (NewBufferedRowGroup + + WriteRecordBatch) honors ArrowWriterProperties::use_threads, so this + opens one buffered row group per `row_group_size` rows and writes the + table's batches for that range into it. The row group layout is the + same as WriteTable's; the difference is that a whole row group is + held in memory until it is complete, instead of one column chunk. + """ + cdef: + CTable* ctable = table.table + FileWriter* writer = self.writer.get() + shared_ptr[CTable] window + shared_ptr[CRecordBatch] batch + unique_ptr[TableBatchReader] reader + int64_t num_rows = ctable.num_rows() + int64_t offset = 0 + int64_t length + + # WriteTable checks these itself; WriteRecordBatch does not. + if not ctable.schema().get().Equals(deref(writer.schema()), False): + raise ValueError( + "Table schema does not match schema used to create file: \n" + f"table:\n{table.schema!s} vs. \n" + f"file:\n{pyarrow_wrap_schema(writer.schema())!s}") + + with nogil: + check_status(ctable.Validate()) + while offset < num_rows: + length = num_rows - offset + if length > row_group_size: + length = row_group_size + window = ctable.Slice(offset, length) + reader.reset(new TableBatchReader(window)) + check_status(writer.NewBufferedRowGroup()) + while True: + check_status(reader.get().ReadNext(&batch)) + if batch.get() == NULL: + break + check_status(writer.WriteRecordBatch(deref(batch))) + offset += length + def add_key_value_metadata(self, key_value_metadata): cdef: shared_ptr[const CKeyValueMetadata] c_metadata diff --git a/python/pyarrow/includes/libparquet.pxd b/python/pyarrow/includes/libparquet.pxd index df353cc7805f..834249ac959c 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -22,8 +22,8 @@ from pyarrow.includes.libarrow cimport (Type, CChunkedArray, CScalar, CSchema, CStatus, CTable, CMemoryPool, CBuffer, CKeyValueMetadata, CRandomAccessFile, COutputStream, CCacheOptions, - TimeUnit, CRecordBatchReader, - CSecureString) + TimeUnit, CRecordBatch, + CRecordBatchReader, CSecureString) cdef extern from "parquet/api/schema.h" namespace "parquet::schema" nogil: @@ -531,6 +531,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil: Builder* disable_compliant_nested_types() Builder* set_engine_version(ArrowWriterEngineVersion version) Builder* set_time_adjusted_to_utc(c_bool adjusted) + Builder* set_use_threads(c_bool use_threads) shared_ptr[ArrowWriterProperties] build() c_bool support_deprecated_int96_timestamps() @@ -620,8 +621,11 @@ cdef extern from "parquet/arrow/writer.h" namespace "parquet::arrow" nogil: const shared_ptr[WriterProperties]& properties, const shared_ptr[ArrowWriterProperties]& arrow_properties) + shared_ptr[CSchema] schema() const CStatus WriteTable(const CTable& table, int64_t chunk_size) CStatus NewRowGroup() + CStatus NewBufferedRowGroup() + CStatus WriteRecordBatch(const CRecordBatch& batch) CStatus Close() CStatus AddKeyValueMetadata(const shared_ptr[const CKeyValueMetadata]& key_value_metadata) diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index ff880fdcf52c..e915bc871809 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -985,6 +985,14 @@ def _sanitize_table(table, new_schema, flavor): - A boolean, with ``True`` indicating that a Bloom filter should be produced with the above mentioned default values of `ndv=1048576` and `fpp=0.05`. This is equivalent to passing an empty dict. +use_threads : bool, default False + Encode the columns of each row group in parallel on the Arrow CPU + thread pool (see :func:`pyarrow.cpu_count`) instead of one after + another. The row groups written are the same either way. While a row + group is being written, all of its column chunks are held in memory + until it is complete, instead of one column chunk at a time; use a + smaller ``row_group_size`` if that is a concern. Do not enable this in + code that itself runs on the Arrow CPU thread pool: it can deadlock. """ _parquet_writer_example_doc = """\ @@ -2017,6 +2025,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_time_adjusted_to_utc=False, max_rows_per_page=None, bloom_filter_options=None, + use_threads=False, **kwargs): # Implementor's note: when adding keywords here / updating defaults, also # update it in write_to_dataset and _dataset_parquet.pyx ParquetFileWriteOptions @@ -2051,6 +2060,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_time_adjusted_to_utc=write_time_adjusted_to_utc, max_rows_per_page=max_rows_per_page, bloom_filter_options=bloom_filter_options, + use_threads=use_threads, **kwargs) as writer: writer.write_table(table, row_group_size=row_group_size) except Exception: diff --git a/python/pyarrow/tests/parquet/test_parquet_writer.py b/python/pyarrow/tests/parquet/test_parquet_writer.py index a49441f09f45..fa75fc0b498b 100644 --- a/python/pyarrow/tests/parquet/test_parquet_writer.py +++ b/python/pyarrow/tests/parquet/test_parquet_writer.py @@ -544,3 +544,92 @@ def test_writer_props_max_rows_per_page_file_size(tempdir): # A smaller maximum rows parameter should produce a larger file assert file_infos[0].size > file_infos[1].size + + +def _threaded_write_sample_table(): + # Several chunks per row group, plus nested and dictionary columns, so + # the parallel path is exercised with more than one batch per row group + # and with columns that have more than one Parquet leaf. + n = 1000 + pieces = [] + for start in range(0, n, 300): + stop = min(start + 300, n) + pieces.append(pa.record_batch({ + 'i': pa.array(range(start, stop), pa.int64()), + 's': pa.array([str(x) if x % 7 else None for x in range(start, stop)]), + 'l': pa.array([[x, x + 1] if x % 5 else None for x in range(start, stop)], + pa.list_(pa.int32())), + 'st': pa.array([{'a': x, 'b': float(x)} for x in range(start, stop)], + pa.struct([('a', pa.int16()), ('b', pa.float64())])), + 'd': pa.array([str(x % 3) for x in range(start, stop)]).dictionary_encode(), + })) + return pa.Table.from_batches(pieces) + + +def _row_group_lengths(path): + metadata = pq.read_metadata(path) + return [metadata.row_group(i).num_rows for i in range(metadata.num_row_groups)] + + +@pytest.mark.parametrize('row_group_size', [None, 1000, 256, 7]) +def test_parquet_writer_use_threads(tempdir, row_group_size): + table = _threaded_write_sample_table() + serial = tempdir / 'serial.parquet' + threaded = tempdir / 'threaded.parquet' + for path, use_threads in ((serial, False), (threaded, True)): + with pq.ParquetWriter(path, table.schema, use_threads=use_threads) as writer: + writer.write_table(table, row_group_size=row_group_size) + # A second call starts new row groups, as it does today. + writer.write_batch(table.slice(0, 10).to_batches()[0]) + + assert _row_group_lengths(threaded) == _row_group_lengths(serial) + threaded_table = pq.read_table(threaded) + threaded_table.validate(full=True) + assert threaded_table.equals(pq.read_table(serial)) + assert threaded_table.schema.equals(table.schema) + + # Statistics are produced on the parallel path too. + serial_meta = pq.read_metadata(serial) + threaded_meta = pq.read_metadata(threaded) + for rg in range(serial_meta.num_row_groups): + for col in range(serial_meta.num_columns): + expected = serial_meta.row_group(rg).column(col).statistics + actual = threaded_meta.row_group(rg).column(col).statistics + assert (actual is None) == (expected is None) + if expected is not None: + assert actual.equals(expected) + + +def test_parquet_writer_use_threads_row_group_latching(tempdir): + # Both paths cap row_group_size at 64Mi rows and default it to the + # smaller of the table size and 1Mi rows. + default_size = 1024 * 1024 + table = pa.table({'x': pa.array(range(default_size + 100), pa.int32())}) + path = tempdir / 'test.parquet' + pq.write_table(table, path, use_threads=True) + assert _row_group_lengths(path) == [default_size, 100] + pq.write_table(table, path, use_threads=True, row_group_size=64 * 1024 * 1024 * 2) + assert _row_group_lengths(path) == [default_size + 100] + assert pq.read_table(path).equals(table) + + +def test_parquet_writer_use_threads_empty_table(tempdir): + table = pa.table({'x': pa.array([], pa.int32())}) + for use_threads in (False, True): + path = tempdir / f'empty_{use_threads}.parquet' + pq.write_table(table, path, use_threads=use_threads) + assert _row_group_lengths(path) == [0] + assert pq.read_table(path).equals(table) + + +def test_parquet_writer_use_threads_schema_mismatch(tempdir): + schema = pa.schema([('x', pa.int32())]) + other = pa.table({'x': pa.array([1], pa.int64())}) + with pq.ParquetWriter(tempdir / 'test.parquet', schema, use_threads=True) as writer: + with pytest.raises(ValueError, match='schema does not match'): + writer.write_table(other) + # The Cython writer checks the schema itself on the parallel path + # (WriteTable does it in C++ on the serial one). + with pytest.raises((ValueError, pa.ArrowInvalid), + match='schema does not match'): + writer.writer.write_table(other)