diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 89c69e721..96bc7752b 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -947,6 +947,15 @@ class AsyncPermutationBuilder: counts: Optional[List[int]] = None, fixed: Optional[int] = None, ) -> "AsyncPermutationBuilder": ... + def bootstrap_sample( + self, + num_samples: int, + *, + ratios: Optional[List[float]] = None, + counts: Optional[List[int]] = None, + fixed: Optional[int] = None, + seed: Optional[int] = None, + ) -> "AsyncPermutationBuilder": ... def split_calculated(self, calculation: str) -> "AsyncPermutationBuilder": ... def shuffle( self, seed: Optional[int], clump_size: Optional[int] diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 103ec0aa4..355ebf798 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -169,6 +169,50 @@ class PermutationBuilder: ) return self + def bootstrap_sample( + self, + num_samples: int, + *, + ratios: Optional[list[float]] = None, + counts: Optional[list[int]] = None, + fixed: Optional[int] = None, + seed: Optional[int] = None, + split_names: Optional[list[str]] = None, + ) -> "PermutationBuilder": + """ + Sample rows uniformly **with replacement**, independently per split. + + Unlike the other split strategies, the number of output rows is + ``num_samples``, not the number of (optionally filtered) input rows: the + same row can appear zero, one, or many times, and ``num_samples`` may be + larger or smaller than the input row count. Useful for bootstrap + resampling, oversampling rare classes, or any training workload that + wants a fixed-size stream of random draws rather than a single + deterministic pass over every row. Combine with :meth:`filter` to + restrict the eligible population. + + One of ratios, counts, or fixed must be provided, with the same meaning + as in :meth:`split_random`, except they divide ``num_samples`` across + splits rather than the input row count: if fixed is 4, each split gets + ``num_samples // 4`` samples. + + Each split draws its own sample using a seed derived independently from + ``(seed, split_index)``, so a split's content does not change if the samples are + later redistributed across a different number of splits. + + The optional split_names can be provided to name the splits. If not provided, + the splits can only be referenced by their index. + """ + self._async.bootstrap_sample( + num_samples, + ratios=ratios, + counts=counts, + fixed=fixed, + seed=seed, + split_names=split_names, + ) + return self + def split_calculated( self, calculation: str, split_names: Optional[list[str]] = None ) -> "PermutationBuilder": diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 91ca6ce28..4a4ffd1f5 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -19,6 +19,7 @@ above. """ import ctypes +import hashlib import heapq import logging import os @@ -39,6 +40,7 @@ from torch.utils.data import DataLoader, IterableDataset, get_worker_info from .permutation import ( Permutation, + PermutationBuilder, Transforms, permutation_builder, _drop_base_version, @@ -244,6 +246,20 @@ class StreamingDataset(IterableDataset): 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. + Ignored when ``split_strategy`` is provided. + split_strategy: + Escape hatch for any ``PermutationBuilder`` split strategy this class + doesn't otherwise expose a dedicated parameter for -- ``split_hash``, + ``split_calculated``, ``bootstrap_sample`` (uniform sampling with + replacement), or any strategy added later. A callable that receives a + fresh ``PermutationBuilder`` for ``table`` and returns a configured one; + it is called and ``.execute()``'d directly, taking over permutation + construction entirely. When provided, ``shuffle``, ``shuffle_seed``, + ``epoch``, and ``shuffle_clump_size`` are ignored -- the callable is + solely responsible for how rows are split and ordered. For example, + ``split_strategy=lambda b: b.bootstrap_sample(10_000, fixed=num_splits, + seed=42)`` draws 10,000 samples with replacement, divided evenly across + ``num_splits`` splits. shuffle_seed: Base seed for the random permutation. Combined with ``epoch`` so each epoch produces a different ordering. Pass ``None`` to generate @@ -393,6 +409,9 @@ class StreamingDataset(IterableDataset): *, num_splits: Optional[int] = None, shuffle: bool = True, + split_strategy: Optional[ + Callable[[PermutationBuilder], PermutationBuilder] + ] = None, shuffle_seed: Optional[int] = 0, epoch: int = 0, rank: int = 0, @@ -583,13 +602,70 @@ class StreamingDataset(IterableDataset): builder = permutation_builder(table) if filter is not None: builder = builder.filter(filter) - if shuffle: + if split_strategy is not None: + # Escape hatch: lets any PermutationBuilder split strategy drive this + # dataset (split_hash, split_calculated, bootstrap_sample, or any future + # strategy) without StreamingDataset needing a bespoke parameter for each + # one. shuffle/shuffle_seed/epoch/shuffle_clump_size are ignored in this + # path -- the caller's split_strategy is solely responsible for ordering. + self._perm_table = split_strategy(builder).execute() + # _rank_splits below assumes split ids are exactly 0..num_splits-1 with no + # gaps (every other path guarantees this by always building with + # fixed=num_splits) -- a split_strategy callable has no such guarantee, and + # a mismatch would otherwise silently drop rows (ids >= num_splits) or leave + # some ranks reading an empty split (missing ids < num_splits). + split_value_counts = ( + self._perm_table.to_arrow().column("split_id").value_counts() + ) + counts_by_split = dict( + zip( + split_value_counts.field("values").to_pylist(), + split_value_counts.field("counts").to_pylist(), + ) + ) + expected_splits = set(range(num_splits)) + if set(counts_by_split) != expected_splits: + raise ValueError( + f"split_strategy produced split ids {sorted(counts_by_split)}, " + f"but num_splits={num_splits} requires exactly " + f"{{0, ..., {num_splits - 1}}}. Pass num_splits matching the " + "strategy's own split count (e.g. `fixed=num_splits` on " + "bootstrap_sample/split_random/etc.)." + ) + # The round-robin iteration below (and in _iter_owned) assumes every + # split has the same row count -- it stops as soon as any one owned + # split is exhausted (see its own comment: "all splits exhaust + # simultaneously"). Every other path guarantees this by always building + # with fixed=num_splits; a split_strategy using ratios/counts/split_hash + # weights has no such guarantee, and unequal sizes would otherwise + # silently drop rows from the larger splits -- or, with world_size > 1 + # where each rank owns different splits, end ranks after different row + # counts and desync or hang a distributed run's end-of-epoch collectives. + if len(set(counts_by_split.values())) > 1: + raise ValueError( + f"split_strategy produced unequal split sizes " + f"{counts_by_split} (split id -> row count). Every split must " + "have the same row count. Use `fixed=num_splits` rather than " + "`ratios`/`counts`/weighted `split_hash`, which can produce " + "unequal sizes, or split unevenly-sized groups into separate " + "StreamingDatasets instead." + ) + elif 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() + # Only the split_strategy path needs a permutation fingerprint in the + # checkpoint (see `_checkpoint_snapshot`): split_random/split_sequential are + # already fully pinned by shuffle_seed/num_splits, and a fingerprint over raw + # row ids would spuriously reject an otherwise-valid resume after those row + # ids get renumbered by compaction (row ids are not stable by default). + self._uses_split_strategy = split_strategy is not None + # Lazily computed and cached by `_perm_fingerprint` -- `_perm_table` never + # changes after construction, so this is correct to compute once. + self._perm_fingerprint_cache: Optional[str] = None if self._blocks_per_epoch == "auto": self._blocks_per_epoch = self._estimate_blocks_per_epoch() @@ -1406,11 +1482,51 @@ class StreamingDataset(IterableDataset): ) return state + @property + def _perm_fingerprint(self) -> str: + """Cheap identity check for the built permutation, computed once and cached + (``_perm_table`` never changes after construction). + + Only meaningful -- and only recorded in a checkpoint -- for the + ``split_strategy`` path (see ``_checkpoint_snapshot``): ``num_splits``/ + ``shuffle_seed`` alone cannot detect a `load_state_dict` resuming saved + per-split progress against a *different* permutation than the one that + produced it, since a ``split_strategy`` callable ignores ``shuffle_seed`` + entirely -- e.g. a ``bootstrap_sample`` with no seed (or any changed + callable) can rebuild a completely different permutation on the next run + while ``shuffle_seed`` keeps its same default value and the existing check + passes anyway. ``split_random``/``split_sequential`` are already fully + pinned by ``shuffle_seed``/``num_splits`` and must *not* additionally be + fingerprinted: this hashes physical row ids, which are not stable across + compaction by default, so checking it on those paths would reject an + otherwise-valid resume after the table gets compacted between save and + load even though the seeded permutation itself is unchanged. + """ + if self._perm_fingerprint_cache is None: + arrow_tbl = self._perm_table.to_arrow() + hasher = hashlib.sha256() + hasher.update(str(arrow_tbl.num_rows).encode()) + for name in ("row_id", "split_id"): + column = arrow_tbl.column(name).combine_chunks() + hasher.update(column.to_numpy(zero_copy_only=False).tobytes()) + self._perm_fingerprint_cache = hasher.hexdigest() + return self._perm_fingerprint_cache + def _checkpoint_snapshot(self) -> dict: + # Only recorded for split_strategy: split_random/split_sequential are + # already fully pinned by shuffle_seed/num_splits below, and row ids are not + # stable across compaction by default, so fingerprinting them unconditionally + # would reject an otherwise-valid resume after the table got compacted. + perm_fingerprint = ( + {"perm_fingerprint": self._perm_fingerprint} + if (self._uses_split_strategy) + else {} + ) if self._pack_sequences is not None: return { "shuffle_seed": self._shuffle_seed, "num_splits": self._num_splits, + **perm_fingerprint, "epoch": self._epoch, "pack_sequences": self._pack_sequences, "eos_id": self._eos_id, @@ -1431,6 +1547,7 @@ class StreamingDataset(IterableDataset): return { "shuffle_seed": self._shuffle_seed, "num_splits": self._num_splits, + **perm_fingerprint, "epoch": self._epoch, "samples_consumed_per_split": samples, "positions_consumed_per_split": positions, @@ -1560,8 +1677,16 @@ class StreamingDataset(IterableDataset): 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. Packed checkpoints - pin ``pack_sequences``, ``eos_id``, ``pad_id``, + makes mid-epoch resumption meaningless. Also raises if the checkpoint + was built with ``split_strategy`` and records a permutation fingerprint + that does not match the current dataset's permutation -- this catches a + changed or non-deterministically-reseeded ``split_strategy`` that + ``num_splits``/``shuffle_seed`` alone cannot detect, since ``shuffle_seed`` + is unused on that path. Checkpoints from ``split_random``/ + ``split_sequential`` never carry this fingerprint (row ids are not stable + across compaction, so it would reject an otherwise-valid resume), and + continue to rely on ``shuffle_seed``/``num_splits`` alone as before. + Packed checkpoints pin ``pack_sequences``, ``eos_id``, ``pad_id``, ``blocks_per_epoch``, and ``epoch``. """ if state["num_splits"] != self._num_splits: @@ -1574,6 +1699,19 @@ class StreamingDataset(IterableDataset): f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, " f"current dataset has {self._shuffle_seed}" ) + checkpoint_fingerprint = state.get("perm_fingerprint") + if ( + checkpoint_fingerprint is not None + and checkpoint_fingerprint != self._perm_fingerprint + ): + raise ValueError( + "permutation mismatch: this checkpoint was saved against a " + "different permutation (e.g. a split_strategy that resamples " + "differently between runs, such as bootstrap_sample with no " + "seed, or a changed callable) than the current dataset built. " + "Resuming would apply saved per-split positions to the wrong " + "rows." + ) self._consumer_checkpoint_requires_uniform = False if "pack_buffers" in state or self._pack_sequences is not None: @@ -1675,6 +1813,10 @@ class StreamingDataset(IterableDataset): first = states[0] packed = "pack_buffers" in first config_keys = ["shuffle_seed", "num_splits", "epoch"] + if "perm_fingerprint" in first: + # Absent only in checkpoints saved before this key existed; skip rather + # than KeyError so old checkpoints keep merging. + config_keys.append("perm_fingerprint") if packed: config_keys.extend( ["pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"] @@ -1684,6 +1826,8 @@ class StreamingDataset(IterableDataset): if ("pack_buffers" in state) != packed: raise ValueError("cannot merge packed and unpacked state dicts") for key in config_keys: + if key not in state: + continue if state[key] != first[key]: raise ValueError( f"{key} mismatch across state dicts: " diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 9d18752f1..5eb3f095c 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -3297,3 +3297,230 @@ def test_streaming_dataset_over_remote_table(): assert len(server.scans) == 1, "the permutation is built with one row-id scan" assert server.takes, "rows must be fetched with row-id takes" assert_server_safe_row_id_requests(server) + + +# --------------------------------------------------------------------------- +# split_strategy escape hatch +# --------------------------------------------------------------------------- + + +def test_split_strategy_bootstrap_sample(lance_table): + """split_strategy lets StreamingDataset use bootstrap_sample directly, producing + more rows than the base table has (sampling with replacement).""" + num_samples = 4 * NUM_ROWS + ds = StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + num_samples, fixed=NUM_SPLITS, seed=SHUFFLE_SEED + ), + num_splits=NUM_SPLITS, + ) + ids = [row["id"] for row in ds] + + assert len(ids) == num_samples + # 4x oversampling a small table must produce at least one repeat. + assert len(set(ids)) < len(ids) + assert set(ids).issubset(set(range(NUM_ROWS))) + + +def test_split_strategy_matches_direct_permutation_builder_use(lance_table): + """StreamingDataset(split_strategy=...), read single-process (no DataLoader + workers, so this one process owns every split and round-robins across all of + them), must yield exactly what calling the same strategy directly on a + PermutationBuilder would read, interleaved one row per split per cycle. + + Reads each split via __getitems__ (explicit indices), not plain iteration: + Permutation's sequential __iter__ returns empty for a permutation with repeated + row_ids (confirmed directly -- len() is correct, __getitems__ returns the right + rows, but `for row in perm` yields nothing). That's a pre-existing limitation of + Permutation's iterator, unrelated to Bootstrap sampling's correctness: it doesn't + affect StreamingDataset or this benchmark, since both only ever read via + explicit-index batched fetches (take_offsets/__getitems__), never the plain + iterator -- flagged separately, not fixed here (out of scope for this change). + """ + from lancedb.permutation import Permutation, permutation_builder + + def strategy(builder): + return builder.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=1) + + direct_perm_tbl = strategy(permutation_builder(lance_table)).execute() + expected_by_split = [] + for split in range(NUM_SPLITS): + perm = Permutation.from_tables(lance_table, direct_perm_tbl, split=split) + rows = perm.__getitems__(list(range(len(perm)))) + expected_by_split.append([row["id"] for row in rows]) + samples_per_split = len(expected_by_split[0]) + expected = [ + expected_by_split[split][cycle] + for cycle in range(samples_per_split) + for split in range(NUM_SPLITS) + ] + + ds = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS) + actual = [row["id"] for row in ds] + assert actual == expected + + +def test_split_strategy_is_reproducible_across_instances(lance_table): + def strategy(builder): + return builder.bootstrap_sample(3 * NUM_ROWS, fixed=NUM_SPLITS, seed=99) + + ds1 = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS) + ds2 = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS) + + assert [row["id"] for row in ds1] == [row["id"] for row in ds2] + + +def test_split_strategy_num_splits_mismatch_raises(lance_table): + """StreamingDataset(split_strategy=...) must reject a strategy whose split + count doesn't match num_splits. `_rank_splits` assumes split ids are exactly + 0..num_splits-1 with no gaps; a mismatch would otherwise silently drop rows + (split ids >= num_splits) or leave some ranks reading an empty split (missing + ids < num_splits).""" + with pytest.raises(ValueError, match="num_splits"): + StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + 2 * NUM_ROWS, fixed=NUM_SPLITS, seed=SHUFFLE_SEED + ), + num_splits=NUM_SPLITS + 1, # the strategy above always builds NUM_SPLITS + ) + + +def test_split_strategy_unequal_split_sizes_raises(lance_table): + """StreamingDataset(split_strategy=...) must reject a strategy whose splits + all exist (ids 0..num_splits-1) but aren't the same size. The round-robin + iteration (see _iter_owned) assumes every owned split has the same row count + and stops as soon as any one is exhausted -- `ratios`/`counts` (unlike + `fixed`) can produce unequal split sizes, which would otherwise silently + drop rows from the larger splits.""" + with pytest.raises(ValueError, match="unequal split sizes"): + StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + 100, ratios=[0.7, 0.3], seed=SHUFFLE_SEED + ), + num_splits=2, + ) + + +def test_split_strategy_permutation_mismatch_on_resume_raises(lance_table): + """load_state_dict must reject resuming against a different permutation than + the one that produced the checkpoint. num_splits/shuffle_seed alone can't catch + this for split_strategy, since that path ignores shuffle_seed entirely: a + bootstrap_sample reseeded differently resamples completely differently while + shuffle_seed keeps whatever value the caller passed either way, so the + pre-existing checks would pass despite the mismatch.""" + + def strategy(seed): + return lambda b: b.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=seed) + + ds1 = StreamingDataset( + lance_table, + split_strategy=strategy(1), + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + it = iter(ds1) + for _ in range(NUM_SPLITS): # one full cycle: every split advances uniformly + next(it) + checkpoint = ds1.state_dict() + + ds2 = StreamingDataset( + lance_table, + split_strategy=strategy(2), # different seed -> a different permutation + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, # unchanged -- irrelevant on this path + ) + with pytest.raises(ValueError, match="permutation mismatch"): + ds2.load_state_dict(checkpoint) + + +def test_split_strategy_resume_with_matching_permutation_succeeds(lance_table): + """The common case: resuming against the same split_strategy (same seed) must + work exactly like the non-split_strategy checkpoint/resume path, continuing + from exactly the same position with no gaps or repeats.""" + + def strategy(b): + return b.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=7) + + ds1 = StreamingDataset( + lance_table, + split_strategy=strategy, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + it = iter(ds1) + for _ in range(NUM_SPLITS): + next(it) + checkpoint = ds1.state_dict() + remaining_original = [row["id"] for row in it] + + ds2 = StreamingDataset( + lance_table, + split_strategy=strategy, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + ds2.load_state_dict(checkpoint) # must not raise + remaining_resumed = [row["id"] for row in ds2] + + assert remaining_original == remaining_resumed + + +def test_default_checkpoint_omits_perm_fingerprint(lance_table): + """split_random/split_sequential checkpoints must not carry a + perm_fingerprint -- see test_default_checkpoint_resumes_after_compaction for + why checking one on this path would be wrong, not just unnecessary.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + assert "perm_fingerprint" not in ds.state_dict() + + +def test_split_strategy_checkpoint_includes_perm_fingerprint(lance_table): + ds = StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + 2 * NUM_ROWS, fixed=NUM_SPLITS, seed=SHUFFLE_SEED + ), + num_splits=NUM_SPLITS, + ) + assert "perm_fingerprint" in ds.state_dict() + + +def test_default_checkpoint_resumes_after_compaction(lance_table): + """split_random/split_sequential checkpoints must resume normally even after + the table is compacted between save and load. Lance row ids are not stable + across compaction by default (confirmed empirically: a 2-fragment table's + row ids differ before and after `optimize()`), so fingerprinting raw row ids + on this path -- instead of relying on shuffle_seed/num_splits alone, as it + does -- would reject this resume even though the seeded permutation itself + is unchanged.""" + lance_table.add(pa.table({"id": [NUM_ROWS]})) # a second fragment to compact + + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + it = iter(ds) + for _ in range(NUM_SPLITS): + next(it) + checkpoint = ds.state_dict() + assert "perm_fingerprint" not in checkpoint + + before_compaction = set( + lance_table.to_lance() + .to_table(columns=[], with_row_id=True)["_rowid"] + .to_pylist() + ) + lance_table.optimize() + after_compaction = set( + lance_table.to_lance() + .to_table(columns=[], with_row_id=True)["_rowid"] + .to_pylist() + ) + assert before_compaction != after_compaction, ( + "test setup didn't actually renumber row ids -- fix the fixture, not " + "the assertion below" + ) + + ds_resumed = StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED + ) + ds_resumed.load_state_dict(checkpoint) # must not raise diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index 5e8bb1d76..ed1894340 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -161,6 +161,76 @@ def test_split_random_with_seed(mem_db): assert data1["split_id"] == data2["split_id"] +def test_bootstrap_sample_fixed(mem_db): + """Output row count is num_samples, not the input row count, and splits divide + it evenly.""" + tbl = mem_db.create_table("test_table", pa.table({"x": range(20)})) + permutation_tbl = ( + permutation_builder(tbl).bootstrap_sample(1000, fixed=4, seed=42).execute() + ) + + assert permutation_tbl.count_rows() == 1000 + + data = permutation_tbl.search(None).to_arrow().to_pydict() + split_ids = data["split_id"] + assert set(split_ids) == {0, 1, 2, 3} + for split_id in range(4): + assert split_ids.count(split_id) == 250 + + # Every sampled row_id must still be one of the 20 real rows. + assert set(data["row_id"]).issubset(set(range(20))) + + +def test_bootstrap_sample_draws_with_replacement(mem_db): + """Sampling more rows than exist must produce repeats.""" + tbl = mem_db.create_table("test_table", pa.table({"x": range(5)})) + permutation_tbl = ( + permutation_builder(tbl).bootstrap_sample(500, fixed=1, seed=1).execute() + ) + + data = permutation_tbl.search(None).to_arrow().to_pydict() + assert len(data["row_id"]) == 500 + # Only 5 distinct rows exist, so with 500 draws every one of them must repeat. + for row_id in range(5): + assert data["row_id"].count(row_id) > 1 + + +def test_bootstrap_sample_with_seed_is_reproducible(mem_db): + tbl = mem_db.create_table("test_table", pa.table({"x": range(50)})) + + perm1 = permutation_builder(tbl).bootstrap_sample(200, fixed=2, seed=42).execute() + perm2 = permutation_builder(tbl).bootstrap_sample(200, fixed=2, seed=42).execute() + + data1 = perm1.search(None).to_arrow().to_pydict() + data2 = perm2.search(None).to_arrow().to_pydict() + assert data1["row_id"] == data2["row_id"] + assert data1["split_id"] == data2["split_id"] + + +def test_bootstrap_sample_respects_filter(mem_db): + tbl = mem_db.create_table("test_table", pa.table({"x": range(20)})) + permutation_tbl = ( + permutation_builder(tbl) + .filter("x >= 15") + .bootstrap_sample(100, fixed=1, seed=7) + .execute() + ) + + data = permutation_tbl.search(None).to_arrow().to_pydict() + assert len(data["row_id"]) == 100 + assert set(data["row_id"]).issubset(set(range(15, 20))) + + +def test_bootstrap_sample_requires_exactly_one_size_arg(mem_db): + tbl = mem_db.create_table("test_table", pa.table({"x": range(10)})) + with pytest.raises(ValueError): + permutation_builder(tbl).bootstrap_sample(10).execute() + with pytest.raises(ValueError): + permutation_builder(tbl).bootstrap_sample( + 10, ratios=[0.5, 0.5], fixed=2 + ).execute() + + def test_split_hash(mem_db): """Test hash-based splitting.""" tbl = mem_db.create_table( diff --git a/python/src/permutation.rs b/python/src/permutation.rs index 24ca2b5a3..4e07d439e 100644 --- a/python/src/permutation.rs +++ b/python/src/permutation.rs @@ -178,6 +178,50 @@ impl PyAsyncPermutationBuilder { }) } + #[pyo3(signature = (num_samples, *, ratios=None, counts=None, fixed=None, seed=None, split_names=None))] + pub fn bootstrap_sample( + slf: PyRefMut<'_, Self>, + num_samples: u64, + ratios: Option>, + counts: Option>, + fixed: Option, + seed: Option, + split_names: Option>, + ) -> PyResult { + // Check that exactly one split type is provided + let split_args_count = [ratios.is_some(), counts.is_some(), fixed.is_some()] + .iter() + .filter(|&&x| x) + .count(); + + if split_args_count != 1 { + return Err(pyo3::exceptions::PyValueError::new_err( + "Exactly one of 'ratios', 'counts', or 'fixed' must be provided", + )); + } + + let sizes = if let Some(ratios) = ratios { + SplitSizes::Percentages(ratios) + } else if let Some(counts) = counts { + SplitSizes::Counts(counts) + } else if let Some(fixed) = fixed { + SplitSizes::Fixed(fixed) + } else { + unreachable!("One of the split arguments must be provided"); + }; + + slf.modify(|builder| { + builder.with_split_strategy( + SplitStrategy::Bootstrap { + num_samples, + sizes, + seed, + }, + split_names, + ) + }) + } + pub fn split_calculated( slf: PyRefMut<'_, Self>, calculation: String, diff --git a/rust/lancedb/src/dataloader/permutation/builder.rs b/rust/lancedb/src/dataloader/permutation/builder.rs index 32e7063a3..15f8c3cd3 100644 --- a/rust/lancedb/src/dataloader/permutation/builder.rs +++ b/rust/lancedb/src/dataloader/permutation/builder.rs @@ -316,6 +316,10 @@ impl PermutationBuilder { } else { self.sort_by_column(rows, ROW_ID).await? }; + // Bootstrap's output row count (num_samples) differs from its input row count + // (num_rows); every other strategy's output matches its input. The shuffler + // must be told the former, not assume it always equals the latter. + let output_rows = splitter.output_row_count(num_rows); let split_data = splitter.apply(rows, num_rows).await?; // Shuffle data if requested @@ -328,7 +332,7 @@ impl PermutationBuilder { temp_dir: self.config.temp_dir.clone(), max_rows_per_file: 10 * 1024 * 1024, }); - shuffler.shuffle(split_data, num_rows).await? + shuffler.shuffle(split_data, output_rows).await? } }; diff --git a/rust/lancedb/src/dataloader/permutation/split.rs b/rust/lancedb/src/dataloader/permutation/split.rs index cfa8e5235..b4e5deac1 100644 --- a/rust/lancedb/src/dataloader/permutation/split.rs +++ b/rust/lancedb/src/dataloader/permutation/split.rs @@ -11,6 +11,7 @@ use arrow_schema::{DataType, Field, Schema}; use datafusion_common::hash_utils::{RandomState, create_hashes}; use futures::{StreamExt, TryStreamExt}; use lance_arrow::SchemaExt; +use rand::Rng; use lance_core::ROW_ID; @@ -19,7 +20,7 @@ use crate::{ arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}, dataloader::{ permutation::shuffle::{Shuffler, ShufflerConfig}, - permutation::util::TemporaryDirectory, + permutation::util::{TemporaryDirectory, non_crypto_rng}, }, query::{Query, QueryBase, Select}, }; @@ -76,6 +77,59 @@ pub enum SplitStrategy { /// /// If this strategy is used then the counts/percentages in the SplitSizes are ignored. Calculated { calculation: String }, + /// Rows are sampled uniformly **with replacement** from the base table (combine with + /// [`PermutationBuilder::with_filter`] to restrict the eligible population), independently + /// per split. Named after "the bootstrap" -- the statistical technique of estimating a + /// dataset's properties by resampling it with replacement -- the same sense `bootstrap=True` + /// has in scikit-learn's bagging classifiers. + /// + /// Every other strategy partitions the table: each row goes to exactly one split, exactly + /// once. `Bootstrap` does not partition -- it resamples. The number of output rows is + /// `num_samples`, not the number of input rows: the same row can appear zero, one, or many + /// times, and `num_samples` may be larger or smaller than the (filtered) input row count. + /// Typical uses: + /// + /// - **Bagging / ensemble training** (e.g. random forests): each split is an independent + /// bootstrap resample of the training set. + /// - **Oversampling a rare class**: filter to the minority class, then sample it up to + /// whatever count balances it against the majority, instead of exhausting its small set + /// once per epoch. + /// - **Fixed-size training steps independent of dataset size**: `num_samples` is set + /// directly rather than being implicitly `len(dataset)`. + /// - **Weighted sampling without a weight column**: filter to "eligible" rows (e.g. a + /// windowed-sampling workload filtering to valid window-start positions), then bootstrap + /// the filtered set -- an item with more eligible rows is naturally drawn proportionally + /// more often, with no explicit per-row weight needed. + /// + /// `sizes` divides `num_samples` across splits (reusing [`SplitSizes`], so `Fixed(n)` means + /// "n splits, each getting `num_samples / n` samples", the same convention `Random` and + /// `Sequential` use). Each split draws its own sample using a seed derived independently + /// from `(seed, split_index)`, so a given split's content does not change if `num_samples` + /// is redistributed across a different number of splits -- the same elastic-determinism + /// property documented for the data loader's other strategies. + /// + /// Example: a 5-row base table (`row_id` 0..4), `Bootstrap { num_samples: 8, sizes: + /// Fixed(2), seed: Some(42) }` produced (an actual run, not a hypothetical one): + /// + /// ```text + /// row_id split_id + /// 0 0 + /// 2 0 + /// 4 0 + /// 4 0 <- row 4 repeated within split 0 + /// 4 1 <- row 4 also drawn in split 1 (independent draw) + /// 4 1 + /// 1 1 + /// 2 1 + /// ``` + /// + /// 8 output rows from a 5-row table, row 4 drawn 4 times total, row 3 never drawn -- + /// exactly the with-replacement behavior the name promises, not an even round-robin. + Bootstrap { + num_samples: u64, + sizes: SplitSizes, + seed: Option, + }, } impl SplitStrategy { @@ -107,6 +161,25 @@ impl SplitStrategy { } Self::Sequential { sizes } => sizes.validate(num_rows), Self::Calculated { .. } => Ok(()), + Self::Bootstrap { + num_samples, sizes, .. + } => { + if num_rows == 0 { + return Err(Error::InvalidInput { + message: "Bootstrap strategy requires at least one eligible row to \ + sample from (check the filter, if one is set)" + .to_string(), + }); + } + if *num_samples == 0 { + return Err(Error::InvalidInput { + message: "Bootstrap strategy requires num_samples > 0".to_string(), + }); + } + // Validated against num_samples, not num_rows: sampling is with replacement, + // so num_samples may exceed the number of eligible rows. + sizes.validate(*num_samples) + } } } } @@ -314,6 +387,95 @@ impl Splitter { ))) } + /// Materializes the full (filtered) row id stream -- random access is required since + /// with-replacement sampling can draw any position, including ones already drawn -- then + /// independently samples each split's share with its own derived RNG seed. + async fn apply_bootstrap( + &self, + source: SendableRecordBatchStream, + num_samples: u64, + sizes: &SplitSizes, + seed: Option, + ) -> Result { + let schema = source.schema(); + let row_id_field = schema.field(0).clone(); + let batches: Vec = source.try_collect().await?; + let row_id_arrays = batches + .iter() + .map(|b| b.column(0).as_ref()) + .collect::>(); + let row_ids = arrow::compute::concat(&row_id_arrays)?; + // Sample against however many rows this scan actually returned, not the + // separately computed `num_rows` count checked by `validate` -- that count and + // this scan are two different queries, not pinned to one snapshot, so a + // concurrent append or delete between them could otherwise push + // `random_range` out of `row_ids`'s real bounds. + let num_eligible_rows = row_ids.len() as u64; + if num_eligible_rows == 0 { + return Err(Error::InvalidInput { + message: "Bootstrap strategy found no eligible rows to sample from in \ + this scan (the table or its filter may have changed since \ + validation)" + .to_string(), + }); + } + + let split_sample_counts = sizes.to_counts(num_samples); + + // `try_reserve_exact` (fallible) rather than `with_capacity` (infallible -- + // aborts the whole process via `handle_alloc_error` on failure, which no + // caller can catch) so an unreasonably large `num_samples` becomes a normal, + // catchable error instead of a crash. This bounds the failure mode, not peak + // memory itself: the output is still one fully materialized RecordBatch, not + // streamed in chunks. + let too_large = || Error::InvalidInput { + message: format!( + "Bootstrap strategy's num_samples ({num_samples}) is too large to \ + allocate in memory" + ), + }; + let mut position_indices = Vec::::new(); + position_indices + .try_reserve_exact(num_samples as usize) + .map_err(|_| too_large())?; + let mut split_ids = Vec::::new(); + split_ids + .try_reserve_exact(num_samples as usize) + .map_err(|_| too_large())?; + for (split_idx, &count) in split_sample_counts.iter().enumerate() { + // Derive an independent seed per split (rather than one RNG shared across all + // splits) so a split's sample doesn't depend on how many other splits there are. + let split_seed = seed.map(|s| { + s.wrapping_add( + (split_idx as u64) + .wrapping_mul(0x9E37_79B9_7F4A_7C15) + .wrapping_add(1), + ) + }); + let mut rng = non_crypto_rng(&split_seed); + for _ in 0..count { + position_indices.push(rng.random_range(0..num_eligible_rows)); + split_ids.push(split_idx as u64); + } + } + + let indices = UInt64Array::from(position_indices); + let sampled_row_ids = arrow::compute::take(row_ids.as_ref(), &indices, None)?; + let split_id_array: Arc = Arc::new(UInt64Array::from(split_ids)); + + let new_schema = Arc::new(Schema::new(vec![ + row_id_field, + Field::new(SPLIT_ID_COLUMN, DataType::UInt64, false), + ])); + let result_batch = + RecordBatch::try_new(new_schema.clone(), vec![sampled_row_ids, split_id_array])?; + + Ok(Box::pin(SimpleRecordBatchStream::new( + futures::stream::once(std::future::ready(Ok(result_batch))), + new_schema, + ))) + } + pub async fn apply( &self, source: SendableRecordBatchStream, @@ -357,6 +519,14 @@ impl Splitter { self.apply_hash(source, split_weights, *discard_weight) .await } + SplitStrategy::Bootstrap { + num_samples, + sizes, + seed, + } => { + self.apply_bootstrap(source, *num_samples, sizes, *seed) + .await + } } } @@ -377,7 +547,13 @@ impl Splitter { pub fn orders_by_split_id(&self) -> bool { match &self.strategy { - SplitStrategy::Hash { .. } | SplitStrategy::Calculated { .. } => true, + // Bootstrap emits samples grouped split-by-split already (see apply_bootstrap), + // rather than relying on the shared pipeline's sort-by-split-id step: within-split + // draw order is meaningful to callers (e.g. a warmup/measured split within one + // split's stream), and sort stability for ties isn't documented behavior to lean on. + SplitStrategy::Hash { .. } + | SplitStrategy::Calculated { .. } + | SplitStrategy::Bootstrap { .. } => true, SplitStrategy::NoSplit | SplitStrategy::Sequential { .. } // It may be strange but for random we shuffle and then assign splits so the result is @@ -385,6 +561,26 @@ impl Splitter { | SplitStrategy::Random { .. } => false, } } + + /// The row count `apply` produces from `num_rows` input rows. Equal to `num_rows` + /// for every strategy except `Bootstrap`, whose output size is `num_samples` + /// regardless of how many (filtered) input rows it draws from. Callers that need + /// to validate or pre-size `apply`'s output (e.g. the shuffler) must use this + /// instead of assuming the input and output row counts match. + /// + /// For `Bootstrap`, this is `sizes.to_counts(num_samples)`'s sum, not + /// `num_samples` itself: `to_counts` can emit less than its input (e.g. + /// `Fixed(n)` when `n` does not divide `num_samples` evenly, or explicit + /// `counts`/`ratios` that do not use every sample) and `apply_bootstrap` only + /// ever emits exactly what `to_counts` returns. + pub fn output_row_count(&self, num_rows: u64) -> u64 { + match &self.strategy { + SplitStrategy::Bootstrap { + num_samples, sizes, .. + } => sizes.to_counts(*num_samples).iter().sum(), + _ => num_rows, + } + } } /// Split configuration - either percentages or absolute counts @@ -807,4 +1003,169 @@ mod tests { assert_eq!(num_in_split_0, 12); // 24% assert_eq!(num_in_split_1, 30); // 60% } + + #[tokio::test] + async fn test_bootstrap_split_schema_and_counts() { + let splitter = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 1000, + sizes: SplitSizes::Fixed(4), + seed: Some(42), + }, + ); + + let split_batches = splitter + .apply(test_data(), 50) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let schema = split_batches[0].schema(); + let split_batch = concat_batches(&schema, &split_batches).unwrap(); + + // Schema stays exactly [row_id, split_id], same as every other strategy. + let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(field_names, vec!["id", SPLIT_ID_COLUMN]); + + // Output row count is num_samples, not num_rows (50) -- sampling is with replacement. + assert_eq!(split_batch.num_rows(), 1000); + + let split_ids = split_batch.column(1).as_primitive::().values(); + for split_id in [0u64, 1, 2, 3] { + assert_eq!(split_ids.iter().filter(|v| **v == split_id).count(), 250); + } + + // With 1000 draws from only 50 distinct rows, some row id must repeat. + let row_ids = split_batch.column(0).as_primitive::().values(); + let mut seen = std::collections::HashSet::new(); + assert!(!row_ids.iter().all(|v| seen.insert(*v))); + } + + #[tokio::test] + async fn test_bootstrap_split_seed_is_reproducible() { + let make_splitter = || { + Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 200, + sizes: SplitSizes::Fixed(2), + seed: Some(7), + }, + ) + }; + + let first = make_splitter() + .apply(test_data(), 50) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let second = make_splitter() + .apply(test_data(), 50) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let schema = first[0].schema(); + let first = concat_batches(&schema, &first).unwrap(); + let second = concat_batches(&schema, &second).unwrap(); + assert_eq!(first.column(0), second.column(0)); + assert_eq!(first.column(1), second.column(1)); + } + + #[tokio::test] + async fn test_bootstrap_split_respects_filter() { + // Bootstrap is driven through a filter the same way every other strategy is -- the + // base table rows reaching Splitter::apply are already filtered by the time they get + // here, so this exercises the same num_rows=0 guard a fully-filtered-out population + // would hit. + let splitter = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 10, + sizes: SplitSizes::Fixed(1), + seed: Some(1), + }, + ); + let Err(err) = splitter.apply(test_data(), 0).await else { + panic!("expected Bootstrap over zero eligible rows to fail"); + }; + assert!(err.to_string().contains("at least one eligible row")); + } + + fn empty_data() -> SendableRecordBatchStream { + lance_datagen::gen_batch() + .with_seed(Seed::from(42)) + .col(ID_COLUMN, lance_datagen::array::step::()) + .into_ldb_stream(RowCount::from(0), BatchCount::from(1)) + } + + #[test] + fn test_output_row_count() { + // Bootstrap's output row count is independent of the input size -- every + // other strategy's output matches its input. The shuffler relies on this + // (see builder.rs) instead of assuming input and output row counts match. + let bootstrap = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 1000, + sizes: SplitSizes::Fixed(4), + seed: Some(42), + }, + ); + assert_eq!(bootstrap.output_row_count(50), 1000); + + // Not num_samples itself: `Fixed(3)` doesn't divide 1000 evenly, so + // `to_counts` emits 333 * 3 = 999, one short of num_samples -- + // `apply_bootstrap` only ever emits exactly what `to_counts` returns. + let bootstrap_uneven = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 1000, + sizes: SplitSizes::Fixed(3), + seed: Some(42), + }, + ); + assert_eq!(bootstrap_uneven.output_row_count(50), 999); + + let sequential = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Sequential { + sizes: SplitSizes::Fixed(2), + }, + ); + assert_eq!(sequential.output_row_count(50), 50); + } + + #[tokio::test] + async fn test_bootstrap_rejects_empty_scan_cleanly() { + // The eligibility count `validate` checks and this scan are two separate + // queries against the table, not pinned to one snapshot -- a concurrent + // delete could shrink the real scan to zero rows after `validate` already + // passed on a nonzero count. Calling `apply_bootstrap` directly (bypassing + // `apply`'s `validate`) reproduces exactly that: a scan that returns zero + // rows regardless of what count was validated upstream. It must return a + // clean `Err`, not panic on an out-of-range `take` or an empty-range + // `random_range`. + let splitter = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 10, + sizes: SplitSizes::Fixed(1), + seed: Some(1), + }, + ); + let Err(err) = splitter + .apply_bootstrap(empty_data(), 10, &SplitSizes::Fixed(1), Some(1)) + .await + else { + panic!("expected Bootstrap over an empty scan to fail cleanly, not panic"); + }; + assert!(err.to_string().contains("no eligible rows")); + } }