From 51771f59ae2c5f7b4452d2f357b072e0c01ec716 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 7 Jul 2026 16:43:41 -0700 Subject: [PATCH 1/8] [SPARK-58019][PYTHON] Convert Arrow list columns to Python rows in bulk 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 769ms -> 507ms (1.5x); list> 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 --- python/benchmarks/bench_arrow.py | 40 +++++++++++ python/pyspark/sql/conversion.py | 66 ++++++++++++++++-- python/pyspark/sql/tests/test_conversion.py | 77 +++++++++++++++++++++ python/pyspark/worker.py | 32 +++++---- 4 files changed, 196 insertions(+), 19 deletions(-) diff --git a/python/benchmarks/bench_arrow.py b/python/benchmarks/bench_arrow.py index 29a95dbcc98b5..81189707016b6 100644 --- a/python/benchmarks/bench_arrow.py +++ b/python/benchmarks/bench_arrow.py @@ -114,3 +114,43 @@ def time_long_with_nulls_to_pandas_ext(self, n_rows, method): 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())), + ) + 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 peakmem_list_of_strings_to_rows(self, n_rows, method): + self.convert(self.list_of_strings) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index bfa0d4a559a51..e3342adce2012 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -980,6 +980,60 @@ class ArrowTableToRowsConversion: Conversion from Arrow Table to Rows. """ + @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. + + ``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 can be removed once the minimum supported PyArrow version includes the fix + for apache/arrow#50326. + """ + import pyarrow as pa + import pyarrow.types as pa_types + + try: + import numpy # noqa: F401 + except ImportError: + return column.to_pylist() + + if isinstance(column, pa.ChunkedArray): + result: List[Any] = [] + for chunk in column.chunks: + result.extend(ArrowTableToRowsConversion._to_pylist(chunk)) + return result + + column_type = column.type + if (pa_types.is_list(column_type) or pa_types.is_large_list(column_type)) and len( + column + ) > 0: + n = len(column) + 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): @@ -1069,9 +1123,9 @@ def convert_struct(value: Any) -> Any: dataType.elementType, none_on_identity=True, binary_as_bytes=binary_as_bytes ) - assert element_conv is not None, ( - f"_need_converter() returned True for ArrayType of {dataType.elementType}" - ) + assert ( + element_conv is not None + ), f"_need_converter() returned True for ArrayType of {dataType.elementType}" def convert_array(value: Any) -> Any: if value is None: @@ -1306,7 +1360,11 @@ def convert( ] 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 dd5c7f44d2818..0944e49664cb8 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -844,6 +844,83 @@ def test_geometry_convert_numpy(self): self.assertEqual(len(result), 0) +class ArrowColumnToPylistTests(unittest.TestCase): + """ + ArrowTableToRowsConversion._to_pylist must return exactly what + column.to_pylist() returns, including exact element types. + """ + + 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 9e629138d5370..cc23eace59bd9 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -1755,9 +1755,9 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record # 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) ] @@ -2616,9 +2616,9 @@ def grouped_func( assert num_udfs == 1, "One GROUPED_MAP_ARROW_ITER UDF expected here." grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] parsed_offsets = extract_key_value_indexes(arg_offsets) - assert len(parsed_offsets) == 1, ( - "Expected one pair of offsets for GROUPED_MAP_ARROW_ITER UDF." - ) + assert ( + len(parsed_offsets) == 1 + ), "Expected one pair of offsets for GROUPED_MAP_ARROW_ITER UDF." arrow_return_type = to_arrow_type( return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types @@ -2752,9 +2752,9 @@ def grouped_func( assert num_udfs == 1, "One GROUPED_MAP_PANDAS_ITER UDF expected here." grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] parsed_offsets = extract_key_value_indexes(arg_offsets) - assert len(parsed_offsets) == 1, ( - "Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF." - ) + assert ( + len(parsed_offsets) == 1 + ), "Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF." key_offsets = parsed_offsets[0][0] value_offsets = parsed_offsets[0][1] @@ -3067,7 +3067,11 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record # --- 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: @@ -3119,9 +3123,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record coerce = ( str if isinstance(udf_return_type, StringType) - else bytes - if isinstance(udf_return_type, BinaryType) - else None + else bytes if isinstance(udf_return_type, BinaryType) else None ) udf_infos.append( ( @@ -3363,9 +3365,9 @@ def process_results(): # See TransformWithStateInPandasExec for how arg_offsets are used to # distinguish between grouping attributes and data attributes parsed_offsets = extract_key_value_indexes(arg_offsets) - assert len(parsed_offsets) == 1, ( - "Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF." - ) + assert ( + len(parsed_offsets) == 1 + ), "Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF." key_offsets = parsed_offsets[0][0] value_offsets = parsed_offsets[0][1] From eebbc4fd05957ca316041ff3682a2473dce5a62b Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 7 Jul 2026 18:11:05 -0700 Subject: [PATCH 2/8] [SPARK-58019][PYTHON] Document the zero-copy offsets invariant Co-authored-by: Isaac --- python/pyspark/sql/conversion.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index e3342adce2012..94ced9f9c58ae 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -1019,6 +1019,9 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: 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( From 691885a81adcba6b6ef806da0c45448b5289c1fc Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 7 Jul 2026 21:00:12 -0700 Subject: [PATCH 3/8] [SPARK-58019][PYTHON] Restore ruff formatting 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 --- python/pyspark/sql/conversion.py | 6 +++--- python/pyspark/worker.py | 22 ++++++++++++---------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index 94ced9f9c58ae..2e4c5bd93500b 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -1126,9 +1126,9 @@ def convert_struct(value: Any) -> Any: dataType.elementType, none_on_identity=True, binary_as_bytes=binary_as_bytes ) - assert ( - element_conv is not None - ), f"_need_converter() returned True for ArrayType of {dataType.elementType}" + assert element_conv is not None, ( + f"_need_converter() returned True for ArrayType of {dataType.elementType}" + ) def convert_array(value: Any) -> Any: if value is None: diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index cc23eace59bd9..83ceeed1bd102 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -2616,9 +2616,9 @@ def grouped_func( assert num_udfs == 1, "One GROUPED_MAP_ARROW_ITER UDF expected here." grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] parsed_offsets = extract_key_value_indexes(arg_offsets) - assert ( - len(parsed_offsets) == 1 - ), "Expected one pair of offsets for GROUPED_MAP_ARROW_ITER UDF." + assert len(parsed_offsets) == 1, ( + "Expected one pair of offsets for GROUPED_MAP_ARROW_ITER UDF." + ) arrow_return_type = to_arrow_type( return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types @@ -2752,9 +2752,9 @@ def grouped_func( assert num_udfs == 1, "One GROUPED_MAP_PANDAS_ITER UDF expected here." grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] parsed_offsets = extract_key_value_indexes(arg_offsets) - assert ( - len(parsed_offsets) == 1 - ), "Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF." + assert len(parsed_offsets) == 1, ( + "Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF." + ) key_offsets = parsed_offsets[0][0] value_offsets = parsed_offsets[0][1] @@ -3123,7 +3123,9 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record coerce = ( str if isinstance(udf_return_type, StringType) - else bytes if isinstance(udf_return_type, BinaryType) else None + else bytes + if isinstance(udf_return_type, BinaryType) + else None ) udf_infos.append( ( @@ -3365,9 +3367,9 @@ def process_results(): # See TransformWithStateInPandasExec for how arg_offsets are used to # distinguish between grouping attributes and data attributes parsed_offsets = extract_key_value_indexes(arg_offsets) - assert ( - len(parsed_offsets) == 1 - ), "Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF." + assert len(parsed_offsets) == 1, ( + "Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF." + ) key_offsets = parsed_offsets[0][0] value_offsets = parsed_offsets[0][1] From 784974ceeceb1605e6586c222b0cbb86c6084fb9 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 7 Jul 2026 23:34:30 -0700 Subject: [PATCH 4/8] [SPARK-58019][PYTHON] Skip Arrow to_pylist tests without PyArrow Co-authored-by: Isaac --- python/pyspark/sql/tests/test_conversion.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index 0944e49664cb8..b559d3ac3081e 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -844,6 +844,7 @@ def test_geometry_convert_numpy(self): self.assertEqual(len(result), 0) +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowColumnToPylistTests(unittest.TestCase): """ ArrowTableToRowsConversion._to_pylist must return exactly what From 1fd52c8e3517278c229ac0120c137d99f0c9feb8 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 8 Jul 2026 13:06:22 -0700 Subject: [PATCH 5/8] [SPARK-58019][PYTHON] Cache NumPy availability; add peakmem and array 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 case plus peakmem variants for the nested cases (1M rows: nested list 862M -> 862M, array 885M -> 921M, ~4% peak increase from the transient flattened pointer list; leaf objects are shared). Co-authored-by: Isaac --- python/benchmarks/bench_arrow.py | 16 ++++++++++++++++ python/pyspark/sql/conversion.py | 19 ++++++++++++++++--- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/python/benchmarks/bench_arrow.py b/python/benchmarks/bench_arrow.py index 81189707016b6..b59a43d420725 100644 --- a/python/benchmarks/bench_arrow.py +++ b/python/benchmarks/bench_arrow.py @@ -141,6 +141,13 @@ def setup(self, n_rows, method): [[[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: @@ -152,5 +159,14 @@ def time_list_of_strings_to_rows(self, n_rows, method): 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 2e4c5bd93500b..bd545e3ab2f61 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -506,6 +506,21 @@ def convert_column( return pa.RecordBatch.from_arrays(arrays, schema.names) +_numpy_available: Optional[bool] = None + + +def _is_numpy_available() -> bool: + global _numpy_available + if _numpy_available is None: + try: + import numpy # noqa: F401 + + _numpy_available = True + except ImportError: + _numpy_available = False + return _numpy_available + + class LocalDataToArrowConversion: """ Conversion from local data (except pandas DataFrame and numpy ndarray) to Arrow. @@ -1003,9 +1018,7 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: import pyarrow as pa import pyarrow.types as pa_types - try: - import numpy # noqa: F401 - except ImportError: + if not _is_numpy_available(): return column.to_pylist() if isinstance(column, pa.ChunkedArray): From 2e405eb6ffbb4dbc3e27f7b5a9b7451157b4fe84 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Thu, 9 Jul 2026 12:14:44 -0700 Subject: [PATCH 6/8] [SPARK-58019][PYTHON] Use native to_pylist when PyArrow converts without 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 --- python/pyspark/sql/conversion.py | 27 ++++++++++++++++++++- python/pyspark/sql/tests/test_conversion.py | 20 +++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index bd545e3ab2f61..08065373fb075 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -506,6 +506,28 @@ def convert_column( return pa.RecordBatch.from_arrays(arrays, schema.names) +# The pure-Python bulk conversion in ArrowTableToRowsConversion._to_pylist is +# a workaround for PyArrow materializing one Scalar per element (see +# apache/arrow#50326). PyArrow releases containing the fix convert natively +# without per-element Scalars, in which case the native conversion is used +# directly. Bump this constant if the fix ships in a different release. +_MIN_PYARROW_NATIVE_TO_PYLIST_VERSION = "25.0.1" + +_pyarrow_native_to_pylist_is_fast: Optional[bool] = None + + +def _has_fast_native_to_pylist() -> bool: + global _pyarrow_native_to_pylist_is_fast + if _pyarrow_native_to_pylist_is_fast is None: + import pyarrow as pa + from pyspark.loose_version import LooseVersion + + _pyarrow_native_to_pylist_is_fast = LooseVersion(pa.__version__) >= LooseVersion( + _MIN_PYARROW_NATIVE_TO_PYLIST_VERSION + ) + return _pyarrow_native_to_pylist_is_fast + + _numpy_available: Optional[bool] = None @@ -1018,7 +1040,10 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: import pyarrow as pa import pyarrow.types as pa_types - if not _is_numpy_available(): + if _has_fast_native_to_pylist() or not _is_numpy_available(): + # Recent PyArrow converts without per-element Scalars natively + # (apache/arrow#50326); without NumPy the bulk paths below are + # unavailable. Either way, use the native conversion. return column.to_pylist() if isinstance(column, pa.ChunkedArray): diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index b559d3ac3081e..12780a1bd4fa7 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -851,6 +851,26 @@ class ArrowColumnToPylistTests(unittest.TestCase): column.to_pylist() returns, including exact element types. """ + def setUp(self): + # Force the 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()). + import pyspark.sql.conversion as conversion_mod + + self._conversion_mod = conversion_mod + self._saved_gate = conversion_mod._pyarrow_native_to_pylist_is_fast + conversion_mod._pyarrow_native_to_pylist_is_fast = False + + def tearDown(self): + self._conversion_mod._pyarrow_native_to_pylist_is_fast = self._saved_gate + + def test_native_to_pylist_gate(self): + import pyarrow as pa + + column = pa.array([[1, None], None], type=pa.list_(pa.int32())) + self._conversion_mod._pyarrow_native_to_pylist_is_fast = True + 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)): From 6828ff7e9f9996bfdb4ea5c164627425d34f54a3 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 7 Jul 2026 17:44:32 -0700 Subject: [PATCH 7/8] [SPARK-58024][PYTHON] Convert Arrow struct and map columns to Python rows in bulk Extend ArrowTableToRowsConversion._to_pylist with bulk paths for struct and map columns: * Struct columns convert each child field in bulk (recursively reusing the list/leaf fast 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. ASV microbenchmark (bench_arrow.ArrowStructMapColumnToRowsBenchmark, 1M rows, 10% nulls): struct 914ms -> 175ms (5.2x); map with 2 entries per row 2.07s -> 303ms (6.8x). Peak memory unchanged. Co-authored-by: Isaac --- python/benchmarks/bench_arrow.py | 43 +++++++++++++ python/pyspark/sql/conversion.py | 69 +++++++++++++++++++-- python/pyspark/sql/tests/test_conversion.py | 46 ++++++++++++++ 3 files changed, 153 insertions(+), 5 deletions(-) diff --git a/python/benchmarks/bench_arrow.py b/python/benchmarks/bench_arrow.py index b59a43d420725..4ed1c91d579c7 100644 --- a/python/benchmarks/bench_arrow.py +++ b/python/benchmarks/bench_arrow.py @@ -170,3 +170,46 @@ def peakmem_nested_ints_with_nulls_to_rows(self, n_rows, method): def peakmem_array_of_structs_to_rows(self, n_rows, method): self.convert(self.array_of_structs) + + +class ArrowStructMapColumnToRowsBenchmark: + """ + Benchmark for converting Arrow struct and map columns to Python rows. + + ``baseline`` measures plain ``column.to_pylist()``; ``bulk`` measures + ``ArrowTableToRowsConversion._to_pylist`` with the struct/map bulk paths. + """ + + params = [ + [100000, 1000000], + ["baseline", "bulk"], + ] + param_names = ["n_rows", "method"] + + def setup(self, n_rows, method): + from pyspark.sql.conversion import ArrowTableToRowsConversion + + self.structs = pa.array( + [{"a": i, "b": f"s{i}"} if i % 10 != 0 else None for i in range(n_rows)], + type=pa.struct([("a", pa.int64()), ("b", pa.string())]), + ) + self.maps = pa.array( + [ + [(f"k{i % 3}", i), (f"q{i % 5}", i + 1)] if i % 10 != 0 else None + for i in range(n_rows) + ], + type=pa.map_(pa.string(), pa.int64()), + ) + if method == "bulk": + self.convert = ArrowTableToRowsConversion._to_pylist + else: + self.convert = lambda column: column.to_pylist() + + def time_structs_to_rows(self, n_rows, method): + self.convert(self.structs) + + def time_maps_to_rows(self, n_rows, method): + self.convert(self.maps) + + def peakmem_structs_to_rows(self, n_rows, method): + self.convert(self.structs) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index 08065373fb075..8b6ed69db543d 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -1020,8 +1020,11 @@ class ArrowTableToRowsConversion: @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. + Equivalent to ``column.to_pylist()``, but converts (nested) list, struct and map + columns in bulk instead of one scalar at a time. Structs become dicts (with + a fallback to ``to_pylist`` for duplicate field names, which raise ``ValueError`` + there) and maps become lists of ``(key, value)`` tuples, matching + ``StructScalar.as_py`` and ``MapScalar.as_py`` exactly. ``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 @@ -1052,10 +1055,46 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: result.extend(ArrowTableToRowsConversion._to_pylist(chunk)) return result + if len(column) == 0: + return [] + column_type = column.type - if (pa_types.is_list(column_type) or pa_types.is_large_list(column_type)) and len( - column - ) > 0: + + if pa_types.is_map(column_type): + # Maps have the same offsets layout as lists; each row becomes a + # list of (key, value) tuples, matching MapScalar.as_py. + n = len(column) + offsets = column.offsets.to_numpy(zero_copy_only=True).tolist() + start = offsets[0] + length = offsets[-1] - start + keys = ArrowTableToRowsConversion._to_pylist(column.keys.slice(start, length)) + items = ArrowTableToRowsConversion._to_pylist(column.items.slice(start, length)) + if column.null_count == 0: + return [ + list( + zip( + keys[offsets[i] - start : offsets[i + 1] - start], + items[offsets[i] - start : offsets[i + 1] - start], + ) + ) + for i in range(n) + ] + valid = column.is_valid().to_numpy(zero_copy_only=False).tolist() + return [ + ( + list( + zip( + keys[offsets[i] - start : offsets[i + 1] - start], + items[offsets[i] - start : offsets[i + 1] - start], + ) + ) + if valid[i] + else None + ) + for i in range(n) + ] + + if pa_types.is_list(column_type) or pa_types.is_large_list(column_type): 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 @@ -1073,6 +1112,26 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: for i in range(n) ] + if pa_types.is_struct(column_type): + n = len(column) + names = [column_type.field(i).name for i in range(column_type.num_fields)] + if len(set(names)) != len(names): + # StructScalar.as_py raises ValueError on duplicate field names; + # let the generic path surface the same error. + return column.to_pylist() + fields = [ + ArrowTableToRowsConversion._to_pylist(column.field(i)) + for i in range(column_type.num_fields) + ] + if column.null_count == 0: + if not names: + return [{} for _ in range(n)] + return [dict(zip(names, row)) for row in zip(*fields)] + valid = column.is_valid().to_numpy(zero_copy_only=False).tolist() + if not names: + return [{} if m else None for m in valid] + return [dict(zip(names, row)) if m else None for row, m in zip(zip(*fields), valid)] + return column.to_pylist() @staticmethod diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index 12780a1bd4fa7..aee5fe888659c 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -15,6 +15,7 @@ # limitations under the License. # import datetime +import decimal import unittest from zoneinfo import ZoneInfo @@ -897,6 +898,37 @@ def test_matches_to_pylist(self): 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)), + # non-list leaves keep as_py semantics (native to_pylist) + pa.array([b"", None, b"\x00\xff"], type=pa.binary()), + pa.array([datetime.date(2020, 1, 2), None], type=pa.date32()), + pa.array([decimal.Decimal("1.23"), None], type=pa.decimal128(10, 2)), + pa.array([[b"x", None], None, [b""]], type=pa.list_(pa.binary())), + pa.array([[True, None], [False]], type=pa.list_(pa.bool_())), + # struct and map bulk paths + pa.array( + [{"a": 1, "b": "x"}, None, {"a": None, "b": None}], + type=pa.struct([("a", pa.int64()), ("b", pa.string())]), + ), + pa.array( + [{"s": {"a": 1}, "l": [1, None]}, None], + type=pa.struct( + [("s", pa.struct([("a", pa.int32())])), ("l", pa.list_(pa.int64()))] + ), + ), + pa.array([{}, None, {}], type=pa.struct([])), + pa.array([None] * 4, type=pa.struct([("a", pa.int32())])), + pa.array( + [[("k1", [1, None]), ("k2", None)], None, []], + type=pa.map_(pa.string(), pa.list_(pa.int32())), + ), + pa.array( + [{"m": [("k", 1)]}, None], + type=pa.struct([("m", pa.map_(pa.string(), pa.int64()))]), + ), + pa.array( + [[{"a": 1}, None], None], + type=pa.list_(pa.struct([("a", pa.int64())])), + ), ] for column in columns: views = [column, column.slice(1), column.slice(0, max(len(column) - 1, 0))] @@ -920,6 +952,20 @@ def test_int_list_with_nulls_stays_int(self): self.assertEqual(result, [[1, None, 3]]) self.assertEqual([type(v) for v in result[0]], [int, type(None), int]) + def test_struct_duplicate_field_names_still_raises(self): + import pyarrow as pa + + dup = pa.StructArray.from_arrays([pa.array([1, 2]), pa.array(["a", "b"])], names=["x", "x"]) + with self.assertRaises(ValueError): + ArrowTableToRowsConversion._to_pylist(dup) + + def test_struct_rows_are_distinct_dicts(self): + import pyarrow as pa + + result = ArrowTableToRowsConversion._to_pylist(pa.array([{}, {}], type=pa.struct([]))) + self.assertEqual(result, [{}, {}]) + self.assertIsNot(result[0], result[1]) + def test_convert_table_with_list_columns(self): import pyarrow as pa From 3074d23fb938f5c62aeb402d589f602c67a94b1e Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 8 Jul 2026 16:21:59 -0700 Subject: [PATCH 8/8] [SPARK-58050][PYTHON] Bulk-assemble Arrow Python UDF results The Arrow Python UDF worker converts every UDF result through a per-row Python converter (defensive list copies for arrays, dict to entry-list for maps, Row/dict to dict for structs) before pa.array. These conversions are Spark's own and independent of PyArrow's conversion performance; pa.array itself is cheap. Add LocalDataToArrowConversion._create_results_to_arrow: when the element/field/value converters are identity, results are assembled in bulk - arrays pass the returned lists to pa.array directly, map results pass dicts directly (PyArrow accepts dicts for map types), and struct results (Rows/tuples) are transposed and assembled via StructArray.from_arrays with a null mask. Any shape or validation mismatch falls back to the per-row path, preserving its error behavior; types whose child converters transform values (e.g. validated integers, timestamps, decimals) keep the per-row path. Microbenchmarks (400k rows, identical outputs): array output 7.5x, map output 4.2x, struct output 4.0x. End-to-end UDF benchmark (6.4M rows, on top of SPARK-58019/58023/58024): array 2.39s -> 1.84s. Arrow-vs-pickle collect() results verified identical with injected nulls across nested types. Co-authored-by: Isaac --- python/pyspark/sql/conversion.py | 106 ++++++++++++++++++++ python/pyspark/sql/tests/test_conversion.py | 73 ++++++++++++++ python/pyspark/worker.py | 28 ++---- 3 files changed, 189 insertions(+), 18 deletions(-) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index 8b6ed69db543d..f47104659d951 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -932,6 +932,112 @@ def convert_other(value: Any) -> Any: else: # pragma: no cover assert False, f"Need converter for {dataType} but failed to find one." + @staticmethod + def _create_results_to_arrow( + dataType: DataType, + arrow_type: "pa.DataType", + *, + safecheck: bool, + int_to_decimal_coercion_enabled: bool, + ) -> Callable[[List[Any]], "pa.Array"]: + """ + Return a function converting a column of UDF results (a list of Python + values) to an Arrow array of ``arrow_type``. + + The generic strategy applies the per-row converter from + ``_create_converter`` and passes the converted values to ``pa.array``. + When the element/field/value converters are identity, the per-row + converter only reshapes values that ``pa.array`` accepts directly, so + results are assembled in bulk instead: arrays are passed as the + returned lists (skipping the per-row defensive copy), map results are + passed as dicts (PyArrow accepts dicts for map types), and struct + results (``Row``/tuples) are transposed and assembled via + ``pa.StructArray.from_arrays``. Any shape or validation mismatch falls + back to the per-row path, preserving its error behavior. + """ + import pyarrow as pa + + result_conv = LocalDataToArrowConversion._create_converter( + dataType, + none_on_identity=True, + int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, + ) + + def _fallback(results: List[Any]) -> "pa.Array": + converted = [result_conv(r) for r in results] if result_conv is not None else results + try: + return pa.array(converted, type=arrow_type) + except pa.lib.ArrowInvalid: + return pa.array(converted).cast(target_type=arrow_type, safe=safecheck) + + if result_conv is None: + return _fallback + + def _child_conv(dt: DataType, nullable: bool) -> Optional[Callable]: + return LocalDataToArrowConversion._create_converter( + dt, + nullable, + none_on_identity=True, + int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, + ) + + if isinstance(dataType, ArrayType) and ( + _child_conv(dataType.elementType, dataType.containsNull) is None + ): + + def _bulk_array(results: List[Any]) -> "pa.Array": + try: + return pa.array(results, type=arrow_type) + except (pa.lib.ArrowInvalid, TypeError): + return _fallback(results) + + return _bulk_array + + if isinstance(dataType, MapType) and ( + _child_conv(dataType.keyType, False) is None + and _child_conv(dataType.valueType, dataType.valueContainsNull) is None + ): + + def _bulk_map(results: List[Any]) -> "pa.Array": + try: + return pa.array(results, type=arrow_type) + except (pa.lib.ArrowInvalid, TypeError): + return _fallback(results) + + return _bulk_map + + if isinstance(dataType, StructType) and all( + _child_conv(f.dataType, f.nullable) is None for f in dataType.fields + ): + names = dataType.fieldNames() + width = len(names) + placeholder = (None,) * width + + def _bulk_struct(results: List[Any]) -> "pa.Array": + # Rows are tuples; fall back for dict/object results. + if any(r is not None and not isinstance(r, tuple) for r in results): + return _fallback(results) + if any(r is not None and len(r) != width for r in results): + return _fallback(results) + mask = [r is None for r in results] + tuples = ( + [r if r is not None else placeholder for r in results] if any(mask) else results + ) + try: + field_arrays = [ + pa.array(list(f), type=arrow_type.field(k).type) + for k, f in enumerate(zip(*tuples)) + ] + return pa.StructArray.from_arrays( + field_arrays, names=names, mask=pa.array(mask, type=pa.bool_()) + ) + except (pa.lib.ArrowInvalid, TypeError): + return _fallback(results) + + return _bulk_struct + + return _fallback + @staticmethod def convert(data: Sequence[Any], schema: StructType, use_large_var_types: bool) -> "pa.Table": require_minimum_pyarrow_version() diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index aee5fe888659c..3cb4a16c94b02 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -988,6 +988,79 @@ def test_convert_table_with_list_columns(self): self.assertEqual(actual[2], Row(arr=[], nested=None)) +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class LocalDataResultsToArrowTests(unittest.TestCase): + """ + LocalDataToArrowConversion._create_results_to_arrow must produce the same + Arrow array as applying the per-row converter and pa.array. + """ + + def _reference(self, results, data_type, arrow_type): + conv = LocalDataToArrowConversion._create_converter( + data_type, none_on_identity=True, int_to_decimal_coercion_enabled=False + ) + converted = [conv(r) for r in results] if conv is not None else results + import pyarrow as pa + + return pa.array(converted, type=arrow_type) + + def _check(self, results, data_type): + from pyspark.sql.pandas.types import to_arrow_type + + arrow_type = to_arrow_type(data_type, timezone="UTC", prefers_large_types=False) + to_arrow = LocalDataToArrowConversion._create_results_to_arrow( + data_type, arrow_type, safecheck=True, int_to_decimal_coercion_enabled=False + ) + actual = to_arrow(results) + expected = self._reference(results, data_type, arrow_type) + self.assertTrue(actual.equals(expected), f"{data_type}: {actual} != {expected}") + + def test_matches_per_row_converter(self): + array_string = ArrayType(StringType()) + map_si = MapType(StringType(), IntegerType()) + struct_t = ( + StructType().add("i", IntegerType()).add("s", StringType()).add("d", DoubleType()) + ) + cases = [ + (array_string, [["a", None], None, []]), + (ArrayType(IntegerType()), [[1, None], None, [2]]), + (map_si, [{"a": 1, "b": None}, None, {}]), + (struct_t, [Row(i=1, s="x", d=0.5), None, Row(i=None, s=None, d=None)]), + # dict results for struct take the fallback path + (struct_t, [{"i": 1, "s": "x", "d": 0.5}, None]), + (StringType(), ["a", None]), + (ArrayType(ArrayType(StringType())), [[["a"], None], None]), + (StructType().add("l", ArrayType(StringType())), [Row(l=["a", None]), None]), + ] + for data_type, results in cases: + with self.subTest(type=str(data_type)): + self._check(results, data_type) + + def test_struct_arity_mismatch_falls_back_to_same_error(self): + from pyspark.sql.pandas.types import to_arrow_type + + struct_t = StructType().add("i", IntegerType()).add("s", StringType()) + arrow_type = to_arrow_type(struct_t, timezone="UTC", prefers_large_types=False) + to_arrow = LocalDataToArrowConversion._create_results_to_arrow( + struct_t, arrow_type, safecheck=True, int_to_decimal_coercion_enabled=False + ) + with self.assertRaises(PySparkValueError): + to_arrow([(1, "x", "extra")]) + + def test_returned_lists_are_not_mutated(self): + from pyspark.sql.pandas.types import to_arrow_type + + data_type = ArrayType(StringType()) + arrow_type = to_arrow_type(data_type, timezone="UTC", prefers_large_types=False) + to_arrow = LocalDataToArrowConversion._create_results_to_arrow( + data_type, arrow_type, safecheck=True, int_to_decimal_coercion_enabled=False + ) + results = [["a", "b"], None] + arr = to_arrow(results) + self.assertEqual(arr.to_pylist(), [["a", "b"], None]) + self.assertEqual(results, [["a", "b"], None]) + + if __name__ == "__main__": from pyspark.testing import main diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 83ceeed1bd102..fd3ef8137608e 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -3025,19 +3025,20 @@ def cogrouped_func( udf_func, udf_args_offsets, udf_kwargs_offsets ) zero_arg = len(args_kwargs_offsets) == 0 + arrow_return_type = to_arrow_type( + udf_return_type, + timezone="UTC", + prefers_large_types=runner_conf.use_large_var_types, + ) udf_infos.append( ( wrapped_func, args_kwargs_offsets or (0,), zero_arg, - to_arrow_type( + LocalDataToArrowConversion._create_results_to_arrow( udf_return_type, - timezone="UTC", - prefers_large_types=runner_conf.use_large_var_types, - ), - LocalDataToArrowConversion._create_converter( - udf_return_type, - none_on_identity=True, + arrow_return_type, + safecheck=runner_conf.safecheck, int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled, ), ) @@ -3079,7 +3080,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record # --- Process: evaluate each UDF row-by-row --- output_arrays = [] - for udf_func, offsets, zero_arg, arrow_return_type, result_conv in udf_infos: + for udf_func, offsets, zero_arg, results_to_arrow in udf_infos: rows = ( [() for _ in range(num_rows)] if zero_arg @@ -3089,16 +3090,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record verify_result_row_count(len(results), num_rows) # --- Output: Python -> Arrow --- - converted = ( - [result_conv(r) for r in results] if result_conv is not None else results - ) - try: - arr = pa.array(converted, type=arrow_return_type) - except pa.lib.ArrowInvalid: - arr = pa.array(converted).cast( - target_type=arrow_return_type, safe=runner_conf.safecheck - ) - output_arrays.append(arr) + output_arrays.append(results_to_arrow(results)) yield pa.RecordBatch.from_arrays(output_arrays, col_names)