mirror of
https://github.com/lancedb/lancedb.git
synced 2026-07-10 22:40:41 +00:00
c6db80dd0b
# 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>
The Multimodal AI Lakehouse
How to Install ✦ Detailed Documentation ✦ Tutorials and Recipes ✦ Contributors
The ultimate multimodal data platform for AI/ML applications.
LanceDB is designed for fast, scalable, and production-ready vector search. It is built on top of the Lance columnar format. You can store, index, and search over petabytes of multimodal data and vectors with ease. LanceDB is a central location where developers can build, train and analyze their AI workloads.
Demo: Multimodal Search by Keyword, Vector or with SQL
Star LanceDB to get updates!
⭐ Click here ⭐ to see how fast we're growing!
Key Features:
- Fast Vector Search: Search billions of vectors in milliseconds with state-of-the-art indexing.
- Comprehensive Search: Support for vector similarity search, full-text search and SQL.
- Multimodal Support: Store, query and filter vectors, metadata and multimodal data (text, images, videos, point clouds, and more).
- Advanced Features: Zero-copy, automatic versioning, manage versions of your data without needing extra infrastructure. GPU support in building vector index.
Products:
- Open Source & Local: 100% open source, runs locally or in your cloud. No vendor lock-in.
- Cloud and Enterprise: Production-scale vector search with no servers to manage. Complete data sovereignty and security.
Ecosystem:
- Columnar Storage: Built on the Lance columnar format for efficient storage and analytics.
- Seamless Integration: Python, Node.js, Rust, and REST APIs for easy integration. Native Python and Javascript/Typescript support.
- Rich Ecosystem: Integrations with LangChain 🦜️🔗, LlamaIndex 🦙, Apache-Arrow, Pandas, Polars, DuckDB and more on the way.
How to Install:
Follow the Quickstart doc to set up LanceDB locally.
API & SDK: We also support Python, Typescript and Rust SDKs
| Interface | Documentation |
|---|---|
| Python SDK | https://lancedb.github.io/lancedb/python/python/ |
| Typescript SDK | https://lancedb.github.io/lancedb/js/globals/ |
| Rust SDK | https://docs.rs/lancedb/latest/lancedb/index.html |
| REST API | https://docs.lancedb.com/api-reference/rest |
Join Us and Contribute
We welcome contributions from everyone! Whether you're a developer, researcher, or just someone who wants to help out.
If you have any suggestions or feature requests, please feel free to open an issue on GitHub or discuss it on our Discord server.
Check out the GitHub Issues if you would like to work on the features that are planned for the future. If you have any suggestions or feature requests, please feel free to open an issue on GitHub.
Contributors
Stay in Touch With Us