feat(python): add backpressure to StreamingDataset post-transform queue (#3897)

Rename prefetch_batches → io_queue_depth and introduce
transform_queue_depth as a symmetric pair: both express "number of
batches to buffer per split at this pipeline stage." The old names are
still accepted as keyword arguments but log a deprecation warning
redirecting callers to the new names.

transform_queue_depth caps how many transform-result batches can
accumulate per split in the post-transform queue. Without this limit a
slow consumer (e.g. a GPU training step) causes cooked rows to pile up
unboundedly. The backpressure check in _try_submit_tx counts both
already-cooked rows and rows expected from in-flight transforms; it
skips proactive transform submission when the combined total reaches the
limit. The reactive _ensure_cooked path bypasses the check so the
consumer never stalls.

---------

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

python/python/tests/test_elastic_dataloader.py

@@ -1374,6 +1374,188 @@ def test_transform_parallelism_must_be_positive(lance_table, transform_paralleli
         )
 
 
+# ---------------------------------------------------------------------------
+# Backpressure / transform_queue_depth tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("transform_queue_depth", [0, -1])
+def test_transform_queue_depth_must_be_positive(lance_table, transform_queue_depth):
+    """transform_queue_depth=0 or negative must raise ValueError."""
+    with pytest.raises(
+        ValueError, match="transform_queue_depth must be greater than 0"
+    ):
+        StreamingDataset(
+            lance_table,
+            num_splits=NUM_SPLITS,
+            transform_queue_depth=transform_queue_depth,
+        )
+
+
+@pytest.mark.parametrize("transform_queue_depth", [1, 2, 4])
+def test_transform_queue_depth_correctness(lance_table, transform_queue_depth):
+    """With backpressure enabled, every row is still yielded exactly once."""
+    ds = StreamingDataset(
+        lance_table,
+        num_splits=NUM_SPLITS,
+        shuffle_seed=SHUFFLE_SEED,
+        transform_queue_depth=transform_queue_depth,
+        read_batch_size=8,
+    )
+    items = list(ds)
+    assert sorted(item["id"] for item in items) == list(range(NUM_ROWS))
+
+
+def test_transform_queue_depth_matches_no_backpressure(lance_table):
+    """With backpressure enabled the same samples are produced as without it."""
+    ds_unlimited = StreamingDataset(
+        lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED
+    )
+    ds_limited = StreamingDataset(
+        lance_table,
+        num_splits=NUM_SPLITS,
+        shuffle_seed=SHUFFLE_SEED,
+        transform_queue_depth=1,
+    )
+    assert [item["id"] for item in ds_unlimited] == [
+        item["id"] for item in ds_limited
+    ], "transform_queue_depth must not affect the sample ordering or set"
+
+
+def test_transform_queue_depth_bounds_cooked_rows(lance_table):
+    """prefetch_queue_depth stays within transform_queue_depth * read_batch_size
+    per split when observed from the main thread during iteration."""
+    n_splits = 4
+    batch_size = 8
+    cooked_depth = 2
+    # max cooked rows across all 4 splits: 4 * 2 * 8 = 64
+    max_allowed = n_splits * cooked_depth * batch_size
+
+    ds = StreamingDataset(
+        lance_table,
+        num_splits=n_splits,
+        shuffle_seed=SHUFFLE_SEED,
+        transform_queue_depth=cooked_depth,
+        read_batch_size=batch_size,
+        transform_parallelism=1,
+        world_size=1,
+    )
+
+    peak = 0
+    for _ in ds:
+        depth = ds.prefetch_queue_depth
+        if depth > peak:
+            peak = depth
+
+    # The main thread observes depth *after* popping a row, so the peak is at
+    # most max_allowed (one row already popped from the split just served).
+    assert peak <= max_allowed, (
+        f"prefetch_queue_depth peaked at {peak}, expected <= {max_allowed}"
+    )
+
+
+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.
+
+    The test intercepts ThreadPoolExecutor.submit to make I/O calls execute
+    synchronously on the main thread.  This ensures all raw batches land in
+    raw_batches (via _drain_io) before _try_submit_tx evaluates the admission
+    predicate for the first time.  Without this, the I/O future for batch N+1
+    might still be in io_pending at the capacity-minus-one transition, leaving
+    raw_batches empty and causing _try_submit_tx to skip the admission check
+    entirely — so both the correct and the broken predicate produce depth=0
+    observations and the test cannot distinguish them.
+
+    With all raw batches pre-loaded in raw_batches the 4→3 cooked transition
+    (consuming one row from a full cooked queue) always triggers _try_submit_tx
+    against a non-empty raw_batches.
+
+    With transform_queue_depth=1 and batch_size=4, max_cooked_rows=4.
+    A transform may only be submitted when in_pipeline + batch_size <= 4, i.e.
+    when in_pipeline == 0 (cooked is completely empty).  Under the old broken
+    predicate (in_pipeline >= max_cooked_rows) the second transform would be
+    admitted with cooked containing batch_size-1 rows still unconsumed.
+    """
+    import concurrent.futures as cf
+    from concurrent.futures import ThreadPoolExecutor
+    from unittest.mock import patch
+
+    db = lancedb.connect(tmp_path)
+    batch_size = 4
+    # Four full batches → four transform submissions to observe.
+    table = db.create_table("t", pa.table({"id": list(range(batch_size * 4))}))
+
+    cooked_at_submit: list[int] = []
+
+    original_submit = ThreadPoolExecutor.submit
+
+    def tracking_submit(self, fn, *args, **kwargs):
+        name = getattr(fn, "__name__", "")
+        if name == "_io_call":
+            # Run I/O synchronously on the calling (main) thread and return an
+            # already-completed Future.  _drain_io checks fut.done(), so a
+            # completed Future is moved to raw_batches immediately on the next
+            # _advance call — making raw-batch readiness deterministic at the
+            # capacity-minus-one transition instead of depending on I/O thread
+            # scheduling.
+            fut = cf.Future()
+            try:
+                fut.set_result(fn(*args, **kwargs))
+            except Exception as exc:
+                fut.set_exception(exc)
+            return fut
+        if name == "_tx_call_guarded":
+            # Capture cooked depth synchronously on the main thread before the
+            # transform worker can drain the queue.
+            ref = ds._cooked_ref
+            cooked_at_submit.append(len(ref[0]) if ref is not None else -1)
+        return original_submit(self, fn, *args, **kwargs)
+
+    with patch.object(ThreadPoolExecutor, "submit", tracking_submit):
+        ds = StreamingDataset(
+            table,
+            num_splits=1,
+            shuffle_seed=42,
+            read_batch_size=batch_size,
+            transform_queue_depth=1,
+            transform_parallelism=1,
+        )
+        list(ds)
+
+    assert len(cooked_at_submit) == 4, (
+        f"Expected 4 transform submissions (one per batch), got {len(cooked_at_submit)}"
+    )
+    # With full-batch backpressure each transform is only admitted when the
+    # cooked queue is completely empty (depth == 0).  The old broken predicate
+    # would admit at depth == batch_size - 1 == 3.
+    assert all(depth == 0 for depth in cooked_at_submit), (
+        "Transform admitted with non-empty cooked queue; full-batch backpressure "
+        "requires in_pipeline + batch_size <= max_cooked_rows before admission. "
+        f"Cooked depths at each submission: {cooked_at_submit}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Deprecated parameter name tests
+# ---------------------------------------------------------------------------
+
+
+def test_prefetch_batches_deprecated_warns(lance_table, caplog):
+    """prefetch_batches logs a deprecation warning and behaves like io_queue_depth."""
+    with caplog.at_level(logging.WARNING, logger="lancedb.streaming"):
+        ds = StreamingDataset(
+            lance_table,
+            num_splits=NUM_SPLITS,
+            shuffle_seed=SHUFFLE_SEED,
+            prefetch_batches=2,
+        )
+    messages = [r.message for r in caplog.records if r.levelno >= logging.WARNING]
+    assert any("deprecated" in m.lower() and "io_queue_depth" in m for m in messages), (
+        f"Expected deprecation warning mentioning io_queue_depth; got: {messages}"
+    )
+    assert sorted(item["id"] for item in ds) == list(range(NUM_ROWS))
+
+
 def test_filter_limits_rows(tmp_path):
     """A filter expression is applied to the permutation so only matching rows
     are yielded.  IDs 0..59 pass ``id < 60``; the other 60 are excluded."""
@@ -1954,7 +2136,7 @@ def test_doc_example_basic(tmp_path):
 
 
 def test_doc_example_prefetch_params(tmp_path):
-    """doc: Prefetching — read_batch_size and prefetch_batches still cover all rows."""
+    """doc: Prefetching — read_batch_size and io_queue_depth still cover all rows."""
     db = lancedb.connect(tmp_path)
     table = db.create_table("t", pa.table({"id": list(range(NUM_ROWS))}))
 
@@ -1963,7 +2145,7 @@ def test_doc_example_prefetch_params(tmp_path):
         num_splits=NUM_SPLITS,
         shuffle_seed=SHUFFLE_SEED,
         read_batch_size=8,
-        prefetch_batches=2,
+        io_queue_depth=2,
     )
     assert sorted(s["id"] for s in ds) == list(range(NUM_ROWS))