spark-test commented on code in PR #60115:
URL: https://github.com/apache/spark/pull/60115#discussion_r4487957754
##########
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
Review Comment:
**Nit (P3):** The reader takes the next item from `iterator`, materializing
it with `list(item)` when it is a lazy group iterator, before `_put` blocks on
a queue that already holds `queue_depth` items. So it can be `queue_depth + 1`
items ahead of the caller. `test_early_exit_stops_reader_blocked_on_full_queue`
shows this at `queue_depth=1`: after the consumer takes one item, the source
has produced three. For grouped UDFs the extra item is a whole materialized
group, so memory sized from this docstring comes out one group short. Counting
the item the reader holds while it waits for room, for example "up to
``queue_depth`` queued items plus the one it is waiting to enqueue", would make
the bound accurate. The `spark.python.udf.pipelined.queueDepth` description
uses the same wording, but this PR does not change it.
**Verification:**
- **Inspection:** The read-ahead bound in the read_in_thread docstring
matches the implementation: a queue holding at most queue_depth items, plus the
one item the reader has already read (and materialized, for a lazy group
iterator) before it blocks waiting for room.
##########
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) or stop_event.is_set():
+ return
+ except Exception as e:
+ error = e
+ finally:
+ _put(_EndOfInput(error))
+
+ def _consume() -> Iterator[Any]:
+ while True:
+ item = item_queue.get()
+ if isinstance(item, _EndOfInput):
Review Comment:
**Nit (P3):** Because the reader creates a new `_EndOfInput` in its
`finally` block, the consumer can recognize it only by type, so every queued
item now pays an `isinstance` call where the base branch's `_queued_iter` did
`item is _SENTINEL`. In paired micro-benchmarks on CPython 3.12 that costs
about 30 ns of GIL-held CPU per item. It is negligible per Arrow batch or
group, but pickled row-at-a-time UDFs (`SQL_BATCHED_UDF`, `SQL_TABLE_UDF`) pay
it once per row in pipelined mode. Results, ordering and error transport are
unaffected, and the default non-pipelined path never reaches this code.
Creating the marker once per call keeps the single end marker that carries
the exception and lets the consumer test it by identity:
```python
end = _EndOfInput(None)
def _read() -> None:
try:
...
except Exception as e:
end.error = e
finally:
_put(end)
def _consume() -> Iterator[Any]:
while True:
item = item_queue.get()
if item is end:
if end.error is not None:
raise end.error
return
yield item
```
In the same measurement this variant runs at the base-branch cost, and it
passes the three `ReadInThreadTests` scenarios.
**Verification:**
- **Regression:** The consumer's per-item end-of-input test is an identity
comparison, and a per-item timing of the consumer loop matches the base-branch
identity test rather than the isinstance test.
- **Behavior:** ReadInThreadTests still pass: a source exception is
re-raised after the preceding items, an early exit stops a reader blocked on a
full queue without another read or a warning, and a reader that does not exit
still triggers 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]
