I Built a Rust-Powered Dataset Cache Because Training Data Keeps Changing

You finish preprocessing your dataset. The expensive part is done.

Then the next experiment needs different class weights.

The next one excludes low-quality samples.

Then you need a tiny development subset. Then you want to prioritize examples the model gets wrong.

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.

One dataset, two connected parts

Imagine a cached record containing an image, several arrays, and some preprocessing results. 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.

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 you want only the training split, excluding samples below a quality threshold:

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

Adding loader processes creates another problem: memory.

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.

In one local benchmark, these were the median second-pass steady-state results across five trials:

The workload used 8,192 records across two caches, with 16 KiB payloads, on an Apple M5. Each delivered payload was hashed once. These were warm filesystem measurements without image decoding or model execution; startup and first-batch latency were excluded from steady throughput.

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.

If you keep revisiting the same expensive records while changing how training sees them, that is the workflow DatasetRT is built for.

Explore DatasetRT on GitHub or install the package:

uv add dataset-rt