From c6db80dd0b1daff2950b33ec886db2e5fae99852 Mon Sep 17 00:00:00 2001 From: Weston Pace Date: Mon, 6 Jul 2026 05:50:45 -0700 Subject: [PATCH] feat: add an elastic dataloader as an iterable dataset (#3509) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Elastic Streaming Dataloader ## Motivation Training large models on LanceDB tables today requires loading the entire dataset into memory or writing bespoke batching logic. This PR introduces `StreamingDataset`, a PyTorch `IterableDataset` that streams directly from a LanceDB table with two hard guarantees that are difficult to achieve together: **elastic determinism** and **resumability**. ## Goals ### Elastic determinism The dataset partitions the table into a fixed number of *splits* (controlled by `num_splits`, `shuffle_seed`, and `epoch`). Samples are yielded by round-robining over splits one sample per split per cycle. Because the split structure is fixed, the set of samples that makes up each global training step is identical regardless of `world_size` or `num_workers`. You can scale your cluster up or down between runs and the model sees the same data in the same order — no re-sharding, no gradient variance from topology changes. ### Resumability `state_dict()` / `load_state_dict()` capture how many samples each split has consumed. Because all splits are the same size and the round-robin design keeps them in lockstep, the state reduces to a single scalar (`samples_consumed_per_split`) that is topology-independent. A checkpoint saved with 8 GPUs can resume correctly on 4 GPUs or 16 GPUs without any adjustment. ### PyTorch `IterableDataset` / streaming `StreamingDataset` implements the standard PyTorch `IterableDataset` interface, so it drops into any existing `DataLoader` pipeline without modification. Data is fetched lazily from Lance in chunks — only the rows needed for the current batch are ever in memory. Compared to the map dataset this takes more work from pytorch and puts it into the dataset itself (e.g. shuffling, filtering, etc.). We do this because we cannot achieve things like elastic determinism or prefiltering otherwise. ### Multi-worker support DataLoader workers are automatically assigned contiguous sub-blocks of splits (the rank's splits are divided evenly across workers). Each worker is independent: no shared state, no inter-process coordination. The only constraint is that `num_splits` must be divisible by `world_size * num_workers`. That being said, multi-worker is highly discouraged as it relies on multiprocessing which is inefficient. Still, we want to support it. ### Filters as prefilters Filters are applied at *permutation-build time* via `PermutationBuilder.filter()`, not re-evaluated on every fetch. The filtered row IDs are stored in the permutation table so that subsequent reads see only the matching rows. This allows us to avoid loading rows that don't match the filter (which is the default pytorch behavior) ### Prefetching Two parameters control the I/O pipeline: - `read_batch_size` (default 64) — number of rows fetched per `take_offsets` call. Larger values amortise per-request overhead, which is critical on object storage where a single round-trip can cost ~100 ms. - `prefetch_batches` (default 4) — number of batches prefetched in parallel per split via a `ThreadPoolExecutor`. While the model processes the current batch, the next several batches are already in flight, hiding storage latency behind compute. If set correctly then you can get good performance even with num_workers=0 (unless you are bottlenecked on transform). ### Transform parallelism The underlying `Permutation` API supports a `with_transform()` callback for decoding, augmentation, and format conversion. Unfortunately, this is not parallelized. Pytorch typically parallelizes this with num_workers which is multiprocessing which is highly inefficient. For simple transforms we should be able to utilize multithreading and Rust based UDFs. For complex python UDFs we could have a dedicated multiprocessing pipeline for just the transform. Or we could just utilize multithreading. In both cases we would exclude the I/O stage from the multiprocessing because that ends up being very memory hungry and inefficient. --------- Co-authored-by: Claude Sonnet 4.6 --- Cargo.lock | 25 +- docs/src/js/interfaces/SplitRandomOptions.md | 8 + nodejs/src/permutation.rs | 8 +- .../benchmarks/bench_streaming_dataloader.py | 135 ++ python/python/lancedb/permutation.py | 5 + python/python/lancedb/streaming.py | 607 ++++++ .../python/tests/test_elastic_dataloader.py | 1682 +++++++++++++++++ python/python/tests/test_remote_db.py | 8 +- python/src/permutation.rs | 12 +- rust/lancedb/Cargo.toml | 5 + .../examples/bench_streaming_dataloader.rs | 272 +++ .../src/dataloader/permutation/builder.rs | 1 + .../src/dataloader/permutation/reader.rs | 12 +- .../src/dataloader/permutation/split.rs | 15 +- rust/lancedb/src/query.rs | 26 +- 15 files changed, 2784 insertions(+), 37 deletions(-) create mode 100644 python/benchmarks/bench_streaming_dataloader.py create mode 100644 python/python/lancedb/streaming.py create mode 100644 python/python/tests/test_elastic_dataloader.py create mode 100644 rust/lancedb/examples/bench_streaming_dataloader.rs diff --git a/Cargo.lock b/Cargo.lock index 69c0de3a3..c2ebeb7f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5279,7 +5279,7 @@ dependencies = [ "criterion", "lance-arrow", "num-traits", - "pprof", + "pprof 0.15.0", "rand 0.9.4", ] @@ -5361,6 +5361,7 @@ dependencies = [ "pin-project", "polars", "polars-arrow", + "pprof 0.14.1", "rand 0.9.4", "random_word", "regex", @@ -7326,6 +7327,28 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "pprof" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "afad4d4df7b31280028245f152d5a575083e2abb822d05736f5e47653e77689f" +dependencies = [ + "aligned-vec", + "backtrace", + "cfg-if 1.0.4", + "findshlibs", + "inferno", + "libc", + "log", + "nix", + "once_cell", + "smallvec", + "spin 0.10.0", + "symbolic-demangle", + "tempfile", + "thiserror 1.0.69", +] + [[package]] name = "pprof" version = "0.15.0" diff --git a/docs/src/js/interfaces/SplitRandomOptions.md b/docs/src/js/interfaces/SplitRandomOptions.md index f2607b98b..1a5507118 100644 --- a/docs/src/js/interfaces/SplitRandomOptions.md +++ b/docs/src/js/interfaces/SplitRandomOptions.md @@ -8,6 +8,14 @@ ## Properties +### clumpSize? + +```ts +optional clumpSize: number; +``` + +*** + ### counts? ```ts diff --git a/nodejs/src/permutation.rs b/nodejs/src/permutation.rs index 43e2e1e8a..5f1fbd73b 100644 --- a/nodejs/src/permutation.rs +++ b/nodejs/src/permutation.rs @@ -16,6 +16,7 @@ pub struct SplitRandomOptions { pub counts: Option>, pub fixed: Option, pub seed: Option, + pub clump_size: Option, pub split_names: Option>, } @@ -125,10 +126,15 @@ impl PermutationBuilder { }; let seed = options.seed.map(|s| s as u64); + let clump_size = options.clump_size.map(|c| c as u64); self.modify(|builder| { builder.with_split_strategy( - SplitStrategy::Random { seed, sizes }, + SplitStrategy::Random { + seed, + sizes, + clump_size, + }, options.split_names.clone(), ) }) diff --git a/python/benchmarks/bench_streaming_dataloader.py b/python/benchmarks/bench_streaming_dataloader.py new file mode 100644 index 000000000..4d0dadcf9 --- /dev/null +++ b/python/benchmarks/bench_streaming_dataloader.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Benchmark for StreamingDataset throughput. + +Sweeps read_batch_size from 1 to 16384 to show how amortising the per-request +overhead scales. Each row at each chunk size is timed via the real +StreamingDataset so the numbers reflect production code. + +Run with: + cd python + uv run --extra tests benchmarks/bench_streaming_dataloader.py + +Optional env vars: + BENCH_NUM_ROWS — total rows in the table (default 49152 = 24 × 2048) + BENCH_NUM_SPLITS — number of splits (default 24) + BENCH_STEPS — round-robin cycles to time per chunk size (default 100) + BENCH_ROW_BYTES — approximate bytes per row padded with a binary column + (default 4096, mimics a small embedding/image patch) +""" + +import os +import time +import tempfile + +import pyarrow as pa +import lancedb + +from lancedb.streaming import StreamingDataset + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +NUM_SPLITS = int(os.environ.get("BENCH_NUM_SPLITS", 24)) +# Default: 2048 rows per split so every chunk size up to 16Ki has ≥1 full +# chunk (except 16Ki itself which gets a single full-split fetch — still valid). +NUM_ROWS = int(os.environ.get("BENCH_NUM_ROWS", NUM_SPLITS * 2048)) +STEPS = int(os.environ.get("BENCH_STEPS", 100)) +ROW_BYTES = int(os.environ.get("BENCH_ROW_BYTES", 4096)) + +assert NUM_ROWS % NUM_SPLITS == 0, "NUM_ROWS must be divisible by NUM_SPLITS" + +CHUNK_SIZES = [1, 4, 16, 64, 256, 1024, 4096, 16384] + + +# --------------------------------------------------------------------------- +# Table helpers +# --------------------------------------------------------------------------- + + +def make_table(db_path: str) -> lancedb.table.Table: + db = lancedb.connect(db_path) + payload = b"x" * ROW_BYTES + data = pa.table( + { + "id": pa.array(range(NUM_ROWS), type=pa.int32()), + "payload": pa.array([payload] * NUM_ROWS, type=pa.large_binary()), + } + ) + return db.create_table("bench", data, mode="overwrite") + + +# --------------------------------------------------------------------------- +# Timing +# --------------------------------------------------------------------------- + + +def bench_chunk(table, chunk_size: int, steps: int) -> tuple[int, float]: + """Return (rows_drained, elapsed_seconds) for one timed run.""" + total_rows = steps * NUM_SPLITS + ds = StreamingDataset( + table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk_size + ) + count = 0 + t0 = time.perf_counter() + for _ in ds: + count += 1 + if count >= total_rows: + break + return count, time.perf_counter() - t0 + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + rows_per_split = NUM_ROWS // NUM_SPLITS + print("Benchmark config:") + print( + f" NUM_ROWS={NUM_ROWS} NUM_SPLITS={NUM_SPLITS} " + f"rows/split={rows_per_split} STEPS={STEPS} ROW_BYTES={ROW_BYTES}" + ) + print(f" ~{NUM_ROWS * ROW_BYTES / 1024 / 1024:.1f} MB total table size") + print() + + with tempfile.TemporaryDirectory() as tmp: + print("Creating table...", flush=True) + table = make_table(tmp) + + cols = ( + f"{'chunk':>6} {'rows':>6} {'elapsed':>8} {'rows/s':>10} {'ms/step':>9}" + ) + print(f"\n{cols}") + print("-" * 52) + + for chunk in CHUNK_SIZES: + # Warm-up pass (one step's worth of rows) + warmup_ds = StreamingDataset( + table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk + ) + warmup_count = 0 + for _ in warmup_ds: + warmup_count += 1 + if warmup_count >= NUM_SPLITS: + break + + drained, elapsed = bench_chunk(table, chunk, STEPS) + rows_per_sec = drained / elapsed if elapsed > 0 else float("inf") + ms_per_step = elapsed / STEPS * 1000 + + print( + f"{chunk:>6} {drained:>6} {elapsed:>7.3f}s " + f"{rows_per_sec:>10.0f} {ms_per_step:>8.1f}ms" + ) + + print() + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 4e1193f56..ac2ebfa4d 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -65,6 +65,7 @@ class PermutationBuilder: counts: Optional[list[int]] = None, fixed: Optional[int] = None, seed: Optional[int] = None, + clump_size: Optional[int] = None, split_names: Optional[list[str]] = None, ) -> "PermutationBuilder": """ @@ -87,6 +88,9 @@ class PermutationBuilder: Rows will be randomly assigned to splits. The optional seed can be provided to make the assignment deterministic. + If clump_size is provided, rows are shuffled as contiguous groups of that size, + preserving I/O locality while still randomising the split assignment. + The optional split_names can be provided to name the splits. If not provided, the splits can only be referenced by their index. """ @@ -95,6 +99,7 @@ class PermutationBuilder: counts=counts, fixed=fixed, seed=seed, + clump_size=clump_size, split_names=split_names, ) return self diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py new file mode 100644 index 000000000..9ffc612d9 --- /dev/null +++ b/python/python/lancedb/streaming.py @@ -0,0 +1,607 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Elastic streaming dataloader for PyTorch. + +Provides StreamingDataset, a PyTorch IterableDataset that guarantees: + +- **Elastic determinism**: for a fixed (num_splits, shuffle_seed, epoch) the set + of samples that forms each global training step is identical regardless of + world_size or num_workers. +- **Resumability**: state_dict / load_state_dict capture per-split consumption + counts so training can resume from an exact mid-epoch position even when the + distributed topology changes between runs. +""" + +import ctypes +import logging +import os +import random +import threading +import time +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from multiprocessing import RawArray +from typing import Any, Callable, Iterator, Optional + +from torch.utils.data import IterableDataset, get_worker_info + +from .permutation import ( + Permutation, + Transforms, + permutation_builder, + _table_from_pickle_state, + _table_to_pickle_state, +) + +logger = logging.getLogger(__name__) + +# Multiplier used to combine shuffle_seed and epoch into a single permutation +# seed. Chosen to be a large prime so different (seed, epoch) pairs produce +# distinct seeds for any practically encountered epoch count. +_EPOCH_PRIME = 100003 + +DEFAULT_READ_BATCH_SIZE = 64 +DEFAULT_PREFETCH_BATCHES = 4 + + +class StreamingDataset(IterableDataset): + """An elastic, resumable PyTorch IterableDataset backed by a LanceDB table. + + The table is partitioned into ``num_splits`` fixed splits using a + deterministic random shuffle controlled by ``shuffle_seed`` and ``epoch``. + Each rank is assigned a contiguous block of splits, and within a rank each + DataLoader worker is assigned a contiguous sub-block. Samples are yielded + by round-robining over the assigned splits, one sample per split per cycle. + + Internally ``__iter__`` runs a two-stage pipeline: + + - **Stage 1 (I/O)**: one thread pool with ``num_splits * prefetch_batches`` + 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 ``os.cpu_count()`` + workers picks up raw batches, applies the transform, and places the + results in a per-split cooked-row queue. + + The main thread round-robins over the cooked queues, yielding one row per + split per cycle. + + Parameters + ---------- + table: + LanceDB table to stream from. + num_splits: + Number of fixed splits to partition the table into. Must be divisible + by ``world_size``. When used with DataLoader workers it must also be + divisible by ``world_size * num_workers``. Defaults to ``world_size``. + If the row count (after any ``filter``) is not evenly divisible by + ``num_splits``, the surplus rows — at most ``num_splits - 1`` per epoch + — are silently dropped to keep all splits the same length. + shuffle: + Whether to randomly assign rows to splits. When ``True`` (the + default) rows are shuffled using ``shuffle_seed`` and ``epoch``. + When ``False`` rows are divided into splits sequentially in storage + order, which can be useful for deterministic debugging or evaluation. + shuffle_seed: + Base seed for the random permutation. Combined with ``epoch`` so + each epoch produces a different ordering. Pass ``None`` to generate + a random seed at construction time. + epoch: + Current training epoch. Combined with ``shuffle_seed`` so that each + epoch produces a different sample ordering. + rank: + This process's rank in the distributed training group. + world_size: + Total number of processes in the distributed training group. + read_batch_size: + Number of rows fetched from each split in a single ``take_offsets`` + 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: + 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). + columns: + Optional list of column names to read. When set, only those columns + are fetched from storage; all others are omitted. ``None`` (the + default) reads every column. + shuffle_clump_size: + When set, rows are shuffled in contiguous groups of this size rather + than individually. Larger clumps improve I/O locality (important on + object storage) at the cost of reduced randomness. ``None`` (the + default) shuffles rows individually. + filter: + Optional SQL filter expression (e.g. ``"label = 'dog'"``). Only rows + that satisfy the predicate are included in the permutation. The filter + is applied during permutation construction so split sizes reflect the + filtered row count. + transform: + Optional callable applied to each ``pyarrow.RecordBatch`` before rows + are yielded. Receives one batch at a time and must return an iterable + whose length equals the number of rows in the batch. When ``None`` + (the default) rows are returned as plain Python dicts. + 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 + that need to simulate multiple workers without spawning real processes. + If both this and the real worker info are non-None a warning is logged + and the override takes precedence. + """ + + def __init__( + self, + table, + *, + num_splits: Optional[int] = None, + shuffle: bool = True, + shuffle_seed: Optional[int] = 0, + epoch: int = 0, + rank: int = 0, + world_size: int = 1, + read_batch_size: int = DEFAULT_READ_BATCH_SIZE, + prefetch_batches: int = DEFAULT_PREFETCH_BATCHES, + columns: Optional[list[str]] = None, + shuffle_clump_size: Optional[int] = None, + filter: Optional[str] = None, + transform: Optional[Callable] = None, + connection_factory: Optional[Callable[[str], Any]] = None, + worker_info_override=None, + ): + super().__init__() + if num_splits is None: + num_splits = world_size + if shuffle_seed is None: + shuffle_seed = random.randrange(2**32) + if num_splits % world_size != 0: + raise ValueError( + f"num_splits ({num_splits}) must be divisible by " + f"world_size ({world_size})" + ) + + self._table = table + self._num_splits = num_splits + self._shuffle = shuffle + self._shuffle_seed = shuffle_seed + self._epoch = epoch + self._rank = rank + self._world_size = world_size + self._read_batch_size = read_batch_size + self._prefetch_batches = prefetch_batches + self._columns = columns + self._shuffle_clump_size = shuffle_clump_size + self._filter = filter + self._transform = transform + self._connection_factory = connection_factory + self._worker_info_override = worker_info_override + + # Live references to pipeline state, set only while __iter__ is running + # in the same process. Used by the observability properties when the + # DataLoader runs with num_workers=0. + self._raw_batches_ref: Optional[list[deque]] = None + self._cooked_ref: Optional[list[deque]] = None + self._fetch_head_ref: Optional[list[int]] = None + self._split_sizes_ref: Optional[list[int]] = None + self._local_consumed_ref: Optional[list[int]] = None + + # Shared-memory counters written by __iter__ (which may run in a + # DataLoader worker process) and read by the observability properties + # in the main process. RawArray is picklable via the forkserver + # reduction protocol so it survives the dataset pickle round-trip. + # Layout: [unscanned_rows, raw_rows, cooked_rows, consumed_rows, + # bytes_loaded, fetch_time_us, transform_time_us] + self._worker_stats: RawArray = RawArray(ctypes.c_int64, 7) + + # Cumulative bytes of Arrow buffer data fetched across all iterations. + self._bytes_loaded: int = 0 + # Cumulative seconds spent in LanceDB I/O and in transform functions. + self._fetch_time: float = 0.0 + self._transform_time: float = 0.0 + + # Number of samples each split has already been consumed. At global + # step boundaries all splits have consumed this many samples, so a + # single scalar captures the topology-independent checkpoint state. + self._resume_offset: int = 0 + + # Build the permutation table once, deterministically. + builder = permutation_builder(table) + if filter is not None: + builder = builder.filter(filter) + if shuffle: + perm_seed = shuffle_seed + epoch * _EPOCH_PRIME + self._perm_table = builder.split_random( + fixed=num_splits, seed=perm_seed, clump_size=shuffle_clump_size + ).execute() + else: + self._perm_table = builder.split_sequential(fixed=num_splits).execute() + + # Contiguous block of global split indices assigned to this rank. + splits_per_rank = num_splits // world_size + rank_start = rank * splits_per_rank + self._rank_splits: list[int] = list( + range(rank_start, rank_start + splits_per_rank) + ) + + def _resolve_my_splits(self) -> list[int]: + """Return the split indices this instance should read in __iter__.""" + torch_worker_info = get_worker_info() + if self._worker_info_override is not None: + if torch_worker_info is not None: + logger.warning( + "worker_info_override is set but get_worker_info() also returned a " + "non-None value; ignoring the real torch worker info and using the " + "override instead. This may lead to duplicated or incorrect data " + "from the dataset." + ) + worker_info = self._worker_info_override + else: + worker_info = torch_worker_info + + if worker_info is None: + return self._rank_splits + + num_workers: int = worker_info.num_workers + worker_id: int = worker_info.id + n_rank_splits = len(self._rank_splits) + if n_rank_splits % num_workers != 0: + raise ValueError( + f"Number of rank splits ({n_rank_splits}) must be divisible by " + f"num_workers ({num_workers})" + ) + splits_per_worker = n_rank_splits // num_workers + start = worker_id * splits_per_worker + return self._rank_splits[start : start + splits_per_worker] + + def __iter__(self) -> Iterator[dict[str, Any]]: + if self._raw_batches_ref is not None: + raise RuntimeError( + "StreamingDataset does not support concurrent iteration. " + "Only one active iterator per dataset instance is allowed." + ) + my_splits = self._resolve_my_splits() + if not my_splits: + return + + # Set identity transform on each Permutation so __getitems__ returns + # the raw RecordBatch. Stage 2 applies the real transform. + permutations: list[Permutation] = [] + for split_idx in my_splits: + perm = Permutation.from_tables( + self._table, self._perm_table, split=split_idx + ) + if self._columns is not None: + perm = perm.select_columns(self._columns) + perm = perm.with_transform(lambda batch: batch) + if self._resume_offset > 0: + perm = perm.with_skip(self._resume_offset) + permutations.append(perm) + + n = len(permutations) + split_sizes = [perm.num_rows for perm in permutations] + initial_offset = self._resume_offset + local_consumed = [0] * n + + batch_size = self._read_batch_size + max_prefetch = self._prefetch_batches + cpu_workers = os.cpu_count() or 1 + final_transform = ( + self._transform if self._transform is not None else Transforms.arrow2python + ) + + # Per-split pipeline state. + fetch_head = [0] * n + io_pending = [deque() for _ in range(n)] # Future[RecordBatch] + raw_batches = [deque() for _ in range(n)] # RecordBatch — fetched, awaiting tx + tx_pending = [deque() for _ in range(n)] # Future[list[Any]] + cooked = [deque() for _ in range(n)] # rows ready to yield + + # Limit simultaneous transforms to cpu_workers across all splits. + tx_semaphore = threading.Semaphore(cpu_workers) + + # ── Stage 1 helpers ─────────────────────────────────────────────────── + + def _io_call(perm, indices): + t0 = time.perf_counter() + batch = perm.__getitems__(indices) + self._bytes_loaded += batch.nbytes + self._fetch_time += time.perf_counter() - t0 + return batch + + def _submit_io(i: int) -> None: + remaining = split_sizes[i] - fetch_head[i] + if remaining <= 0: + return + fetch = min(batch_size, remaining) + start = fetch_head[i] + fetch_head[i] += fetch + perm_i = permutations[i] + indices = list(range(start, start + fetch)) + io_pending[i].append(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]: + _submit_io(i) + + def _drain_io(i: int) -> None: + """Move completed I/O futures into raw_batches non-blockingly.""" + while io_pending[i] and io_pending[i][0].done(): + raw_batches[i].append(io_pending[i].popleft().result()) + + # ── Stage 2 helpers ─────────────────────────────────────────────────── + + def _tx_call_guarded(batch): + try: + t0 = time.perf_counter() + result = final_transform(batch) + self._transform_time += time.perf_counter() - t0 + return result + finally: + tx_semaphore.release() + + 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): + batch = raw_batches[i].popleft() + tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch)) + + def _drain_tx(i: int) -> None: + """Move completed transform futures into cooked non-blockingly.""" + while tx_pending[i] and tx_pending[i][0].done(): + cooked[i].extend(tx_pending[i].popleft().result()) + + # ── Combined advance ────────────────────────────────────────────────── + + def _advance(i: int) -> None: + """Non-blocking pipeline pump for split i.""" + _drain_io(i) + _drain_tx(i) + _try_submit_tx(i) + _fill_io(i) + + def _ensure_cooked(i: int) -> None: + """Ensure cooked[i] has at least one row, blocking if necessary.""" + _advance(i) + while not cooked[i]: + if tx_pending[i]: + # Wait for the oldest in-flight transform. + cooked[i].extend(tx_pending[i].popleft().result()) + _advance(i) + elif raw_batches[i]: + # Acquire a transform slot (may block briefly if all + # cpu_workers are busy with other splits). + tx_semaphore.acquire() + batch = raw_batches[i].popleft() + tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch)) + elif io_pending[i]: + # Block on the oldest in-flight I/O fetch. + raw_batches[i].append(io_pending[i].popleft().result()) + _advance(i) + else: + break # split exhausted + + # ── Main loop ───────────────────────────────────────────────────────── + + with ThreadPoolExecutor(max_workers=n * max_prefetch) as io_pool: + with ThreadPoolExecutor(max_workers=cpu_workers) as tx_pool: + self._raw_batches_ref = raw_batches + self._cooked_ref = cooked + self._fetch_head_ref = fetch_head + self._split_sizes_ref = split_sizes + self._local_consumed_ref = local_consumed + try: + for i in range(n): + _fill_io(i) + + while True: + # Stop when any split is exhausted (all exhaust + # simultaneously: equal split sizes + round-robin). + if any(local_consumed[i] >= split_sizes[i] for i in range(n)): + break + + for i in range(n): + _ensure_cooked(i) + row = cooked[i].popleft() + local_consumed[i] += 1 + _advance(i) + + # After the last split in each cycle: update the + # global offset and refresh the shared-memory stats + # so the main process can observe pipeline depth + # even when __iter__ runs in a worker process. + if i == n - 1: + self._resume_offset = initial_offset + local_consumed[i] + ws = self._worker_stats + ws[0] = sum( + split_sizes[j] - fetch_head[j] for j in range(n) + ) + ws[1] = sum( + batch.num_rows for q in raw_batches for batch in q + ) + ws[2] = sum(len(q) for q in cooked) + ws[3] = sum(local_consumed) + ws[4] = self._bytes_loaded + ws[5] = int(self._fetch_time * 1_000_000) + ws[6] = int(self._transform_time * 1_000_000) + + yield row + finally: + self._raw_batches_ref = None + self._cooked_ref = None + self._fetch_head_ref = None + self._split_sizes_ref = None + self._local_consumed_ref = None + + @property + def bytes_loaded(self) -> int: + """Cumulative bytes of raw Arrow buffer data fetched from storage. + + Measured on the ``RecordBatch`` before any transform is applied, so + the value reflects actual I/O rather than the size of transformed + output. Accumulates across multiple iterations of the same dataset + instance and is never reset automatically. + """ + if self._raw_batches_ref is not None: + return self._bytes_loaded + return int(self._worker_stats[4]) + + @property + def fetch_time(self) -> float: + """Cumulative seconds spent waiting for data from LanceDB. + + Measured per batch in the Stage 1 I/O threads as the total elapsed + time of the ``take_offsets`` call. Accumulates across all splits and + all iterations. + """ + if self._raw_batches_ref is not None: + return self._fetch_time + return self._worker_stats[5] / 1_000_000 + + @property + def transform_time(self) -> float: + """Cumulative seconds spent applying the transform. + + Measured per batch in the Stage 2 transform threads as the elapsed + time inside the transform callable (or the default ``arrow2python`` + conversion when no transform is set). Accumulates across all splits + and all iterations. + """ + if self._raw_batches_ref is not None: + return self._transform_time + return self._worker_stats[6] / 1_000_000 + + @property + def raw_queue_depth(self) -> int: + """Number of raw rows waiting for a transform thread across all splits. + + A persistently non-zero value means Stage 2 (transform) is the + bottleneck: I/O is completing faster than transforms can consume + batches. Returns 0 when not iterating. + """ + if self._raw_batches_ref is not None: + return sum(batch.num_rows for q in self._raw_batches_ref for batch in q) + return int(self._worker_stats[1]) + + @property + def prefetch_queue_depth(self) -> int: + """Number of rows transformed and ready to yield across all splits. + + Counts rows whose transform has completed and are sitting in memory + waiting for the main thread — rows that can be handed off with no + I/O or CPU wait. Returns 0 when not iterating. + """ + if self._cooked_ref is not None: + return sum(len(q) for q in self._cooked_ref) + return int(self._worker_stats[2]) + + @property + def unscanned_rows(self) -> int: + """Number of rows not yet submitted to the I/O stage across all splits. + + Decreases as the I/O stage submits fetch requests. When this reaches + zero all data has been requested from storage (though it may not have + arrived yet). Returns 0 when not iterating. + """ + if self._fetch_head_ref is not None: + return sum( + size - head + for size, head in zip(self._split_sizes_ref, self._fetch_head_ref) + ) + return int(self._worker_stats[0]) + + @property + def consumed_rows(self) -> int: + """Number of rows already yielded to the caller across all splits. + + Monotonically increases throughout iteration. Returns 0 when not + iterating. + """ + if self._local_consumed_ref is not None: + return sum(self._local_consumed_ref) + return int(self._worker_stats[3]) + + def __getstate__(self): + """Support pickling for multi-worker DataLoader (forkserver / spawn). + + The live LanceDB table object contains non-picklable connection state + (sockets, Rust-backed PyO3 objects). If a ``connection_factory`` was + supplied only the table name is serialised; the factory is called in + the worker to reopen the connection without embedding any credentials. + Without a factory the table's own picklable reopen state is captured + via ``_table_to_pickle_state`` (mirrors the ``Permutation`` approach). + """ + state = self.__dict__.copy() + # _table: replace with reconnect info (credentials must not be embedded). + state["_table_name"] = self._table.name + if self._connection_factory is not None: + state["_table"] = None + else: + state["_table"] = _table_to_pickle_state(self._table) + # _perm_table: always in-memory; serialise as Arrow data (mirrors + # how Permutation.__getstate__ handles its permutation_table). + state["_perm_table"] = ( + self._perm_table.name, + self._perm_table.to_arrow(), + ) + for key in ( + "_raw_batches_ref", + "_cooked_ref", + "_fetch_head_ref", + "_split_sizes_ref", + "_local_consumed_ref", + ): + state[key] = None + return state + + def __setstate__(self, state): + """Reconnect to LanceDB after unpickling in a worker process.""" + from . import connect as _connect + + table_name = state.pop("_table_name") + table_state = state.pop("_table") + perm_name, perm_data = state.pop("_perm_table") + self.__dict__.update(state) + if self._connection_factory is not None: + self._table = self._connection_factory(table_name) + else: + self._table = _table_from_pickle_state(table_state) + self._perm_table = _connect("memory://").create_table(perm_name, perm_data) + + def state_dict(self) -> dict: + """Snapshot the dataset's consumption state. + + The returned dict is topology-independent: at global step boundaries + every split has been consumed the same number of times (by the + round-robin design), so the per-split count is a single uniform value + that is identical across all ranks and DataLoader workers. + """ + return { + "shuffle_seed": self._shuffle_seed, + "num_splits": self._num_splits, + "epoch": self._epoch, + "samples_consumed_per_split": [self._resume_offset] * self._num_splits, + } + + def load_state_dict(self, state: dict) -> None: + """Resume from a previously snapshotted state. + + Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ + from the checkpoint, since a different split structure or shuffle order + makes mid-epoch resumption meaningless. + """ + if state["num_splits"] != self._num_splits: + raise ValueError( + f"num_splits mismatch: checkpoint has {state['num_splits']}, " + f"current dataset has {self._num_splits}" + ) + if state["shuffle_seed"] != self._shuffle_seed: + raise ValueError( + f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, " + f"current dataset has {self._shuffle_seed}" + ) + consumed = state["samples_consumed_per_split"] + # All entries are equal at step boundaries; use the first. + if isinstance(consumed, list): + self._resume_offset = consumed[0] if consumed else 0 + else: + self._resume_offset = int(consumed) diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py new file mode 100644 index 000000000..6a98bed0c --- /dev/null +++ b/python/python/tests/test_elastic_dataloader.py @@ -0,0 +1,1682 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Tests for elastic dataloader properties. + +Two properties are verified: + +1. **Elastic determinism** — for a fixed (num_splits, shuffle_seed, epoch), the + set of samples that forms each global training step is identical regardless of + the training topology (world_size). This mirrors the MDS guarantee described + in dataloader.md: "the global batch at a given step will always contain the + same samples." + +2. **Resumability** — after calling state_dict() and load_state_dict(), a fresh + dataset continues from exactly where the previous run stopped: no sample is + skipped, no sample is repeated. The checkpoint is topology-independent (it + captures per-split consumption counts), so resumption works even when + world_size changes between runs. + +Tests use explicit rank / world_size constructor parameters rather than +torch.distributed, so they run in a single process without a distributed process +group. The import guard below causes the whole module to be skipped when +lancedb.streaming does not yet exist, which keeps CI green while the +implementation is in progress. + +Parameters used throughout: + NUM_ROWS = 120 (divisible by all tested world_sizes and num_splits) + NUM_SPLITS = 12 (divides cleanly by world_sizes 1,2,3,4,6,12) + GLOBAL_BATCH_SIZE = 12 (= num_splits → each split contributes 1 sample/step) + STEPS_PER_EPOCH = 10 (= NUM_ROWS / GLOBAL_BATCH_SIZE) +""" + +import dataclasses +import logging +from unittest.mock import patch + +import lancedb +import pyarrow as pa +import pytest + +torch = pytest.importorskip("torch") +streaming = pytest.importorskip("lancedb.streaming") +StreamingDataset = streaming.StreamingDataset + +# --------------------------------------------------------------------------- +# Dataset parameters +# --------------------------------------------------------------------------- + +NUM_ROWS = 120 +NUM_SPLITS = 12 +GLOBAL_BATCH_SIZE = NUM_SPLITS # one sample per split per step +STEPS_PER_EPOCH = NUM_ROWS // GLOBAL_BATCH_SIZE # = 10 +SHUFFLE_SEED = 42 + +# world_sizes compatible with NUM_SPLITS (i.e. NUM_SPLITS % world_size == 0) +COMPATIBLE_WORLD_SIZES = [1, 2, 3, 4, 6, 12] + +# Parameters for the "large global batch" tests where global_batch_size is a +# multiple of num_splits rather than equal to it. Using 3× gives 3 samples per +# split per step, which exercises the path where each rank drains more than one +# sample per split per global step. +LARGE_GLOBAL_BATCH_SIZE = 36 # = 3 * NUM_SPLITS +LARGE_NUM_ROWS = 360 # must be divisible by LARGE_GLOBAL_BATCH_SIZE (360/36=10) + +# (world_size, num_workers) pairs used in multi-worker tests. +# Constraint: num_splits % (world_size * num_workers) == 0, i.e. world_size * +# num_workers must divide NUM_SPLITS (12). +MULTI_WORKER_TOPOLOGIES = [ + (1, 2), + (1, 3), + (1, 4), + (2, 2), + (2, 3), +] + + +@dataclasses.dataclass +class FakeWorkerInfo: + """Stand-in for torch.utils.data.WorkerInfo, accepted by worker_info_override. + + Mirrors the attributes that StreamingDataset reads from the real WorkerInfo: + id — this worker's index (0-based) + num_workers — total workers for this rank's DataLoader + """ + + id: int + num_workers: int + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def lance_table(tmp_path): + """A simple LanceDB table with integer IDs 0 .. NUM_ROWS-1.""" + db = lancedb.connect(tmp_path) + return db.create_table("data", pa.table({"id": list(range(NUM_ROWS))})) + + +@pytest.fixture +def lance_table_large(tmp_path): + """A larger table for tests where global_batch_size > num_splits. + + Uses LARGE_NUM_ROWS rows. + """ + db = lancedb.connect(tmp_path) + return db.create_table("data", pa.table({"id": list(range(LARGE_NUM_ROWS))})) + + +# --------------------------------------------------------------------------- +# Simulation helpers +# --------------------------------------------------------------------------- + + +def _make_dataset( + table, + rank: int, + world_size: int, + *, + num_splits: int = NUM_SPLITS, + shuffle_seed: int = SHUFFLE_SEED, + epoch: int = 0, + worker_info_override: FakeWorkerInfo | None = None, +) -> StreamingDataset: + """Create a StreamingDataset configured for a specific rank in a simulated + distributed run. rank and world_size are passed explicitly so we do not + need an actual torch.distributed process group. + + Pass worker_info_override to simulate a specific DataLoader worker without + spawning real worker processes. + """ + return StreamingDataset( + table, + num_splits=num_splits, + shuffle_seed=shuffle_seed, + epoch=epoch, + rank=rank, + world_size=world_size, + worker_info_override=worker_info_override, + ) + + +def _collect_global_batches( + table, + world_size: int, + *, + num_splits: int = NUM_SPLITS, + global_batch_size: int = GLOBAL_BATCH_SIZE, + shuffle_seed: int = SHUFFLE_SEED, + epoch: int = 0, +) -> list[frozenset[int]]: + """Drain a full epoch and return one frozenset of sample IDs per global step. + + Each rank is represented by one StreamingDataset. A "global step" consumes + (global_batch_size // world_size) samples from every rank, so the global batch + at step k is the union of all per-rank micro-batches at that step. frozensets + are used so comparisons are order-independent within a batch. + + global_batch_size must be a multiple of num_splits (each split contributes + global_batch_size // num_splits samples per step). + + Raises AssertionError if the per-rank iterators exhaust at different times, + which would indicate a padding/truncation bug in the implementation. + """ + assert num_splits % world_size == 0, ( + f"world_size={world_size} does not divide num_splits={num_splits}" + ) + assert global_batch_size % num_splits == 0, ( + f"global_batch_size={global_batch_size} is not a multiple of " + f"num_splits={num_splits}" + ) + micro = global_batch_size // world_size # samples per rank per global step + + datasets = [ + _make_dataset( + table, + rank, + world_size, + num_splits=num_splits, + shuffle_seed=shuffle_seed, + epoch=epoch, + ) + for rank in range(world_size) + ] + iters = [iter(ds) for ds in datasets] + + _STOP = object() + global_batches: list[frozenset[int]] = [] + while True: + step_samples: set[int] = set() + exhausted = 0 + for it in iters: + for _ in range(micro): + val = next(it, _STOP) + if val is _STOP: + exhausted += 1 + break + step_samples.add(val["id"]) + if exhausted == len(iters): + # All iterators exhausted at the same step — end of epoch. + break + assert exhausted == 0, ( + "Rank iterators exhausted at different steps. " + "StreamingDataset must produce equal-length sequences across ranks." + ) + global_batches.append(frozenset(step_samples)) + + return global_batches + + +def _advance_and_checkpoint( + table, + world_size: int, + steps: int, + *, + num_splits: int = NUM_SPLITS, + global_batch_size: int = GLOBAL_BATCH_SIZE, + shuffle_seed: int = SHUFFLE_SEED, + epoch: int = 0, +) -> tuple[list[frozenset[int]], dict]: + """Run `steps` global steps, then snapshot the dataset state. + + Returns (batches_consumed, checkpoint) where checkpoint is the state_dict + from rank 0's dataset. Because state is per-split (not per-rank), any + rank's state_dict can be used to resume on any compatible topology. + """ + assert num_splits % world_size == 0 + assert global_batch_size % num_splits == 0 + micro = global_batch_size // world_size + + datasets = [ + _make_dataset( + table, + rank, + world_size, + num_splits=num_splits, + shuffle_seed=shuffle_seed, + epoch=epoch, + ) + for rank in range(world_size) + ] + iters = [iter(ds) for ds in datasets] + + seen: list[frozenset[int]] = [] + for _ in range(steps): + step_samples: set[int] = set() + for it in iters: + for _ in range(micro): + step_samples.add(next(it)["id"]) + seen.append(frozenset(step_samples)) + + checkpoint = datasets[0].state_dict() + return seen, checkpoint + + +def _resume_and_collect( + table, + world_size: int, + checkpoint: dict, + *, + num_splits: int = NUM_SPLITS, + global_batch_size: int = GLOBAL_BATCH_SIZE, + shuffle_seed: int = SHUFFLE_SEED, + epoch: int = 0, +) -> list[frozenset[int]]: + """Create fresh datasets, load checkpoint, and drain to epoch end. + + world_size may differ from the run that produced the checkpoint, + exercising the elastic-resume path. + """ + assert num_splits % world_size == 0 + assert global_batch_size % num_splits == 0 + micro = global_batch_size // world_size + + datasets = [ + _make_dataset( + table, + rank, + world_size, + num_splits=num_splits, + shuffle_seed=shuffle_seed, + epoch=epoch, + ) + for rank in range(world_size) + ] + for ds in datasets: + ds.load_state_dict(checkpoint) + + iters = [iter(ds) for ds in datasets] + _STOP2 = object() + remaining: list[frozenset[int]] = [] + while True: + step_samples: set[int] = set() + exhausted = 0 + for it in iters: + for _ in range(micro): + val = next(it, _STOP2) + if val is _STOP2: + exhausted += 1 + break + step_samples.add(val["id"]) + if exhausted == len(iters): + break + assert exhausted == 0 + remaining.append(frozenset(step_samples)) + + return remaining + + +# --------------------------------------------------------------------------- +# Elastic determinism tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_elastic_det_full_coverage(lance_table, world_size): + """Every sample ID appears exactly once across all steps of an epoch.""" + batches = _collect_global_batches(lance_table, world_size) + all_seen = sorted(sid for b in batches for sid in b) + assert all_seen == list(range(NUM_ROWS)), ( + f"world_size={world_size}: expected IDs 0..{NUM_ROWS - 1} but got {all_seen}" + ) + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_elastic_det_correct_step_count(lance_table, world_size): + """Number of global steps equals NUM_ROWS // GLOBAL_BATCH_SIZE.""" + batches = _collect_global_batches(lance_table, world_size) + assert len(batches) == STEPS_PER_EPOCH, ( + f"world_size={world_size}: expected {STEPS_PER_EPOCH} steps, got {len(batches)}" + ) + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_elastic_det_no_intra_batch_duplicates(lance_table, world_size): + """Within a single global batch, every sample ID is distinct.""" + batches = _collect_global_batches(lance_table, world_size) + for step, batch in enumerate(batches): + assert len(batch) == GLOBAL_BATCH_SIZE, ( + f"world_size={world_size} step {step}: " + f"expected {GLOBAL_BATCH_SIZE} samples, got {len(batch)}" + ) + + +def test_elastic_det_same_batches_across_world_sizes(lance_table): + """The global batch at each step is identical for every compatible world_size. + + This is the core elastic-determinism guarantee from the design doc. + """ + reference = _collect_global_batches(lance_table, world_size=1) + for ws in COMPATIBLE_WORLD_SIZES[1:]: + batches = _collect_global_batches(lance_table, world_size=ws) + assert len(batches) == len(reference), ( + f"world_size={ws} produced {len(batches)} steps; " + f"world_size=1 produced {len(reference)}" + ) + for step, (ref_batch, batch) in enumerate(zip(reference, batches)): + assert ref_batch == batch, ( + f"Global batch at step {step} differs between world_size=1 and " + f"world_size={ws}.\n" + f" world_size=1 → {sorted(ref_batch)}\n" + f" world_size={ws} → {sorted(batch)}" + ) + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_elastic_det_reproducible(lance_table, world_size): + """Two independent runs with the same parameters produce identical epochs.""" + batches_a = _collect_global_batches(lance_table, world_size) + batches_b = _collect_global_batches(lance_table, world_size) + assert batches_a == batches_b, ( + f"world_size={world_size}: second run produced different batches" + ) + + +@pytest.mark.parametrize("epoch_a,epoch_b", [(0, 1), (1, 2), (0, 2)]) +def test_elastic_det_different_epochs_differ(lance_table, epoch_a, epoch_b): + """Different epochs use different shuffles and therefore produce different + step sequences. (Both cover all samples; they just appear in different order.) + """ + batches_a = _collect_global_batches(lance_table, world_size=1, epoch=epoch_a) + batches_b = _collect_global_batches(lance_table, world_size=1, epoch=epoch_b) + # Same set of samples overall... + assert sorted(sid for b in batches_a for sid in b) == list(range(NUM_ROWS)) + assert sorted(sid for b in batches_b for sid in b) == list(range(NUM_ROWS)) + # ...but a different step-by-step ordering. + assert batches_a != batches_b, ( + f"Epoch {epoch_a} and epoch {epoch_b} produced identical step sequences; " + "shuffle_seed should incorporate the epoch number." + ) + + +def test_elastic_det_different_seeds_differ(lance_table): + """Different shuffle seeds produce different sample orderings.""" + batches_seed0 = _collect_global_batches(lance_table, world_size=1, shuffle_seed=0) + batches_seed1 = _collect_global_batches(lance_table, world_size=1, shuffle_seed=1) + assert batches_seed0 != batches_seed1, ( + "Different shuffle_seed values must produce different sample orderings." + ) + + +# --------------------------------------------------------------------------- +# Resumability tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "checkpoint_at_step", [0, 1, 4, STEPS_PER_EPOCH - 1, STEPS_PER_EPOCH] +) +def test_resumability_same_world_size(lance_table, checkpoint_at_step): + """Resuming on the same world_size produces the same global batches as a + reference full-epoch run, split at the checkpoint step. + + Combined (before + after resume): + - No sample is repeated. + - No sample is skipped. + - Each global batch matches the reference exactly. + """ + world_size = 2 + reference = _collect_global_batches(lance_table, world_size) + + seen_before, checkpoint = _advance_and_checkpoint( + lance_table, world_size, checkpoint_at_step + ) + seen_after = _resume_and_collect(lance_table, world_size, checkpoint) + + # Correct total number of steps. + assert len(seen_before) + len(seen_after) == STEPS_PER_EPOCH, ( + f"Expected {STEPS_PER_EPOCH} total steps; " + f"got {len(seen_before)} + {len(seen_after)}" + ) + + # Each pre-checkpoint batch matches the reference. + for step, (ref, got) in enumerate(zip(reference, seen_before)): + assert ref == got, ( + f"Pre-checkpoint batch {step} doesn't match reference.\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + # Each post-resume batch matches the reference (continuing from + # checkpoint_at_step). + for offset, (ref, got) in enumerate( + zip(reference[checkpoint_at_step:], seen_after) + ): + step = checkpoint_at_step + offset + assert ref == got, ( + f"Post-resume batch {step} doesn't match reference.\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + # Full coverage: every sample seen exactly once. + all_seen = sorted(sid for b in seen_before + seen_after for sid in b) + assert all_seen == list(range(NUM_ROWS)) + + +@pytest.mark.parametrize( + "initial_world_size,resume_world_size", + [ + (1, 2), + (1, 4), + (1, 12), + (2, 1), + (2, 4), + (4, 2), + (4, 1), + (3, 6), + (6, 3), + (12, 1), + ], +) +def test_resumability_elastic_world_size_change( + lance_table, initial_world_size, resume_world_size +): + """Resuming on a *different* world_size (elastic resume) produces the same + global batches as the reference epoch for both the pre- and post-checkpoint + phases. No sample is skipped or repeated. + + This is the central elastic-resume guarantee: because state is stored + per-split (not per-rank), the checkpoint is topology-independent. + """ + checkpoint_at_step = 4 + reference = _collect_global_batches(lance_table, world_size=initial_world_size) + + seen_before, checkpoint = _advance_and_checkpoint( + lance_table, initial_world_size, checkpoint_at_step + ) + seen_after = _resume_and_collect(lance_table, resume_world_size, checkpoint) + + # Pre-checkpoint batches match reference. + for step, (ref, got) in enumerate(zip(reference, seen_before)): + assert ref == got, ( + f"Pre-checkpoint batch {step} differs " + f"(initial_world_size={initial_world_size}).\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + # Post-resume batches match reference (elastic determinism also holds + # for the resumed phase with the new world_size). + for offset, (ref, got) in enumerate( + zip(reference[checkpoint_at_step:], seen_after) + ): + step = checkpoint_at_step + offset + assert ref == got, ( + f"Post-resume batch {step} differs " + f"(resume_world_size={resume_world_size}).\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + # No overlap and full coverage. + before_ids = {sid for b in seen_before for sid in b} + after_ids = {sid for b in seen_after for sid in b} + assert before_ids.isdisjoint(after_ids), ( + f"Samples seen before and after resume overlap: " + f"{sorted(before_ids & after_ids)}" + ) + assert before_ids | after_ids == set(range(NUM_ROWS)), ( + f"Not all samples covered. Missing: " + f"{sorted(set(range(NUM_ROWS)) - (before_ids | after_ids))}" + ) + + +def test_resumability_state_dict_is_topology_independent(lance_table): + """state_dict() captures per-split consumption, so it must be identical + across all ranks at the same training step. + + If ranks produce different state dicts, resumption with a different + world_size would give inconsistent results depending on which rank's + checkpoint was saved. + """ + checkpoint_at_step = 4 + world_size = 4 + micro = NUM_SPLITS // world_size + + datasets = [ + _make_dataset(lance_table, rank, world_size) for rank in range(world_size) + ] + iters = [iter(ds) for ds in datasets] + + for _ in range(checkpoint_at_step): + for it in iters: + for _ in range(micro): + next(it) + + state_dicts = [ds.state_dict() for ds in datasets] + for rank in range(1, world_size): + assert state_dicts[rank] == state_dicts[0], ( + f"state_dict from rank {rank} differs from rank 0.\n" + f" rank 0: {state_dicts[0]}\n" + f" rank {rank}: {state_dicts[rank]}" + ) + + +def test_resumability_round_trip_is_deterministic(lance_table): + """Loading the same checkpoint twice produces identical post-resume sequences.""" + world_size = 2 + checkpoint_at_step = 3 + + _, checkpoint = _advance_and_checkpoint(lance_table, world_size, checkpoint_at_step) + remaining_a = _resume_and_collect(lance_table, world_size, checkpoint) + remaining_b = _resume_and_collect(lance_table, world_size, checkpoint) + + assert remaining_a == remaining_b, ( + "Same checkpoint produced different sequences on two separate resumes." + ) + + +def test_resumability_at_epoch_start(lance_table): + """A checkpoint taken before any samples are consumed resumes as a full epoch.""" + world_size = 2 + reference = _collect_global_batches(lance_table, world_size) + + seen_before, checkpoint = _advance_and_checkpoint(lance_table, world_size, steps=0) + seen_after = _resume_and_collect(lance_table, world_size, checkpoint) + + assert seen_before == [] + assert seen_after == reference + + +def test_resumability_at_epoch_end(lance_table): + """A checkpoint taken after all steps are consumed yields an empty resume.""" + world_size = 2 + + _, checkpoint = _advance_and_checkpoint(lance_table, world_size, STEPS_PER_EPOCH) + remaining = _resume_and_collect(lance_table, world_size, checkpoint) + + assert remaining == [], ( + f"Expected no remaining samples after epoch-end checkpoint, " + f"but got {len(remaining)} batches: {remaining}" + ) + + +def test_resumability_state_dict_contains_required_keys(lance_table): + """state_dict() must contain the keys the design doc mandates. + + From the design doc: + - shuffle_seed (must match on resume) + - num_splits (must match on resume) + - epoch (current epoch) + - samples_consumed_per_split (per-split counter, topology-independent) + """ + _, checkpoint = _advance_and_checkpoint(lance_table, world_size=1, steps=3) + required_keys = { + "shuffle_seed", + "num_splits", + "epoch", + "samples_consumed_per_split", + } + missing = required_keys - checkpoint.keys() + assert not missing, ( + f"state_dict() is missing required keys: {missing}\n" + f"Got keys: {set(checkpoint.keys())}" + ) + + +def test_resumability_mismatched_num_splits_raises(lance_table): + """Loading a checkpoint with a different num_splits must raise an error. + + Changing num_splits invalidates the split partitioning and mid-epoch resume + is no longer meaningful. + """ + _, checkpoint = _advance_and_checkpoint(lance_table, world_size=1, steps=3) + + different_splits = NUM_SPLITS * 2 + ds = StreamingDataset( + lance_table, + num_splits=different_splits, + shuffle_seed=SHUFFLE_SEED, + rank=0, + world_size=1, + ) + with pytest.raises((ValueError, RuntimeError)): + ds.load_state_dict(checkpoint) + + +def test_resumability_mismatched_shuffle_seed_raises(lance_table): + """Loading a checkpoint with a different shuffle_seed must raise an error. + + A different seed produces a different sample ordering, so the per-split + consumption counts in the checkpoint no longer refer to the same samples. + """ + _, checkpoint = _advance_and_checkpoint(lance_table, world_size=1, steps=3) + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED + 1, + rank=0, + world_size=1, + ) + with pytest.raises((ValueError, RuntimeError)): + ds.load_state_dict(checkpoint) + + +# --------------------------------------------------------------------------- +# Large global batch size tests (global_batch_size is a multiple of num_splits) +# --------------------------------------------------------------------------- +# These tests use LARGE_GLOBAL_BATCH_SIZE = 36 = 3 * NUM_SPLITS, so each split +# contributes 3 samples per global step instead of 1. The elastic-determinism +# and resumability properties must hold regardless of this multiplier. + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_large_batch_elastic_det_full_coverage(lance_table_large, world_size): + """Every sample appears exactly once per epoch regardless of global_batch_size.""" + batches = _collect_global_batches( + lance_table_large, + world_size, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + all_seen = sorted(sid for b in batches for sid in b) + assert all_seen == list(range(LARGE_NUM_ROWS)) + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_large_batch_elastic_det_correct_step_count(lance_table_large, world_size): + """Steps per epoch = LARGE_NUM_ROWS // LARGE_GLOBAL_BATCH_SIZE.""" + expected_steps = LARGE_NUM_ROWS // LARGE_GLOBAL_BATCH_SIZE + batches = _collect_global_batches( + lance_table_large, + world_size, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + assert len(batches) == expected_steps, ( + f"world_size={world_size}: expected {expected_steps} steps, got {len(batches)}" + ) + + +@pytest.mark.parametrize("world_size", COMPATIBLE_WORLD_SIZES) +def test_large_batch_elastic_det_correct_batch_size(lance_table_large, world_size): + """Each global batch contains exactly LARGE_GLOBAL_BATCH_SIZE distinct samples.""" + batches = _collect_global_batches( + lance_table_large, + world_size, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + for step, batch in enumerate(batches): + assert len(batch) == LARGE_GLOBAL_BATCH_SIZE, ( + f"world_size={world_size} step {step}: " + f"expected {LARGE_GLOBAL_BATCH_SIZE} samples, got {len(batch)}" + ) + + +def test_large_batch_elastic_det_same_across_topologies(lance_table_large): + """Global batch k is identical for every compatible world_size when + global_batch_size (36) is a 3× multiple of num_splits (12).""" + reference = _collect_global_batches( + lance_table_large, + world_size=1, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + for ws in COMPATIBLE_WORLD_SIZES[1:]: + batches = _collect_global_batches( + lance_table_large, + world_size=ws, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + assert len(batches) == len(reference) + for step, (ref, got) in enumerate(zip(reference, batches)): + assert ref == got, ( + f"Global batch at step {step} differs between world_size=1 and " + f"world_size={ws} (global_batch_size={LARGE_GLOBAL_BATCH_SIZE}).\n" + f" world_size=1 → {sorted(ref)}\n" + f" world_size={ws} → {sorted(got)}" + ) + + +@pytest.mark.parametrize( + "initial_world_size,resume_world_size", + [(1, 4), (2, 4), (4, 2), (4, 1), (3, 6), (6, 3)], +) +def test_large_batch_resumability_elastic_world_size_change( + lance_table_large, initial_world_size, resume_world_size +): + """Elastic resume with global_batch_size=36 (3× num_splits): combined + samples before and after the checkpoint match the reference epoch exactly.""" + checkpoint_at_step = 4 + reference = _collect_global_batches( + lance_table_large, + world_size=initial_world_size, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + + seen_before, checkpoint = _advance_and_checkpoint( + lance_table_large, + initial_world_size, + checkpoint_at_step, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + seen_after = _resume_and_collect( + lance_table_large, + resume_world_size, + checkpoint, + global_batch_size=LARGE_GLOBAL_BATCH_SIZE, + ) + + for step, (ref, got) in enumerate(zip(reference, seen_before)): + assert ref == got, ( + f"Pre-checkpoint batch {step} differs " + f"(initial_world_size={initial_world_size}, " + f"global_batch_size={LARGE_GLOBAL_BATCH_SIZE}).\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + for offset, (ref, got) in enumerate( + zip(reference[checkpoint_at_step:], seen_after) + ): + step = checkpoint_at_step + offset + assert ref == got, ( + f"Post-resume batch {step} differs " + f"(resume_world_size={resume_world_size}, " + f"global_batch_size={LARGE_GLOBAL_BATCH_SIZE}).\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + before_ids = {sid for b in seen_before for sid in b} + after_ids = {sid for b in seen_after for sid in b} + assert before_ids.isdisjoint(after_ids), ( + f"Samples overlap across checkpoint: {sorted(before_ids & after_ids)}" + ) + assert before_ids | after_ids == set(range(LARGE_NUM_ROWS)) + + +# --------------------------------------------------------------------------- +# Multi-worker tests (num_workers > 1) +# --------------------------------------------------------------------------- +# PyTorch DataLoader collects a full batch from each worker before moving to the +# next (batch-level, not item-level, round-robin). For the dataset's batches to +# be consistent across num_workers, the dataset must assign each worker a +# contiguous block of splits. Worker k owns splits in the range: +# [k * (num_splits // num_workers), (k+1) * (num_splits // num_workers)) +# (within a single rank's split allocation). +# +# Tests in this section use worker_info_override to simulate multiple workers +# without spawning real DataLoader worker processes. + + +def _collect_global_batches_multi_worker( + table, + world_size: int, + num_workers: int, + *, + num_splits: int = NUM_SPLITS, + global_batch_size: int = GLOBAL_BATCH_SIZE, + shuffle_seed: int = SHUFFLE_SEED, + epoch: int = 0, +) -> list[frozenset[int]]: + """Drain a full epoch using multiple simulated workers per rank. + + Creates one StreamingDataset per (rank, worker_id) pair, each configured + via worker_info_override. Samples from all (rank, worker_id) pairs at the + same global step are combined into one frozenset, mirroring the contiguous + batch-level collection that PyTorch DataLoader performs. + + Constraint: num_splits % (world_size * num_workers) == 0 + """ + assert num_splits % (world_size * num_workers) == 0, ( + f"num_splits={num_splits} must be divisible by " + f"world_size * num_workers = {world_size * num_workers}" + ) + assert global_batch_size % (world_size * num_workers) == 0 + samples_per_worker_per_step = global_batch_size // (world_size * num_workers) + + # Collect all samples from every (rank, worker_id) pair. + worker_samples: dict[tuple[int, int], list[int]] = {} + for rank in range(world_size): + for worker_id in range(num_workers): + ds = _make_dataset( + table, + rank, + world_size, + num_splits=num_splits, + shuffle_seed=shuffle_seed, + epoch=epoch, + worker_info_override=FakeWorkerInfo( + id=worker_id, num_workers=num_workers + ), + ) + worker_samples[(rank, worker_id)] = [item["id"] for item in ds] + + lengths = {len(s) for s in worker_samples.values()} + assert len(lengths) == 1, ( + f"Workers produced unequal sample counts: " + f"{dict(zip(worker_samples, (len(s) for s in worker_samples.values())))}" + ) + n_per_worker = lengths.pop() + assert n_per_worker % samples_per_worker_per_step == 0 + n_steps = n_per_worker // samples_per_worker_per_step + + global_batches: list[frozenset[int]] = [] + for step in range(n_steps): + batch: set[int] = set() + for samples in worker_samples.values(): + start = step * samples_per_worker_per_step + batch.update(samples[start : start + samples_per_worker_per_step]) + global_batches.append(frozenset(batch)) + + return global_batches + + +def _advance_and_checkpoint_multi_worker( + table, + world_size: int, + num_workers: int, + steps: int, + *, + num_splits: int = NUM_SPLITS, + global_batch_size: int = GLOBAL_BATCH_SIZE, + shuffle_seed: int = SHUFFLE_SEED, + epoch: int = 0, +) -> tuple[list[frozenset[int]], dict]: + """Run `steps` global steps with multiple workers per rank, then checkpoint.""" + assert num_splits % (world_size * num_workers) == 0 + assert global_batch_size % (world_size * num_workers) == 0 + samples_per_worker_per_step = global_batch_size // (world_size * num_workers) + + datasets: dict[tuple[int, int], StreamingDataset] = {} + iters: dict[tuple[int, int], object] = {} + for rank in range(world_size): + for worker_id in range(num_workers): + ds = _make_dataset( + table, + rank, + world_size, + num_splits=num_splits, + shuffle_seed=shuffle_seed, + epoch=epoch, + worker_info_override=FakeWorkerInfo( + id=worker_id, num_workers=num_workers + ), + ) + datasets[(rank, worker_id)] = ds + iters[(rank, worker_id)] = iter(ds) + + seen: list[frozenset[int]] = [] + for _ in range(steps): + batch: set[int] = set() + for it in iters.values(): + for _ in range(samples_per_worker_per_step): + batch.add(next(it)["id"]) + seen.append(frozenset(batch)) + + # State is per-split → same across all rank/worker combinations. + checkpoint = datasets[(0, 0)].state_dict() + return seen, checkpoint + + +# ── Multi-worker correctness ────────────────────────────────────────────────── + + +@pytest.mark.parametrize("world_size,num_workers", MULTI_WORKER_TOPOLOGIES) +def test_multi_worker_full_coverage(lance_table, world_size, num_workers): + """Every sample is seen exactly once across all workers and all steps.""" + batches = _collect_global_batches_multi_worker(lance_table, world_size, num_workers) + all_seen = sorted(sid for b in batches for sid in b) + assert all_seen == list(range(NUM_ROWS)), ( + f"world_size={world_size} num_workers={num_workers}: " + f"expected IDs 0..{NUM_ROWS - 1}" + ) + + +@pytest.mark.parametrize("world_size,num_workers", MULTI_WORKER_TOPOLOGIES) +def test_multi_worker_correct_step_count(lance_table, world_size, num_workers): + """Step count equals NUM_ROWS // GLOBAL_BATCH_SIZE regardless of topology.""" + batches = _collect_global_batches_multi_worker(lance_table, world_size, num_workers) + assert len(batches) == STEPS_PER_EPOCH, ( + f"world_size={world_size} num_workers={num_workers}: " + f"expected {STEPS_PER_EPOCH} steps, got {len(batches)}" + ) + + +@pytest.mark.parametrize("world_size,num_workers", MULTI_WORKER_TOPOLOGIES) +def test_multi_worker_no_cross_worker_overlap(lance_table, world_size, num_workers): + """Workers within the same rank have disjoint sample sets.""" + for rank in range(world_size): + per_worker: list[set[int]] = [] + for worker_id in range(num_workers): + ds = _make_dataset( + lance_table, + rank, + world_size, + worker_info_override=FakeWorkerInfo( + id=worker_id, num_workers=num_workers + ), + ) + per_worker.append({item["id"] for item in ds}) + + for i in range(num_workers): + for j in range(i + 1, num_workers): + overlap = per_worker[i] & per_worker[j] + assert not overlap, ( + f"rank={rank} workers {i} and {j} share samples: {sorted(overlap)}" + ) + + +# ── Elastic determinism with num_workers ───────────────────────────────────── + + +@pytest.mark.parametrize("world_size,num_workers", MULTI_WORKER_TOPOLOGIES) +def test_multi_worker_same_global_batches_as_single_worker( + lance_table, world_size, num_workers +): + """Global batches produced with num_workers > 1 match those from num_workers=1. + + PyTorch DataLoader collects a complete batch from each worker before moving + to the next. With contiguous split assignment, worker k's batch at step s + contains the same samples as the corresponding segment of the single-worker + batch at step s. The union across all workers at step s is therefore the + same global batch. + """ + reference = _collect_global_batches(lance_table, world_size) + batches = _collect_global_batches_multi_worker(lance_table, world_size, num_workers) + + assert len(batches) == len(reference) + for step, (ref, got) in enumerate(zip(reference, batches)): + assert ref == got, ( + f"Global batch at step {step} differs between num_workers=1 and " + f"num_workers={num_workers} (world_size={world_size}).\n" + f" num_workers=1 → {sorted(ref)}\n" + f" num_workers={num_workers} → {sorted(got)}" + ) + + +def test_multi_worker_elastic_det_across_worker_counts(lance_table): + """Global batches match for every compatible world_size / num_workers pair.""" + reference = _collect_global_batches(lance_table, world_size=1) + for world_size, num_workers in MULTI_WORKER_TOPOLOGIES: + batches = _collect_global_batches_multi_worker( + lance_table, world_size, num_workers + ) + assert batches == reference, ( + f"Mismatch for world_size={world_size}, num_workers={num_workers}" + ) + + +# ── Resumability with num_workers ───────────────────────────────────────────── + + +def test_multi_worker_resumability_same_topology(lance_table): + """Checkpoint with num_workers=2, resume with num_workers=2: exact continuation.""" + world_size = 1 + num_workers = 2 + checkpoint_at_step = 4 + + reference = _collect_global_batches_multi_worker( + lance_table, world_size, num_workers + ) + seen_before, checkpoint = _advance_and_checkpoint_multi_worker( + lance_table, world_size, num_workers, checkpoint_at_step + ) + + # Resume: create new datasets, load checkpoint, collect remaining. + remaining: list[frozenset[int]] = [] + samples_per_worker_per_step = GLOBAL_BATCH_SIZE // (world_size * num_workers) + datasets = {} + iters = {} + for rank in range(world_size): + for worker_id in range(num_workers): + ds = _make_dataset( + lance_table, + rank, + world_size, + worker_info_override=FakeWorkerInfo( + id=worker_id, num_workers=num_workers + ), + ) + ds.load_state_dict(checkpoint) + datasets[(rank, worker_id)] = ds + iters[(rank, worker_id)] = iter(ds) + + _STOP3 = object() + while True: + batch: set[int] = set() + exhausted = 0 + for it in iters.values(): + for _ in range(samples_per_worker_per_step): + val = next(it, _STOP3) + if val is _STOP3: + exhausted += 1 + break + batch.add(val["id"]) + if exhausted == len(iters): + break + assert exhausted == 0 + remaining.append(frozenset(batch)) + + for step, (ref, got) in enumerate(zip(reference, seen_before)): + assert ref == got, f"Pre-checkpoint step {step} doesn't match reference" + for offset, (ref, got) in enumerate(zip(reference[checkpoint_at_step:], remaining)): + step = checkpoint_at_step + offset + assert ref == got, f"Post-resume step {step} doesn't match reference" + + all_seen = sorted(sid for b in seen_before + remaining for sid in b) + assert all_seen == list(range(NUM_ROWS)) + + +def test_multi_worker_resumability_worker_count_change(lance_table): + """Checkpoint with num_workers=1, resume with num_workers=2. + + Because state is per-split, it is independent of num_workers just as it is + independent of world_size. + """ + world_size = 1 + checkpoint_at_step = 4 + reference = _collect_global_batches(lance_table, world_size=1) + + seen_before, checkpoint = _advance_and_checkpoint( + lance_table, world_size=1, steps=checkpoint_at_step + ) + + # Resume with num_workers=2 using the checkpoint. + num_workers = 2 + samples_per_worker_per_step = GLOBAL_BATCH_SIZE // (world_size * num_workers) + iters = {} + for worker_id in range(num_workers): + ds = _make_dataset( + lance_table, + rank=0, + world_size=world_size, + worker_info_override=FakeWorkerInfo(id=worker_id, num_workers=num_workers), + ) + ds.load_state_dict(checkpoint) + iters[worker_id] = iter(ds) + + _STOP4 = object() + remaining: list[frozenset[int]] = [] + while True: + batch: set[int] = set() + exhausted = 0 + for it in iters.values(): + for _ in range(samples_per_worker_per_step): + val = next(it, _STOP4) + if val is _STOP4: + exhausted += 1 + break + batch.add(val["id"]) + if exhausted == len(iters): + break + assert exhausted == 0 + remaining.append(frozenset(batch)) + + for offset, (ref, got) in enumerate(zip(reference[checkpoint_at_step:], remaining)): + step = checkpoint_at_step + offset + assert ref == got, ( + f"Post-resume step {step} doesn't match reference when switching " + f"from num_workers=1 to num_workers=2.\n" + f" reference: {sorted(ref)}\n" + f" got: {sorted(got)}" + ) + + before_ids = {sid for b in seen_before for sid in b} + after_ids = {sid for b in remaining for sid in b} + assert before_ids.isdisjoint(after_ids) + assert before_ids | after_ids == set(range(NUM_ROWS)) + + +# ── worker_info_override warning behaviour ──────────────────────────────────── + + +def test_worker_info_override_logs_warning_when_torch_worker_active( + lance_table, caplog +): + """When worker_info_override is set but get_worker_info() also returns a + value (i.e. the dataset is being iterated inside a real DataLoader worker), + the dataset must log a warning that the override takes precedence and that + this may produce incorrect or duplicated data. + + The warning is important because silently ignoring the real worker assignment + would cause all workers to read the same data. + """ + override = FakeWorkerInfo(id=0, num_workers=2) + ds = _make_dataset(lance_table, rank=0, world_size=1, worker_info_override=override) + + fake_torch_info = FakeWorkerInfo(id=0, num_workers=4) + with patch("lancedb.streaming.get_worker_info", return_value=fake_torch_info): + with caplog.at_level(logging.WARNING, logger="lancedb.streaming"): + list(ds) + + warning_messages = [ + r.message for r in caplog.records if r.levelno >= logging.WARNING + ] + assert warning_messages, ( + "Expected at least one WARNING when worker_info_override is active inside " + "a real DataLoader worker, but no warnings were logged." + ) + combined = " ".join(warning_messages).lower() + assert "override" in combined or "worker_info_override" in combined, ( + f"Warning message did not mention the override. Messages: {warning_messages}" + ) + + +def test_worker_info_override_no_warning_in_main_process(lance_table, caplog): + """When get_worker_info() returns None (main process / num_workers=0), + using worker_info_override is the normal testing path and must not warn. + """ + override = FakeWorkerInfo(id=0, num_workers=2) + ds = _make_dataset(lance_table, rank=0, world_size=1, worker_info_override=override) + + with patch("lancedb.streaming.get_worker_info", return_value=None): + with caplog.at_level(logging.WARNING, logger="lancedb.streaming"): + list(ds) + + warning_messages = [ + r.message for r in caplog.records if r.levelno >= logging.WARNING + ] + assert not warning_messages, ( + f"Unexpected warnings when override used in main process: {warning_messages}" + ) + + +# --------------------------------------------------------------------------- +# Prefetch queue depth tests +# --------------------------------------------------------------------------- + + +def test_concurrent_iteration_raises(lance_table): + """Starting a second iterator while one is already active must raise.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + it1 = iter(ds) + next(it1) # advance it1 so the pipeline is live + + it2 = iter(ds) + with pytest.raises(RuntimeError, match="concurrent"): + next(it2) + + +def test_raw_queue_depth_zero_when_not_iterating(lance_table): + """raw_queue_depth is 0 before iteration starts and after it ends.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + assert ds.raw_queue_depth == 0 + + list(ds) # drain the whole epoch + + assert ds.raw_queue_depth == 0 + + +def test_prefetch_queue_depth_zero_when_not_iterating(lance_table): + """prefetch_queue_depth is 0 before iteration starts and after it ends.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + assert ds.prefetch_queue_depth == 0 + + list(ds) # drain the whole epoch + + assert ds.prefetch_queue_depth == 0 + + +def test_prefetch_queue_depth_positive_during_iteration(lance_table): + """prefetch_queue_depth is > 0 while rows are being yielded.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + it = iter(ds) + next(it) # advance past the first yield; pipeline is now primed + # The other splits still have their initial futures in flight. + assert ds.prefetch_queue_depth > 0 + + list(it) # exhaust the remaining rows + + assert ds.prefetch_queue_depth == 0 + + +# --------------------------------------------------------------------------- +# Transform tests +# --------------------------------------------------------------------------- + + +def test_fetch_and_transform_time_zero_before_iteration(lance_table): + """fetch_time and transform_time start at 0.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + assert ds.fetch_time == 0.0 + assert ds.transform_time == 0.0 + + +def test_fetch_and_transform_time_positive_after_iteration(lance_table): + """Both timers are positive after a full epoch.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + list(ds) + assert ds.fetch_time > 0.0 + assert ds.transform_time > 0.0 + + +def test_fetch_time_excludes_transform(lance_table): + """fetch_time does not include transform time: fetch + transform < total wall time, + and neither counter bleeds into the other.""" + import pyarrow as pa + import time + + def slow_transform(batch: pa.RecordBatch) -> list: + time.sleep(0.01) # 10 ms of artificial transform work + return batch.to_pydict()["id"] + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform=slow_transform, + ) + list(ds) + + # The slow transform should dominate transform_time. + assert ds.transform_time > ds.fetch_time + + +def test_bytes_loaded_increases_after_iteration(lance_table): + """bytes_loaded is 0 before iteration and positive after.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + assert ds.bytes_loaded == 0 + + list(ds) + + assert ds.bytes_loaded > 0 + + +def test_bytes_loaded_measured_before_transform(lance_table): + """bytes_loaded measures raw Arrow size even when transform discards everything.""" + import pyarrow as pa + + # This transform throws away every value. If bytes_loaded were measured + # after the transform, it would see no Arrow data and stay at 0. + def discard_everything(batch: pa.RecordBatch) -> list: + return [None] * batch.num_rows + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform=discard_everything, + ) + list(ds) + + assert ds.bytes_loaded > 0 + + +def test_transform_is_applied(lance_table): + """A custom transform passed to StreamingDataset is forwarded to the + underlying Permutation and applied to every yielded item.""" + import pyarrow as pa + + def id_only(batch: pa.RecordBatch) -> list[int]: + return batch.column("id").to_pylist() + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform=id_only, + ) + items = list(ds) + + assert len(items) == NUM_ROWS + assert all(isinstance(item, int) for item in items), ( + f"Expected ints from transform, got {type(items[0])}" + ) + assert sorted(items) == list(range(NUM_ROWS)) + + +def test_transform_none_yields_dicts(lance_table): + """With no transform (the default), items are plain Python dicts.""" + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + items = list(ds) + + assert len(items) == NUM_ROWS + assert all(isinstance(item, dict) for item in items) + assert all("id" in item for item in items) + + +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.""" + db = lancedb.connect(tmp_path) + table = db.create_table("data", pa.table({"id": list(range(NUM_ROWS))})) + + ds = StreamingDataset( + table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + filter="id < 60", + ) + items = list(ds) + + ids = [item["id"] for item in items] + assert sorted(ids) == list(range(60)), f"Expected ids 0-59, got {sorted(ids)}" + + +def test_filter_too_few_rows_raises(tmp_path): + """A filter that leaves fewer rows than num_splits raises ValueError at + construction time because each split must have at least one row.""" + db = lancedb.connect(tmp_path) + table = db.create_table("data", pa.table({"id": list(range(NUM_ROWS))})) + + with pytest.raises(ValueError, match="at least 1 row per split"): + StreamingDataset( + table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + filter="id < 0", + ) + + +def test_columns_limits_output_columns(tmp_path): + """Only the requested columns are present in each yielded row.""" + db = lancedb.connect(tmp_path) + table = db.create_table( + "data", + pa.table({"id": list(range(NUM_ROWS)), "val": list(range(NUM_ROWS))}), + ) + + ds = StreamingDataset( + table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + columns=["id"], + ) + items = list(ds) + + assert len(items) == NUM_ROWS + assert all(list(item.keys()) == ["id"] for item in items), ( + "Expected only 'id' column in each row" + ) + assert sorted(item["id"] for item in items) == list(range(NUM_ROWS)) + + +def test_columns_invalid_column_raises(tmp_path): + """Requesting a column that does not exist raises an error at iteration time.""" + db = lancedb.connect(tmp_path) + table = db.create_table("data", pa.table({"id": list(range(NUM_ROWS))})) + + ds = StreamingDataset( + table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + columns=["nonexistent"], + ) + with pytest.raises(ValueError): + list(ds) + + +def test_shuffle_clump_size_yields_all_rows(lance_table): + """shuffle_clump_size still produces a complete epoch with no duplicates or + omissions — clumping affects I/O locality, not correctness.""" + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + shuffle_clump_size=4, + ) + items = list(ds) + + assert sorted(item["id"] for item in items) == list(range(NUM_ROWS)), ( + "Expected every row exactly once with shuffle_clump_size set" + ) + + +def test_num_splits_defaults_to_world_size(lance_table): + """Omitting num_splits gives world_size splits (one per rank).""" + ds = StreamingDataset( + lance_table, + shuffle_seed=SHUFFLE_SEED, + ) + assert ds._num_splits == 1 # world_size defaults to 1 + + ds_ws2 = StreamingDataset( + lance_table, + world_size=2, + shuffle_seed=SHUFFLE_SEED, + ) + assert ds_ws2._num_splits == 2 + + +def test_shuffle_false_sequential_and_deterministic(lance_table): + """shuffle=False produces identical ordering across two fresh instances.""" + ds1 = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle=False) + ds2 = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle=False) + first = [item["id"] for item in ds1] + second = [item["id"] for item in ds2] + assert first == second, "shuffle=False must be deterministic" + assert sorted(first) == list(range(NUM_ROWS)), "All rows must be present" + + +def test_shuffle_false_vs_true_differ(lance_table): + """shuffle=True and shuffle=False produce different orderings.""" + ds_shuf = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=True, + shuffle_seed=SHUFFLE_SEED, + ) + ds_seq = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + ) + shuffled = [item["id"] for item in ds_shuf] + sequential = [item["id"] for item in ds_seq] + assert shuffled != sequential, "Shuffled and sequential orderings should differ" + + +def test_shuffle_seed_none_generates_stable_seed(lance_table): + """shuffle_seed=None resolves to a concrete integer at construction time. + + A second dataset built with the same resolved seed must produce the same ordering. + """ + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=None, + ) + assert isinstance(ds._shuffle_seed, int), ( + "shuffle_seed=None must resolve to an integer" + ) + first = [item["id"] for item in ds] + + ds2 = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=ds._shuffle_seed, + ) + second = [item["id"] for item in ds2] + assert first == second, "Same resolved seed must produce the same ordering" + + +# --------------------------------------------------------------------------- +# Doc examples — each test mirrors the code snippet in index.mdx so that +# broken doc examples are caught before they ship. +# --------------------------------------------------------------------------- + + +def test_doc_example_basic(tmp_path): + """doc: Basic Data loading — StreamingDataset with default params.""" + db = lancedb.connect(tmp_path) + table = db.create_table( + "some_table", + pa.table({"feature": [float(i) for i in range(24)], "label": ["cat"] * 24}), + ) + + ds = StreamingDataset(table, shuffle_seed=42) + samples = list(ds) + + assert len(samples) == 24 + assert all(isinstance(s, dict) for s in samples) + assert all("feature" in s and "label" in s for s in samples) + + +def test_doc_example_prefetch_params(tmp_path): + """doc: Prefetching — read_batch_size and prefetch_batches still cover all rows.""" + db = lancedb.connect(tmp_path) + table = db.create_table("t", pa.table({"id": list(range(NUM_ROWS))})) + + ds = StreamingDataset( + table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + read_batch_size=8, + prefetch_batches=2, + ) + assert sorted(s["id"] for s in ds) == list(range(NUM_ROWS)) + + +def test_doc_example_transform(tmp_path): + """doc: Transformation — normalize scales values into [0, 1].""" + db = lancedb.connect(tmp_path) + table = db.create_table("t", pa.table({"value": list(range(NUM_ROWS))})) + + def normalize(batch: pa.RecordBatch) -> list[dict]: + rows = batch.to_pylist() + for row in rows: + row["value"] = row["value"] / 255.0 + return rows + + ds = StreamingDataset( + table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED, transform=normalize + ) + samples = list(ds) + + assert len(samples) == NUM_ROWS + assert all(0.0 <= s["value"] <= 1.0 for s in samples) + + +def test_doc_example_observability(lance_table): + """doc: Observability — counters are zero before and positive after iteration.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + + assert ds.unscanned_rows == 0 + assert ds.raw_queue_depth == 0 + assert ds.prefetch_queue_depth == 0 + assert ds.consumed_rows == 0 + assert ds.bytes_loaded == 0 + assert ds.fetch_time == 0.0 + assert ds.transform_time == 0.0 + + list(ds) + + assert ds.bytes_loaded > 0 + assert ds.fetch_time >= 0.0 + assert ds.transform_time >= 0.0 + + +def test_doc_example_columns_and_filter(tmp_path): + """doc: Filtering data — columns + filter reduce both dimensions independently.""" + db = lancedb.connect(tmp_path) + table = db.create_table( + "t", + pa.table( + { + "id": list(range(NUM_ROWS)), + "value": list(range(NUM_ROWS)), + "category": ["train"] * 60 + ["val"] * 60, + } + ), + ) + + ds = StreamingDataset( + table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + columns=["id"], + filter="category = 'train'", + ) + samples = list(ds) + + assert all(list(s.keys()) == ["id"] for s in samples), "Only 'id' column expected" + assert len(samples) == 60, "Only train rows expected" + assert all(s["id"] < 60 for s in samples) + + +def test_doc_example_epoch_shuffle(lance_table): + """doc: Shuffling rows — different epochs produce different orderings.""" + ids_e0 = [ + s["id"] + for s in StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED, epoch=0 + ) + ] + ids_e1 = [ + s["id"] + for s in StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED, epoch=1 + ) + ] + + assert sorted(ids_e0) == list(range(NUM_ROWS)) + assert sorted(ids_e1) == list(range(NUM_ROWS)) + assert ids_e0 != ids_e1, "Different epochs must produce different orderings" + + +def test_doc_example_shuffle_false_eval(lance_table): + """doc: Shuffling rows — shuffle=False gives deterministic sequential order.""" + ids_a = [ + s["id"] + for s in StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle=False) + ] + ids_b = [ + s["id"] + for s in StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle=False) + ] + + assert ids_a == ids_b, "shuffle=False must produce identical orderings" + assert sorted(ids_a) == list(range(NUM_ROWS)) + + +def test_doc_example_shuffle_clump_size(lance_table): + """doc: Shuffling rows (Note) — shuffle_clump_size still covers every row.""" + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + shuffle_clump_size=16, + ) + assert sorted(s["id"] for s in ds) == list(range(NUM_ROWS)) + + +def test_doc_example_elastic_ddp(lance_table): + """doc: Data splits and elasticity — rank/world_size/num_splits covers all rows.""" + # Use a highly composite num_splits that works across many world sizes + ELASTIC_SPLITS = 12 # works with world_size 1, 2, 3, 4, 6, 12 + + all_ids: set[int] = set() + for rank in range(4): + ds = StreamingDataset( + lance_table, + num_splits=ELASTIC_SPLITS, + shuffle_seed=SHUFFLE_SEED, + rank=rank, + world_size=4, + ) + all_ids.update(s["id"] for s in ds) + + assert all_ids == set(range(NUM_ROWS)), "All ranks together must cover every row" + + +def test_doc_example_checkpoint(lance_table): + """doc: state_dict/load_state_dict checkpointing resumes without gaps or repeats.""" + STEPS_BEFORE_CHECKPOINT = 4 # consume 4 full cycles then checkpoint + + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + it = iter(ds) + consumed = [next(it)["id"] for _ in range(STEPS_BEFORE_CHECKPOINT * NUM_SPLITS)] + checkpoint = ds.state_dict() + remaining_original = [s["id"] for s in it] # drain the rest + + # Resume from checkpoint on a fresh dataset + ds_resumed = StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED + ) + ds_resumed.load_state_dict(checkpoint) + remaining_resumed = [s["id"] for s in ds_resumed] + + assert remaining_original == remaining_resumed, ( + "Resumed dataset must continue from exactly the same position" + ) + assert sorted(consumed + remaining_original) == list(range(NUM_ROWS)), ( + "Consumed + remaining must cover every row exactly once" + ) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 1c1b26295..4f54f5961 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -412,10 +412,12 @@ def test_remote_permutation_is_picklable(): content_len = int(request.headers.get("Content-Length")) body = json.loads(request.rfile.read(content_len)) if "filter" in body: - match = re.search(r"_rowoffset in \((.*?)\)", body["filter"]) - offsets = [int(offset.strip()) for offset in match.group(1).split(",")] + match = re.search( + r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE + ) + offsets = [int(o.strip()) for o in match.group(1).split(",")] else: - offsets = rows + offsets = list(range(len(rows))) table = pa.table({"a": [rows[offset] for offset in offsets]}) request.send_response(200) diff --git a/python/src/permutation.rs b/python/src/permutation.rs index 75e1fe1b7..4dc49cfd3 100644 --- a/python/src/permutation.rs +++ b/python/src/permutation.rs @@ -79,13 +79,14 @@ impl PyAsyncPermutationBuilder { #[pymethods] impl PyAsyncPermutationBuilder { - #[pyo3(signature = (*, ratios=None, counts=None, fixed=None, seed=None, split_names=None))] + #[pyo3(signature = (*, ratios=None, counts=None, fixed=None, seed=None, clump_size=None, split_names=None))] pub fn split_random( slf: PyRefMut<'_, Self>, ratios: Option>, counts: Option>, fixed: Option, seed: Option, + clump_size: Option, split_names: Option>, ) -> PyResult { // Check that exactly one split type is provided @@ -111,7 +112,14 @@ impl PyAsyncPermutationBuilder { }; slf.modify(|builder| { - builder.with_split_strategy(SplitStrategy::Random { seed, sizes }, split_names) + builder.with_split_strategy( + SplitStrategy::Random { + seed, + sizes, + clump_size, + }, + split_names, + ) }) } diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index ba5fd4bcd..761d6e55d 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -108,6 +108,8 @@ http-body = "1" # Matching reqwest rstest = "0.23.0" test-log = "0.2" serial_test = "3" +[target.'cfg(unix)'.dev-dependencies] +pprof = { version = "0.14", features = ["flamegraph"] } [features] @@ -164,6 +166,9 @@ required-features = ["sentence-transformers"] name = "bedrock" required-features = ["bedrock"] +[[example]] +name = "bench_streaming_dataloader" + [[example]] name = "simple" diff --git a/rust/lancedb/examples/bench_streaming_dataloader.rs b/rust/lancedb/examples/bench_streaming_dataloader.rs new file mode 100644 index 000000000..087268ff8 --- /dev/null +++ b/rust/lancedb/examples/bench_streaming_dataloader.rs @@ -0,0 +1,272 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Benchmark + CPU profiler for the PermutationReader used by the elastic +//! streaming dataloader. +//! +//! Normal sweep: +//! cargo run --release --example bench_streaming_dataloader +//! +//! Flamegraph (self-contained, no perf/dtrace needed): +//! BENCH_PROFILE=1 BENCH_CHUNK=64 cargo run --release \ +//! --example bench_streaming_dataloader +//! # writes flamegraph.svg in the current directory +//! +//! Environment variables: +//! BENCH_NUM_ROWS – total rows (default 49152 = 24 × 2048) +//! BENCH_NUM_SPLITS – number of splits (default 24) +//! BENCH_STEPS – round-robin cycles per chunk-size trial (default 200) +//! BENCH_ROW_BYTES – bytes of payload per row (default 4096) +//! BENCH_CHUNK – restrict sweep to this single chunk size +//! BENCH_PROFILE – if set to "1", capture a pprof flamegraph SVG + +use std::{sync::Arc, time::Instant}; + +use arrow_array::{Int32Array, LargeBinaryArray, RecordBatch}; +use arrow_schema::{DataType, Field, Schema}; +use lancedb::{ + Result, Table, + arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}, + connect, + dataloader::permutation::{ + builder::{PermutationBuilder, ShuffleStrategy}, + reader::PermutationReader, + split::{SplitSizes, SplitStrategy}, + }, + query::Select, +}; + +fn env_usize(key: &str, default: usize) -> usize { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +// --------------------------------------------------------------------------- +// Table creation +// --------------------------------------------------------------------------- + +async fn make_base_table(num_rows: usize, row_bytes: usize) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("payload", DataType::LargeBinary, false), + ])); + let payload = vec![0u8; row_bytes]; + let ids: Int32Array = (0..num_rows as i32).collect(); + let payloads: LargeBinaryArray = (0..num_rows).map(|_| Some(payload.as_slice())).collect(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(payloads)])?; + let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream::new( + futures::stream::once(std::future::ready(Ok(batch))), + schema, + )); + let db = connect("memory:///").execute().await?; + db.create_table("base", stream).execute().await +} + +async fn make_permutation_table(base: &Table, num_splits: usize) -> Result
{ + PermutationBuilder::new(base.clone()) + .with_split_strategy( + SplitStrategy::Random { + seed: Some(42), + sizes: SplitSizes::Fixed(num_splits as u64), + clump_size: None, + }, + None, + ) + .with_shuffle_strategy(ShuffleStrategy::Random { + seed: Some(42), + clump_size: None, + }) + .build() + .await +} + +// --------------------------------------------------------------------------- +// Round-robin hot loop (mirrors StreamingDataset.__iter__) +// --------------------------------------------------------------------------- + +async fn run_hot_loop( + readers: &[PermutationReader], + chunk_size: usize, + steps: usize, +) -> Result<(usize, f64)> { + let n = readers.len(); + let split_sizes: Vec = readers.iter().map(|r| r.count_rows() as usize).collect(); + + struct SplitBuf { + batch: Option, + row_in_batch: usize, + consumed: usize, + } + let mut bufs: Vec = (0..n) + .map(|_| SplitBuf { + batch: None, + row_in_batch: 0, + consumed: 0, + }) + .collect(); + + // Pre-fill + for i in 0..n { + let fetch = chunk_size.min(split_sizes[i]); + if fetch > 0 { + let offsets: Vec = (0..fetch as u64).collect(); + bufs[i].batch = Some(readers[i].take_offsets(&offsets, Select::All).await?); + } + } + + let mut total_rows = 0usize; + let t0 = Instant::now(); + + 'outer: for _step in 0..steps { + for i in 0..n { + if bufs[i].consumed >= split_sizes[i] { + break 'outer; + } + let need_refill = bufs[i] + .batch + .as_ref() + .map(|b| bufs[i].row_in_batch >= b.num_rows()) + .unwrap_or(true); + if need_refill { + let start = bufs[i].consumed as u64; + let remaining = (split_sizes[i] - bufs[i].consumed) as u64; + let fetch = chunk_size.min(remaining as usize); + let offsets: Vec = (start..start + fetch as u64).collect(); + bufs[i].batch = Some(readers[i].take_offsets(&offsets, Select::All).await?); + bufs[i].row_in_batch = 0; + } + bufs[i].row_in_batch += 1; + bufs[i].consumed += 1; + total_rows += 1; + } + } + + Ok((total_rows, t0.elapsed().as_secs_f64())) +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +#[tokio::main] +async fn main() -> Result<()> { + let num_splits = env_usize("BENCH_NUM_SPLITS", 24); + let num_rows = env_usize("BENCH_NUM_ROWS", num_splits * 2048); + let steps = env_usize("BENCH_STEPS", 200); + let row_bytes = env_usize("BENCH_ROW_BYTES", 4096); + let single_chunk: Option = std::env::var("BENCH_CHUNK") + .ok() + .and_then(|v| v.parse().ok()); + let do_profile = std::env::var("BENCH_PROFILE") + .map(|v| v == "1") + .unwrap_or(false); + + assert_eq!( + num_rows % num_splits, + 0, + "NUM_ROWS must be divisible by NUM_SPLITS" + ); + + println!("Benchmark config:"); + println!( + " num_rows={} num_splits={} rows/split={} steps={} row_bytes={}", + num_rows, + num_splits, + num_rows / num_splits, + steps, + row_bytes, + ); + println!( + " ~{:.1} MB total", + (num_rows * row_bytes) as f64 / (1024.0 * 1024.0) + ); + println!(); + + print!("Building base table... "); + let _ = std::io::Write::flush(&mut std::io::stdout()); + let base = make_base_table(num_rows, row_bytes).await?; + println!("done"); + + print!("Building permutation table... "); + let _ = std::io::Write::flush(&mut std::io::stdout()); + let perm = make_permutation_table(&base, num_splits).await?; + println!("done"); + + print!("Building {} PermutationReaders... ", num_splits); + let _ = std::io::Write::flush(&mut std::io::stdout()); + let base_inner = base.base_table().clone(); + let perm_inner = perm.base_table().clone(); + let mut readers = Vec::with_capacity(num_splits); + for split in 0..num_splits { + readers.push( + PermutationReader::try_from_tables( + base_inner.clone(), + perm_inner.clone(), + split as u64, + ) + .await?, + ); + } + println!("done ({} rows/split)", readers[0].count_rows()); + println!(); + + let chunk_sizes: Vec = if let Some(c) = single_chunk { + vec![c] + } else { + vec![1, 4, 16, 64, 256, 1024, 4096, 16384] + }; + + if do_profile { + #[cfg(unix)] + { + let chunk = chunk_sizes[0]; + println!("Profiling chunk={chunk} for {steps} steps..."); + // Warm-up outside the profiler window + let _ = run_hot_loop(&readers, chunk, 1).await?; + + let guard = pprof::ProfilerGuardBuilder::default() + .frequency(1000) + .build() + .unwrap(); + + let (rows, elapsed) = run_hot_loop(&readers, chunk, steps).await?; + + if let Ok(report) = guard.report().build() { + let svg_path = "flamegraph.svg"; + let file = std::fs::File::create(svg_path).unwrap(); + report.flamegraph(file).unwrap(); + println!("Flamegraph written to {svg_path}"); + } + + let rows_per_sec = rows as f64 / elapsed; + println!("chunk={chunk} {rows} rows {elapsed:.3}s {rows_per_sec:.0} rows/s"); + } + #[cfg(not(unix))] + { + println!("Flamegraph profiling (BENCH_PROFILE=1) is not supported on this platform."); + println!("Run without BENCH_PROFILE to get throughput numbers."); + } + } else { + println!( + "{:>6} {:>7} {:>8} {:>11} {:>10}", + "chunk", "rows", "elapsed", "rows/s", "ms/step" + ); + println!("{}", "-".repeat(52)); + + for &chunk in &chunk_sizes { + let _ = run_hot_loop(&readers, chunk, 1).await?; + let (rows, elapsed) = run_hot_loop(&readers, chunk, steps).await?; + let rows_per_sec = rows as f64 / elapsed; + let ms_per_step = elapsed / steps as f64 * 1000.0; + println!( + "{:>6} {:>7} {:>7.3}s {:>11.0} {:>9.1}ms", + chunk, rows, elapsed, rows_per_sec, ms_per_step, + ); + } + } + + println!("\nDone."); + Ok(()) +} diff --git a/rust/lancedb/src/dataloader/permutation/builder.rs b/rust/lancedb/src/dataloader/permutation/builder.rs index 377712e98..6b4ae2303 100644 --- a/rust/lancedb/src/dataloader/permutation/builder.rs +++ b/rust/lancedb/src/dataloader/permutation/builder.rs @@ -391,6 +391,7 @@ mod tests { SplitStrategy::Random { seed: Some(42), sizes: SplitSizes::Percentages(vec![0.05, 0.30]), + clump_size: None, }, None, ) diff --git a/rust/lancedb/src/dataloader/permutation/reader.rs b/rust/lancedb/src/dataloader/permutation/reader.rs index 65d065db7..afe79b0ad 100644 --- a/rust/lancedb/src/dataloader/permutation/reader.rs +++ b/rust/lancedb/src/dataloader/permutation/reader.rs @@ -21,6 +21,7 @@ use arrow::compute::concat_batches; use arrow::datatypes::UInt64Type; use arrow_array::{RecordBatch, UInt64Array}; use arrow_schema::SchemaRef; +use datafusion_expr::{Expr, col, lit}; use futures::{StreamExt, TryStreamExt}; use lance::dataset::scanner::DatasetRecordBatchStream; use lance::io::RecordBatchStream; @@ -196,17 +197,10 @@ impl PermutationReader { .expect_ok()? .values(); - let filter = format!( - "_rowid in ({})", - row_ids - .iter() - .map(|o| o.to_string()) - .collect::>() - .join(",") - ); + let in_list: Vec = row_ids.iter().map(|id| lit(*id)).collect(); let base_query = QueryRequest { - filter: Some(QueryFilter::Sql(filter)), + filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))), select: selection, with_row_id: true, ..Default::default() diff --git a/rust/lancedb/src/dataloader/permutation/split.rs b/rust/lancedb/src/dataloader/permutation/split.rs index bb6f59d34..ca9920eff 100644 --- a/rust/lancedb/src/dataloader/permutation/split.rs +++ b/rust/lancedb/src/dataloader/permutation/split.rs @@ -35,9 +35,13 @@ pub enum SplitStrategy { /// Rows will be randomly assigned to splits /// /// A seed can be provided to make the assignment deterministic. + /// + /// A clump_size can be provided to shuffle contiguous groups of rows together, + /// preserving I/O locality while still randomising the split assignment. Random { seed: Option, sizes: SplitSizes, + clump_size: Option, }, /// Rows will be assigned to splits based on the values in the specified columns. /// @@ -323,13 +327,17 @@ impl Splitter { self.apply_sequential(source, num_rows, &SplitSizes::Counts(vec![num_rows])) .await } - SplitStrategy::Random { seed, sizes } => { + SplitStrategy::Random { + seed, + sizes, + clump_size, + } => { let shuffler = Shuffler::new(ShufflerConfig { seed: *seed, // In this case we are only shuffling row ids so we can use a large max_rows_per_file max_rows_per_file: 10 * 1024 * 1024, temp_dir: self.temp_dir.clone(), - clump_size: None, + clump_size: *clump_size, }); let shuffled = shuffler.shuffle(source, num_rows).await?; @@ -692,6 +700,7 @@ mod tests { SplitStrategy::Random { seed: Some(42), sizes: SplitSizes::Fixed(3), + clump_size: None, }, ); @@ -718,6 +727,7 @@ mod tests { SplitStrategy::Random { seed: Some(42), sizes: SplitSizes::Counts(vec![5, 15, 10]), + clump_size: None, }, ); @@ -744,6 +754,7 @@ mod tests { SplitStrategy::Random { seed: Some(42), sizes: SplitSizes::Percentages(vec![0.217, 0.168, 0.17]), + clump_size: None, }, ); diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index d00ca5836..f3dbca1ab 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -7,7 +7,7 @@ use std::{future::Future, time::Duration}; use arrow::compute::concat_batches; use arrow_array::{Array, Float16Array, Float32Array, Float64Array, RecordBatch, make_array}; use arrow_schema::{DataType, SchemaRef}; -use datafusion_expr::Expr; +use datafusion_expr::{Expr, col, lit}; use datafusion_physical_plan::ExecutionPlan; use futures::{FutureExt, TryFutureExt, TryStreamExt, stream, try_join}; use half::f16; @@ -1468,18 +1468,13 @@ impl TakeQuery { /// /// See [`crate::Table::take_offsets`] for more details. pub fn from_offsets(parent: Arc, offsets: Vec) -> Self { - let filter = format!( - "_rowoffset in ({})", - offsets - .iter() - .map(|o| o.to_string()) - .collect::>() - .join(",") - ); + let in_list: Vec = offsets.iter().map(|o| lit(*o)).collect(); Self { parent, request: QueryRequest { - filter: Some(QueryFilter::Sql(filter)), + filter: Some(QueryFilter::Datafusion( + col("_rowoffset").in_list(in_list, false), + )), ..Default::default() }, } @@ -1489,18 +1484,11 @@ impl TakeQuery { /// /// See [`crate::Table::take_row_ids`] for more details. pub fn from_row_ids(parent: Arc, row_ids: Vec) -> Self { - let filter = format!( - "_rowid in ({})", - row_ids - .iter() - .map(|o| o.to_string()) - .collect::>() - .join(",") - ); + let in_list: Vec = row_ids.iter().map(|id| lit(*id)).collect(); Self { parent, request: QueryRequest { - filter: Some(QueryFilter::Sql(filter)), + filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))), ..Default::default() }, }