spark-test commented on code in PR #59845:
URL: https://github.com/apache/spark/pull/59845#discussion_r4395510658
##########
python/pyspark/eval_handlers/_arrow.py:
##########
@@ -879,13 +879,18 @@ def run(self, split_index: int, data:
Iterator[pa.RecordBatch]) -> Iterator[pa.R
phase_timers = None
@fail_on_stopiteration
- def _evaluate_batch_udf(udf_func: Any, rows: list[tuple]) -> list[Any]:
+ def _evaluate_batch_udf(udf_func: Any, args: list[list[Any]],
num_rows: int) -> list[Any]:
if runner_conf.arrow_concurrency_level <= 0:
- return [udf_func(*row) for row in rows]
+ if not args:
+ return [udf_func() for _ in range(num_rows)]
+ # Feed the columns straight to map() to avoid building a tuple
per row.
+ return list(map(udf_func, *args))
Review Comment:
**Non-blocking (P2):** `list(map(udf_func, *args))` calls the UDF inside
`map.__next__`. When the UDF raises StopIteration, `map` passes it on and
`list()` takes it as the end of the input, so it returns the rows computed so
far without raising. This handler sets `_fail_on_stopiteration = False` and
relies on the `@fail_on_stopiteration` around `_evaluate_batch_udf`, which
never sees the exception. `verify_result_row_count` then fails the task with
`RESULT_ROWS_MISMATCH` ("The number of output rows (k) must match the number of
input rows (n)."). The user's exception text and any mention of StopIteration
are gone. Before this change, the list comprehension let the exception
propagate, and users got `STOP_ITERATION_OCCURRED` ("Caught StopIteration
thrown from user's code; failing the task: ..."). I checked the
`list(map(...))` behavior on CPython 3.12.16 with standalone snippets.
This hits every Arrow-optimized `@udf` with at least one argument in the
default configuration: Arrow on, non-legacy conversion,
`arrow.concurrency.level` unset. The task still fails, so no rows are lost, but
the error now blames the engine. A common trigger is calling `next()` on an
exhausted iterator inside a UDF. The zero-argument branch and the thread-pool
path behave as before. The legacy pandas-conversion handler and
`ArrowElementwiseUDFHandler` still report `STOP_ITERATION_OCCURRED`, so the
same UDF now fails differently depending on the legacy flag.
`test_stop_iteration_raises_runtime_error` only checks
`assertRaises(PySparkRuntimeError)`, and `RESULT_ROWS_MISMATCH` satisfies that,
so it stays green. The inherited end-to-end `test_raise_stop_iteration` expects
"StopIteration" in the `PythonException`. It should fail in the non-legacy
Arrow UDF suites (classic and Connect parity); I did not run it.
Calling the UDF in a comprehension body keeps the exception visible to
`fail_on_stopiteration`, for example `[udf_func(v) for v in args[0]]` for one
argument and `[udf_func(*row) for row in zip(*args)]` otherwise. Neither builds
a materialized list of row tuples. In a quick standalone timing on CPython
3.12.16 (10,000 rows, trivial UDF), the one-argument comprehension was slightly
faster than `map()`. The three-argument `zip()` form kept about two-thirds of
the saving over the old code, and like `map()` it triggered no GC collections.
It would also help to tighten `test_stop_iteration_raises_runtime_error` to
check the error condition, as the elementwise handler test does.
**Verification:**
- **Regression:** With the concurrency level unset, a one-argument UDF and a
multi-argument UDF that raise StopIteration partway through a batch each fail
with error condition STOP_ITERATION_OCCURRED, and the message includes the
user's exception text.
- **Behavior:** Default-path results for one-argument, multi-argument,
keyword-argument, and zero-argument UDFs keep the same values and order across
multiple batches, and the thread-pool path keeps its current results for one-,
multi-, and zero-argument UDFs.
- **Compatibility:** With Arrow-optimized Python UDFs on and the default
non-legacy conversion, a UDF that raises StopIteration for some rows fails with
a Python worker error that names StopIteration, as the legacy pandas-conversion
variant does.
- **Inspection:** The default serial branch calls the UDF directly in a
comprehension body rather than inside an iterator consumer such as
`list(map(...))`, and it holds no materialized list of per-row argument tuples.
--
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]
