fix(python): bound StreamingDataset's raw_batches queue against io_queue_depth (#4530)

## Summary
- `StreamingDataset`'s Stage 1 (I/O) queue (`raw_batches`) had no cap:
`_fill_io` only limited *in-flight* fetch futures (`io_pending`) to
`io_queue_depth`, but every completed fetch was unconditionally drained
into `raw_batches` and the freed slot immediately refilled, so I/O could
race arbitrarily far ahead of a slow transform stage with no bound on
buffered memory.
- `transform_queue_depth`'s admission check (`_try_submit_tx`) couldn't
compensate either, since it only counts `cooked` rows and in-flight
transform futures — it never looks at `raw_batches`.
- Fix: `_fill_io` now counts both `io_pending` and `raw_batches` against
the `io_queue_depth` budget, matching the parameter's documented
"batches to keep in flight per split" contract.

## Test plan
- [x] Added
`test_raw_queue_depth_bounded_when_transform_is_the_bottleneck`
(deterministic, via synchronous I/O) asserting `raw_queue_depth <=
io_queue_depth * read_batch_size` even when the transform stage is the
bottleneck.
- [x] Full `python/python/tests/test_elastic_dataloader.py` suite passes
(190/190), confirming no change to sample ordering, counts, or
checkpoint/resumability behavior.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

python/python/lancedb/streaming.py

@@ -261,10 +261,15 @@ class StreamingDataset(IterableDataset):
         storage) at the cost of higher memory usage per split buffer.  Defaults
         to ``DEFAULT_READ_BATCH_SIZE`` (64).
     io_queue_depth:
-        Number of I/O batches to keep in flight per split.  Higher values
-        overlap storage latency with transform and training compute at the cost
-        of more memory and threads.  Must be greater than zero.  Defaults to
-        ``DEFAULT_PREFETCH_BATCHES`` (4).
+        Number of fetched batches to keep ahead of the transform stage per
+        split: this bounds the combined count of in-flight fetches and
+        already-fetched batches still waiting for a transform slot.  Higher
+        values overlap storage latency with transform and training compute
+        at the cost of more memory and threads, but I/O stops running ahead
+        once this many batches are buffered, applying backpressure when the
+        transform stage (or ``transform_queue_depth``) is the bottleneck.
+        Must be greater than zero.  Defaults to ``DEFAULT_PREFETCH_BATCHES``
+        (4).
     columns:
         Optional list of column names to read.  When set, only those columns
         are fetched from storage; all others are omitted.  ``None`` (the
@@ -836,8 +841,17 @@ class StreamingDataset(IterableDataset):
             io_pending[i].append((abs_start, io_pool.submit(_io_call, perm_i, indices)))
 
         def _fill_io(i: int) -> None:
+            # Budget covers both in-flight fetches (io_pending) and batches
+            # that have already landed but are still waiting for a transform
+            # slot (raw_batches).  Without raw_batches in the count, a slow
+            # transform stage never applies backpressure to Stage 1: each
+            # completed fetch is drained into raw_batches and immediately
+            # replaced by a new one, so I/O keeps racing arbitrarily far
+            # ahead of the transform stage instead of merely overlapping
+            # storage latency with it.
             while (
-                len(io_pending[i]) < io_queue_depth and fetch_head[i] < split_sizes[i]
+                len(io_pending[i]) + len(raw_batches[i]) < io_queue_depth
+                and fetch_head[i] < split_sizes[i]
             ):
                 _submit_io(i)
 

python/python/tests/test_elastic_dataloader.py

@@ -2036,6 +2036,92 @@ def test_transform_queue_depth_bounds_cooked_rows(lance_table):
     )
 
 
+def test_raw_queue_depth_bounded_when_transform_is_the_bottleneck(tmp_path):
+    """io_queue_depth bounds raw_batches even when transform is the bottleneck.
+
+    Regression test for a bug where Stage 1 (I/O) applied no backpressure at
+    all: _fill_io only capped io_pending (in-flight fetches) at
+    io_queue_depth, but every completed fetch was drained into raw_batches
+    unconditionally (_drain_io) and immediately replaced by a new fetch
+    (_fill_io again), so I/O kept racing arbitrarily far ahead of a slow
+    transform stage -- raw_batches had no cap of its own.
+    transform_queue_depth did not help either: its admission check in
+    _try_submit_tx only counts cooked rows and in-flight transform futures
+    (`len(cooked[i]) + len(tx_pending[i]) * batch_size`); it never looks at
+    len(raw_batches[i]), so it cannot throttle Stage 1.
+
+    The fix makes _fill_io count both io_pending and raw_batches against the
+    io_queue_depth budget, so raw_queue_depth should never exceed
+    io_queue_depth * read_batch_size per split, regardless of how slow the
+    transform stage is.
+
+    I/O is forced to run synchronously (same technique as
+    test_transform_queue_depth_does_not_admit_at_capacity_minus_one) so the
+    assertion is deterministic regardless of machine load, rather than
+    relying on a real-time race between I/O and transform threads.  Only the
+    transform runs on a real background thread (with a short sleep) so the
+    consumer genuinely blocks waiting for it, letting Stage 1 attempt to run
+    ahead between blocking waits.
+    """
+    import concurrent.futures as cf
+    from concurrent.futures import ThreadPoolExecutor
+    from unittest.mock import patch
+    import pyarrow as pa
+    import time
+
+    n_rows = 200
+    read_batch_size = 10
+    io_queue_depth = 2
+    transform_queue_depth = 2
+    max_allowed = io_queue_depth * read_batch_size
+
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("t", pa.table({"id": list(range(n_rows))}))
+
+    def slow_transform(batch: pa.RecordBatch) -> list:
+        time.sleep(0.005)
+        return batch.column("id").to_pylist()
+
+    original_submit = ThreadPoolExecutor.submit
+
+    def synchronous_io_submit(self, fn, *args, **kwargs):
+        if getattr(fn, "__name__", "") == "_io_call":
+            fut = cf.Future()
+            try:
+                fut.set_result(fn(*args, **kwargs))
+            except Exception as exc:
+                fut.set_exception(exc)
+            return fut
+        return original_submit(self, fn, *args, **kwargs)
+
+    with patch.object(ThreadPoolExecutor, "submit", synchronous_io_submit):
+        ds = StreamingDataset(
+            table,
+            num_splits=1,
+            shuffle_seed=SHUFFLE_SEED,
+            read_batch_size=read_batch_size,
+            io_queue_depth=io_queue_depth,
+            transform_queue_depth=transform_queue_depth,
+            transform_parallelism=1,
+            transform=slow_transform,
+        )
+
+        peak_raw = 0
+        ids = []
+        for item in ds:
+            ids.append(item)  # slow_transform yields bare ids, not dicts
+            if ds.raw_queue_depth > peak_raw:
+                peak_raw = ds.raw_queue_depth
+
+    assert peak_raw <= max_allowed, (
+        f"raw_queue_depth peaked at {peak_raw}, expected <= {max_allowed} "
+        "(io_queue_depth * read_batch_size)"
+    )
+    # The backpressure fix must not come at the cost of dropping or
+    # duplicating rows.
+    assert sorted(ids) == list(range(n_rows))
+
+
 def test_transform_queue_depth_does_not_admit_at_capacity_minus_one(tmp_path):
     """Admission requires a full read_batch_size of free space, not just one slot.