diff --git a/python/benchmarks/bench_arrow.py b/python/benchmarks/bench_arrow.py index 29a95dbcc98b5..92c63aa609c81 100644 --- a/python/benchmarks/bench_arrow.py +++ b/python/benchmarks/bench_arrow.py @@ -114,3 +114,102 @@ 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())), + ) + 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) + + +class ArrowLeafColumnToRowsBenchmark: + """ + Benchmark for converting flat (leaf) Arrow columns to Python rows. + + ``baseline`` measures plain ``column.to_pylist()``; ``bulk`` measures + ``ArrowTableToRowsConversion._to_pylist`` with the string/binary/numeric + fast paths. + """ + + params = [ + [100000, 1000000], + ["baseline", "bulk"], + ] + param_names = ["n_rows", "method"] + + def setup(self, n_rows, method): + from pyspark.sql.conversion import ArrowTableToRowsConversion + + self.strings = pa.array( + [f"s{i}" if i % 10 != 0 else None for i in range(n_rows)], type=pa.string() + ) + self.longs_with_nulls = pa.array( + [i if i % 10 != 0 else None for i in range(n_rows)], type=pa.int64() + ) + self.doubles = pa.array([float(i) for i in range(n_rows)], type=pa.float64()) + if method == "bulk": + self.convert = ArrowTableToRowsConversion._to_pylist + else: + self.convert = lambda column: column.to_pylist() + + def time_strings_with_nulls_to_rows(self, n_rows, method): + self.convert(self.strings) + + def time_longs_with_nulls_to_rows(self, n_rows, method): + self.convert(self.longs_with_nulls) + + def time_doubles_to_rows(self, n_rows, method): + self.convert(self.doubles) + + def peakmem_strings_with_nulls_to_rows(self, n_rows, method): + self.convert(self.strings) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index bfa0d4a559a51..512dd92e86ab5 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -506,6 +506,43 @@ 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 + + +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. @@ -980,6 +1017,97 @@ 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, with fast paths for string, binary, integral, + floating point and boolean leaves. + + ``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). Results + are exactly identical to ``to_pylist``: ``None`` stays ``None`` and values inside + numeric lists stay Python ints, unlike a pandas round trip which would coerce + them to floats/NaN. In particular the leaf fast paths cannot coerce: string and + binary columns use Arrow's object-dtype conversion, which only produces ``str`` / + ``bytes`` / ``None``; nullable numeric and boolean columns are converted from the + original values (nulls filled with a placeholder and restored to ``None`` from the + validity bitmap afterwards), so ints are materialized from the int buffer, never + via a float representation. Types whose ``as_py`` returns non-primitive objects + (dates, timestamps, decimals, ...) keep using ``to_pylist``. + + 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 + + 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): + result: List[Any] = [] + for chunk in column.chunks: + result.extend(ArrowTableToRowsConversion._to_pylist(chunk)) + return result + + if len(column) == 0: + return [] + + column_type = column.type + + if ( + pa_types.is_string(column_type) + or pa_types.is_large_string(column_type) + or pa_types.is_binary(column_type) + or pa_types.is_large_binary(column_type) + ): + # The object-dtype conversion produces exactly str/bytes and None. + return column.to_numpy(zero_copy_only=False).tolist() + + if ( + pa_types.is_integer(column_type) + or pa_types.is_float32(column_type) + or pa_types.is_float64(column_type) + or pa_types.is_boolean(column_type) + ): + # Booleans are bit-packed, so their conversion to NumPy is never zero-copy. + zero_copy = not pa_types.is_boolean(column_type) + if column.null_count == 0: + return column.to_numpy(zero_copy_only=zero_copy).tolist() + import pyarrow.compute as pc + + valid = column.is_valid().to_numpy(zero_copy_only=False).tolist() + fill_value = False if pa_types.is_boolean(column_type) else 0 + values = pc.fill_null(column, fill_value).to_numpy(zero_copy_only=zero_copy).tolist() + return [v if m else None for v, m in zip(values, valid)] + + 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 + # 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 +1434,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..49808f977144f 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 @@ -844,6 +845,125 @@ 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 + 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)): + 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)), + # leaf fast paths: exact str/bytes/int/float/bool types, None for nulls + pa.array(["", None, "日本語", "\N{GRINNING FACE}", "x" * 40], type=pa.string()), + pa.array(["a", None, ""], type=pa.large_string()), + pa.array([b"", None, b"\x00\xff"], type=pa.binary()), + pa.array([b"a", None], type=pa.large_binary()), + pa.array([1, None, -(2**62), 3], type=pa.int64()), + pa.array([0, None, 2**63 + 7], type=pa.uint64()), + pa.array([-128, 127, None], type=pa.int8()), + pa.array([1.5, None, float("nan"), float("inf")], type=pa.float64()), + pa.array([1.5, None], type=pa.float32()), + pa.array([True, None, False], type=pa.bool_()), + pa.array([True, False] * 5, type=pa.bool_()), + pa.array(list(range(10)), type=pa.int32()), + # non-primitive leaves must keep as_py semantics (fallback path) + pa.array([datetime.date(2020, 1, 2), None], type=pa.date32()), + pa.array([datetime.datetime(2020, 1, 2, 3, 4, 5), None], type=pa.timestamp("us")), + pa.array([decimal.Decimal("1.23"), None], type=pa.decimal128(10, 2)), + # lists of fast-path leaves + pa.array([[b"x", None], None, [b""]], type=pa.list_(pa.binary())), + pa.array([[True, None], [False]], type=pa.list_(pa.bool_())), + pa.array([[1.5, None], None], type=pa.list_(pa.float32())), + ] + 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..83ceeed1bd102 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) ] @@ -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: