Spenserrrr commented on code in PR #58903:
URL: https://github.com/apache/spark/pull/58903#discussion_r4066719302
##########
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(
Review Comment:
I don’t expect this transformation to be reused by other eval types. It
depends on element-wise nesting depth, aligned inputs, and first-input shape
metadata. I moved it out of the shared transformer and kept it local as
_elementwise_flatten_inputs; flattened pandas inputs still go through the
batch-level to_pandas API.
--
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]
