Loading Data for Model Training

LanceDB provides an excellent data backend for training machine learning models. For basic use cases, a Table can be fed directly into a data loader. For a more complete solution, a StreamingDataset adds prefetching, elastic determinism, resumability, multithreaded transformations, and more.

Basic Data Loading

Most model training frameworks feed data into a model by sequencing through successive batches, a process often referred to as data loading. The simplest way to perform data loading is to iterate through a LanceDB table in a loop.

import lancedb

db = lancedb.connect("file://some/db/path")
table = db.open_table("some_table")

for batch in table:
    print(batch.to_pydict())

Advanced Data Loading

Advanced data loading scenarios can be handled by the StreamingDataset. This object allows you to apply preprocessing before training, pipelined reads, and more. Use this when the training data does not fit in memory, or when you need deterministic global batches across cluster sizes.

import lancedb
import torch
from lancedb.streaming import StreamingDataset

db = lancedb.connect("file://some/db/path")
table = db.open_table("some_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)

The StreamingDataset wraps a LanceDB Table and yields plain Python dictionaries by default. Whether iterating directly or through a PyTorch DataLoader, only one iterator can be simultaneously active on a StreamingDataset instance.

StreamingDataset is built on the permutation API and works with both LanceDB OSS and LanceDB Enterprise. The underlying table data can live on local disk or object storage.

To understand the wide range of scenarios handled by StreamingDataset, 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 fix the following terminology:

World size - The number of GPUs being loaded simultaneously. For example, if we have $2$ servers with $4$ GPUs each, the world size is $8$ .
Rank - The identifier of an individual GPU. It is an integer in the range [0, world_size). Each rank is allocated a disjoint portion of the source data.
Batch size - The number of rows processed by one GPU in a single training step. For example, $128$ .
Global batch size - The total number of rows processed across all GPUs in a training step (sometimes called a “minibatch” in machine learning literature). For example, if we have $8$ GPUs with batch size $128$ , the global batch size is $1024$ .

Other concepts, such as read batch size and num_workers, are introduced in the relevant sections below.

Filtering

The StreamingDataset object allows you to train models on only the subset of rows that satisfy a user-defined filter, specified using the filter parameter. This can be useful for removing large rows, badly-sourced data, and a wide range of other applications. You can also use the columns parameter to load only a subset of columns.

ds = StreamingDataset(
    table,
    shuffle_seed=42,
    columns=["image", "label"],       # skip all other columns
    filter="category = 'train'",      # only training rows
)

Shuffling Rows

In model training, model quality can be hurt by unintentionally learning patterns in the ordering of the source data. Shuffling the dataset before training alleviates this problem. However, fetching rows with scattered sparse reads across object storage can be dramatically slower than fetching contiguous batches. StreamingDataset offers two different implementations to read data in shuffled order, toggled by the presence of the block_size parameter.

StreamingDataset shuffle seeds combine shuffle_seed with epoch; each epoch produces a different permutation while runs with the same inputs remain reproducible. Set shuffle_seed = None to generate a random seed upon pipeline initialization.

One-Phase Shuffle

Setting block_size = None corresponds to the naive

1

-phase shuffled read: the rows of the source dataset are shuffled according to a pseudorandom permutation, then fed in batches to the accelerator. This gives the best possible randomness, but can exhibit low throughput.

# 1-phase shuffled read: rows shuffled individually.
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)

Two-Phase Shuffle

Setting block_size to an integer corresponds to selecting

2

-phase shuffled read. In this implementation, the source dataset is first partitioned into fixed-size blocks, which are then assigned randomly to splits. Within each split, rows are then further shuffled, with each row moving at most max_shuffle_distance distance away from its original position. Reading blocks allows storage I/O to constitute mostly large, contiguous reads rather than single row fetches. This yields much higher throughput while slightly compromising randomness.

# 2-phase shuffled read: rows are shuffled block-by-block instead of
# individually, trading a bit of randomness for much cheaper I/O.
for epoch in range(num_epochs):
    ds = StreamingDataset(
        table,
        shuffle_seed=42,
        epoch=epoch,               
        block_size=1024,           # enables 2-phase
        window_blocks=4,           # blocks resident per split at once
        max_shuffle_distance=256,  # how far a row may move from its
                                    # original position
    )
    for sample in ds:
        train_step(sample)

# Evaluation: deterministic sequential order, no shuffle.
eval_ds = StreamingDataset(table, shuffle=False, block_size=1024)
for sample in eval_ds:
    eval_step(sample)

Parallelism & Elasticity

When loading data, StreamingDataset partitions the the source dataset among disjoint parallel threads called splits. Each rank is allocated a contiguous group of splits, and each DataLoader worker associated to a rank gets a contiguous subgroup of its rank’s splits. If using shuffle = false, data is allocated to splits naively in ordered contiguous blocks of equal size. If shuffle = true, in one-phase shuffling, the rows of the source dataset are distributed randomly among splits such that each split receives the same number of rows. In two-phase shuffling, the dataset is first partitioned into blocks, which are then distributed randomly among splits such that each split receives the same number of blocks. The number of splits defaults to world_size if num_splits is omitted during initialization. This is not recommended, as a later run using a different world size will then construct a different layout and structure of 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. 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 common multiple such as

24

for num_splits. Highly composite values can support many layouts; for example, num_splits = 48 is compatible with

10

distinct possible GPU counts. Remember to include num_workers when checking divisibility.

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 many integers,
# so this dataset works unchanged as you add or remove 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)

Fetching

Traditional datasets have often been designed to yield a single sample at a time upon iteration. However, in modern use cases, accessing data on object storage one row at a time introduces too much per-call overhead. Thus, StreamingDataset fetches data in read batches, whose size is controlled by the read_batch_size parameter (defaults = 64). In addition to batching requests, the prefetching mechanism reads ahead in the background. While one batch is being processed, StreamingDataset reads several subsequent batches. The prefetch_batches parameter specifies the number of batches simultaneously kept in flight per split (default = 4). A larger value provides more buffering against jittery workloads, but also increases memory use and I/O concurrency.

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
)

Transforming

Many model training workloads require performing a transformation to the data after it is loaded from object storage and before it is passed to the GPU. For example, we may need to decode images, tokenize text, or normalize vectors. A transformation function can be provided using the transform parameter. A transformation function must receive a PyArrow RecordBatch and return an iterable. For example, the default transform function returns a list of Python dictionaries. Furthermore, the input expected by PyTorch must match the collation function used in the PyTorch data loader. PyTorch’s default collation function accepts a Python dictionary and several other types, and is thus compatible with the default StreamingDataset transform function. If your custom transform function returns something different (for example, Arrow record batches), you may need to provide a compatible PyTorch collation function.

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)

StreamingDataset manages expensive transforms through a ThreadPoolExecutor whose worker count equals the number of available CPUs. Furthermore, transformations are applied to batches (instead of individual samples) to amortize per-batch overhead.

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.

Python

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=batch_size,
    num_workers=2,
    multiprocessing_context="forkserver",
    persistent_workers=True,
)

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

Optimizing data loader performance can be tricky, as runtime bottlenecks are often difficult to locate across various stages of the I/O and CPU pipelines. To help distinguish bottlenecks, StreamingDataset exposes several pipeline counters: raw_queue_depth shows the number of loaded rows waiting to be transformed, and 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 likely 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.

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"
)

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.

Checkpointing

StreamingDataset provides checkpointing in case your model training fails partway through a run. Use StreamingDataset.state_dict() to capture the current number of samples consumed from each split in a Python dictionary, and load_state_dict() to restore the last saved position. A stored StreamingDataset state reflects only round-robin cycles that have been completed across all splits; samples completed partway through a cycle at the time of saving can be replayed. To ensure efficient saving, make sure global_batch_size is a multiple of num_splits, as in the example below.

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)

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.

Checkpointing with Multiple DataLoader Workers

Python’s built-in torch.utils.data.DataLoader produces a safe checkpoint when num_workers = 0. With num_workers > 0, calling state_dict() often raises RuntimeError. In this case, use StreamingDataLoader for multi-worker prefetch and exact resumability. It carries a state snapshot alongside every internal batch and commits it to the parent StreamingDataset only when that batch is returned to the trainer. The trainer receives the same collated batch it would receive from a standard DataLoader.

from lancedb.streaming import StreamingDataset, StreamingDataLoader

dataset = StreamingDataset(
    table,
    num_splits=48,
    shuffle_seed=42,
    epoch=current_epoch,
    rank=rank,
    world_size=world_size,
)
dataloader = StreamingDataLoader(
    dataset,
    batch_size=batch_size,
    num_workers=4,
)

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",
        )

StreamingDataLoader accepts the same arguments as torch.utils.data.DataLoader, with a few restrictions:

dataset must be a StreamingDataset (subclasses that override __iter__ are not supported).
in_order = True is required so that consumer-committed checkpoints stay deterministic.
persistent_workers = True is not supported, because prefetched worker copies cannot be restored from parent-committed state.
drop_last = True is not supported, because incomplete tails discarded by worker replicas cannot be checkpointed topology-independently.

With more than one worker, call state_dict() at a complete logical step boundary where every split assigned to the rank has the same consumed-sample count. Calling it mid-step raises RuntimeError asking you to consume more batches first.

Resuming Across Different Topologies

When training across ranks, each rank owns its own subset of splits and only its own splits have exact progress. To resume on a different world_size, collect the state_dict() from every rank of the previous run and merge them with StreamingDataset.merge_state_dicts before calling load_state_dict() on the new run. The merge is topology-agnostic: pass the full list of per-rank states in, and hand the identical merged dict to every rank of the resumed job, regardless of whether the topology grew, shrank, or stayed the same.

# On each rank of the previous run:
state = dataset.state_dict()
# ...gather `state` from every rank into a single list `states` (for example
# via torch.distributed all-gather or by saving one file per rank).

merged = StreamingDataset.merge_state_dicts(states)

# On each rank of the resumed run (world_size may differ):
dataset = StreamingDataset(
    table,
    num_splits=merged["num_splits"],
    shuffle_seed=merged["shuffle_seed"],
    epoch=merged["epoch"],
    rank=rank,
    world_size=new_world_size,
)
dataset.load_state_dict(merged)

Permutations

In specific use cases, you may need the flexibility to shuffle, split, and select data without using the full StreamingDataset object. In these cases, use Permutation, the lower-level functionality on which StreamingDataset is built. A Permutation defines a customizable row ordering, and supports map-style access through __getitem__() and batched access through __getitems__().

Version Pinning

A Permutation is pinned to the version of its base table at the time it was constructed, and every read (including future reads from a StreamingDataset that wraps it, and every DataLoader worker after a fork) resolves against that pinned version. This makes iteration deterministic across compactions and worker forks: the permutation only addresses rows that existed when it was built, and rows appended to the base table afterwards are not visible through the existing permutation. To include newly appended rows, build a new Permutation (and any StreamingDataset that depends on it) against the updated table. Permutations written before this behavior shipped continue to read as they did before.

Get started

Model training

Guides

Feature Engineering (Geneva)

Support

Loading Data for Model Training

Basic Data Loading

Advanced Data Loading

Filtering

Shuffling Rows

One-Phase Shuffle

Two-Phase Shuffle

Parallelism & Elasticity

Fetching

Transforming

DataLoader Workers

Observability

Checkpointing

Checkpointing with Multiple DataLoader Workers

Resuming Across Different Topologies

Permutations

Version Pinning

​Basic Data Loading

​Advanced Data Loading

​Filtering

​Shuffling Rows

​One-Phase Shuffle

​Two-Phase Shuffle

​Parallelism & Elasticity

​Fetching

​Transforming

DataLoader Workers

​Observability

​Checkpointing

Checkpointing with Multiple DataLoader Workers

Resuming Across Different Topologies

​Permutations

Version Pinning

Basic Data Loading

Advanced Data Loading

Filtering

Shuffling Rows

One-Phase Shuffle

Two-Phase Shuffle

Parallelism & Elasticity

Fetching

Transforming

Observability

Checkpointing

Permutations