[SPARK-58019][PYTHON] Convert Arrow list columns to Python rows in bulk

[viirya]

What changes were proposed in this pull request?

Add ArrowTableToRowsConversion._to_pylist, which converts Arrow list-typed columns to Python values in bulk: the flattened child values are converted with a single recursive to_pylist call, and the resulting Python list is sliced per row using the offsets and the validity bitmap. It is used in the Arrow-to-rows conversion paths (Spark Connect collect(), Arrow batch UDF inputs, Arrow UDTF inputs). Non-list columns, map columns and environments without NumPy fall back to plain column.to_pylist().

Leaf values are still converted by Arrow's own to_pylist, so results are exactly identical to column.to_pylist(): None stays None and values inside numeric lists stay Python ints. NumPy is used only for the offsets (non-null integers) and the validity bitmap (booleans), never for the values, so the type coercion problems of a pandas round trip (list<int32> of [1, None, 3] becoming [1.0, nan, 3.0]) cannot occur.

This is an interim measure for a PyArrow-side inefficiency: Array.to_pylist() materializes one Scalar per element, and for list types each row additionally allocates a C++ scalar, a Python Scalar wrapper and a Python Array wrapper before converting elements one by one (root-caused in apache/arrow#50326, fix proposed in apache/arrow#50327). The helper documents this and can be removed once the minimum supported PyArrow version includes the upstream fix.

Why are the changes needed?

Converting Arrow list columns to Python rows is the hot path of Arrow-optimized Python UDF inputs and Spark Connect collect(). With plain to_pylist() it is several times slower than necessary, which caused a performance regression on array columns when Arrow serialization became the default for regular Python UDFs.

ASV microbenchmark (python/benchmarks/bench_arrow.py::ArrowListColumnToRowsBenchmark, added in this PR; 1M rows, macOS arm64, PyArrow 24.0.0):

case	to_pylist()	this PR	speedup
list<string>	769 ms	507 ms	1.5x
list<list<int32>> with nulls	1.86 s	537 ms	3.5x

Peak memory is unchanged.

Does this PR introduce any user-facing change?

No. Only performance; conversion results are byte-identical (covered by exact-type tests).

How was this patch tested?

New ArrowColumnToPylistTests in python/pyspark/sql/tests/test_conversion.py, comparing _to_pylist against column.to_pylist() with exact element-type assertions across list/large_list/nested/struct/map/fixed-size-list columns, sliced and chunked variants, plus a dedicated test that list<int32> of [1, None, 3] stays ints, and an end-to-end ArrowTableToRowsConversion.convert test with array columns. Full test_conversion.py passes. The new ASV benchmark class parametrizes baseline vs bulk, so the comparison above is reproducible with ./python/asv run --python=same --quick -b 'bench_arrow.ArrowListColumnToRowsBenchmark'.

Was this patch authored or co-authored using generative AI tooling?

Yes. This pull request and its description were written by Isaac (Claude Code).

[HyukjinKwon]

cc @Yicong-Huang FYI

[dbtsai]

Reviewed the full diff. Nice, well-scoped optimization. I traced the correctness-critical _to_pylist path against to_pylist() semantics and it holds up on the edges that matter:

• Sliced arrays: column.values returns the full unsliced child and column.offsets are absolute into it, so reading start = offsets[0], slicing values.slice(start, offsets[-1]-start), and rebasing every per-row slice by -start is correct. Using .values rather than .flatten() is the right call here -- .flatten() would drop null-slot ranges and break the rebasing.
• Null slots: Arrow does not guarantee a null list slot has an empty offset range, and keying None-ness off the validity bitmap (valid[i]) rather than off an empty range handles that correctly.
• Exact types: leaf values still go through Arrow's own to_pylist; NumPy only touches offsets and the validity bitmap, so list<int32> [1, None, 3] stays ints -- the pandas round-trip float/NaN coercion is genuinely avoided.
• Fallbacks: map / fixed_size_list / list_view don't match is_list/is_large_list and fall through to to_pylist(). list_view is the important one (non-monotonic offsets would break the slicing), and it correctly does not take the fast path.

The test matrix -- comparing _to_pylist vs to_pylist() with exact element-type assertions across list/large_list/nested/struct/map/fixed_size, and across sliced and chunked views -- is exactly what's needed for a function that reimplements Arrow semantics.

Approve from my side. Two minor comments inline, both non-blocking.

[dbtsai]

Minor: zero_copy_only=True here is load-bearing and works only because list offset buffers never carry a validity bitmap. That invariant is correct, but non-obvious -- a one-line comment would stop a future reader from "fixing" it to zero_copy_only=False (and would flag the assumption if some new Arrow list variant ever violated it).

[viirya]

Good point — added the comment in eebbc4f. zero_copy_only=True also doubles as an assertion: if a future Arrow list variant ever produced offsets with a validity bitmap, this would fail loudly instead of silently copying.

[dbtsai]

Minor (no change needed): the per-row slicing here is still O(rows) Python (offsets.tolist() + a comprehension slicing flat), so the win is fewer object allocations, not vectorization -- consistent with the benchmark numbers. Fine as the documented interim tradeoff; noting it so the "bulk" framing isn't read as fully vectorized.

[viirya]

Right — row assembly is still an O(rows) Python loop; the win is that each row costs one list slice instead of a C++ scalar + Python Scalar wrapper + Python Array wrapper + generator. "Bulk" refers to the child-values conversion being a single call, not to vectorized row assembly. The fully vectorized fix is the upstream one (apache/arrow#50327); this stays intentionally simple as the interim.

[dongjoon-hyun]

We need skipIf like the other arrow-based test cases.

Suggested change

	@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)

[viirya]

Good catch — added in 784974c, thanks!

[dongjoon-hyun]

+1, LGTM (with one comment about skipIf).

[viirya]

Added one more commit (2e405eb) after discussion: since Spark 4.2 ships with this code frozen, _to_pylist now checks the installed PyArrow version and short-circuits to the native column.to_pylist() when it contains the apache/arrow#50326 fix (planned for PyArrow 25.0.1 per the dev-list discussion). This way Spark 4.2 users automatically pick up the native fast path — and any future PyArrow-side improvements — just by upgrading PyArrow, while older PyArrow keeps the bulk conversion here. The version constant is documented and can be bumped in a patch release if the fix ships in a different version. Tests pin the gate off so the bulk paths stay covered regardless of the CI's PyArrow.

[viirya]

I will go to merge this. @gaogaotiantian Do you want to do a review before that?

[gaogaotiantian]

Please give me some time to review this. Will do this asap.

[huaxingao]

Thanks @viirya for the fix! I will cut RC6 as soon as this is in.

[Yicong-Huang]

change to has_numpy and np to align with other files? for reference: #57163

[viirya]

Renamed the cached flag to has_numpy in 8e24c2f to align with #57163. I didn't cache an np module global though: unlike stateful_processor_api_client.py, this file never references the numpy module directly — numpy is only needed because pyarrow.Array.to_numpy() requires it — so np would have no user here.

[gaogaotiantian]

Overall I think this approach is good and the code is maintainable (also discardable in the future). The ROI of the patch is nice.

I have a few suggestions.

1. Move all the numpy and pyarrow version check into a staticmethod of LocalDataToArrowConversion. Name it something like _should_manual_bulk() (or whatever you think is proper). Do a @functools.cache on the function so the actual import path is only triggered once. Let's leave no garbage in the module level block so it's easier to cleanly remove it in the future.
2. Comment wise, for both method, let's emphasize that this is an internal function that should not be used externally, and it should (not can) be removed when the minumum version contains the patch.

[gaogaotiantian]

I don't think this is really necessary.

[viirya]

It is actually needed: ArrowTableToRowsConversion.convert passes table.columns, and pa.Table.columns are ChunkedArrays (the Spark Connect collect path). Without this branch a chunked list column falls into the list path and fails with AttributeError: 'pyarrow.lib.ChunkedArray' object has no attribute 'offsets'. The worker.py call sites pass RecordBatch columns (plain Arrays), so only the convert path exercises it; the chunked variants in ArrowColumnToPylistTests cover it.

[gaogaotiantian]

I meant assigning column.type to column_type. I think it's also some optimization that tries to save attribute access cost.

[viirya]

Ah, I misread your comment as being about the ChunkedArray block above. Removed in a96c82b — column.type is used inline now.

[viirya]

Thanks @huaxingao! Addressing review comments now — will keep this moving so RC6 isn't blocked.

[viirya]

@gaogaotiantian Both done in 5ce205e:

1. The module-level block is gone — the version and NumPy checks are folded into a single ArrowTableToRowsConversion._should_manual_bulk() staticmethod with @functools.cache (I put it on ArrowTableToRowsConversion since that's the class that owns _to_pylist). Tests patch the method directly instead of poking module globals, so removing the whole thing later is a clean two-site deletion.
2. Both docstrings now state the helpers are internal-only and that they (together with the manual bulk paths) should be removed once the minimum supported PyArrow version contains the fix.

[viirya]

Merge Summary:

• merged into master 484342a
• merged into branch-4.x ad0eebf
• merged into branch-4.2 d0856c7

Posted by merge_spark_pr.py

[viirya]

Thanks @dbtsai @dongjoon-hyun @Yicong-Huang @huaxingao @gaogaotiantian