Loading Data for Model Training

import lancedb

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

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

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)

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
)

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)

ds = StreamingDataset(
    table,
    shuffle_seed=42,
    transform=normalize,
    transform_parallelism=4,
)

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

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

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

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

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)

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)