mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-24 04:32:27 +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:
@@ -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