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/lancedb/streaming.py

@@ -61,7 +61,7 @@ class StreamingDataset(IterableDataset):
 
     Internally ``__iter__`` runs a two-stage pipeline:
 
-    - **Stage 1 (I/O)**: one thread pool with ``num_splits * prefetch_batches``
+    - **Stage 1 (I/O)**: one thread pool with ``num_splits * io_queue_depth``
       workers fetches raw ``RecordBatch`` objects from LanceDB in parallel
       across all splits and places them in a per-split raw-batch queue.
     - **Stage 2 (transform)**: a second thread pool with
@@ -104,11 +104,11 @@ class StreamingDataset(IterableDataset):
         call.  Larger values amortise per-request overhead (critical on object
         storage) at the cost of higher memory usage per split buffer.  Defaults
         to ``DEFAULT_READ_BATCH_SIZE`` (64).
-    prefetch_batches:
+    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.  Defaults to ``DEFAULT_PREFETCH_BATCHES``
-        (4).
+        of more memory and threads.  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
@@ -175,6 +175,16 @@ class StreamingDataset(IterableDataset):
         Prefer the ``filter`` parameter when bad rows can be expressed as a
         SQL predicate (e.g. ``"col IS NOT NULL"``) — filtering happens before
         splits are built, so every guarantee is fully preserved.
+    transform_queue_depth:
+        Number of transform-result batches to buffer per split in the
+        post-transform queue before backpressure is applied to the transform
+        stage.  When the combined count of in-flight transform futures and
+        already-buffered rows for a split reaches
+        ``transform_queue_depth * read_batch_size``, no new transforms are
+        submitted for that split until the consumer catches up.  Useful for
+        capping peak memory when the consumer (e.g. a GPU training step) is
+        slower than the transform stage.  Must be greater than zero.
+        ``None`` (the default) imposes no limit.
     worker_info_override:
         If set, used in place of ``torch.utils.data.get_worker_info()`` to
         determine the DataLoader worker assignment.  Intended for unit tests
@@ -194,17 +204,26 @@ class StreamingDataset(IterableDataset):
         rank: int = 0,
         world_size: int = 1,
         read_batch_size: int = DEFAULT_READ_BATCH_SIZE,
-        prefetch_batches: int = DEFAULT_PREFETCH_BATCHES,
+        io_queue_depth: int = DEFAULT_PREFETCH_BATCHES,
         columns: Optional[list[str]] = None,
         shuffle_clump_size: Optional[int] = None,
         filter: Optional[str] = None,
         transform: Optional[Callable] = None,
         transform_parallelism: Optional[int] = None,
         on_transform_error: Union[str, Callable[[Exception], bool]] = "raise",
+        transform_queue_depth: Optional[int] = None,
         connection_factory: Optional[Callable[[str], Any]] = None,
         worker_info_override=None,
+        # Deprecated; use io_queue_depth instead.
+        prefetch_batches: Optional[int] = None,
     ):
         super().__init__()
+        if prefetch_batches is not None:
+            logger.warning(
+                "prefetch_batches is deprecated and will be removed in a future "
+                "version; use io_queue_depth instead"
+            )
+            io_queue_depth = prefetch_batches
         if num_splits is None:
             num_splits = world_size
         if shuffle_seed is None:
@@ -214,6 +233,8 @@ class StreamingDataset(IterableDataset):
                 f"num_splits ({num_splits}) must be divisible by "
                 f"world_size ({world_size})"
             )
+        if io_queue_depth <= 0:
+            raise ValueError("io_queue_depth must be greater than 0")
         if transform_parallelism is not None and transform_parallelism <= 0:
             raise ValueError("transform_parallelism must be greater than 0")
         if on_transform_error not in ("raise", "skip", "warn") and not callable(
@@ -223,6 +244,8 @@ class StreamingDataset(IterableDataset):
                 "on_transform_error must be 'raise', 'skip', 'warn', or a "
                 f"callable, got {on_transform_error!r}"
             )
+        if transform_queue_depth is not None and transform_queue_depth <= 0:
+            raise ValueError("transform_queue_depth must be greater than 0")
 
         self._table = table
         self._num_splits = num_splits
@@ -232,13 +255,14 @@ class StreamingDataset(IterableDataset):
         self._rank = rank
         self._world_size = world_size
         self._read_batch_size = read_batch_size
-        self._prefetch_batches = prefetch_batches
+        self._io_queue_depth = io_queue_depth
         self._columns = columns
         self._shuffle_clump_size = shuffle_clump_size
         self._filter = filter
         self._transform = transform
         self._transform_parallelism = transform_parallelism
         self._on_transform_error = on_transform_error
+        self._transform_queue_depth = transform_queue_depth
         self._connection_factory = connection_factory
         self._worker_info_override = worker_info_override
 
@@ -365,7 +389,7 @@ class StreamingDataset(IterableDataset):
         pos_consumed = list(initial_positions)
 
         batch_size = self._read_batch_size
-        max_prefetch = self._prefetch_batches
+        io_queue_depth = self._io_queue_depth
         transform_workers = (
             self._transform_parallelism
             if self._transform_parallelism is not None
@@ -374,6 +398,13 @@ class StreamingDataset(IterableDataset):
         final_transform = (
             self._transform if self._transform is not None else Transforms.arrow2python
         )
+        # None means no limit; otherwise cap rows per split to
+        # transform_queue_depth batches worth (including in-flight transforms).
+        max_cooked_rows = (
+            self._transform_queue_depth * batch_size
+            if self._transform_queue_depth is not None
+            else None
+        )
 
         # Per-split pipeline state.  Batches are paired with the absolute
         # permutation position of their first row so that skipped rows can be
@@ -409,7 +440,9 @@ class StreamingDataset(IterableDataset):
             io_pending[i].append((abs_start, io_pool.submit(_io_call, perm_i, indices)))
 
         def _fill_io(i: int) -> None:
-            while len(io_pending[i]) < max_prefetch and fetch_head[i] < split_sizes[i]:
+            while (
+                len(io_pending[i]) < io_queue_depth and fetch_head[i] < split_sizes[i]
+            ):
                 _submit_io(i)
 
         def _drain_io(i: int) -> None:
@@ -487,7 +520,19 @@ class StreamingDataset(IterableDataset):
 
         def _try_submit_tx(i: int) -> None:
             """Submit transforms for raw_batches[i] up to available capacity."""
-            while raw_batches[i] and tx_semaphore.acquire(blocking=False):
+            while raw_batches[i]:
+                # Backpressure: only submit a new transform when there is room
+                # for a full batch in the post-transform queue.  Checking for
+                # a full batch prevents submitting a transform that would
+                # overflow the limit mid-batch (e.g. 990 rows queued with a
+                # capacity of 1000 and a batch_size of 128 must wait until
+                # 128 rows have been consumed, not just 1).
+                if max_cooked_rows is not None:
+                    in_pipeline = len(cooked[i]) + len(tx_pending[i]) * batch_size
+                    if in_pipeline + batch_size > max_cooked_rows:
+                        break
+                if not tx_semaphore.acquire(blocking=False):
+                    break
                 abs_start, batch = raw_batches[i].popleft()
                 tx_pending[i].append(tx_pool.submit(_tx_call_guarded, abs_start, batch))
 
@@ -531,7 +576,7 @@ class StreamingDataset(IterableDataset):
 
         # ── Main loop ─────────────────────────────────────────────────────────
 
-        with ThreadPoolExecutor(max_workers=n * max_prefetch) as io_pool:
+        with ThreadPoolExecutor(max_workers=n * io_queue_depth) as io_pool:
             with ThreadPoolExecutor(max_workers=transform_workers) as tx_pool:
                 self._raw_batches_ref = raw_batches
                 self._cooked_ref = cooked

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))