Spenserrrr commented on code in PR #58903:
URL: https://github.com/apache/spark/pull/58903#discussion_r4067459522
##########
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:
Makes sense. I changed `concat_batches` to accept `Iterable` and materialize
it once inside the helper, so callers can pass iterators directly. I also
updated the test to cover an iterator input.
##########
python/pyspark/worker.py:
##########
@@ -1733,38 +1733,10 @@ def mapper(_, it):
return mapper, ser
-def _elementwise_renest(flat_values, shape_lengths, is_large):
- """Re-nest a flat Array of per-element results into an ``array<R>`` column.
-
- ``flat_values`` holds the results for every non-null element in order;
``shape_lengths`` is
- the per-array element count of the iterated argument (``None`` for a null
array, which stays
- null and consumes no elements). ``is_large`` preserves the input's list
width (``ListArray``
- with int32 offsets vs. ``LargeListArray`` with int64).
-
- Shared by the vectorized element-wise worker paths (scalar pandas / Arrow
and their iterator
- variants) that back Python UDFs inside higher-order function lambdas. See
- ``ExtractPythonUDFFromLambda``.
- """
- import pyarrow as pa
-
- offsets = [0]
- running = 0
- mask = []
- for n in shape_lengths:
- mask.append(n is None)
- if n is not None:
- running += n
- offsets.append(running)
- list_cls = pa.LargeListArray if is_large else pa.ListArray
- offsets_arr = pa.array(offsets, type=pa.int64() if is_large else
pa.int32())
- null_mask = pa.array(mask, type=pa.bool_())
- return list_cls.from_arrays(offsets_arr, flat_values, mask=null_mask)
-
-
-def _elementwise_leaf_type(data_type, depth):
+def _elementwise_udf_input_type(data_type, depth):
Review Comment:
Done. I added parameter and return type annotations to the modified
element-wise helpers and conversion adapters.
--
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]
