dongjoon-hyun commented on PR #57899:
URL: https://github.com/apache/spark/pull/57899#issuecomment-5247035382
Reviewed the change. The design faithfully extends the 102 path: every place
that enumerates `SQL_ARROW_ELEMENTWISE_UDF` (`PythonRunner`,
`PythonUDF.SCALAR_TYPES`, `ArrowEvalPythonExec`, `ArrowPythonRunner.evalConf`,
`ExtractPythonUDFs`, the `worker.py` serializer / `read_single_udf` branches,
`util.py`, `_typing.pyi`) is covered with no omissions, the worker-side
verification mirrors the base vectorized paths (`__len__` / struct-DataFrame
checks, `OUTPUT_EXCEEDS_INPUT_ROWS` fail-fast, final row-count match,
`verify_iterator_exhausted`), and the fusion constraints for the iterator
flavors match the base iterator types. Nice catch on taking the iterator buffer
type from the first chunk to avoid the session-timezone `pa.concat_arrays`
mismatch.
One correctness edge case and a few minor points:
### 1. Iterator path can silently drop rows when a whole partition has only
null/empty arrays
In the new `SQL_SCALAR_*_ITER_ELEMENTWISE_UDF` path in `worker.py`:
```python
def emit_ready():
nonlocal buffered
while pending_shapes:
lengths = pending_shapes[0]
needed = sum(n for n in lengths if n is not None)
if buffered is None or len(buffered) < needed:
break
```
`buffered` stays `None` until the UDF yields its first chunk. If every row
in a partition has a null or empty array, the flattened inputs are all
zero-length, and a common defensive iterator UDF such as
```python
def f(it):
for s in it:
if len(s) == 0:
continue
yield s + 1
```
yields no chunks at all. Then `num_output_elements == num_input_elements ==
0`, so `verify_result_row_count` passes, but the final `if pending_shapes:
yield from emit_ready()` breaks immediately on `buffered is None` even though
`needed == 0`, and no output batch is emitted. On the JVM side the join is
driven by the output iterator (`outputRowIterator.map {
resultProj(joined(queue.remove(), outputRow)) }` in
`EvalPythonEvaluatorFactory`), so the unmatched input rows are dropped
**silently** — wrong results rather than an error.
A narrow fix in `emit_ready`:
```python
if needed and (buffered is None or len(buffered) < needed):
break
flat = pa.nulls(0, type=arrow_element_type) if buffered is None else
buffered.slice(0, needed)
buffered = None if buffered is None else buffered.slice(needed)
```
(A stream that reaches this state has never emitted a non-empty batch, so
the UTC-typed empty array cannot cause a mixed stream schema.) The adjacent
comment "Any residual shapes are all-null arrays" is also slightly off —
residual shapes can be empty (non-null) arrays too; those happen to work today,
the `buffered is None` case is the one that doesn't. A regression test with
`transform(values, x -> iter_udf(x))` over a partition of only `[]`/`NULL`
arrays and a skip-empty iterator UDF would pin this down.
### Minor
- **Repeated `pa.concat_arrays`**: if the UDF yields many small chunks,
`buffered = pa.concat_arrays([buffered, chunk])` re-copies the whole buffer
each time (worst case O(n^2)). Keeping a list of chunks plus a running length
and concatenating once at emit time would avoid it. Low priority since typical
UDFs yield once per input batch.
- **`_elementwise_renest` second return value**: both call sites discard
`running` (`nested, _ = ...`); the non-iter path computes `total_elements`
separately and the iter path guarantees the length via slicing. The docstring
says it exists "so callers can verify the flat result length", but no caller
does — it could return just the nested array.
- **Test message nit**: `s"UDF still inside a lambda for eval type
${udf.func}"` in `ExtractPythonUDFFromLambdaSuite` prints the function object,
not an eval type.
- **Test gaps**: no HOF test for a struct-element return type (pandas UDF
returning a `pd.DataFrame`), and the timestamp test only asserts iter ==
non-iter consistency, so a timezone bug common to both paths would go unnoticed.
Also worth noting (not a regression — inherent to the 102 design too):
flattening amplifies batch size, so a UDF receives `rows x avg array length`
elements in one Series/Array and `arrowMaxRecordsPerBatch` no longer bounds the
element count.
Overall LGTM once item 1 is addressed.
--
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]
