Yicong-Huang commented on code in PR #58903:
URL: https://github.com/apache/spark/pull/58903#discussion_r4063947904
##########
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
Review Comment:
let's keep the comment's scope about the method itself to make it general
enough. the call site usage information (e.g., Element-wise iterator UDFs)
should be omitted IMO.
##########
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":
Review Comment:
this is a question not a required change: do you think if allowing
`Iterable` is better? Iterable could be Sequence or Iterator.
##########
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:
ditto.
Also I feel there are a few operations bundled/fused together in this method?
##########
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:
hmm I am a bit on the fence for this one. from the signature it seems to be
a specific method for a specific eval type (elementwise input). If we were to
extract a helper method in this shared utility `ArrowBatchTransformer`, the
naming and purpose should be general so that other eval types could potentially
reuse it. So two questions
1. will it likely to be reused by other eval types?
2. if the answer above is yes, can we revise the function name to be general
and describe what it does only?
If not, I would prefer to keep this method local to the element wise eval
type handlers.
##########
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
Review Comment:
let's make them class method?
--
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]
