spark-test commented on code in PR #60115:
URL: https://github.com/apache/spark/pull/60115#discussion_r4476879938
##########
python/pyspark/worker_util.py:
##########
@@ -265,6 +265,95 @@ def send_accumulator_updates(outfile: IO) -> None:
pickleSer._write_with_length((aid, accum._value), outfile)
+class _EndOfInput:
+ """Last item ``read_in_thread`` queues, carrying the exception that ended
the source."""
+
+ __slots__ = ("error",)
+
+ def __init__(self, error: Optional[Exception]) -> None:
+ self.error = error
+
+
+@contextmanager
+def read_in_thread(
+ iterator: Iterable[Any], queue_depth: int
+) -> Generator[Iterator[Any], None, None]:
+ """
+ Read ``iterator`` in a background thread, keeping up to ``queue_depth``
items ahead of the
+ caller, and yield an iterator over the same items.
+
+ This lets reading and deserializing the input overlap with the caller's
work. Items that are
+ themselves iterators (e.g. lazy per-group iterators that still read the
input stream) are
+ materialized into lists on the reader thread, so the caller never reads
the stream. An
+ exception raised by ``iterator`` is re-raised to the caller once the items
before it are
+ consumed.
+
+ On exit the reader is stopped and joined for up to 5 seconds. The input is
never closed, as
+ the stream it reads may be reused by the next task of a reused worker; if
the reader is still
+ blocked reading it, a RuntimeWarning is emitted instead.
+ """
+ import queue
+ import threading
+
+ item_queue: queue.Queue = queue.Queue(maxsize=queue_depth)
+ stop_event = threading.Event()
+
+ def _put(item: Any) -> bool:
+ # Block until there is room, giving up once the consumer has stopped.
+ while not stop_event.is_set():
+ try:
+ item_queue.put(item, timeout=0.1)
+ return True
+ except queue.Full:
+ continue
+ return False
+
+ def _read() -> None:
+ error = None
+ try:
+ for item in iterator:
+ if hasattr(item, "__next__"):
+ item = list(item)
+ if not _put(item):
Review Comment:
**Non-blocking (P2):** When the consumer leaves the `read_in_thread` context
early, the exit sets `stop_event` and drains the queue, which lets a reader
that is waiting in `_put` complete its put. `_put` then returns `True`, so this
loop calls `next(iterator)` again and reads one more item from the worker input
stream before the next `_put` sees the stop. The reader removed from
`worker.py` re-checked `stop_event.is_set()` after every successful put and
returned, so it never read after the stop.
This affects only pipelined mode
(`spark.python.udf.pipelined.enabled=true`), when the consumer exits before the
end of input while the reader is waiting on a full queue, which is the usual
state for a UDF that is slower than its input. Examples are a UDF or
verification error, or a `mapInPandas`/`mapInArrow` function that returns
early. Each such exit now reads and deserializes an extra Arrow batch, or
materializes a whole extra group, that nothing uses; for pickled-row UDFs the
extra item is a row, which reads the stream only when it starts a new pickled
batch. If that item is slow to arrive, the exit waits for it, up to the 5s
join, before the error is reported or the task finishes. Past 5s it also emits
the "did not exit within 5s" `RuntimeWarning` and leaves the reader blocked on
the same socket that `invoke_udf` then reads the finish signal from.
A JVM-free reproduction that runs the base-branch reader and this one
against a stub source reads the source 3 times on the base branch and 4 times
here at queue depth 1 (one extra read at every depth tried). With the next item
8s away, the exit takes 0.05s on the base branch and 5.05s plus the warning
here. The new `test_early_exit_stops_reader_blocked_on_full_queue` allows the
extra read (`assertLessEqual(len(produced), 4)`).
Re-checking the stop after a successful put restores the base-branch
behavior:
```suggestion
if not _put(item) or stop_event.is_set():
```
**Verification:**
- **Regression:** When the consumer leaves read_in_thread while the reader
waits to enqueue on a full queue, the source is not read after the stop, at any
queue depth and also when the items are lazy group iterators the reader
materializes. For example, at queue depth 1 a consumer that takes one item and
then leaves causes three source reads in total, the base-branch count, rather
than four.
- **Regression:** When the source's next read after the stop would block for
longer than the 5 second join, leaving the context early returns promptly,
emits no RuntimeWarning, and leaves no reader thread running.
- **Behavior:** A source exception is still re-raised to the consumer after
the items that preceded it, and a reader that is already blocked inside a read
when the consumer stops still triggers the bounded join and the RuntimeWarning.
--
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]
