mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-14 06:36:54 +00:00
feat: add Bootstrap split strategy for sampling with replacement (#4597)
## 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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user