[python] Group write rows by partition/bucket in Arrow to avoid GIL-bound per-row extraction

[yugan95]

Purpose

Linked issue: #9043

TableWrite.write_arrow_batch groups the input rows by (partition, bucket) in pure Python, converting every row of the RecordBatch into Python objects via per-row .as_py() before any data is written. CPython's GIL lets only one thread run bytecode at a time, so multi-threaded embedders (Daft's native runner, Ray actors sharing a process, Spark/Flink PyArrow UDFs) get no write parallelism — throughput stays flat at ~1 core no matter how many threads write.

Measured on an unaware table partitioned by (ds, batch), 8 threads each with its own TableWrite (independent files, so the ceiling is the GIL, not snapshot-commit contention): speedup 1.16x / 8. This matches a production 16-thread worker at ~1.6 / 16 effective cores.

This is the write-side analogue of the read-side GIL problem fixed in #8406.

Root Cause

RowKeyExtractor._extract_partitions_batch materializes partition values one row at a time:

for row_idx in range(data.num_rows):
    partition_values = tuple(col[row_idx].as_py() for col in partition_columns)

For N rows and P partition columns this is N*P Arrow scalar __getitem__ + .as_py() conversions plus N tuple/dict operations, all under the GIL. The only part of the write that releases the GIL is the final file_io.write_parquet (pyarrow C++). cProfile of one write_arrow_batch + prepare_commit (200k rows): the per-row partition extraction is 72% of the time; the per-row .as_py() generator alone is 48%. The cost is paid even when every row is in the same partition (the common case after where(ds==X).where(batch==Y)): all N rows are materialized before the code discovers there is a single group. A bulk col.to_pylist() does not help — it still allocates N Python objects and stays GIL-bound (~1.0x thread speedup).

Changes

• Add RowKeyExtractor.extract_partition_bucket_groups(data), which returns [(partition, bucket, row_indices)]. Grouping runs in Arrow (group_by([...]).aggregate([("__idx", "list")]), C++/Acero, GIL released), so only the K distinct group keys are materialized into Python, not N rows.
• Buckets are computed once, in row order, via _extract_buckets_batch before grouping, so stateful extractors (dynamic bucket) keep their exact assignment sequence and notify_new_record side effects regardless of which grouping path runs.
• The row index fed to the aggregation is built with np.arange (C speed, releases the GIL); a Python range() there makes pyarrow iterate it element-by-element under the GIL and caps multi-thread scaling.
• row_indices is None when the whole batch is a single group, so write_arrow_batch can pass the original batch through without gathering rows — important for BLOB columns, which would otherwise be copied through take.
• write_arrow_batch keeps the contiguous-group zero-copy slice fast path, but determines contiguity from pa.compute.min_max(row_indices) rather than the first/last positions. Arrow's grouped list aggregation runs multi-threaded and does not guarantee within-group order; using positions could mistake a non-contiguous group for a contiguous one and slice the wrong rows. Distinct indices with max - min + 1 == count are provably contiguous regardless of order.
• A _group_indices_python fallback reproduces the legacy per-row grouping, reusing the already-computed buckets (no double-notify). It is only reached on Arrow's own ArrowNotImplementedError / ArrowInvalid (a partition key type Arrow cannot group); any other exception propagates instead of silently degrading to the GIL-bound path, and the fallback is logged.

Performance

Same workload after the change:

metric	before	after
one write_arrow_batch + prepare_commit (200k rows)	0.276s	~0.033s (~8x)
8-thread speedup (independent writers)	1.16x / 8	4.76x / 8
isolated Arrow group_by aggregation	—	9.3x / 8

cProfile hot spot moves from per-row .as_py() (72%) to write_parquet (pyarrow C++, GIL released). Remaining non-linearity is in prepare_commit (per-file stats via per-column .as_py(), GenericRow, DataFileMeta) and, for non-unaware tables, per-row bucket hashing in _extract_buckets_batch — good follow-ups toward near-linear scaling.

Tests

table_write_test.py:

• test_write_arrow_batch_handles_unsorted_row_indices — a non-contiguous group delivered out of order whose endpoints span exactly len(group); asserts the correct rows are written (guards the min_max contiguity fix; fails against a first/last positional check).
• test_write_arrow_batch_contiguous_group_detected_despite_unsorted — a contiguous group delivered shuffled still takes the zero-copy slice path (take not called).
• test_extract_partition_bucket_groups_multi_partition / _single_group_is_none / _arrow_matches_fallback — interleaved partitions map to correct row indices, a single-partition batch yields row_indices=None, and the Arrow and Python grouping paths agree on membership.
• Existing write_arrow_batch fast-path tests (full-batch reuse, contiguous zero-copy slice, non-contiguous take) updated to the new grouping contract.

API and Format

N/A — extract_partition_bucket_groups is a new internal method on RowKeyExtractor; no public API, wire format, or on-disk layout changes.

Documentation

N/A

[JingsongLi]

I found two correctness/compatibility regressions in the new grouping path. Both inline comments include a concrete trigger and a suggested fix. The focused tests pass on PyArrow 19.0.1, but they do not cover these cases.

[JingsongLi]

Could we preserve the original row order before this take? hash_list with the default threaded group-by can return a group's indices out of order. For contiguous groups the slice above restores input order, but this path writes rows in aggregation order. KeyValueDataWriter._add_system_fields then assigns sequence numbers in that order, so an interleaved group containing repeated primary keys can make an earlier input row receive the highest sequence and win deduplication or partial update. I reproduced this with PyArrow 19.0.1: a group whose last input index was 2,999,988 ended with 1,048,575. Please sort the row-index values before gathering (or otherwise make grouping stable) and assert row order/latest-wins in the regression test.

[yugan95]

Good catch, thanks — confirmed and fixed in d851cc3.

Rather than sorting at the take call site, I made grouping stable at the source: _group_indices_arrow now sorts each group's row indices back to ascending input order (np.sort on the aggregated list, so it stays in C and releases the GIL). Both grouping paths therefore return indices in input order — the per-row fallback already appended them ascending.

With that invariant, write_arrow_batch no longer needs min_max; it derives the contiguous span from the first/last positions directly.

Regression test test_group_indices_arrow_sorts_unordered_aggregation forces the aggregation to report a group out of order ([4,0,2] / [3,1]) and asserts the output comes back [0,2,4] / [1,3], and the existing grouping tests now assert order-sensitively instead of by membership.

[JingsongLi]

Table.group_by alone is not a sufficient capability check. PyArrow 7 has it, but hash_list was added in Arrow 8, while the Python 3.7 dependency range still permits pyarrow>=7,<13. On PyArrow 7 this aggregate lookup raises ArrowKeyError, which is not caught above, so every write_arrow_batch fails instead of using the fallback. Could we probe for hash_list when initializing the capability flag (or explicitly handle only the missing-kernel case) and add a PyArrow 7 regression test?

[yugan95]

You're right, hasattr(pa.Table, "group_by") isn't sufficient — fixed in d851cc3.

The flag now comes from _probe_arrow_group_by(), which actually runs a tiny group_by(...).aggregate([("__idx", "list")]) once at import. That way pyarrow 7 (has group_by, but hash_list only landed in Arrow 8 → ArrowKeyError) is detected as unsupported and falls through to per-row grouping, same as pyarrow < 7 which has no group_by at all. I went with probing the kernel rather than catching ArrowKeyError per-batch so there's no repeated warning/exception on the hot path.

Regression test test_probe_arrow_group_by_false_when_hash_list_missing simulates the missing kernel by making the aggregate raise ArrowKeyError and asserts the probe returns False.

[JingsongLi]

+1