rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-23 20:22:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Bootstrap split strategy for sampling with replacement (#4597)

Browse Source 
## Context

`StreamingDataset`'s shipped split strategies (`split_random`,
`split_sequential`, `split_hash`, `split_calculated`) all share one
trait: every row in the base table is assigned to exactly one split,
exactly once. They only differ in *how* that one-to-one assignment is
chosen -- none of them can express "give me a resampled stream of N
rows, allowing repeats."

That's a distinct, common need in ML training: **the bootstrap** -- the
statistical technique (Efron, 1979) of estimating a dataset's properties
by resampling it *with replacement*. It's the same sense scikit-learn's
bagging classifiers mean by `bootstrap=True`. `SplitStrategy::Bootstrap`
adds exactly this to the data loader: sample rows uniformly with
replacement from the (optionally filtered) base table, independently per
split.

## Main use cases

- **Bagging / ensemble training** (e.g. random forests): each split
trains on 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**: want
exactly N batches per epoch regardless of whether the table has 50 rows
or 50 million? Set `num_samples` directly instead of it 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.

## Example (an actual run, not a hypothetical one)

A 5-row base table (`row_id` 0..4), sampled via `Bootstrap {
num_samples: 8, sizes: Fixed(2), seed: Some(42) }`:

```
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 -- not 5. Row 4 drawn 4 times total,
row 3 never drawn. That's with-replacement sampling actually behaving
like with-replacement sampling, not an even round-robin dressed up with
a new name.

## What changed

Adds `SplitStrategy::Bootstrap` (uniform sampling with replacement,
optionally filtered, independently seeded per split -- a given split's
content doesn't depend on how many other splits there are, the same
elastic-determinism property documented for the loader's other
strategies). No new `PermutationBuilder` Rust method was needed:
`with_split_strategy` already accepts any `SplitStrategy` value
generically. Python binding:
`PermutationBuilder.bootstrap_sample(num_samples, *, ratios=None,
counts=None, fixed=None, seed=None, split_names=None)`, mirroring
`split_random`'s existing shape.

Also adds a generic `split_strategy` parameter to `StreamingDataset`'s
constructor (a callable taking/returning a `PermutationBuilder`), since
it previously could only reach `split_random`/`split_sequential` even
though `split_hash`/`split_calculated` already shipped. This unblocks
`Bootstrap` and any future strategy without further constructor changes.
This commit is contained in:
Jack Ye
2026-10-21 23:18:30 -07:00
committed by GitHub
parent a4544c5472
commit 1564ec10c8
8 changed files with 909 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+9
View File
@@ -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]
python/python/lancedb/permutation.py
+44
View File
@@ -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":
python/python/lancedb/streaming.py
+147 -3
View File
@@ -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: "
python/python/tests/test_elastic_dataloader.py
+227
View File
@@ -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
python/python/tests/test_permutation.py
+70
View File
@@ -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(
python/src/permutation.rs
+44
View File
@@ -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<Vec<f64>>,
counts: Option<Vec<u64>>,
fixed: Option<u64>,
seed: Option<u64>,
split_names: Option<Vec<String>>,
) -> PyResult<Self> {
// 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,
rust/lancedb/src/dataloader/permutation/builder.rs
+5 -1
View File
@@ -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?
}
};
rust/lancedb/src/dataloader/permutation/split.rs
+363 -2
View File
@@ -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<u64>,
},
}
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<u64>,
) -> Result<SendableRecordBatchStream> {
let schema = source.schema();
let row_id_field = schema.field(0).clone();
let batches: Vec<RecordBatch> = source.try_collect().await?;
let row_id_arrays = batches
.iter()
.map(|b| b.column(0).as_ref())
.collect::<Vec<_>>();
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::<u64>::new();
position_indices
.try_reserve_exact(num_samples as usize)
.map_err(|_| too_large())?;
let mut split_ids = Vec::<u64>::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<dyn Array> = 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::<Vec<_>>()
.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::<UInt64Type>().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::<Int32Type>().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::<Vec<_>>()
.await
.unwrap();
let second = make_splitter()
.apply(test_data(), 50)
.await
.unwrap()
.try_collect::<Vec<_>>()
.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::<Int32Type>())
.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"));
}
}