Yicong-Huang opened a new pull request, #60115: URL: https://github.com/apache/spark/pull/60115
### What changes were proposed in this pull request? Simplify the pipelined UDF mode (SPARK-56642) in `invoke_udf`: - Keep one processing path. `process()` and `pipelined_process()` duplicated `func` -> `dump_stream` -> `out_iter.close()`; now pipelined mode only decides whether the input from `load_stream` is read ahead in a background thread. - Move the background reader into a `read_in_thread(iterator, queue_depth)` context manager in `worker_util.py`. Thread start, stop, queue drain, the bounded join and the leak warning live in its exit. The env vars are still read in `worker.py`. - Inside it, merge the two identical "put until stopped" loops into one `_put`, and replace the `reader_error` list + sentinel pair with a single end marker that carries the reader exception. - Expose `ArrowStreamSerializer.flush_per_batch` (already a constructor argument) instead of setting the private `_flush_per_batch` from `worker.py`. Behavior is unchanged: the input stream is never closed (worker reuse), lazy per-group iterators are still materialized on the reader thread, shutdown order and the warning are the same, reader exceptions are re-raised in the main thread, and the non-pipelined path is unchanged. ### Why are the changes needed? Removes the duplicated processing path and makes the reader logic testable on its own. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - New `ReadInThreadTests` in `pyspark.tests.test_worker` (no JVM): ordering, reads stay on the reader thread, lazy items materialized there, reader error re-raised, early consumer exit stops a reader blocked on a full queue, join timeout warns. - `pyspark.sql.tests.pandas.test_pipelined_udf`, `pyspark.sql.tests.test_pipelined_shuffle`, `pyspark.sql.tests.test_python_worker_env`, `pyspark.tests.test_worker` pass. - ASV `bench_pipelined_udf.py` before/after: results to follow. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 5.5 This pull request and its description were written by Isaac. -- 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]
