DatasetRT: Stop Rebuilding Datasets to Rebalance Training

You’re running another training experiment on the same dataset. The images are
already processed, the features are computed, or the text is tokenized. Millions
of records are sitting on disk, ready to use.

This time, you want to rebalance a rare class. Or exclude low-quality samples.
Or run on 1% of the data while debugging a model change.

You also want your PyTorch loading workers to keep the GPU fed without filling
RAM with copies of dataset metadata.

This article is for ML engineers and researchers who reuse large preprocessed
datasets across training experiments and need to change how those samples are
selected. The examples use Python, Polars, and PyTorch; the cache itself is
independent of the training framework.

The problem becomes familiar once your dataset contains more than just an input
and a target. You might have an image with annotations and preprocessing
outputs, or a tokenized document with several arrays. Loading that whole record
just to check its label or quality score is work your experiment shouldn’t need.

The payloads haven’t changed. Your decisions about which records to train on
have.

That distinction is the starting point for DatasetRT: a Rust-powered dataset cache for Python that keeps serialized samples and their sampling metadata separate.

It lets you change the training population without rebuilding the payload cache.
The useful part is the combination: Python records you already know how to
serialize, separately queryable metadata, native sample reads, and shared mapped
buffers for multiprocessing loaders.

One dataset, two connected parts

Imagine a training record containing an image, its class label, and
arrays produced during preprocessing. Your Python source packs that structure
with pickle or dill.

That record might be large. Its label, split, quality score, and sampling
weight are usually much smaller. Keeping those attributes queryable lets you
ask for more examples of an underrepresented class without unpacking every
image first. The same pattern applies to source categories in tokenized corpora.

DatasetRT keeps both:

A (cache_id, sample_id) identity connects a metadata row to its stored payload.

Polars operates on the metadata. Your pickled records stay in the payload store.

That makes a surprisingly useful question cheap to answer:

“How many examples of each class do I have?”

metadata = dataset.get_metadata()

print(
    metadata.group_by("label")
    .agg(pl.len().alias("samples"))
)

You don’t need to unpickle every training record to count its label.

Changing an experiment shouldn’t mean copying a dataset

Suppose your next training run should exclude low-quality records while keeping
the original cache available for other experiments. With label, split, and
quality stored as metadata columns, the selection is:

selected = (
    dataset.get_metadata()
    .filter(
        (pl.col("split") == "train")
        & (pl.col("quality") >= 0.8)
    )
)

dataset.update_metadata(selected)

The cache files stay unchanged. Future dataset iterators use the selected population.

For a quick development run:

dataset.update_metadata(selected.head(100))

You now have a smaller active dataset backed by the same immutable payloads.

That matters when preprocessing is expensive, storage is large, or experiments keep changing the selection rules. You can express those rules directly instead of maintaining another collection of copied records.

Analyze the metadata. Change the population. Reuse the bytes.

Class balancing becomes a table operation

Consider a dataset with 9,000 examples of one class and 1,000 of another.

With uniform per-record sampling, their expected shares are 90% and 10%.

Assign each record the inverse of its class count:

balanced = selected.with_columns(
    (1.0 / pl.len().over("label")).alias("weight")
)

dataset.update_metadata(balanced)

Each class now has total weight 1, giving them equal expected sampling probability.

Rust validates the accepted identities and weights before applying the update. Payloads remain untouched.

The same mechanism can support hard-example sampling: compute scores in your training workflow, turn them into positive sampling weights, and update the population before creating the next iterator or loader.

One detail deserves explicit attention: shuffled sampling uses replacement. A sample can appear several times while another is absent. Equal expected class probabilities do not guarantee exactly equal class counts in every epoch.

Multiprocessing shouldn’t require another full metadata copy per worker

Once your Python decode or augmentation step becomes a bottleneck, adding
PyTorch loading processes is a reasonable experiment. It also creates another
budget to watch: RAM.

If each worker reconstructs a private metadata table and decoded index, worker count becomes a multiplier on those allocations.

DatasetRT uses read-only, file-backed mmap for the loader’s immutable Arrow metadata snapshot. Workers map the same snapshot, letting the operating system share its backing pages on the same host.

Cache indexes are also mapped from their original files.

This works with spawn: workers open their mappings independently, so sharing does not depend on inheriting a live parent runtime.

There are still private allocations. Weights, sampling state, decoded records, and prefetched batches consume memory per consumer. Separate metadata exports also have separate backing files.

The saving is specific and useful: workers can reference shared metadata and index pages instead of each retaining a complete private copy of those buffers.

Rust handles the sample reads

After selecting a sample, DatasetRT’s native runtime resolves its index entry and reads the payload using native workers and bounded prefetching.

Python receives the serialized bytes and performs application-specific unpacking:

record = pickle.loads(sample.data)

Your source chooses the serialization format. DatasetRT doesn’t need to understand your nested dictionary, tensor layout, or domain object to store and retrieve its bytes.

Keeping that boundary small lets the native reader focus on indexed access, ordering, and bounded work.

Payload size matters when comparing readers: 10,000 small records and 10,000
large images represent very different amounts of work. I measured random
weighted reads plus unpickling, and report the reconstructed image bytes per
second so the amount of work is explicit.

Each cached payload is a pickled dictionary containing a 256 × 256 RGB NumPy
array (uint8), a label, and sample identity. Each reconstructed image occupies
192 KiB. There are 4096 image records across two caches: 768 MiB of image data,
plus pickle headers and the other dictionary fields.

This benchmark uses shuffle=True: weighted multinomial sampling with
replacement, not sequential reading. Four equally populated classes have
per-record weights 1:2:4:8, giving expected draw shares of approximately
6.7%, 13.3%, 26.7%, and 53.3%. Each pass draws 4096 samples; repeats count as
additional reads and reconstructions.

Every delivered record goes through pickle.loads() inside the timed consumer
or loading worker. The benchmark validates the reconstructed image's shape,
dtype, sample identity, and identifying pixels. It does not cache unpacked
records between draws.

These are medians of five trials per configuration, with case order rotated:

Measurements used an Apple M5 with 24 GiB RAM, macOS 27.0.1, Python 3.11.15,
PyTorch 2.14.1, and DatasetRT 0.4.6. Each consumer has two native reader threads;
batch size and native prefetch capacity are both 64. Loading workers use spawn,
replicated weighted populations, two prefetched batches per worker, and automatic
persistence between passes. Fresh readers with the same seed and topology
reproduced their sample traces; successive passes continued the sampling stream.

The timer covers iterator creation through exhaustion: sample reading,
unpickling, validation, batching, and consumption. The first-pass figures include
loading-worker startup; the second pass reuses persistent workers. Dataset and
loader construction, imports, fixture creation, and metadata-weight setup are
outside these timers. MiB/s is reconstructed image bytes divided by total pass
time, including first-batch latency; 1 MiB is 1,048,576 bytes. The raw report also
retains serialized-payload byte counts for the steady interval after the first
batch.

The images are synthetic RGB arrays, not photographs or encoded JPEG/PNG
files. Unpickling reconstructs the dictionary and NumPy array; this benchmark
does not include codec decoding, augmentation, GPU transfer, or model execution.
Loading workers send compact validation results back to the parent rather than
full image arrays, so the table does not measure full-image IPC throughput.

The dataset fits in RAM and the filesystem cache is warm. These are logical
read-and-reconstruction rates, including replacement draws—not physical disk
bandwidth. A benchmark of compressed images or a complete training pipeline
would measure a different workload.

The one-worker result is worth keeping in the table. It was slower than the zero-worker baseline. More processes also increased startup time and memory.

Fast native reads give you a useful foundation. Your actual decoding workload determines whether additional Python workers help.

The workflow I wanted

The combination is what makes DatasetRT useful:

1. Preprocess and serialize records once.
2. Keep their sampling attributes in a separate metadata table.
3. Analyze, filter, and reweight that table as experiments change.
4. Read selected payloads through the Rust runtime.
5. Share mapped metadata and indexes across loader workers where possible.

There is one deliberate responsibility left with the application: cache reuse follows the source name. If preprocessing or source contents change, use a new source name or cache root. DatasetRT does not automatically fingerprint your preprocessing code.

DatasetRT is worth evaluating when preprocessing is expensive, cached records
are large, and your experiment loop repeatedly changes subsets or sampling
weights. It is especially relevant when multiple loading processes make private
metadata and index copies costly.

If your data is small enough to keep in memory and a simple dataset plus sampler
already does the job, the extra cache layer may buy you little. The intended
user is the person who keeps revisiting the same expensive records and wants the
next experiment to start with a metadata change rather than another data-preparation
job.

Explore DatasetRT on GitHub or install the package:

uv add dataset-rt