GH-50866: [Python][Parquet] Add use_threads to ParquetWriter to encode a row group's columns in parallel

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

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)
 

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:

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)