Yicong-Huang commented on code in PR #58903:
URL: https://github.com/apache/spark/pull/58903#discussion_r4067145915
##########
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:
if we are modifying the methods, let's take the chance to fill in the type
annotations.
--
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]
