[SPARK-55336][PYTHON] Let createDF use create_batch logic for decoupling

Yicong-Huang · HyukjinKwon · commit b91d4079036a · 2026-02-04T07:32:39.000+09:00
### What changes were proposed in this pull request? This PR duplicates the pandas-to-Arrow batch conversion logic in `ArrowStreamPandasSerializer` to decouple it. - `create_arrow_array_from_pandas()` - converts a pandas Series to Arrow Array - `create_arrow_batch_from_pandas()` - converts a list of (series, spark_type) tuples to Arrow RecordBatch Both `_create_from_pandas_with_arrow` (classic Spark) and `createDataFrame` (Spark Connect) now use these standalone functions directly with `ArrowStreamSerializer`, instead of depending on `ArrowStreamPandasSerializer`. ### Why are the changes needed? For better decoupling. Previously, `createDataFrame` had to instantiate `ArrowStreamPandasSerializer` just to call its `_create_batch` method. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #54111 from Yicong-Huang/SPARK-55336/refactor/factor-out-create-batch-logic. Authored-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py
@@ -69,7 +69,7 @@
 from pyspark.sql.connect.readwriter import DataFrameReader
 from pyspark.sql.connect.streaming.readwriter import DataStreamReader
 from pyspark.sql.connect.streaming.query import StreamingQueryManager
-from pyspark.sql.pandas.serializers import ArrowStreamPandasSerializer
+from pyspark.sql.pandas.conversion import create_arrow_batch_from_pandas
 from pyspark.sql.pandas.types import (
     to_arrow_schema,
     _deduplicate_field_names,
@@ -621,17 +621,12 @@ def createDataFrame(
 
             safecheck = configs["spark.sql.execution.pandas.convertToArrowArraySafely"]
 
-            ser = ArrowStreamPandasSerializer(
-                timezone=cast(str, timezone),
-                safecheck=safecheck == "true",
-                int_to_decimal_coercion_enabled=False,
-                prefers_large_types=prefers_large_types,
-            )
-
             _table = pa.Table.from_batches(
                 [
-                    ser._create_batch(
+                    create_arrow_batch_from_pandas(
                         [(c, st) for (_, c), st in zip(data.items(), spark_types)],
+                        timezone=cast(str, timezone),
+                        safecheck=safecheck == "true",
                         prefers_large_types=prefers_large_types,
                     )
                 ]
diff --git a/python/pyspark/sql/pandas/conversion.py b/python/pyspark/sql/pandas/conversion.py
@@ -18,9 +18,11 @@
 from typing import (
     Any,
     Callable,
+    Iterable,
     List,
     Optional,
     Sequence,
+    Tuple,
     Union,
     cast,
     no_type_check,
@@ -48,13 +50,132 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    import pandas as pd
     import pyarrow as pa
     from py4j.java_gateway import JavaObject
 
     from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike
     from pyspark.sql import DataFrame
 
 
+def create_arrow_array_from_pandas(
+    series: "pd.Series",
+    spark_type: Optional[DataType],
+    *,
+    timezone: Optional[str] = None,
+    safecheck: bool = False,
+    prefers_large_types: bool = False,
+) -> "pa.Array":
+    """
+    Create an Arrow Array from the given pandas.Series and Spark type.
+
+    Parameters
+    ----------
+    series : pandas.Series
+        A single series
+    spark_type : DataType, optional
+        The Spark return type. If None, pyarrow's inferred type will be used.
+    timezone : str, optional
+        The timezone to use for timestamp conversions.
+    safecheck : bool, optional
+        Whether to enable safe type checking during conversion.
+    prefers_large_types : bool, optional
+        Whether to prefer large Arrow types (e.g., large_string instead of string).
+
+    Returns
+    -------
+    pyarrow.Array
+    """
+    import pyarrow as pa
+    import pandas as pd
+    from pyspark.sql.pandas.types import to_arrow_type, _create_converter_from_pandas
+
+    if isinstance(series.dtype, pd.CategoricalDtype):
+        series = series.astype(series.dtype.categories.dtype)
+
+    # Derive arrow_type from spark_type
+    arrow_type = (
+        to_arrow_type(spark_type, timezone=timezone, prefers_large_types=prefers_large_types)
+        if spark_type is not None
+        else None
+    )
+
+    if spark_type is not None:
+        conv = _create_converter_from_pandas(
+            spark_type,
+            timezone=timezone,
+            error_on_duplicated_field_names=False,
+        )
+        series = conv(series)
+
+    if hasattr(series.array, "__arrow_array__"):
+        mask = None
+    else:
+        mask = series.isnull()
+    try:
+        return pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck)
+    except TypeError as e:
+        error_msg = (
+            "Exception thrown when converting pandas.Series (%s) "
+            "with name '%s' to Arrow Array (%s)."
+        )
+        raise PySparkTypeError(error_msg % (series.dtype, series.name, arrow_type)) from e
+    except ValueError as e:
+        error_msg = (
+            "Exception thrown when converting pandas.Series (%s) "
+            "with name '%s' to Arrow Array (%s)."
+        )
+        if safecheck:
+            error_msg = error_msg + (
+                " It can be caused by overflows or other "
+                "unsafe conversions warned by Arrow. Arrow safe type check "
+                "can be disabled by using SQL config "
+                "`spark.sql.execution.pandas.convertToArrowArraySafely`."
+            )
+        raise PySparkValueError(error_msg % (series.dtype, series.name, arrow_type)) from e
+
+
+def create_arrow_batch_from_pandas(
+    series_with_types: Iterable[Tuple["pd.Series", Optional[DataType]]],
+    *,
+    timezone: Optional[str] = None,
+    safecheck: bool = False,
+    prefers_large_types: bool = False,
+) -> "pa.RecordBatch":
+    """
+    Create an Arrow record batch from the given iterable of (series, spark_type) tuples.
+
+    Parameters
+    ----------
+    series_with_types : iterable
+        Iterable of (series, spark_type) tuples.
+    timezone : str, optional
+        The timezone to use for timestamp conversions.
+    safecheck : bool, optional
+        Whether to enable safe type checking during conversion.
+    prefers_large_types : bool, optional
+        Whether to prefer large Arrow types (e.g., large_string instead of string).
+
+    Returns
+    -------
+    pyarrow.RecordBatch
+        Arrow RecordBatch
+    """
+    import pyarrow as pa
+
+    arrs = [
+        create_arrow_array_from_pandas(
+            s,
+            spark_type,
+            timezone=timezone,
+            safecheck=safecheck,
+            prefers_large_types=prefers_large_types,
+        )
+        for s, spark_type in series_with_types
+    ]
+    return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))])
+
+
 def _convert_arrow_table_to_pandas(
     arrow_table: "pa.Table",
     schema: "StructType",
@@ -807,7 +928,7 @@ def _create_from_pandas_with_arrow(
 
         assert isinstance(self, SparkSession)
 
-        from pyspark.sql.pandas.serializers import ArrowStreamPandasSerializer
+        from pyspark.sql.pandas.serializers import ArrowStreamSerializer
         from pyspark.sql.types import TimestampType
         from pyspark.sql.pandas.types import (
             from_arrow_type,
@@ -877,20 +998,20 @@ def _create_from_pandas_with_arrow(
         step = step if step > 0 else len(pdf)
         pdf_slices = (pdf.iloc[start : start + step] for start in range(0, len(pdf), step))
 
-        # Create list of (columns, spark_type) for serializer dump_stream
-        arrow_data = [
-            [(c, t) for (_, c), t in zip(pdf_slice.items(), spark_types)]
+        # Create Arrow batches directly using the standalone function
+        arrow_batches = [
+            create_arrow_batch_from_pandas(
+                [(c, t) for (_, c), t in zip(pdf_slice.items(), spark_types)],
+                timezone=timezone,
+                safecheck=safecheck,
+                prefers_large_types=prefers_large_var_types,
+            )
             for pdf_slice in pdf_slices
         ]
 
         jsparkSession = self._jsparkSession
 
-        ser = ArrowStreamPandasSerializer(
-            timezone=timezone,
-            safecheck=safecheck,
-            int_to_decimal_coercion_enabled=False,
-            prefers_large_types=prefers_large_var_types,
-        )
+        ser = ArrowStreamSerializer()
 
         @no_type_check
         def reader_func(temp_filename):
@@ -901,7 +1022,7 @@ def create_iter_server():
             return self._jvm.ArrowIteratorServer()
 
         # Create Spark DataFrame from Arrow stream file, using one batch per partition
-        jiter = self._sc._serialize_to_jvm(arrow_data, ser, reader_func, create_iter_server)
+        jiter = self._sc._serialize_to_jvm(arrow_batches, ser, reader_func, create_iter_server)
         assert self._jvm is not None
         jdf = self._jvm.PythonSQLUtils.toDataFrame(jiter, schema.json(), jsparkSession)
         df = DataFrame(jdf, self)
