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() }, }