mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-14 06:36:54 +00:00
43bccbadc574a8b1f13114f571f3266c501dc026
10
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Jack Ye
|
1564ec10c8 |
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.
|
||
  Lei XuandClaude Opus 5.5
|
1772f4ef83 |
fix: don't panic building permutations over tables with schema metadata (#4602)
## Problem
`permutation_builder(table).execute()` and `StreamingDataset(table,
...)` (which builds a permutation at construction) panic with
`RustPanic` on any table whose Arrow schema has schema-level metadata:
```python
schema = pa.schema([pa.field("x", pa.int64())], metadata={b"huggingface": b"{}"})
t = lancedb.connect(tmp).create_table("t", pa.table({"x": list(range(100))}, schema=schema))
permutation_builder(t).split_random(fixed=2, seed=1).execute() # RustPanic
```
Every Hugging Face Lance dataset triggers this, e.g.
`hf://datasets/lance-format/lerobot-xvla-soft-fold/data` table `frames`
carries `source-column-name-map`.
## Root cause
`PermutationBuilder::add_config_metadata` replaced the stream schema's
metadata with the permutation config and then called
`batch.with_schema(schema).unwrap()`. The base-table scan's batches
carry the base schema metadata, and `RecordBatch::with_schema` requires
the target to be a superset of the current schema, metadata included. It
returned `SchemaError("target schema is not superset of current schema
...")`, and the `unwrap` turned that into a panic on a tokio worker.
## Fix
- The permutation table deliberately does **not** inherit base-table
schema metadata. That metadata describes columns the permutation table
(`row_id`, `split_id`) does not have, and its keys can collide with
config keys such as `split_names`. A merge would let a base table's
`split_names` leak into, or override, the permutation's split names.
Batches are rebuilt with `RecordBatch::try_new(schema, columns)`.
- Errors are propagated (`map_ok` → `and_then`) instead of `unwrap`.
- The same panic-on-worker pattern in the splitter's streaming closures
(`split.rs`: `try_new`/`filter(...).unwrap()`) now propagates errors
too. Nothing else in `dataloader/` used `with_schema(...).unwrap()`.
## Tests
- Rust: `test_permutation_over_table_with_schema_metadata` in
`builder.rs` builds a random split + shuffled permutation over a table
with `huggingface` and a colliding `split_names` key. It panics with the
original `SchemaError` on the old code and passes now.
- Python: `test_permutation_over_table_with_schema_metadata` in
`test_permutation.py` covers the repro and reads both splits through
`Permutations`/`Permutation.from_tables`.
`test_streaming_dataset_over_table_with_schema_metadata` in
`test_elastic_dataloader.py` covers `StreamingDataset` construction and
iteration.
- `cargo fmt`, `cargo clippy --features remote --tests --examples`, the
dataloader Rust tests (42), `test_permutation.py` (72), and
`test_elastic_dataloader.py` all pass.
`test_torch.py::test_remote_permutation_dataloader_multiprocessing`
fails locally with `take query result did not include ordering column
'__lancedb_take_row_offset'`. It uses `Permutation.identity`, which does
not go through `PermutationBuilder`, so this change does not cause it.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
|
||
 Weston PaceandClaude Sonnet 5
|
ec53c82ef8 |
fix(python): bound StreamingDataset's raw_batches queue against io_queue_depth (#4530)
## Summary - `StreamingDataset`'s Stage 1 (I/O) queue (`raw_batches`) had no cap: `_fill_io` only limited *in-flight* fetch futures (`io_pending`) to `io_queue_depth`, but every completed fetch was unconditionally drained into `raw_batches` and the freed slot immediately refilled, so I/O could race arbitrarily far ahead of a slow transform stage with no bound on buffered memory. - `transform_queue_depth`'s admission check (`_try_submit_tx`) couldn't compensate either, since it only counts `cooked` rows and in-flight transform futures — it never looks at `raw_batches`. - Fix: `_fill_io` now counts both `io_pending` and `raw_batches` against the `io_queue_depth` budget, matching the parameter's documented "batches to keep in flight per split" contract. ## Test plan - [x] Added `test_raw_queue_depth_bounded_when_transform_is_the_bottleneck` (deterministic, via synchronous I/O) asserting `raw_queue_depth <= io_queue_depth * read_batch_size` even when the transform stage is the bottleneck. - [x] Full `python/python/tests/test_elastic_dataloader.py` suite passes (190/190), confirming no change to sample ordering, counts, or checkpoint/resumability behavior. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> |
||
  lancedb-gatefixer[bot]andXuanwo
|
105fd73bc6 |
fix(python): commit streaming worker checkpoints on consumption (#4023)
## Summary - add `StreamingDataLoader`, which transports worker snapshots with prefetched batches and commits them to the parent dataset only when the trainer receives each batch - preserve exact non-uniform per-split progress and resume lagging splits without replaying already-consumed rows - reject stale parent checkpoints after a standard multi-process `DataLoader` has started, with guidance to use the consumer-aware loader - document the new public loader and merge non-uniform state across ranks ## Root cause PyTorch runs `StreamingDataset.__iter__` in private worker-process copies, while callers invoke `state_dict()` on the parent dataset. Sharing producer counters would still be incorrect because DataLoader prefetch can advance workers beyond batches returned to the trainer. ## Validation - `uv run --extra tests pytest python/tests/test_elastic_dataloader.py -q` (154 passed) - focused non-uniform merge regression (1 passed) - `uv run --project python --extra tests --extra dev ruff format .` - `uv run --project python --extra tests --extra dev ruff check .` - `cd docs && PYTHONPATH=. ../python/.venv/bin/mkdocs build` Fixes #3967 <!-- lance-gatekeeper-fix:v1 agent=572be272619660b97e87fd5c85188341 generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io> |
||
  Ayush ChaurasiaandOpenAI Codex
|
242ade8017 |
feat(python): support sequence packing in streaming dataset (#3920)
## How packing works
Consider four tokenized documents:
[1]
[2]
[10, 11, 12, 13, 14, 15, 16, 17]
[20]
With:
```
StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
num_splits=2,
pack_sequences=5,
eos_id=9,
pad_id=0,
blocks_per_epoch=6,
)
```
the documents are assigned to two fixed logical splits. Each split
maintains an independent token buffer, appends eos_id after every
document, and emits blocks of five tokens.
Because blocks_per_epoch=6, each split emits exactly three blocks:
Cycl/e 1:
Split 0: [1, 9, 2, 9, 0] # 9 is eos, 0 is padding
Split 1: [10, 11, 12, 13, 14]
Cycle 2:
Split 0: [0, 0, 0, 0, 0]
Split 1: [15, 16, 17, 9, 20]
Cycle 3:
Split 0: [0, 0, 0, 0, 0]
Split 1: [9, 0, 0, 0, 0]
If a split runs out of tokens early, it emits padded blocks through the
fixed budget. This prevents one rank from finishing before another.
Logical splits are independent of rank and worker ownership. A
checkpoint records each split’s consumed-document count, emitted-block
count, remaining tokens, and document boundaries. Merging
those per-split states allows the same packed stream to resume after the
topology changes.
doc_ids identifies document segments, including continuations across
block boundaries. It is not a padding mask: padding retains the
preceding document ID, so callers must mask padding using a
reserved pad_id.
blocks_per_epoch="auto" is also available. It estimates the budget from
a deterministic bounded sample and warns that the result is approximate.
WIP pre-training tests:
```
┌────────────────────────────────────┬────────────────────────┬─────────────────────────────────────┐
│ │ GPT-2 124M │ GPT-2 medium 354M │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Corpus │ 2.4M docs / 12GB table │ 9.67M docs / 45GB table │
│ Tokens (Chinchilla) │ 2.43B │ 7.0B │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Data prep (ingest→curate→tokenize) │ ~12 min │ ~51 min │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Training wall time │ ~50 min │ 3h 06m │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Throughput / MFU │ 1.60M tok/s / 35% │ 684k tok/s / 42.0%, │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Final val loss │ 3.230 │ 2.840 │
└────────────────────────────────────┴────────────────────────┴─────────────────────────────────────┘
```
---------
Co-authored-by: OpenAI Codex <codex@openai.com>
|
||
|
Weston PaceandClaude Sonnet 4.6
|
1f1d03f306 |
feat(python): add backpressure to StreamingDataset post-transform queue (#3897)
Rename prefetch_batches → io_queue_depth and introduce transform_queue_depth as a symmetric pair: both express "number of batches to buffer per split at this pipeline stage." The old names are still accepted as keyword arguments but log a deprecation warning redirecting callers to the new names. transform_queue_depth caps how many transform-result batches can accumulate per split in the post-transform queue. Without this limit a slow consumer (e.g. a GPU training step) causes cooked rows to pile up unboundedly. The backpressure check in _try_submit_tx counts both already-cooked rows and rows expected from in-flight transforms; it skips proactive transform submission when the combined total reaches the limit. The reactive _ensure_cooked path bypasses the check so the consumer never stalls. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
|
Jack Ye
|
f39a7a4dd9 |
feat: support remote tables in the data loader (#3981)
`StreamingDataset`, `PermutationBuilder`, and `Permutation` now work
against a `RemoteTable` (LanceDB Cloud and Enterprise), which unblocks
benchmarking the loader against the enterprise cluster cache.
```python
db = lancedb.connect("db://my-db", api_key=..., host_override=...)
ds = StreamingDataset(db.open_table("training"), world_size=8, rank=r)
```
Rows are addressed by `_rowid` exactly as before —
`PermutationReader::load_batch` already built the same `_rowid IN (...)`
filter that `Table::take_row_ids` sends, so the loader's fetch was
always the take path. It just was never allowed to run.
### The guard
`PermutationBuilder.__init__` rejected anything without `_inner`, so a
`RemoteTable` raised `TypeError` before reaching the PyO3 layer — which
already unwraps one via `_table._inner`.
### A bounded schema lookup
`PermutationReader::output_schema` reads the schema off a query plan,
and building a plan on a remote table *executes* the query
(`create_plan` → `execute_query`). With no limit that is `k =
isize::MAX`, so asking a remote table for its output schema pulled the
whole table over HTTP and threw it away — once per assigned split, on
every epoch, since `StreamingDataset.__iter__` constructs a
`Permutation` per split.
One row rather than zero, deliberately: lance gates its limit node on
`self.limit.unwrap_or(0) > 0`, so `Some(0)` means *no limit*.
### Tables with an LSM write spec are refused
A permutation references rows by row id, and rows that have not been
flushed to the base table do not have one yet. The loader could read
around them, but they would then be missing from training with nothing
said about it, so the build refuses such a table up front instead of
half supporting it.
### Fallible identity construction
`PermutationReader::identity` resolved `inner_new` with `unwrap`. That
was near total against a local dataset, but construction counts the base
table — an HTTP round trip for a remote one — so a transient network or
auth failure became a panic across the PyO3 boundary.
### Tests
End-to-end `permutation_builder` and `StreamingDataset` runs against a
mock server, the former torch-free so it runs wherever the suite does,
plus a test that a build succeeds without an LSM write spec and is
refused once one is installed.
|
||
 Sravan AvvaruandClaude Fable 5
|
a615306f39 |
feat(python): add on_transform_error fault tolerance to StreamingDataset (#3763)
Closes #3704 ## Problem Transforms can fail on bad data (e.g. nulls/NaNs from incomplete user surveys). Today any transform exception aborts iteration, and there is no way to skip invalid rows during loading. ## Solution New `on_transform_error` parameter on `StreamingDataset`: - `"raise"` (default, matches current behavior and the convention in tf.data / WebDataset / Ray Data) - `"skip"` — drop the failing rows and continue - `"warn"` — like skip, plus a logged warning per failing batch - a WebDataset-style callable `handler(exc) -> bool`, so users can skip only expected error types Key design points: - **Row-granular skipping**: when a batch fails, the transform is re-run on single-row slices so only the rows that actually fail are dropped (avoids Ray-style whole-block loss). Skips are counted in a new `rows_skipped` property. - **No crash on uneven skips**: the round-robin loop now ends the epoch at the last cycle where every split still has a row, instead of hitting `IndexError` when a split runs dry early. - **Exact resumability under skips**: checkpoints are now position-based. `state_dict` gains `positions_consumed_per_split` (exact for owned splits), and a new `merge_state_dicts` static method combines per-rank states via elementwise max for elastic resume across topology changes. Old checkpoints without the new key still load. Positions equal sample counts when nothing is skipped, so existing behavior is unchanged. - **Guardrail**: transforms returning the wrong number of rows now raise a clear `ValueError` instead of silently corrupting split accounting. ### Answers to the issue's open questions - *Can we do this?* Yes — all transforms funnel through one guarded call in the Stage 2 pipeline. - *What do other libraries do?* tf.data `ignore_errors()`, WebDataset `handler=`, Ray `max_errored_blocks`; MosaicML StreamingDataset offers nothing (skipping conflicts with its determinism model). This design follows the common conventions: raise by default, opt-in skipping, count/log drops. - *Error handling or pre-filtering?* Both: the existing `filter=` remains the recommended tool for predictable bad data (splits are built post-filter, so all guarantees hold — now documented); `on_transform_error` covers failures not expressible as a predicate. - *Impact on splits / elastic determinism?* Per-split sample sequences stay deterministic (skips are data-dependent, not topology-dependent). With unequal bad-row counts across splits the last few global steps of an epoch can differ across topologies (bounded by the skew), which is documented on the parameter. With equal counts per split, full determinism is preserved — covered by a test. ## Testing 15 new tests in `test_elastic_dataloader.py` covering: default raise, invalid values, uniform and uneven skips (including epoch-end truncation), warn logging, selective callable handlers, wrong-row-count guardrail, determinism across runs and across world sizes (1/2/3/4) with skips, exact mid-epoch resume with skips on the same topology, elastic resume via `merge_state_dicts` (ws=2 → ws=1), merge validation, and backward-compat loading of old checkpoints. Note: relying on CI for the test run — my local machine OOMs during the final link of the native extension. The change itself is pure Python. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
  buduoqiuandbuduoqiu
|
98a52267a2 |
feat(python): configure streaming transform parallelism (#3699)
## Summary - add a keyword-only `transform_parallelism` option to `StreamingDataset` - preserve CPU auto-detection by default and fall back to one worker when unavailable - apply the configured limit to both the transform executor and concurrency semaphore - document and test explicit, default, fallback, and invalid values ## Testing - `uv run --extra tests --with torch pytest python/tests/test_elastic_dataloader.py -q` (`136 passed`) - `uvx ruff check python/lancedb/streaming.py python/tests/test_elastic_dataloader.py` - `uvx ruff format --check python/lancedb/streaming.py python/tests/test_elastic_dataloader.py` - `git diff --check origin/main...HEAD` Closes #3695 Co-authored-by: buduoqiu <yaodong-shen@users.noreply.github.com> |
||
|
Weston PaceandClaude Sonnet 4.6
|
c6db80dd0b |
feat: add an elastic dataloader as an iterable dataset (#3509)
# Elastic Streaming Dataloader ## Motivation Training large models on LanceDB tables today requires loading the entire dataset into memory or writing bespoke batching logic. This PR introduces `StreamingDataset`, a PyTorch `IterableDataset` that streams directly from a LanceDB table with two hard guarantees that are difficult to achieve together: **elastic determinism** and **resumability**. ## Goals ### Elastic determinism The dataset partitions the table into a fixed number of *splits* (controlled by `num_splits`, `shuffle_seed`, and `epoch`). Samples are yielded by round-robining over splits one sample per split per cycle. Because the split structure is fixed, the set of samples that makes up each global training step is identical regardless of `world_size` or `num_workers`. You can scale your cluster up or down between runs and the model sees the same data in the same order — no re-sharding, no gradient variance from topology changes. ### Resumability `state_dict()` / `load_state_dict()` capture how many samples each split has consumed. Because all splits are the same size and the round-robin design keeps them in lockstep, the state reduces to a single scalar (`samples_consumed_per_split`) that is topology-independent. A checkpoint saved with 8 GPUs can resume correctly on 4 GPUs or 16 GPUs without any adjustment. ### PyTorch `IterableDataset` / streaming `StreamingDataset` implements the standard PyTorch `IterableDataset` interface, so it drops into any existing `DataLoader` pipeline without modification. Data is fetched lazily from Lance in chunks — only the rows needed for the current batch are ever in memory. Compared to the map dataset this takes more work from pytorch and puts it into the dataset itself (e.g. shuffling, filtering, etc.). We do this because we cannot achieve things like elastic determinism or prefiltering otherwise. ### Multi-worker support DataLoader workers are automatically assigned contiguous sub-blocks of splits (the rank's splits are divided evenly across workers). Each worker is independent: no shared state, no inter-process coordination. The only constraint is that `num_splits` must be divisible by `world_size * num_workers`. That being said, multi-worker is highly discouraged as it relies on multiprocessing which is inefficient. Still, we want to support it. ### Filters as prefilters Filters are applied at *permutation-build time* via `PermutationBuilder.filter()`, not re-evaluated on every fetch. The filtered row IDs are stored in the permutation table so that subsequent reads see only the matching rows. This allows us to avoid loading rows that don't match the filter (which is the default pytorch behavior) ### Prefetching Two parameters control the I/O pipeline: - `read_batch_size` (default 64) — number of rows fetched per `take_offsets` call. Larger values amortise per-request overhead, which is critical on object storage where a single round-trip can cost ~100 ms. - `prefetch_batches` (default 4) — number of batches prefetched in parallel per split via a `ThreadPoolExecutor`. While the model processes the current batch, the next several batches are already in flight, hiding storage latency behind compute. If set correctly then you can get good performance even with num_workers=0 (unless you are bottlenecked on transform). ### Transform parallelism The underlying `Permutation` API supports a `with_transform()` callback for decoding, augmentation, and format conversion. Unfortunately, this is not parallelized. Pytorch typically parallelizes this with num_workers which is multiprocessing which is highly inefficient. For simple transforms we should be able to utilize multithreading and Rust based UDFs. For complex python UDFs we could have a dedicated multiprocessing pipeline for just the transform. Or we could just utilize multithreading. In both cases we would exclude the I/O stage from the multiprocessing because that ends up being very memory hungry and inefficient. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |