[SPARK-58050][PYTHON] Bulk-assemble Arrow Python UDF results

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

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.
@@ -895,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()
@@ -980,6 +1123,123 @@ 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, 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
+        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
+
+        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_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
+            # 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)
+            ]
+
+        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
     def _need_converter(dataType: DataType) -> bool:
         if isinstance(dataType, NullType):
@@ -1306,7 +1566,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)
             ]