[SPARK-58023][PYTHON] Fast paths for converting Arrow string, binary, numeric and boolean columns to Python rows

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)

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)
             ]
 

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