[SPARK-58023][PYTHON] Fast paths for converting Arrow string, binary, numeric and boolean columns to Python rows#57104
[SPARK-58023][PYTHON] Fast paths for converting Arrow string, binary, numeric and boolean columns to Python rows#57104viirya wants to merge 7 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
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
af56dd0 to
136e63c
Compare
Co-authored-by: Isaac
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
0641520 to
936960b
Compare
…out per-element Scalars Spark 4.2 ships with this code frozen, so gate the pure-Python bulk conversion on the installed PyArrow version: releases containing the apache/arrow#50326 fix (planned for 25.0.1) convert natively without per-element Scalars and keep improving (apache/arrow#50448), so _to_pylist short-circuits to column.to_pylist() there and only uses the bulk paths on older PyArrow. The version constant can be bumped in a patch release if the fix ships elsewhere. Tests force the gate off so the bulk paths stay covered on any PyArrow, plus a gate pass-through test. 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
936960b to
a947ef0
Compare
|
Closing this PR. Rationale: During review of the sibling PRs, concerns were raised about the maintainability of routing data values through NumPy (behavior across the supported arrow/numpy version matrix, and keeping the "no type coercion" invariant guarded by tests over time). While the conversions here are exactness-tested — the object-dtype conversion for strings/binary can only produce With the version gate added in SPARK-58019 (#57099), Spark automatically uses PyArrow's native conversion once the installed PyArrow contains that fix — so the leaf-level speedup reaches users through a PyArrow upgrade rather than through NumPy-based code maintained here. The remaining Spark-side PRs are scoped to what Arrow cannot do for Spark: SPARK-58019/SPARK-58024 (bulk offsets/validity slicing; no data values through NumPy) and SPARK-58050 (bulk result assembly; no NumPy at all). For the record: this PR gave flat string 196ms→20ms and int64 99ms→28ms per 1M rows, and its removal lowers the struct/map gains in SPARK-58024 from 5.2x/6.8x to 2.0x/2.7x until the PyArrow fix ships. We accept that interim gap in exchange for the smaller maintenance surface. |
…rows in bulk ### 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.** (Originally stacked on SPARK-58023 (#57104), which has been closed in favor of the upstream apache/arrow#50326 fix; the numbers below are measured without it.) Extend `ArrowTableToRowsConversion._to_pylist` with bulk paths for struct and map columns: - **Struct** columns convert each child field in bulk (recursively reusing the bulk list paths), then zip the field values into one dict per row, masked by the validity bitmap. Duplicate field names fall back to `to_pylist`, so they raise the same `ValueError` that `StructScalar.as_py` raises. - **Map** columns share the list offsets layout: the flattened keys and items children are each converted in bulk, and every row becomes a list of `(key, value)` tuples, matching `MapScalar.as_py` exactly. ### Why are the changes needed? Struct and map columns still convert one Scalar per row; per-row Scalar/wrapper allocation dominates (apache/arrow#50326), and maps are the worst case since every row also wraps its keys/items in per-row Arrays. ASV microbenchmark (`bench_arrow.ArrowStructMapColumnToRowsBenchmark`, 1M rows, 10% nulls, PyArrow 24.0.0): | case | `to_pylist()` | this PR | speedup | |---|---|---|---| | `struct<int64,string>` | 970 ms | 474 ms | 2.0x | | `map<string,int64>` (2 entries/row) | 2.16 s | 799 ms | 2.7x | Peak memory unchanged. ### 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 `ArrowColumnToPylistTests` with struct (incl. nested struct/list children, empty struct, all-null), map (incl. list values), struct-of-map, list-of-struct, sliced and chunked views — all compared against `column.to_pylist()` with exact element-type assertions — plus dedicated tests that duplicate struct field names still raise `ValueError` and that empty-struct rows are distinct dict objects. New ASV benchmark class `ArrowStructMapColumnToRowsBenchmark`. ### Was this patch authored or co-authored using generative AI tooling? Yes. This pull request and its description were written by Isaac (Claude Code). Closes #57105 from viirya/arrow-to-pylist-maps-structs. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
…rows in bulk ### 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.** (Originally stacked on SPARK-58023 (#57104), which has been closed in favor of the upstream apache/arrow#50326 fix; the numbers below are measured without it.) Extend `ArrowTableToRowsConversion._to_pylist` with bulk paths for struct and map columns: - **Struct** columns convert each child field in bulk (recursively reusing the bulk list paths), then zip the field values into one dict per row, masked by the validity bitmap. Duplicate field names fall back to `to_pylist`, so they raise the same `ValueError` that `StructScalar.as_py` raises. - **Map** columns share the list offsets layout: the flattened keys and items children are each converted in bulk, and every row becomes a list of `(key, value)` tuples, matching `MapScalar.as_py` exactly. ### Why are the changes needed? Struct and map columns still convert one Scalar per row; per-row Scalar/wrapper allocation dominates (apache/arrow#50326), and maps are the worst case since every row also wraps its keys/items in per-row Arrays. ASV microbenchmark (`bench_arrow.ArrowStructMapColumnToRowsBenchmark`, 1M rows, 10% nulls, PyArrow 24.0.0): | case | `to_pylist()` | this PR | speedup | |---|---|---|---| | `struct<int64,string>` | 970 ms | 474 ms | 2.0x | | `map<string,int64>` (2 entries/row) | 2.16 s | 799 ms | 2.7x | Peak memory unchanged. ### 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 `ArrowColumnToPylistTests` with struct (incl. nested struct/list children, empty struct, all-null), map (incl. list values), struct-of-map, list-of-struct, sliced and chunked views — all compared against `column.to_pylist()` with exact element-type assertions — plus dedicated tests that duplicate struct field names still raise `ValueError` and that empty-struct rows are distinct dict objects. New ASV benchmark class `ArrowStructMapColumnToRowsBenchmark`. ### Was this patch authored or co-authored using generative AI tooling? Yes. This pull request and its description were written by Isaac (Claude Code). Closes #57105 from viirya/arrow-to-pylist-maps-structs. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit 9c6e57b) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
…rows in bulk ### 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.** (Originally stacked on SPARK-58023 (#57104), which has been closed in favor of the upstream apache/arrow#50326 fix; the numbers below are measured without it.) Extend `ArrowTableToRowsConversion._to_pylist` with bulk paths for struct and map columns: - **Struct** columns convert each child field in bulk (recursively reusing the bulk list paths), then zip the field values into one dict per row, masked by the validity bitmap. Duplicate field names fall back to `to_pylist`, so they raise the same `ValueError` that `StructScalar.as_py` raises. - **Map** columns share the list offsets layout: the flattened keys and items children are each converted in bulk, and every row becomes a list of `(key, value)` tuples, matching `MapScalar.as_py` exactly. ### Why are the changes needed? Struct and map columns still convert one Scalar per row; per-row Scalar/wrapper allocation dominates (apache/arrow#50326), and maps are the worst case since every row also wraps its keys/items in per-row Arrays. ASV microbenchmark (`bench_arrow.ArrowStructMapColumnToRowsBenchmark`, 1M rows, 10% nulls, PyArrow 24.0.0): | case | `to_pylist()` | this PR | speedup | |---|---|---|---| | `struct<int64,string>` | 970 ms | 474 ms | 2.0x | | `map<string,int64>` (2 entries/row) | 2.16 s | 799 ms | 2.7x | Peak memory unchanged. ### 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 `ArrowColumnToPylistTests` with struct (incl. nested struct/list children, empty struct, all-null), map (incl. list values), struct-of-map, list-of-struct, sliced and chunked views — all compared against `column.to_pylist()` with exact element-type assertions — plus dedicated tests that duplicate struct field names still raise `ValueError` and that empty-struct rows are distinct dict objects. New ASV benchmark class `ArrowStructMapColumnToRowsBenchmark`. ### Was this patch authored or co-authored using generative AI tooling? Yes. This pull request and its description were written by Isaac (Claude Code). Closes #57105 from viirya/arrow-to-pylist-maps-structs. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit 9c6e57b) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
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).