This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new ad0eebf8c3aa [SPARK-58019][PYTHON] Convert Arrow list columns to
Python rows in bulk
ad0eebf8c3aa is described below
commit ad0eebf8c3aa3fe34f86fc38082c6ecc731bb9f7
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Fri Jul 10 16:24:22 2026 -0700
[SPARK-58019][PYTHON] Convert Arrow list columns to Python rows in bulk
### 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 `co [...]
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 [...]
### 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 #57099 from viirya/arrow-to-pylist-shim.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
(cherry picked from commit 484342ab403271283b2298c43e5e93c3d6d0bc87)
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
python/benchmarks/bench_arrow.py | 56 ++++++++++++++++
python/pyspark/sql/conversion.py | 91 ++++++++++++++++++++++++-
python/pyspark/sql/tests/test_conversion.py | 100 ++++++++++++++++++++++++++++
python/pyspark/worker.py | 10 ++-
4 files changed, 253 insertions(+), 4 deletions(-)
diff --git a/python/benchmarks/bench_arrow.py b/python/benchmarks/bench_arrow.py
index 29a95dbcc98b..b59a43d42072 100644
--- a/python/benchmarks/bench_arrow.py
+++ b/python/benchmarks/bench_arrow.py
@@ -114,3 +114,59 @@ class NullableLongArrowToPandasBenchmark:
def peakmem_long_with_nulls_to_pandas_ext(self, n_rows, method):
self.run_long_with_nulls_to_pandas_ext(n_rows, method)
+
+
+class ArrowListColumnToRowsBenchmark:
+ """
+ Benchmark for converting Arrow list-typed columns to Python rows, the hot
+ path of Arrow-optimized Python UDF inputs and Spark Connect collect().
+
+ ``baseline`` measures plain ``column.to_pylist()``; ``bulk`` measures
+ ``ArrowTableToRowsConversion._to_pylist`` (see apache/arrow#50326).
+ """
+
+ params = [
+ [100000, 1000000],
+ ["baseline", "bulk"],
+ ]
+ param_names = ["n_rows", "method"]
+
+ def setup(self, n_rows, method):
+ from pyspark.sql.conversion import ArrowTableToRowsConversion
+
+ self.list_of_strings = pa.array(
+ [[f"s{i}", f"t{i}"] for i in range(n_rows)],
type=pa.list_(pa.string())
+ )
+ self.nested_ints_with_nulls = pa.array(
+ [[[i, i + 1], None, [i + 2]] if i % 10 != 0 else None for i in
range(n_rows)],
+ type=pa.list_(pa.list_(pa.int32())),
+ )
+ self.array_of_structs = pa.array(
+ [
+ [{"i": i, "s": f"a{i}"}, {"i": i + 1, "s": f"b{i}"}] if i % 10
!= 0 else None
+ for i in range(n_rows)
+ ],
+ type=pa.list_(pa.struct([("i", pa.int32()), ("s", pa.string())])),
+ )
+ if method == "bulk":
+ self.convert = ArrowTableToRowsConversion._to_pylist
+ else:
+ self.convert = lambda column: column.to_pylist()
+
+ def time_list_of_strings_to_rows(self, n_rows, method):
+ self.convert(self.list_of_strings)
+
+ def time_nested_ints_with_nulls_to_rows(self, n_rows, method):
+ self.convert(self.nested_ints_with_nulls)
+
+ def time_array_of_structs_to_rows(self, n_rows, method):
+ self.convert(self.array_of_structs)
+
+ def peakmem_list_of_strings_to_rows(self, n_rows, method):
+ self.convert(self.list_of_strings)
+
+ def peakmem_nested_ints_with_nulls_to_rows(self, n_rows, method):
+ self.convert(self.nested_ints_with_nulls)
+
+ def peakmem_array_of_structs_to_rows(self, n_rows, method):
+ self.convert(self.array_of_structs)
diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py
index bfa0d4a559a5..fddb42c63d8c 100644
--- a/python/pyspark/sql/conversion.py
+++ b/python/pyspark/sql/conversion.py
@@ -18,6 +18,7 @@
import array
import datetime
import decimal
+import functools
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence,
Union, overload
import pyspark
@@ -980,6 +981,90 @@ class ArrowTableToRowsConversion:
Conversion from Arrow Table to Rows.
"""
+ @staticmethod
+ @functools.cache
+ def _should_manual_bulk() -> bool:
+ """
+ Whether ``_to_pylist`` should convert nested columns manually in bulk.
+
+ Internal helper for ``_to_pylist`` only; do not use externally.
Returns True
+ when the installed PyArrow still materializes one Scalar per element in
+ ``to_pylist`` (apache/arrow#50326, fix expected in PyArrow 25.0.1 —
adjust the
+ version below if it ships in a different release) and NumPy (used for
the
+ offsets and validity buffers) is available.
+
+ This method and the manual bulk paths in ``_to_pylist`` should be
removed once
+ the minimum supported PyArrow version contains the fix.
+ """
+ import pyarrow as pa
+ from pyspark.loose_version import LooseVersion
+
+ if LooseVersion(pa.__version__) >= LooseVersion("25.0.1"):
+ # Native to_pylist converts without per-element Scalars.
+ return False
+ try:
+ import numpy # noqa: F401
+ except ImportError:
+ return False
+ return True
+
+ @staticmethod
+ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]:
+ """
+ Equivalent to ``column.to_pylist()``, but converts (nested) list
columns in bulk
+ instead of one scalar at a time.
+
+ Internal helper for the worker and ``convert`` call sites; do not use
+ externally.
+
+ ``Array.to_pylist()`` materializes one Scalar per element; for list
types each row
+ additionally allocates a C++ scalar, a Python Scalar wrapper and a
Python Array
+ wrapper for the row's values before converting elements one by one,
which is
+ several times slower than converting the flattened child values in a
single pass
+ and slicing the resulting Python list per row (see
apache/arrow#50326). The values
+ themselves 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 would coerce them to floats/NaN.
NumPy is used
+ only for the offsets (non-null integers) and the validity bitmap
(booleans), so no
+ value coercion can occur.
+
+ This method should be removed (its call sites reverting to plain
+ ``column.to_pylist()``) once the minimum supported PyArrow version
includes the
+ fix for apache/arrow#50326.
+ """
+ import pyarrow as pa
+
+ if not ArrowTableToRowsConversion._should_manual_bulk():
+ return column.to_pylist()
+
+ if isinstance(column, pa.ChunkedArray):
+ result = []
+ for chunk in column.chunks:
+ result.extend(ArrowTableToRowsConversion._to_pylist(chunk))
+ return result
+
+ if (pa.types.is_list(column.type) or
pa.types.is_large_list(column.type)) and len(
+ column
+ ) > 0:
+ n = len(column)
+ # List offset buffers never carry a validity bitmap, so this
conversion is
+ # always zero-copy; zero_copy_only=True asserts that invariant and
would
+ # fail loudly if a future Arrow list variant ever violated it.
+ offsets = column.offsets.to_numpy(zero_copy_only=True).tolist()
+ start = offsets[0]
+ flat = ArrowTableToRowsConversion._to_pylist(
+ column.values.slice(start, offsets[-1] - start)
+ )
+ if column.null_count == 0:
+ return [flat[offsets[i] - start : offsets[i + 1] - start] for
i in range(n)]
+ valid = column.is_valid().to_numpy(zero_copy_only=False).tolist()
+ return [
+ flat[offsets[i] - start : offsets[i + 1] - start] if valid[i]
else None
+ for i in range(n)
+ ]
+
+ return column.to_pylist()
+
@staticmethod
def _need_converter(dataType: DataType) -> bool:
if isinstance(dataType, NullType):
@@ -1306,7 +1391,11 @@ class ArrowTableToRowsConversion:
]
columnar_data = [
- [conv(v) for v in column.to_pylist()] if conv is not None else
column.to_pylist()
+ (
+ [conv(v) for v in
ArrowTableToRowsConversion._to_pylist(column)]
+ if conv is not None
+ else ArrowTableToRowsConversion._to_pylist(column)
+ )
for column, conv in zip(table.columns, field_converters)
]
diff --git a/python/pyspark/sql/tests/test_conversion.py
b/python/pyspark/sql/tests/test_conversion.py
index dd5c7f44d281..7ce2dfe01411 100644
--- a/python/pyspark/sql/tests/test_conversion.py
+++ b/python/pyspark/sql/tests/test_conversion.py
@@ -16,6 +16,7 @@
#
import datetime
import unittest
+import unittest.mock
from zoneinfo import ZoneInfo
from pyspark.errors import PySparkRuntimeError, PySparkTypeError,
PySparkValueError
@@ -844,6 +845,105 @@ class
ArrowArrayToPandasConversionTests(unittest.TestCase):
self.assertEqual(len(result), 0)
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowColumnToPylistTests(unittest.TestCase):
+ """
+ ArrowTableToRowsConversion._to_pylist must return exactly what
+ column.to_pylist() returns, including exact element types.
+ """
+
+ def setUp(self):
+ # Force the manual bulk paths so they stay covered regardless of the
+ # installed PyArrow version (with a fast native PyArrow the method
+ # short-circuits to column.to_pylist()).
+ self._gate_patcher = unittest.mock.patch.object(
+ ArrowTableToRowsConversion, "_should_manual_bulk", lambda: True
+ )
+ self._gate_patcher.start()
+
+ def tearDown(self):
+ self._gate_patcher.stop()
+
+ def test_native_to_pylist_gate(self):
+ import pyarrow as pa
+
+ column = pa.array([[1, None], None], type=pa.list_(pa.int32()))
+ with unittest.mock.patch.object(
+ ArrowTableToRowsConversion, "_should_manual_bulk", lambda: False
+ ):
+ self.assertEqual(ArrowTableToRowsConversion._to_pylist(column),
[[1, None], None])
+
+ def _assert_identical_types(self, actual, expected):
+ self.assertIs(type(actual), type(expected))
+ if isinstance(actual, (list, tuple)):
+ self.assertEqual(len(actual), len(expected))
+ for a, e in zip(actual, expected):
+ self._assert_identical_types(a, e)
+
+ def test_matches_to_pylist(self):
+ import pyarrow as pa
+
+ columns = [
+ pa.array([[1, None, 3], None, [], [4]], type=pa.list_(pa.int32())),
+ pa.array([["a", None], None, [], ["bcd", ""]],
type=pa.list_(pa.string())),
+ pa.array([["a", None], None, ["b"]],
type=pa.large_list(pa.string())),
+ pa.array([[[1], None, [2, None]], None],
type=pa.list_(pa.list_(pa.int32()))),
+ pa.array(
+ [[{"a": 1, "b": "x"}, None], None],
+ type=pa.list_(pa.struct([("a", pa.int32()), ("b",
pa.string())])),
+ ),
+ pa.array([[("k1", 1), ("k2", None)], None, []],
type=pa.map_(pa.string(), pa.int32())),
+ pa.array([[1.5, None], [float("nan")]],
type=pa.list_(pa.float64())),
+ pa.array([1, None, 3], type=pa.int64()),
+ pa.array(["x", None], type=pa.string()),
+ pa.array([], type=pa.list_(pa.int32())),
+ pa.array([None, None], type=pa.list_(pa.string())),
+ pa.array([[1, 2], None], type=pa.list_(pa.int64(), 2)),
+ ]
+ for column in columns:
+ views = [column, column.slice(1), column.slice(0, max(len(column)
- 1, 0))]
+ views.append(pa.chunked_array([column, column.slice(1)],
type=column.type))
+ for view in views:
+ with self.subTest(type=str(column.type), length=len(view)):
+ actual = ArrowTableToRowsConversion._to_pylist(view)
+ expected = view.to_pylist()
+ # NaN != NaN; compare via repr for the float case
+ self.assertEqual(repr(actual), repr(expected))
+ self._assert_identical_types(actual, expected)
+
+ def test_int_list_with_nulls_stays_int(self):
+ # The exact case that makes a pandas round trip unusable: ints must not
+ # become floats/NaN when the list contains nulls.
+ import pyarrow as pa
+
+ result = ArrowTableToRowsConversion._to_pylist(
+ pa.array([[1, None, 3]], type=pa.list_(pa.int32()))
+ )
+ self.assertEqual(result, [[1, None, 3]])
+ self.assertEqual([type(v) for v in result[0]], [int, type(None), int])
+
+ def test_convert_table_with_list_columns(self):
+ import pyarrow as pa
+
+ schema = (
+ StructType()
+ .add("arr", ArrayType(IntegerType()))
+ .add("nested", ArrayType(ArrayType(StringType())))
+ )
+ tbl = pa.table(
+ {
+ "arr": pa.array([[1, None], None, []],
type=pa.list_(pa.int32())),
+ "nested": pa.array(
+ [[["a"], None], [[]], None],
type=pa.list_(pa.list_(pa.string()))
+ ),
+ }
+ )
+ actual = ArrowTableToRowsConversion.convert(tbl, schema)
+ self.assertEqual(actual[0], Row(arr=[1, None], nested=[["a"], None]))
+ self.assertEqual(actual[1], Row(arr=None, nested=[[]]))
+ self.assertEqual(actual[2], Row(arr=[], nested=None))
+
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 9e629138d537..83ceeed1bd10 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -1755,9 +1755,9 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
# then call eval once per input row.
pylist = [
(
- [conv(v) for v in column.to_pylist()]
+ [conv(v) for v in
ArrowTableToRowsConversion._to_pylist(column)]
if conv is not None
- else column.to_pylist()
+ else ArrowTableToRowsConversion._to_pylist(column)
)
for column, conv in zip(batch.columns, converters)
]
@@ -3067,7 +3067,11 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# --- Input: Arrow -> Python columns ---
columns = [
- [conv(v) for v in col.to_pylist()] if conv is not None
else col.to_pylist()
+ (
+ [conv(v) for v in
ArrowTableToRowsConversion._to_pylist(col)]
+ if conv is not None
+ else ArrowTableToRowsConversion._to_pylist(col)
+ )
for col, conv in zip(input_batch.itercolumns(),
arrow_to_py_converters)
]
if not columns:
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
