mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +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:
@@ -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]
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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: "
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user