dataloader-rs

PyTorch's DataLoader model — dataset, sampler, collator, loader — as a typed Rust library, with a Python binding that runs the same pipeline on threads instead of worker processes.

• Familiar model, typed end to end: DataLoader<Dataset, Sampler, Collator> is generic and monomorphized; batches come out in epoch order as Result<Batch>.
• Threads, not processes: workers share the dataset through Arc, so samples are handed over by reference instead of being pickled across processes. Against torch.utils.data.DataLoader on 4 cores, the Python binding is 4.9–6.6× faster on the loader's own overhead, 3.2–3.9× on network-bound datasets, 2.5–3.5× on array samples with free-threaded Python, and starts an epoch in milliseconds (BENCH.md).
• Two explicit levels of parallelism: num_workers batches in parallel and items_in_parallel items of each batch — on one thread pool built with the loader and reused by every epoch.
• Bounded, ordered prefetch: at most num_workers + prefetch_depth batches are in flight or waiting, even behind one slow batch; they come out in epoch order.
• Failures surface: a dataset or collator error is returned for its batch, a panic is re-raised on the consumer, and Python exceptions keep their type.
• Python binding: PyDataloader mirrors the PyTorch arguments (batch_size, shuffle, sampler, collate_fn, drop_last, seed, __getitems__); wheels for CPython 3.13 to 3.15 and free-threaded 3.14t / 3.15t on Linux (x86_64, arm64) and macOS (arm64).
• No unsafe code: #![forbid(unsafe_code)].

When PyTorch is faster. The workers are threads in one process, not processes. On regular CPython (the GIL build, such as 3.13), only one thread runs Python code at a time: a __getitem__ that spends its time in pure Python (tokenising text, parsing records, Python loops) gets no faster with more workers, while PyTorch's worker processes each run their own interpreter and are 2.4–2.6× faster on such datasets. Threads win when the work releases the GIL (file and network I/O, NumPy, image decoding), and on free-threaded CPython 3.14t / 3.15t. Even there, threads that all read one shared Python dict or list contend on its reference count, and PyTorch is 2.5–2.9× faster: give each thread its own copy (threading.local). Measurements: BENCH.md.

Contents

• Quick start: Rust
• Quick start: Python
• Concepts
• Configuration
• Feature flags
• Performance
• Documentation
• Project layout
• Building and testing
• Minimum supported Rust version
• License

Quick start: Rust

[dependencies]
dataloader-rs = "0.1"

use dataloader_rs::{DataLoader, Dataset, RandomSampler, error::Result};

/// One sample per file; `get` is called from worker threads.
struct Files(Vec<std::path::PathBuf>);

impl Dataset for Files {
    type Item = Vec<u8>;

    fn get(&self, index: usize) -> Result<Vec<u8>> {
        Ok(std::fs::read(&self.0[index])?)
    }

    fn len(&self) -> usize {
        self.0.len()
    }
}

fn main() -> Result<()> {
    let paths = std::fs::read_dir("data")?
        .map(|entry| entry.map(|e| e.path()))
        .collect::<std::io::Result<Vec<_>>>()?;

    let mut loader = DataLoader::builder(Files(paths))
        .batch_size(32)
        .sampler(RandomSampler::new(42)) // reproducible shuffle
        .num_workers(4)                  // 4 batches in flight
        .prefetch_depth(8)               // run up to 8 batches ahead
        .drop_last(true)
        .build();

    for epoch in 0..3 {
        for batch in &mut loader {
            let batch: Vec<Vec<u8>> = batch?; // VecCollator: a Vec of samples
            println!("epoch {epoch}: {} samples", batch.len());
        }
    }
    Ok(())
}

cargo run --example basic runs a complete example.

Quick start: Python

pip install dataloader-rs

from dataloader_rs import PyDataloader, PyDataset


class Files(PyDataset):
    def __init__(self, paths: list[str]) -> None:
        super().__init__()
        self.paths = paths

    def __len__(self) -> int:
        return len(self.paths)

    def __getitem__(self, index: int) -> bytes:
        with open(self.paths[index], "rb") as f:
            return f.read()


loader = PyDataloader(Files(paths), batch_size=32, shuffle=True, seed=0, num_workers=4)
for batch in loader:  # a list of 32 `bytes`, or whatever collate_fn returns
    ...

The full Python guide — arguments, samplers, collate_fn, __getitems__, errors, threads and the GIL — is python/README.md.

Concepts

One epoch: the sampler produces indices, BatchSampler cuts them into batches of batch_size, each batch is fetched (Dataset::get_batch, or get per item in parallel) and collated, and the iterator yields the batches in order. How the threads, the bounded prefetch and the reorder buffer fit together is in ARCHITECTURE.md.

Configuration

Every setting is a DataLoaderBuilder method; the Python argument of the same name maps to it.

At most max(num_workers, 1) × items_in_parallel fetches run at once, on that many threads.

Rules of thumb: num_workers = 0 is the fastest path for cheap samples (no thread hand-off); use workers once a batch costs well more than its hand-off, about 0.5 µs in Rust (a few microseconds through Python); for blocking I/O, items_in_parallel so that num_workers × items_in_parallel requests are in flight, far more than cores; the default prefetch_depth 2 is enough when loading keeps up, beyond that it only holds memory (BENCH.md).

Feature flags

Performance

Measured on an Intel Core i9-13900H pinned to 4 cores, the size of the CI runners, on CPython 3.13, 3.14t and 3.15t; method, every table and the caveats are in BENCH.md. Continuous results per commit: dashboards for Rust and Python 3.14t.

Documentation

Project layout

Building and testing

Prerequisites: Rust ≥ 1.95, uv. The development environment uses free-threaded CPython 3.14t (.python-version); for a GIL build prefix commands with uv run --python 3.13.

# Rust
cargo build
cargo test --features ndarray,python          # add torch-rs to test tch (downloads libtorch)

# Python
uv sync --group dev                            # maturin, pytest, numpy, torch
uv run maturin develop --features python       # build the extension into the venv
uv run pytest                                  # tests/ (benchmarks are deselected)

# Lint, as CI does
CI=1 prek run --all-files                      # ruff, ty, cargo fmt, clippy -D warnings

Benchmarks: BENCH.md, section 9.

Commits follow Conventional Commits; releases, the changelog, the crate and the wheels are published by semantic-release from main.

Minimum supported Rust version

Rust 1.95 (rust-version in Cargo.toml), checked in CI by the Rust tests (MSRV) job.

License

MIT — see LICENSE.