[SPARK-59624][PYTHON] Move element-wise UDF conversion behind Arrow batch interfaces

python/pyspark/sql/conversion.py

@@ -180,6 +180,115 @@ def select_columns(cls, batch: "pa.RecordBatch", column_indices: list[int]) -> "
             [batch.schema.names[i] for i in column_indices],
         )
 
+    @staticmethod
+    def concat_batches(batches: Sequence["pa.RecordBatch"]) -> "pa.RecordBatch":
+        """Concatenate same-schema RecordBatches by row.
+
+        A single batch is returned unchanged. PyArrow before 19.0.0 has no ``concat_batches``;
+        the fallback concatenates the equivalent StructArrays and converts the result back to a
+        RecordBatch. Element-wise iterator UDFs use this when one input batch's flattened result
+        spans multiple output chunks.
+        """
+        import pyarrow as pa
+
+        assert batches
+        if len(batches) == 1:
+            return batches[0]
+        if hasattr(pa, "concat_batches"):
+            return pa.concat_batches(batches)
+        return pa.RecordBatch.from_struct_array(
+            pa.concat_arrays([batch.to_struct_array() for batch in batches])
+        )
+
+    @staticmethod
+    def flatten_elementwise_inputs(
+        batch: "pa.RecordBatch", input_column_indices: Sequence[int], depth: int
+    ) -> tuple["pa.RecordBatch", list[list[Optional[int]]], list[bool]]:
+        """Flatten ``depth`` list levels from an element-wise UDF's input columns.
+
+        Returns ``(flat_input_batch, shape_levels, is_large_levels)``. ``flat_input_batch``
+        contains each selected input's fully flattened leaf Array under a positional ``_N`` name.
+        ``shape_levels[k]`` contains the per-slot list length at level ``k`` (0 is outermost),
+        using ``None`` for a null list. ``is_large_levels[k]`` records whether that level uses
+        ``LargeListArray`` and therefore requires int64 rather than int32 offsets when rebuilt.
+
+        Only the first selected column supplies shape and list-width metadata. The other inputs are
+        aligned to it by ``ExtractPythonUDFFromLambda``, so recording their shapes would repeat
+        the ``list_value_length(...).to_pylist()`` work without changing re-nesting. ``depth`` is 1
+        for a UDF in one higher-order-function lambda and greater for nested lambdas.
+
+        Shared by the row, scalar pandas / Arrow, and iterator element-wise worker paths. See
+        ``ExtractPythonUDFFromLambda``.
+        """
+        import pyarrow as pa
+        import pyarrow.compute as pc
+
+        assert input_column_indices
+        assert depth > 0
+
+        flat_inputs = []
+        shape_levels = []
+        is_large_levels = []
+        for input_index, column_index in enumerate(input_column_indices):
+            current = batch.column(column_index)
+            for _ in range(depth):
+                if input_index == 0:
+                    shape_levels.append(pc.list_value_length(current).to_pylist())
+                    is_large_levels.append(pa.types.is_large_list(current.type))
+                current = current.flatten()
+            flat_inputs.append(current)
+
+        return (
+            pa.RecordBatch.from_arrays(
+                flat_inputs, names=[f"_{index}" for index in range(len(flat_inputs))]
+            ),
+            shape_levels,
+            is_large_levels,
+        )
+
+    @staticmethod
+    def renest_elementwise_outputs(
+        flat_outputs: Sequence[tuple["pa.RecordBatch", list[list[Optional[int]]], list[bool]]],
+        column_names: Sequence[str],
+    ) -> "pa.RecordBatch":
+        """Rebuild nested list columns from flattened element-wise UDF result batches.
+
+        Each input tuple contains a one-column flat result batch plus the ``shape_levels`` and
+        ``is_large_levels`` returned by ``flatten_elementwise_inputs`` for that UDF. Levels are
+        rebuilt from innermost to outermost. A ``None`` length creates a null list and consumes no
+        flat values; a zero length creates an empty, non-null list. ``is_large_levels`` preserves
+        each input level's int32 ``ListArray`` versus int64 ``LargeListArray`` offset width.
+
+        Different fused UDFs may carry different shapes, so every result batch is rebuilt with its
+        own metadata before the columns are assembled into one output RecordBatch. This is the
+        batch-level inverse of ``flatten_elementwise_inputs``.
+        """
+        import pyarrow as pa
+
+        assert len(flat_outputs) == len(column_names)
+        nested_columns = []
+        for flat_batch, shape_levels, is_large_levels in flat_outputs:
+            assert flat_batch.num_columns == 1
+            result = flat_batch.column(0)
+            for shape_lengths, is_large in zip(reversed(shape_levels), reversed(is_large_levels)):
+                offsets = [0]
+                running = 0
+                nulls = []
+                for length in shape_lengths:
+                    nulls.append(length is None)
+                    if length is not None:
+                        running += length
+                    offsets.append(running)
+                list_type = pa.LargeListArray if is_large else pa.ListArray
+                result = list_type.from_arrays(
+                    pa.array(offsets, type=pa.int64() if is_large else pa.int32()),
+                    result,
+                    mask=pa.array(nulls, type=pa.bool_()),
+                )
+            nested_columns.append(result)
+
+        return pa.RecordBatch.from_arrays(nested_columns, names=column_names)
+
     @staticmethod
     def wrap_struct(batch: "pa.RecordBatch") -> "pa.RecordBatch":
         """

python/pyspark/sql/tests/test_conversion.py

@@ -115,6 +115,58 @@ def test_flatten_struct_empty_batch(self):
         self.assertEqual(flattened.num_rows, 0)
         self.assertEqual(flattened.num_columns, 2)
 
+    def test_concat_batches(self):
+        import pyarrow as pa
+
+        batches = [
+            pa.RecordBatch.from_arrays([pa.array([1, 2])], ["x"]),
+            pa.RecordBatch.from_arrays([pa.array([3])], ["x"]),
+        ]
+        result = ArrowBatchTransformer.concat_batches(batches)
+        self.assertEqual(result.column(0).to_pylist(), [1, 2, 3])
+        self.assertIs(ArrowBatchTransformer.concat_batches(batches[:1]), batches[0])
+
+    def test_flatten_elementwise_inputs_and_renest_outputs(self):
+        import pyarrow as pa
+
+        int_values = pa.array(
+            [[[1, 2], []], None, [[3], None]],
+            type=pa.large_list(pa.list_(pa.int64())),
+        )
+        string_values = pa.array(
+            [[["a", "b"], []], None, [["c"], None]],
+            type=pa.large_list(pa.list_(pa.string())),
+        )
+        batch = pa.RecordBatch.from_arrays([int_values, string_values], ["ints", "strings"])
+
+        flat, shape_levels, is_large_levels = ArrowBatchTransformer.flatten_elementwise_inputs(
+            batch, [0, 1], depth=2
+        )
+        self.assertEqual(flat.schema.names, ["_0", "_1"])
+        self.assertEqual(flat.column(0).to_pylist(), [1, 2, 3])
+        self.assertEqual(flat.column(1).to_pylist(), ["a", "b", "c"])
+        self.assertEqual(shape_levels, [[2, None, 2], [2, 0, 1, None]])
+        self.assertEqual(is_large_levels, [True, False])
+
+        restored = ArrowBatchTransformer.renest_elementwise_outputs(
+            [
+                (
+                    pa.RecordBatch.from_arrays([flat.column(0)], ["_0"]),
+                    shape_levels,
+                    is_large_levels,
+                ),
+                (
+                    pa.RecordBatch.from_arrays([flat.column(1)], ["_0"]),
+                    shape_levels,
+                    is_large_levels,
+                ),
+            ],
+            ["ints", "strings"],
+        )
+        self.assertEqual(restored.schema.names, ["ints", "strings"])
+        self.assertTrue(restored.column(0).equals(int_values))
+        self.assertTrue(restored.column(1).equals(string_values))
+
     def test_wrap_struct_basic(self):
         """Test wrapping columns into a struct."""
         import pyarrow as pa