lit-deltalake

Bring Delta Lake to Pytorch Lightning

Features

• Stream immutable Delta table snapshots into PyTorch IterableDataset instances.
• Shard Delta-rs source fragments across initialized DDP ranks and DataLoader workers.
• Use Delta-rs, an existing PySpark session, or an existing Sail Spark Connect session.
• Select, filter, transform, decode, and rename columns before training.
• Pin Delta table versions for reproducible training data.

Installation

Install the default Delta-rs and PyArrow reader:

uv add lit-deltalake

Add an optional reader when needed:

uv add "lit-deltalake[spark]"  # existing PySpark session
uv add "lit-deltalake[sail]"   # existing Sail Spark Connect session

With pip:

python -m pip install lit-deltalake

Usage

from lit_deltalake.dataloaders import create_pytorch_dataloader
from lit_deltalake.readers.types import Filter
import lightning
import torch
from torch import nn
from torch.utils.data import DataLoader


loader = create_pytorch_dataloader(
    "/path/to/training-table",
    columns=("feature_a", "feature_b", "label"),
    filters=(Filter("split", lambda column, value: column == value, "train"),),
    batch_size=128,
    num_workers=4,
    multiprocessing_context="spawn",
)


class EventsDataModule(lightning.LightningDataModule):
    def __init__(self) -> None:
        super().__init__()
        self.loader = loader

    def train_dataloader(self) -> DataLoader:
        return self.loader


class Classifier(lightning.LightningModule):
    def __init__(self, learning_rate: float = 1e-3) -> None:
        super().__init__()
        self.model = nn.Linear(2, 2)
        self.learning_rate = learning_rate
        self.loss_function = nn.CrossEntropyLoss()

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.model(features)

    def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:
        del batch_idx
        features = torch.stack((batch["feature_a"], batch["feature_b"]), dim=1).float()
        labels = batch["label"]
        loss = self.loss_function(self(features), labels)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self) -> torch.optim.Optimizer:
        return torch.optim.AdamW(self.parameters(), lr=self.learning_rate)


trainer = lightning.Trainer(max_epochs=10)
trainer.fit(Classifier(), datamodule=EventsDataModule())

Examples

See examples and notebooks for more usage examples.

Docs

uv run mkdocs build -f ./mkdocs.yml -d ./_build/

Update template

copier update --trust -A --vcs-ref=HEAD

Credits

• Thanks for deltatorch for inspiration and code

This project was generated with