diff --git a/docs/training/index.mdx b/docs/training/index.mdx
index 01ac5e6..a36142a 100644
--- a/docs/training/index.mdx
+++ b/docs/training/index.mdx
@@ -1,15 +1,17 @@
---
title: "Loading Data for Model Training"
sidebarTitle: "Data loading and shuffles"
-description: "Introduction to loading data, shuffling, and creating permutations for model training in LanceDB."
+description: "Stream, shuffle, transform, and resume model training data from LanceDB."
icon: boxes-stacked
---
-LanceDB makes an excellent data backend for training machine learning models. This section will describe the
-`Permutation` API in LanceDB that's designed to facilitate the
-process of training a model and explain how to use LanceDB as a data backend for training.
+LanceDB makes an excellent data backend for training machine learning models. A `Table` can be used directly as input
+to a data loader, but this is typically limited. A `Permutation` gives you control over which rows are accessed and in
+what order. For a more complete solution, LanceDB also provides a streaming data loader through `StreamingDataset`.
+This PyTorch `IterableDataset` adapts the lower-level `Permutation` API and adds prefetching, elastic determinism,
+resumability, and multithreaded transformations.
-## Data loading
+## Basic data loading
Most model training frameworks iterate through data in batches and feed this data into the model. This process is
often referred to as **data loading**. The simplest way to load data into a model is to iterate a LanceDB table in
@@ -29,116 +31,387 @@ for batch in table:
In practice, this is too simplistic for effective training. We may not want to load all the data, or we may want
to load the data in a different order, or we may need to apply some sort of processing to the data before training.
-To achieve this, we will often want to create a `Permutation` of the table.
+To achieve this, we can use the `StreamingDataset`.
-## Selecting rows
+
+```py Python icon=Python
+import lancedb
+import torch
+from lancedb.streaming import StreamingDataset
-When training a model, we might not want to load all of the data. For example, we might filter out columns that
-are not needed for training. We might also divide the data into training and validation sets. Or we could divide
-the data into multiple sets for cross-validation.
+db = lancedb.connect("file://some/db/path")
+table = db.open_table("some_table")
-Whenever we create a permutation of the table, we have to first decide which rows we want to include (and in what
-order). This is stored in a **permutation table** which marks out the row ids that make up our data. Other decisions,
-such as which columns to include, and what transformations to apply, can be defined at read time and don't require
-a separate permutation table.
+dataset = StreamingDataset(table, shuffle_seed=42)
+dataloader = torch.utils.data.DataLoader(
+ dataset,
+ batch_size=128,
+ num_workers=0,
+)
+
+for batch in dataloader:
+ train_step(batch)
+```
+
+
+`StreamingDataset` yields plain Python dictionaries by default. PyTorch's default collation function combines those
+samples into batches. You can also iterate the dataset directly when you do not need collation.
-Permutation tables are tables, just like any other table in LanceDB. By default, they are
-stored in memory but they can be persisted to storage as well. This is useful when you want to share a permutation
-table across processes or nodes.
+`StreamingDataset` is built on the permutation API and works with both embedded LanceDB tables and remote tables
+accessed through the LanceDB Cloud API. The underlying table data can live on local disk or cloud object storage.
-## Selecting all rows
+Use the streaming data loader when the training data does not fit in memory, when you need deterministic global batches
+across cluster sizes, or when you want filtering and prefetching to happen before PyTorch requests individual samples.
+Use `Permutation` directly when you need map-style random access instead. Only one iterator can be active on a
+`StreamingDataset` instance at a time; create a separate instance for each concurrent consumer.
-To select all rows, we can use the `Permutation.identity` method. This gives us a `Permutation` without requiring
-us to create a separate permutation table. This allows us to refine our columns and apply transformations and can
-be useful when the data loader itself is responsible for handling sampling and shuffling.
+## Advanced data loading
-
-```py Python icon=Python
-from lancedb.permutation import Permutation
+The `StreamingDataset` wraps a LanceDB `Table` and, by default, adds prefetching and conversion from Arrow to Python.
+It can also handle more advanced scenarios. To explain these, consider a model trained with stochastic gradient
+descent (SGD) and distributed data parallelism (DDP). In this example, we need to load batches onto multiple GPUs
+across multiple servers. After each batch is processed, the GPUs exchange weights and the next batch is loaded. We
+will use the following PyTorch terms:
-# We can create an identity permutation without needing any separate permutation table.
-permutation = Permutation.identity(table)
+- **World size** - The number of GPUs that we are loading. For example, if we have 2 servers and each server has 4
+ GPUs, the world size is 8.
+- **Rank** - The identifier of the process loading data for a GPU. It is an integer in the range `[0, world_size)`.
+ Each rank gets its own portion of the data.
+- **Global batch size** - The number of rows processed across all GPUs in each step of the SGD algorithm. For
+ example, if we have 8 GPUs and a global batch size of 1024, we load 128 rows onto each GPU for each step.
+- **Batch size** - The number of rows processed by one GPU in each step. In the preceding example, the batch size
+ is 128.
-# This allows us to refine our columns and apply transformations
-permutation = permutation.select_columns(["id", "prompt"])
-```
-
+Other concepts, such as read batch size and `num_workers`, are introduced in the relevant sections below.
+
+### Prefetching
-## Filtering rows
+PyTorch datasets were originally built around in-memory structures like a Pandas DataFrame. When they are iterated,
+they yield a single sample at a time. This makes sense for a simple in-memory structure, but accessing data on object
+storage one row at a time introduces too much per-call overhead. To avoid this, `StreamingDataset` fetches and
+transforms data in batches. The `read_batch_size` parameter controls how many rows are fetched from each split per
+storage request and defaults to `64`.
-If we only want to select a subset of rows, then we can use a filter. This will require us to create a permutation
-table which identifies which rows we want to include.
+In addition to batching requests, the prefetching mechanism reads ahead in the background. While one batch is being
+transformed and processed by the GPU, `StreamingDataset` reads subsequent batches. The `prefetch_batches` parameter
+controls how many batches are kept in flight per split and defaults to `4`. A larger value can provide more buffering
+against jittery workloads, but it also increases memory use and I/O concurrency.
-```py Python icon=Python
-from lancedb.permutation import Permutation, permutation_builder
+```py Python icon=Python
+ds = StreamingDataset(
+ table,
+ shuffle_seed=42,
+ read_batch_size=256, # rows fetched per LanceDB call per split
+ prefetch_batches=8, # batches to keep in flight per split
+)
+```
+
+
+### Transformation
-# We can create a permutation table which identifies which rows we want to include.
-permutation_tbl = permutation_builder(table).filter("category = 'cat'").execute()
+Many model training workloads require a transformation step between loading the data and training the model. For
+example, we may need to decode images, tokenize text, or normalize data. A transformation function can be provided
+using the `transform` parameter. Transformations can be expensive, so `StreamingDataset` applies them with a
+`ThreadPoolExecutor` whose worker count equals the number of available CPUs.
-# We can then use this permutation table to create a Permutation object
-permutation = Permutation.from_tables(table, permutation_tbl)
+Transformations are applied to batches, not individual samples, to amortize per-batch overhead. A transformation
+function receives a PyArrow `RecordBatch` and must return an iterable with exactly one output sample for every input
+row. The sample format should match what your data loader expects. For example, PyTorch's default collation function
+accepts several sample types, with a Python dictionary being one of the most common. When no transform is provided,
+the default transform converts the Arrow record batch into Python dictionaries without further processing.
+
+
+```py Python icon=Python
+import pyarrow as pa
+
+def normalize(batch: pa.RecordBatch) -> list[dict]:
+ # This pure-Python loop holds the GIL and is shown for illustration only.
+ # In practice, prefer a library like torchvision or numpy that releases the
+ # GIL so the ThreadPoolExecutor can run transforms in parallel.
+ rows = batch.to_pylist()
+ for row in rows:
+ row["image"] = [v / 255.0 for v in row["image"]]
+ return rows
+
+ds = StreamingDataset(table, shuffle_seed=42, transform=normalize)
```
-## Creating splits
+#### DataLoader workers
+
+The thread-based transformation model that `StreamingDataset` uses by default is only effective when the transform
+function releases the GIL. This is true for many Python scientific libraries, including NumPy, PyArrow, and
+TorchVision, but not for pure-Python transforms. PyTorch can launch multiple DataLoader worker processes per rank,
+and `StreamingDataset` uses `get_worker_info()` to give each worker a non-overlapping group of splits.
+
+Multiprocessing adds pickling and transfer overhead and increases memory use. Start with `num_workers=0`, which keeps
+loading in the rank's main process, and add workers only after confirming an unavoidable GIL bottleneck. If you use
+workers, `num_splits` must be divisible by `world_size * num_workers`. Use the `forkserver` or `spawn` multiprocessing
+context because LanceDB uses internal threads.
+
+```py Python icon=Python
+dataloader = torch.utils.data.DataLoader(
+ dataset,
+ batch_size=batch_size,
+ num_workers=2,
+ multiprocessing_context="forkserver",
+ persistent_workers=True,
+)
+```
-LanceDB also provides several different methods for creating splits. These allow us to divide our dataset into
-smaller non-overlapping sets. The split can then be specified when creating the `Permutation` object to view
-only a subset of the data.
+By default, `StreamingDataset` serializes enough table state to reopen the table in each worker. For custom connection
+setup, pass a picklable `connection_factory` callable that accepts a table name and returns an open table. This avoids
+serializing connection credentials into worker state.
+
+### Observability & performance
+
+Optimizing data loader performance is tricky because it can be difficult to locate the bottleneck. What is often
+blamed on I/O can be a CPU bottleneck in the transform stage, or vice versa. To help distinguish them,
+`StreamingDataset` exposes pipeline counters. `raw_queue_depth` is the number of loaded rows waiting to be transformed,
+while `prefetch_queue_depth` is the number of transformed rows ready to be consumed. `unscanned_rows` and
+`consumed_rows` show how far the current iterator has progressed.
+
+If the `prefetch_queue_depth` is consistently zero but the `raw_queue_depth` is not then you have a CPU
+transformation bottleneck. You should investigate GIL bottlenecks or look for ways to optimize your transformation.
+This can often be done by batching the compute work. If both `prefetch_queue_depth` and `raw_queue_depth` are
+consistently zero while the consumer is waiting, I/O is the likely bottleneck. A larger read batch size or clumped
+shuffling could help.
-```py Python icon=Python
-from lancedb.permutation import Permutation, permutation_builder
+```py Python icon=Python
+import threading, time
+
+ds = StreamingDataset(table, shuffle_seed=42)
+
+def log_pipeline_health():
+ while True:
+ print(
+ f"unscanned={ds.unscanned_rows} "
+ f"raw={ds.raw_queue_depth} "
+ f"cooked={ds.prefetch_queue_depth} "
+ f"consumed={ds.consumed_rows}"
+ )
+ time.sleep(1.0)
+
+monitor = threading.Thread(target=log_pipeline_health, daemon=True)
+monitor.start()
+
+for sample in ds:
+ train_step(sample)
+
+print(
+ f"bytes loaded: {ds.bytes_loaded} "
+ f"fetch time: {ds.fetch_time:.2f}s "
+ f"transform time: {ds.transform_time:.2f}s"
+)
+```
+
-# Here we create two splits, one for training and one for validation. By default, splits have no
-# name and are accessed by index.
-permutation_tbl = permutation_builder(table).split_random(ratios=[0.95, 0.05]).execute()
+`bytes_loaded` measures raw Arrow buffer bytes before transformation. `bytes_loaded`, `fetch_time`, and `transform_time`
+are cumulative across iterations of the same dataset instance; because work runs concurrently, the summed stage times
+can exceed wall-clock time. The queue depths report rows currently waiting in the pipeline, and the progress counters
+reflect the latest iterator snapshot.
-# Let's create a permutation object which views only the training data.
-permutation = Permutation.from_tables(table, permutation_tbl, split=0)
+### Filtering data
-# Splits can also be given names. The names can then be used later to access the split instead of
-# requiring us to know the index.
-permutation_tbl = permutation_builder(table).split_random(ratios=[0.95, 0.05], split_names=["train", "test"]).execute()
-permutation = Permutation.from_tables(table, permutation_tbl, split="train")
+By default, the streaming data loader includes all rows and columns. LanceDB is a columnar database that also supports
+efficient random access. Reducing the number of columns you load has a direct impact on I/O performance. Reducing the
+number of rows can also help, especially with a selective filter, large values, local data, or the LanceDB Enterprise
+cache.
+
+Use the `columns` parameter to specify which columns to load and the `filter` parameter to specify which rows to load.
+The filter is evaluated once when the permutation is constructed, before the rows are divided into splits. Filtered-out
+rows are not loaded from storage during iteration, and split sizes reflect the filtered row count.
+
+
+```py Python icon=Python
+ds = StreamingDataset(
+ table,
+ shuffle_seed=42,
+ columns=["image", "label"], # skip all other columns
+ filter="category = 'train'", # only training rows
+)
```
-## Shuffling rows
+### Shuffling rows
+
+By default, `StreamingDataset` sets `shuffle=True` and randomly assigns rows to splits. This helps prevent the model
+from learning artifacts from storage order. Set `shuffle=False` to divide rows into splits sequentially, which is
+useful for deterministic evaluation and debugging.
-By default, permutations will access the data in the order the data is stored in the table. This can cause our
-model to learn artifacts specific to the order of the data. This is one of many ways we can "overfit" our model
-to our data. To avoid this, we typically want to shuffle the data before training. Model training frameworks
-(like PyTorch) will often provide a way to shuffle the data. If you are not using one of these frameworks, or if
-you want to shuffle the data with LanceDB, you can shuffle the rows when you create a permutation table.
+The effective shuffle seed combines `shuffle_seed` with `epoch`, so each epoch has a different ordering while runs
+with the same inputs remain reproducible. Keep `epoch=0` if you want the same ordering across iterations. The default
+`shuffle_seed` is `0`; set it to `None` to generate a random seed when the dataset is constructed.
-```py Python icon=Python
-from lancedb.permutation import Permutation, permutation_builder
+```py Python icon=Python
+# Training loop: each epoch gets a different shuffled ordering
+for epoch in range(num_epochs):
+ ds = StreamingDataset(
+ table,
+ shuffle_seed=42,
+ epoch=epoch, # changes the permutation each epoch
+ )
+ for sample in ds:
+ train_step(sample)
+
+# Evaluation: deterministic sequential order, no shuffle
+eval_ds = StreamingDataset(table, shuffle=False)
+for sample in eval_ds:
+ eval_step(sample)
+```
+
-# We can shuffle the rows when we create the permutation table.
-permutation_tbl = permutation_builder(table).shuffle().execute()
+
+Shuffling can have significant impacts on I/O performance, especially if you are loading data from cloud storage.
+In many cases the GPU pipeline is slow enough that this penalty will not be noticeable. However,
+you can use the `shuffle_clump_size` parameter to shuffle the data in clumps (small contiguous batches that get
+shuffled together). This will give some penalty to the randomness of the shuffle, but will significantly improve
+I/O performance.
+
-# We can then use this permutation table to create a Permutation object, this will now
-# access the data in a random order.
-permutation = Permutation.from_tables(table, permutation_tbl)
+
+```py Python icon=Python
+# Clumped shuffle: groups of 16 contiguous rows are shuffled together,
+# preserving read locality while still randomising the global ordering.
+ds = StreamingDataset(
+ table,
+ shuffle_seed=42,
+ shuffle_clump_size=16,
+)
```
-## Selecting columns
+### Data splits and elasticity
+
+`StreamingDataset` partitions the permutation into a fixed number of equal-sized groups called splits. Each rank gets
+a contiguous group of splits, and each DataLoader worker gets a contiguous subgroup of its rank's splits. Samples are
+then yielded by round-robining over the assigned splits.
+
+If `num_splits` is omitted, it defaults to `world_size`. This works when `num_workers=0`, but it ties the split layout
+to the current number of ranks. A run resumed with a different world size would then construct a different layout and
+would not see the same global batches.
+
+To get **elastic determinism**, choose a fixed `num_splits` that is compatible with every topology you plan to use.
+For example, a run with `num_splits=8` and 4 GPUs assigns 2 splits to each rank. The ranks pull from their splits in
+round-robin order, producing the same global batches as a run with 8 GPUs and the same `num_splits`, `shuffle_seed`,
+and `epoch`.
-By default, permutations will return all columns in the table. If you only need a subset of the columns, you can
-significantly reduce your I/O requirements by selecting only the columns you need. This can be done on the
-permutation object itself, and does not require us to create a separate permutation table.
+The following constraints apply:
+
+- `num_splits` must be divisible by `world_size`.
+- With DataLoader worker processes, `num_splits` must be divisible by `world_size * num_workers`.
+- For the same samples to form each global training step across topologies, `global_batch_size` must be a multiple of
+ `num_splits`.
+- The filtered row count must be at least `num_splits`. If it is not evenly divisible by `num_splits`, up to
+ `num_splits - 1` surplus rows are dropped so that every split has the same length.
+
+For example, to switch between 8 and 6 GPUs, use a `num_splits` value divisible by both, such as 24. Highly composite
+values can support several layouts: 48 supports 1, 2, 3, 4, 6, 8, 12, 16, 24, and 48 GPUs, while 60 supports 1, 2, 3,
+4, 5, 6, 10, 12, 15, 20, 30, and 60 GPUs. Remember to include `num_workers` when checking divisibility.
-```py Python icon=Python
-from lancedb.permutation import Permutation
+```py Python icon=Python
+import torch.distributed as dist
+
+dist.init_process_group("nccl")
+rank = dist.get_rank()
+world_size = dist.get_world_size()
+
+# num_splits=48 is divisible by 1, 2, 3, 4, 6, 8, 12, 16, 24, 48
+# so this dataset works unchanged as you scale up or down GPUs.
+ds = StreamingDataset(
+ table,
+ num_splits=48,
+ shuffle_seed=42,
+ epoch=current_epoch,
+ rank=rank,
+ world_size=world_size,
+)
+
+for sample in ds:
+ train_step(sample)
+```
+
+
+### Checkpointing and resumability
+
+Model training is expensive, and failures can occur partway through a run. A model checkpoint is not enough for an
+exact resume: the streaming data loader must also continue from the same position. `StreamingDataset.state_dict()`
+captures the number of samples consumed from every split in a plain Python dictionary, and `load_state_dict()` restores
+that position.
-# We can select only the columns we need.
-permutation = Permutation.identity(table).select_columns(["id", "prompt"])
+Save this state at a global-step boundary where every split has contributed equally. The easiest way to guarantee this
+is to make `global_batch_size` a multiple of `num_splits`, as in the example below. If you checkpoint partway through a
+round-robin cycle, the state reflects the preceding complete cycle and those partial-cycle samples can be replayed.
+
+
+```py Python icon=Python
+import torch
+
+global_batch_size = 48 # one sample per split in each global step
+batch_size = global_batch_size // world_size
+
+dataset = StreamingDataset(
+ table,
+ num_splits=48,
+ shuffle_seed=42,
+ epoch=current_epoch,
+ rank=rank,
+ world_size=world_size,
+)
+dataloader = torch.utils.data.DataLoader(
+ dataset,
+ batch_size=batch_size,
+ num_workers=0,
+)
+
+for step, batch in enumerate(dataloader):
+ train_step(batch)
+
+ if (step + 1) % checkpoint_interval == 0:
+ torch.save(
+ {"model": model.state_dict(), "dataset": dataset.state_dict()},
+ f"checkpoint_{step + 1}.pt",
+ )
+
+# --- resuming after a crash ---
+checkpoint = torch.load("checkpoint_100.pt")
+model.load_state_dict(checkpoint["model"])
+
+dataset_state = checkpoint["dataset"]
+dataset = StreamingDataset(
+ table,
+ num_splits=dataset_state["num_splits"],
+ shuffle_seed=dataset_state["shuffle_seed"],
+ epoch=dataset_state["epoch"],
+ rank=rank,
+ world_size=world_size, # may differ from the run that saved the checkpoint
+)
+dataset.load_state_dict(dataset_state)
+dataloader = torch.utils.data.DataLoader(
+ dataset,
+ batch_size=global_batch_size // world_size,
+ num_workers=0,
+)
+
+for batch in dataloader:
+ train_step(batch)
```
-
\ No newline at end of file
+
+
+`load_state_dict()` rejects a checkpoint whose `num_splits` or `shuffle_seed` differs from the new dataset. For an
+exact resume, also use the same table snapshot, `epoch`, `shuffle`, `filter`, and other data-selection settings. The
+`world_size` and number of DataLoader workers may change as long as the split and batch-size divisibility constraints
+still hold.
+
+## Permutations
+
+In more complicated scenarios, you may want the flexibility to shuffle, split, and select data without using the full
+iterable streaming data loader. In these cases, use `Permutation`, the lower-level class on which `StreamingDataset`
+is built. A `Permutation` defines a custom ordering of the data and supports map-style access through `__getitem__()`
+and batched access through `__getitems__()`.