[SPARK-58023][PYTHON] Fast paths for converting Arrow string, binary, numeric and boolean columns to Python rows#57104
Open
viirya wants to merge 6 commits into
Open
[SPARK-58023][PYTHON] Fast paths for converting Arrow string, binary, numeric and boolean columns to Python rows#57104viirya wants to merge 6 commits into
viirya wants to merge 6 commits into
Conversation
PyArrow's Array.to_pylist() materializes one Scalar per element; for list-typed columns each row additionally allocates a C++ scalar, a Python Scalar wrapper and a Python Array wrapper before converting elements one by one. This makes Arrow-optimized Python UDF inputs and Spark Connect collect() several times slower on array columns than necessary (apache/arrow#50326; a fix is proposed upstream in apache/arrow#50327 but will only be available in a future PyArrow release). Add ArrowTableToRowsConversion._to_pylist, which converts the flattened child values of a list column in a single pass and slices the resulting Python list per row using the offsets and the validity bitmap, and use it in the Arrow-to-rows conversion paths (Spark Connect collect, Arrow batch UDF inputs, Arrow UDTF inputs). Leaf values are still converted by Arrow's own to_pylist, so results are exactly identical - None stays None and values inside numeric lists stay Python ints, unlike a pandas round trip which coerces them to floats/NaN. NumPy is only used for the offsets and validity bitmap, never for the values. ASV microbenchmark (python/benchmarks/bench_arrow.py, 1M rows): list<string> 769ms -> 507ms (1.5x); list<list<int32>> with nulls 1.86s -> 537ms (3.5x). Peak memory unchanged. The helper can be removed once the minimum supported PyArrow version includes the upstream fix. Co-authored-by: Isaac
Co-authored-by: Isaac
viirya
force-pushed
the
arrow-to-pylist-leaf
branch
from
5c4488f to
af56dd0
Compare
The files were formatted with black 26.3.1 (dev/requirements.txt), but the CI linter checks with ruff format, which disagrees on hunks unrelated to this change. Reformat with ruff 0.14.8 to match CI. Co-authored-by: Isaac
viirya
force-pushed
the
arrow-to-pylist-leaf
branch
from
af56dd0 to
136e63c
Compare
Co-authored-by: Isaac
viirya
force-pushed
the
arrow-to-pylist-leaf
branch
from
136e63c to
0641520
Compare
…<struct> benchmarks Address review: cache the NumPy availability check in a module-level helper instead of re-running try/import on every (recursive) invocation, and extend the ASV benchmark with an array<struct> case plus peakmem variants for the nested cases (1M rows: nested list 862M -> 862M, array<struct> 885M -> 921M, ~4% peak increase from the transient flattened pointer list; leaf objects are shared). Co-authored-by: Isaac
… numeric and boolean columns to Python rows Extend ArrowTableToRowsConversion._to_pylist with leaf fast paths: * string/large_string/binary/large_binary columns use Arrow's object-dtype NumPy conversion, which produces exactly str/bytes and None - no other type can come out of it. * integer/float32/float64/boolean columns without nulls convert via a zero-copy (except bit-packed booleans) NumPy view and ndarray.tolist, which materializes exact Python ints/floats/bools. With nulls, the values are filled with a placeholder first (pc.fill_null) and nulls are restored to None from the validity bitmap, so ints are always materialized from the original int buffer, never through a float representation. Types whose as_py returns non-primitive objects (dates, timestamps, decimals, ...) keep using to_pylist. Since list columns convert their flattened child values through _to_pylist, list-typed columns get the leaf speedup on top of the bulk offsets slicing. ASV microbenchmark (bench_arrow.ArrowLeafColumnToRowsBenchmark, 1M rows): string with nulls 196ms -> 20ms (9.7x); int64 with nulls 99ms -> 28ms (3.6x); float64 without nulls 100ms -> 9ms (11x). End-to-end list<string> conversion (ArrowListColumnToRowsBenchmark, 1M rows) improves from 507ms to 118ms on top of SPARK-58019. Co-authored-by: Isaac
viirya
force-pushed
the
arrow-to-pylist-leaf
branch
from
0641520 to
936960b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
Follow-up of SPARK-58019 (#57099); only the last commit is new — this PR is stacked on #57099 and will be rebased once it merges.
Extend
ArrowTableToRowsConversion._to_pylistwith fast paths for leaf (non-nested) columns:str/bytes/None— no other type can come out of it.ndarray.tolist(), which materializes exact Pythonint/float/bool. With nulls, values are filled with a placeholder first (pc.fill_null) and nulls are restored toNonefrom the validity bitmap, so ints are always materialized from the original int buffer, never through a float representation.Types whose
as_pyreturns non-primitive objects (dates, timestamps, decimals, ...) keep usingto_pylist. Since list columns convert their flattened child values through_to_pylist, list-typed columns get the leaf speedup on top of the SPARK-58019 bulk offsets slicing.Why are the changes needed?
After SPARK-58019, leaf values are still converted one Scalar at a time by PyArrow's
to_pylist()(apache/arrow#50326). ASV microbenchmark (bench_arrow.ArrowLeafColumnToRowsBenchmark, 1M rows, PyArrow 24.0.0):to_pylist()End-to-end
list<string>conversion (ArrowListColumnToRowsBenchmark, 1M rows) improves from 507 ms (SPARK-58019) to 118 ms.Does this PR introduce any user-facing change?
No. Only performance; conversion results are identical (covered by exact-type tests).
How was this patch tested?
Extended
ArrowColumnToPylistTestswith flat string/large_string/binary/large_binary, int8/int64/uint64 (including values beyond 2^62), float32/float64 (NaN/inf), bool, no-null and all-null variants, sliced and chunked views — all compared againstcolumn.to_pylist()with exact element-type assertions — plus non-primitive leaves (date/timestamp/decimal) exercising the fallback, and lists of fast-path leaves. New ASV benchmark classArrowLeafColumnToRowsBenchmark.Was this patch authored or co-authored using generative AI tooling?
Yes. This pull request and its description were written by Isaac (Claude Code).