[SPARK-59455][PYTHON] Move Arrow-to-pandas batch conversion into ArrowToPandasConversion

python/benchmarks/bench_arrow_to_pandas.py

@@ -16,15 +16,15 @@
 #
 
 """
-Microbenchmarks for ``ArrowBatchTransformer.to_pandas``, the hot path of pandas
+Microbenchmarks for ``ArrowToPandasConversion.to_pandas``, the hot path of pandas
 UDF inputs: every pandas UDF eval type calls it once per batch to build the
 Series it passes to the user's function.
 
 Part of the per-batch cost is fixed per COLUMN and does not scale with row
 count, so ``n_cols`` is swept alongside ``n_rows``: wide batches and small
 batches are the shapes where that fixed cost dominates.
 
-``ArrowArrayToPandasConversion.convert`` routes each column by type. ``long`` and
+``ArrowToPandasConversion.convert`` routes each column by type. ``long`` and
 ``timestamp`` are in the ``_prefer_convert_numpy`` allowlist and take
 ``convert_numpy``; ``string`` is not, and takes ``convert_legacy``. The two
 allowlist types differ by an order of magnitude in conversion cost -- a timestamp
@@ -37,7 +37,7 @@
 
 
 class ArrowBatchToPandasBenchmark:
-    """Benchmark ``ArrowBatchTransformer.to_pandas`` over a whole RecordBatch."""
+    """Benchmark ``ArrowToPandasConversion.to_pandas`` over a whole RecordBatch."""
 
     params = [
         [128, 10000],
@@ -47,7 +47,7 @@ class ArrowBatchToPandasBenchmark:
     param_names = ["n_rows", "n_cols", "col_type"]
 
     def setup(self, n_rows, n_cols, col_type):
-        from pyspark.sql.conversion import ArrowBatchTransformer
+        from pyspark.sql.conversion import ArrowToPandasConversion
         from pyspark.sql.types import (
             LongType,
             StringType,
@@ -74,7 +74,7 @@ def setup(self, n_rows, n_cols, col_type):
         names = [f"c{i}" for i in range(n_cols)]
         self.batch = pa.RecordBatch.from_arrays([column] * n_cols, names)
         self.schema = StructType([StructField(name, spark_type) for name in names])
-        self.to_pandas = ArrowBatchTransformer.to_pandas
+        self.to_pandas = ArrowToPandasConversion.to_pandas
 
     def time_batch_to_pandas(self, n_rows, n_cols, col_type):
         self.to_pandas(self.batch, timezone="UTC", schema=self.schema)

python/pyspark/sql/conversion.py

@@ -74,7 +74,7 @@
 
 class ArrowBatchTransformer:
     """
-    Pure functions that transform RecordBatch -> RecordBatch.
+    Pure functions that transform Arrow RecordBatches and Tables.
     They should have no side effects (no I/O, no writing to streams).
     """
 
@@ -244,64 +244,6 @@ def enforce_schema(
             return pa.Table.from_arrays(coerced_arrays, names=output_names)
         return pa.RecordBatch.from_arrays(coerced_arrays, names=output_names)
 
-    @classmethod
-    def to_pandas(
-        cls,
-        batch: Union["pa.RecordBatch", "pa.Table"],
-        timezone: str,
-        schema: Optional["StructType"] = None,
-        struct_in_pandas: str = "dict",
-        ndarray_as_list: bool = False,
-        prefer_int_ext_dtype: bool = False,
-        df_for_struct: bool = False,
-    ) -> List[Union["pd.Series", "pd.DataFrame"]]:
-        """
-        Convert a RecordBatch or Table to a list of pandas Series.
-
-        Parameters
-        ----------
-        batch : pa.RecordBatch or pa.Table
-            The Arrow RecordBatch or Table to convert.
-        timezone : str
-            Timezone for timestamp conversion.
-        schema : StructType, optional
-            Spark schema for type conversion. If None, types are inferred from Arrow.
-        struct_in_pandas : str
-            How to represent struct in pandas ("dict", "row", etc.)
-        ndarray_as_list : bool
-            Whether to convert ndarray as list.
-        prefer_int_ext_dtype : bool, optional
-            Whether to convert integers to Pandas ExtensionDType.
-        df_for_struct : bool
-            If True, convert struct columns to DataFrame instead of Series.
-
-        Returns
-        -------
-        List[Union[pd.Series, pd.DataFrame]]
-            List of pandas Series (or DataFrame if df_for_struct=True), one for each column.
-        """
-        import pandas as pd
-
-        if batch.num_columns == 0:
-            return [pd.Series([pyspark._NoValue] * batch.num_rows)]
-
-        if schema is None:
-            schema = from_arrow_schema(batch.schema)
-
-        return [
-            ArrowArrayToPandasConversion.convert(
-                batch.column(i),
-                schema[i].dataType,
-                ser_name=schema[i].name,
-                timezone=timezone,
-                struct_in_pandas=struct_in_pandas,
-                ndarray_as_list=ndarray_as_list,
-                prefer_int_ext_dtype=prefer_int_ext_dtype,
-                df_for_struct=df_for_struct,
-            )
-            for i in range(batch.num_columns)
-        ]
-
 
 class PandasToArrowConversion:
     """
@@ -1750,16 +1692,68 @@ def convert_func(arr: pa.Array) -> pa.Array:
         )
 
 
-class ArrowArrayToPandasConversion:
+class ArrowToPandasConversion:
+    """
+    Conversion utilities from Arrow batches and arrays to pandas for UDF execution.
     """
-    Conversion utilities for converting PyArrow Arrays and ChunkedArrays to pandas.
 
-    This class provides methods to convert PyArrow columnar data structures to pandas
-    Series or DataFrames, with support for Spark-specific type handling and conversions.
+    @classmethod
+    def to_pandas(
+        cls,
+        batch: Union["pa.RecordBatch", "pa.Table"],
+        timezone: str,
+        schema: Optional["StructType"] = None,
+        struct_in_pandas: str = "dict",
+        ndarray_as_list: bool = False,
+        prefer_int_ext_dtype: bool = False,
+        df_for_struct: bool = False,
+    ) -> List[Union["pd.Series", "pd.DataFrame"]]:
+        """
+        Convert a RecordBatch or Table to a list of pandas Series.
 
-    The class is primarily used by PySpark's Arrow-based serializers for UDF execution,
-    where Arrow data needs to be converted to pandas for Python UDF processing.
-    """
+        Parameters
+        ----------
+        batch : pa.RecordBatch or pa.Table
+            The Arrow RecordBatch or Table to convert.
+        timezone : str
+            Timezone for timestamp conversion.
+        schema : StructType, optional
+            Spark schema for type conversion. If None, types are inferred from Arrow.
+        struct_in_pandas : str
+            How to represent struct in pandas ("dict", "row", etc.)
+        ndarray_as_list : bool
+            Whether to convert ndarray as list.
+        prefer_int_ext_dtype : bool, optional
+            Whether to convert integers to Pandas ExtensionDType.
+        df_for_struct : bool
+            If True, convert struct columns to DataFrame instead of Series.
+
+        Returns
+        -------
+        List[Union[pd.Series, pd.DataFrame]]
+            List of pandas Series (or DataFrame if df_for_struct=True), one for each column.
+        """
+        import pandas as pd
+
+        if batch.num_columns == 0:
+            return [pd.Series([pyspark._NoValue] * batch.num_rows)]
+
+        if schema is None:
+            schema = from_arrow_schema(batch.schema)
+
+        return [
+            cls.convert(
+                batch.column(i),
+                schema[i].dataType,
+                ser_name=schema[i].name,
+                timezone=timezone,
+                struct_in_pandas=struct_in_pandas,
+                ndarray_as_list=ndarray_as_list,
+                prefer_int_ext_dtype=prefer_int_ext_dtype,
+                df_for_struct=df_for_struct,
+            )
+            for i in range(batch.num_columns)
+        ]
 
     @classmethod
     def convert(

python/pyspark/sql/tests/test_conversion.py

@@ -23,9 +23,9 @@
 from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
 from pyspark.sql.conversion import (
     ArrowArrayConversion,
-    ArrowArrayToPandasConversion,
     ArrowBatchTransformer,
     ArrowTableToRowsConversion,
+    ArrowToPandasConversion,
     LocalDataToArrowConversion,
     PandasToArrowConversion,
 )
@@ -719,7 +719,7 @@ def test_arrow_array_localize_tz(self):
 
 
 @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
-class ArrowArrayToPandasConversionTests(unittest.TestCase):
+class ArrowToPandasConversionTests(unittest.TestCase):
     def test_convert_numpy_ser_name_survives_preprocess_time(self):
         # convert_numpy reads the Arrow field name before preprocess_time, because the
         # pa.compute kernels it runs for timestamps return a new array with no field name.
@@ -731,7 +731,7 @@ def test_convert_numpy_ser_name_survives_preprocess_time(self):
             )
             col = pa.RecordBatch.from_arrays([ts], ["tscol"]).column(0)
             spark_type = TimestampType() if pa_type.tz is not None else TimestampNTZType()
-            result = ArrowArrayToPandasConversion.convert_numpy(col, spark_type, timezone="UTC")
+            result = ArrowToPandasConversion.convert_numpy(col, spark_type, timezone="UTC")
             self.assertEqual(result.name, "tscol", f"name lost for {pa_type}")
 
     def test_udt_convert_numpy(self):
@@ -741,21 +741,21 @@ def test_udt_convert_numpy(self):
 
         # basic conversion with nulls
         arr = pa.array([[1.0, 2.0], None, [3.0, 4.0]], type=pa.list_(pa.float64()))
-        result = ArrowArrayToPandasConversion.convert_numpy(arr, udt, ser_name="my_point")
+        result = ArrowToPandasConversion.convert_numpy(arr, udt, ser_name="my_point")
         self.assertIsInstance(result.iloc[0], ExamplePoint)
         self.assertEqual(result.iloc[0], ExamplePoint(1.0, 2.0))
         self.assertIsNone(result.iloc[1])
         self.assertEqual(result.iloc[2], ExamplePoint(3.0, 4.0))
         self.assertEqual(result.name, "my_point")
 
         # empty
-        result = ArrowArrayToPandasConversion.convert_numpy(
+        result = ArrowToPandasConversion.convert_numpy(
             pa.array([], type=pa.list_(pa.float64())), udt
         )
         self.assertEqual(len(result), 0)
 
         # PythonOnlyUDT
-        result = ArrowArrayToPandasConversion.convert_numpy(
+        result = ArrowToPandasConversion.convert_numpy(
             pa.array([[5.0, 6.0]], type=pa.list_(pa.float64())), PythonOnlyUDT()
         )
         self.assertIsInstance(result.iloc[0], PythonOnlyPoint)
@@ -767,7 +767,7 @@ def test_udt_chunked_array(self):
         chunk1 = pa.array([[1.0, 2.0]], type=pa.list_(pa.float64()))
         chunk2 = pa.array([[3.0, 4.0]], type=pa.list_(pa.float64()))
         chunked = pa.chunked_array([chunk1, chunk2])
-        result = ArrowArrayToPandasConversion.convert_numpy(chunked, ExamplePointUDT())
+        result = ArrowToPandasConversion.convert_numpy(chunked, ExamplePointUDT())
         self.assertEqual(result.iloc[0], ExamplePoint(1.0, 2.0))
         self.assertEqual(result.iloc[1], ExamplePoint(3.0, 4.0))
 
@@ -790,7 +790,7 @@ def test_variant_convert_numpy(self):
             ],
             type=variant_type,
         )
-        result = ArrowArrayToPandasConversion.convert_numpy(arr, VariantType(), ser_name="v")
+        result = ArrowToPandasConversion.convert_numpy(arr, VariantType(), ser_name="v")
         self.assertIsInstance(result.iloc[0], VariantVal)
         self.assertEqual(result.iloc[0].value, b"\x01")
         self.assertEqual(result.iloc[0].metadata, b"\x02")
@@ -800,7 +800,7 @@ def test_variant_convert_numpy(self):
         self.assertEqual(result.name, "v")
 
         # empty
-        result = ArrowArrayToPandasConversion.convert_numpy(
+        result = ArrowToPandasConversion.convert_numpy(
             pa.array([], type=variant_type), VariantType()
         )
         self.assertEqual(len(result), 0)
@@ -832,14 +832,14 @@ def test_geography_convert_numpy(self):
             ],
             type=geography_type,
         )
-        result = ArrowArrayToPandasConversion.convert_numpy(arr, GeographyType(4326), ser_name="g")
+        result = ArrowToPandasConversion.convert_numpy(arr, GeographyType(4326), ser_name="g")
         self.assertEqual(result.iloc[0], Geography(wkb1, 4326))
         self.assertIsNone(result.iloc[1])
         self.assertEqual(result.iloc[2], Geography(wkb2, 4326))
         self.assertEqual(result.name, "g")
 
         # empty
-        result = ArrowArrayToPandasConversion.convert_numpy(
+        result = ArrowToPandasConversion.convert_numpy(
             pa.array([], type=geography_type), GeographyType(4326)
         )
         self.assertEqual(len(result), 0)
@@ -871,14 +871,14 @@ def test_geometry_convert_numpy(self):
             ],
             type=geometry_type,
         )
-        result = ArrowArrayToPandasConversion.convert_numpy(arr, GeometryType(0), ser_name="g")
+        result = ArrowToPandasConversion.convert_numpy(arr, GeometryType(0), ser_name="g")
         self.assertEqual(result.iloc[0], Geometry(wkb1, 0))
         self.assertIsNone(result.iloc[1])
         self.assertEqual(result.iloc[2], Geometry(wkb2, 0))
         self.assertEqual(result.name, "g")
 
         # empty
-        result = ArrowArrayToPandasConversion.convert_numpy(
+        result = ArrowToPandasConversion.convert_numpy(
             pa.array([], type=geometry_type), GeometryType(0)
         )
         self.assertEqual(len(result), 0)

python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py

@@ -239,7 +239,7 @@ class PyArrowArrayToPandasZeroCopyTests(_PyArrowToPandasTestBase):
     asks pandas to keep pointing at the Arrow buffers instead of materializing them
     into NumPy -- avoiding the copy is the point of that backend.  Its own golden file
     records the result, so the two can be read side by side.  PySpark takes this path
-    in ``ArrowArrayToPandasConversion.convert_numpy``
+    in ``ArrowToPandasConversion.convert_numpy``
     (``python/pyspark/sql/conversion.py``).
     """
 
@@ -479,7 +479,7 @@ class PyArrowArrayToPandasIntegerObjectNullsTests(_PyArrowToPandasTestBase):
     large value silently changes.  ``integer_object_nulls=True`` keeps ``object``
     dtype (Python ``int`` and ``None``) instead, preserving the values.
 
-    PySpark passes it in ``ArrowArrayToPandasConversion.convert_legacy``
+    PySpark passes it in ``ArrowToPandasConversion.convert_legacy``
     (``python/pyspark/sql/conversion.py``), bundled with ``date_as_object`` and
     ``coerce_temporal_nanoseconds``, then narrows the object Series to a nullable
     extension dtype (``Int8Dtype`` .. ``Int64Dtype``) -- the only bridge from Arrow to