[SPARK-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace manual type coercion logic

### What changes were proposed in this pull request?

Replace manual column-by-column type coercion with `ArrowBatchTransformer.enforce_schema` in three places:

1. `ArrowStreamArrowUDTFSerializer.apply_type_coercion` in serializers.py
2. `ArrowStreamArrowUDFSerializer.create_batch` in serializers.py
3. `process_results` in worker.py (scalar Arrow iter UDF path)

Also:
- Add `arrow_cast` parameter to `enforce_schema` for strict type matching mode
- Add `KeyError` handling in `enforce_schema` for missing columns with user-friendly error
- Remove now-unused `coerce_arrow_array` imports from serializers.py and worker.py

### Why are the changes needed?

These three places duplicated the same coerce-and-reassemble logic that `enforce_schema` already provides. Consolidating reduces code duplication and ensures consistent error handling.

### Does this PR introduce _any_ user-facing change?

Error messages for type/schema mismatches in Arrow UDTFs are slightly changed to be consistent with other Arrow UDF error messages.

### How was this patch tested?

Existing tests in `test_arrow_udtf.py` and `test_arrow_udf_scalar.py`.

### Was this patch authored or co-authored using generative AI tooling?

Yes.

Author: Yicong-Huang

python/pyspark/sql/conversion.py

@@ -113,6 +113,7 @@ def enforce_schema(
         batch: "pa.RecordBatch",
         arrow_schema: "pa.Schema",
         safecheck: bool = True,
+        arrow_cast: bool = True,
     ) -> "pa.RecordBatch":
         """
         Enforce target schema on a RecordBatch by reordering columns and coercing types.
@@ -126,6 +127,9 @@ def enforce_schema(
             to_arrow_schema() to avoid repeated conversion.
         safecheck : bool, default True
             If True, use safe casting (fails on overflow/truncation).
+        arrow_cast : bool, default True
+            If True, cast mismatched types to the target type.
+            If False, raise an error on type mismatch instead of casting.
 
         Returns
         -------
@@ -149,8 +153,19 @@ def enforce_schema(
 
         coerced_arrays = []
         for i, field in enumerate(arrow_schema):
-            arr = batch.column(i) if use_index else batch.column(field.name)
+            try:
+                arr = batch.column(i) if use_index else batch.column(field.name)
+            except KeyError:
+                raise PySparkTypeError(
+                    f"Result column '{field.name}' does not exist in the output. "
+                    f"Expected schema: {arrow_schema}, got: {batch.schema}."
+                )
             if arr.type != field.type:
+                if not arrow_cast:
+                    raise PySparkTypeError(
+                        f"Result type of column '{field.name}' does not match "
+                        f"the expected type. Expected: {field.type}, got: {arr.type}."
+                    )
                 try:
                     arr = arr.cast(target_type=field.type, safe=safecheck)
                 except (pa.ArrowInvalid, pa.ArrowTypeError) as e:

python/pyspark/sql/pandas/serializers.py

@@ -23,7 +23,7 @@
 from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple
 
 import pyspark
-from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
+from pyspark.errors import PySparkRuntimeError, PySparkValueError
 from pyspark.serializers import (
     Serializer,
     read_int,
@@ -37,7 +37,6 @@
     ArrowTableToRowsConversion,
     ArrowBatchTransformer,
     PandasToArrowConversion,
-    coerce_arrow_array,
 )
 from pyspark.sql.pandas.types import (
     from_arrow_schema,
@@ -299,40 +298,9 @@ def apply_type_coercion():
                 assert isinstance(arrow_return_type, pa.StructType), (
                     f"Expected pa.StructType, got {type(arrow_return_type)}"
                 )
-
-                # Handle empty struct case specially
-                if batch.num_columns == 0:
-                    coerced_batch = batch  # skip type coercion
-                else:
-                    expected_field_names = [field.name for field in arrow_return_type]
-                    actual_field_names = batch.schema.names
-
-                    if expected_field_names != actual_field_names:
-                        raise PySparkTypeError(
-                            "Target schema's field names are not matching the record batch's "
-                            "field names. "
-                            f"Expected: {expected_field_names}, but got: {actual_field_names}."
-                        )
-
-                    coerced_arrays = []
-                    for i, field in enumerate(arrow_return_type):
-                        try:
-                            coerced_arrays.append(
-                                coerce_arrow_array(
-                                    batch.column(i),
-                                    field.type,
-                                    safecheck=True,
-                                )
-                            )
-                        except (pa.ArrowInvalid, pa.ArrowTypeError) as e:
-                            raise PySparkTypeError(
-                                f"Result type of column '{field.name}' does not "
-                                f"match the expected type. Expected: {field.type}, "
-                                f"got: {batch.column(i).type}."
-                            ) from e
-                    coerced_batch = pa.RecordBatch.from_arrays(
-                        coerced_arrays, names=expected_field_names
-                    )
+                coerced_batch = ArrowBatchTransformer.enforce_schema(
+                    batch, pa.schema(arrow_return_type), safecheck=True
+                )
                 yield coerced_batch, arrow_return_type
 
         return super().dump_stream(apply_type_coercion(), stream)
@@ -615,13 +583,15 @@ def dump_stream(self, iterator, stream):
         def create_batch(
             arr_tuples: List[Tuple["pa.Array", "pa.DataType"]],
         ) -> "pa.RecordBatch":
-            arrs = [
-                coerce_arrow_array(
-                    arr, arrow_type, safecheck=self._safecheck, arrow_cast=self._arrow_cast
-                )
-                for arr, arrow_type in arr_tuples
-            ]
-            return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))])
+            names = ["_%d" % i for i in range(len(arr_tuples))]
+            arrs = [arr for arr, _ in arr_tuples]
+            batch = pa.RecordBatch.from_arrays(arrs, names)
+            target_schema = pa.schema(
+                [pa.field(n, t) for n, (_, t) in zip(names, arr_tuples)]
+            )
+            return ArrowBatchTransformer.enforce_schema(
+                batch, target_schema, safecheck=self._safecheck, arrow_cast=self._arrow_cast
+            )
 
         def normalize(packed):
             if len(packed) == 2 and isinstance(packed[1], pa.DataType):

python/pyspark/sql/tests/arrow/test_arrow_udtf.py

@@ -211,7 +211,9 @@ def eval(self) -> Iterator["pa.Table"]:
 
         with self.assertRaisesRegex(
             PythonException,
-            "Target schema's field names are not matching the record batch's field names",
+            r"(?s)Result column 'x' does not exist in the output\. "
+            r"Expected schema: x: int32\ny: string, "
+            r"got: wrong_col: int32\nanother_wrong_col: double\.",
         ):
             result_df = MismatchedSchemaUDTF()
             result_df.collect()
@@ -373,8 +375,8 @@ def eval(self) -> Iterator["pa.Table"]:
         # Should fail with Arrow cast exception since string cannot be cast to int
         with self.assertRaisesRegex(
             PythonException,
-            "PySparkTypeError: Result type of column 'id' does not match the expected type. "
-            "Expected: int32, got: string.",
+            "Result type of column 'id' does not match "
+            "the expected type. Expected: int32, got: string.",
         ):
             result_df = StringToIntUDTF()
             result_df.collect()

python/pyspark/worker.py

@@ -51,7 +51,6 @@
     LocalDataToArrowConversion,
     ArrowTableToRowsConversion,
     ArrowBatchTransformer,
-    coerce_arrow_array,
 )
 from pyspark.sql.functions import SkipRestOfInputTableException
 from pyspark.sql.pandas.serializers import (
@@ -2909,13 +2908,15 @@ def extract_args(batch: pa.RecordBatch):
             # Call UDF and verify result type (iterator of pa.Array)
             verified_iter = verify_result(pa.Array)(udf_func(args_iter))
 
-            # Process results: coerce type and assemble into RecordBatch
+            # Process results: enforce schema and assemble into RecordBatch
+            target_schema = pa.schema([pa.field("_0", arrow_return_type)])
+
             def process_results():
                 for result in verified_iter:
-                    result = coerce_arrow_array(
-                        result, arrow_return_type, safecheck=True, arrow_cast=True
+                    batch = pa.RecordBatch.from_arrays([result], ["_0"])
+                    yield ArrowBatchTransformer.enforce_schema(
+                        batch, target_schema, safecheck=True
                     )
-                    yield pa.RecordBatch.from_arrays([result], ["_0"])
 
             # Apply row limit check (fail-fast)
             # TODO(SPARK-55579): Create Arrow-specific error class (e.g., ARROW_UDF_OUTPUT_EXCEEDS_INPUT_ROWS)