Spenserrrr commented on code in PR #58903:
URL: https://github.com/apache/spark/pull/58903#discussion_r4066728462
##########
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":
Review Comment:
Yeah. I also moved this logic local and separated the functionality.
_elementwise_renest_output rebuilds one output Array, and the worker extracts
each UDF output column and assembles the rebuilt Arrays into the final
RecordBatch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
