Yicong-Huang commented on code in PR #60115:
URL: https://github.com/apache/spark/pull/60115#discussion_r4486618268
##########
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:
Fixed in 1f893e1: the reader now re-checks the stop after a successful put,
and `test_early_exit_stops_reader_blocked_on_full_queue` asserts exactly three
reads and no warning when the next read would block.
--
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]
