loaderx 0.9.4

A compact, high-performance persistent record store with zero-copy batch gathering and transparent per-record compression, designed for single-machine AI training and serving pipelines

Loaderx

A compact, high-performance persistent record store with zero-copy batch gathering and transparent per-record compression, designed for single-machine AI training and serving pipelines

pip install loaderx

Wheels are published for Linux (glibc ≥ 2.17 and musl, x86-64 and arm64), macOS (≥ 11.0, Intel and Apple Silicon) and Windows (x64 and arm64). The bindings use cffi in ABI mode, so nothing links against the CPython ABI and one wheel per platform serves every supported Python.

Design Philosophy

loaderx is built around several core principles:

1. A pragmatic approach that prioritizes minimal memory overhead and minimal dependencies.
2. A strong focus on single-machine training workflows.
3. We implement based on NumPy semantics, persisted through the Zrecord storage runtime.
4. An immortal (endless) step-based data loader, rather than the traditional epoch-based design—better aligned with modern ML training practices.
5. Dense and ragged are separate contracts; the loader is dense-only. zrecord still speaks one language underneath — records are packed byte ranges located by offsets, and every read/write is a batch operation where the batch size 1 is the degenerate case. But the Python layer does not force equal length to be a special case of variable length: :class:DenseDataset reads fixed-shape records with the shape inferred straight from the schema and no per-record metadata, while :class:RaggedDataset reads variable-length records into a list. Two kinds, one store engine — and a training pipeline only ever needs the dense one.

Quick Start

from loaderx.converter import from_numpy
from loaderx.dataset import DenseDataset
from loaderx.dataloader import DataLoader

from_numpy('train_data', np.load('data.npy', mmap_mode='r'))
from_numpy('train_label', np.load('label.npy', mmap_mode='r'))

loader = DataLoader({'data': DenseDataset('train_data'), 'label': DenseDataset('train_label')},
                    transform=lambda batch: batch)

for i, batch in enumerate(loader):
    if i >= 256:
        break

print(batch['data'].shape)
print(batch['label'].shape)

A batch is a dict {name: values}: each value is the stacked (batch_size, *item_shape) array for that stream. Every stream is gathered at the same indices, so record i lines up across them. The transform callback is the collate step — reshape, cast, stack — where values is the plain dense batch ready for the model.

Converting a NumPy tensor

import numpy as np
from loaderx.converter import from_numpy
from loaderx.dataset import DenseDataset

from_numpy('train_data', np.load('data.npy', mmap_mode='r'))
from_numpy('train_label', np.load('label.npy', mmap_mode='r'))

# many small, similar records compress far better against a trained dictionary:
from_numpy('train_data', images, codec='zstd_dict')

One record per slice along axis 0; a 1-D array (the usual shape of a label set) becomes a dataset of scalars. The conversion streams in bounded chunks, so an mmapped array is never fully materialized. Open the result with :class:DenseDataset:

ds = DenseDataset('train_data')

The store carries a single schema.msgpack recording the kind ("dense" or "ragged"), per-sample dtype and row shape — that is the only loaderx-level metadata. Everything beneath it is plain Zrecord, reachable through loaderx.zrecord.Zrecord when you want raw byte records (ragged samples, pre-encoded images) instead of tensors.

Records

One store engine, two kinds of view. from_numpy writes a dense store where every record is exactly one row of the recorded item_shape; reads are fixed-stride gathers and the batch shape follows from the schema, so no per-record metadata is touched:

import numpy as np
from loaderx.converter import from_numpy
from loaderx.dataset import DenseDataset

from_numpy('data', np.arange(64, dtype=np.float32).reshape(8, 2, 4))
ds = DenseDataset('data')
ds[[0, 5, 2]]                       # (3, 2, 4) — shape straight from the schema

from_iterator writes a ragged store of variable-length records. It is a separate contract: :class:RaggedDataset hands back a list of arrays, so a loader never has to carry row_splits around. Padding a list into a dense batch is a plain function, wherever you need it:

from loaderx.converter import from_iterator
from loaderx.dataset import RaggedDataset, pad

seqs = [np.arange(L, dtype=np.int32) for L in (3, 1, 4, 1, 5)]
from_iterator('tokens', seqs)
rs = RaggedDataset('tokens')

records = rs[[0, 2, 4]]              # list of ndarray — one per record
padded, lengths = pad(records)       # dense (B, max_len) + row counts

A DataLoader works with dense streams, so collation is just the transform — a batch dict in, a batch dict out:

def collate(batch):
    return {'input_ids': batch['tokens'], 'label': batch['label']}

loader = DataLoader({'tokens': dense_tokens, 'label': labelset}, batch_size=32,
                    transform=collate)

Codec notes

"zstd" compresses each record independently with plain zstd (level 3). Use it for general-purpose compression — it is fast and the default.

"zstd_dict" trains a shared dictionary on a sample of the data before writing any record, then compresses every record against it at level 19. The dictionary captures structure shared across records that per-record compression cannot see — a large win for many small, similar records (image tiles, token sequences). Dictionary size adapts to the training data automatically; pass dict_size to cap it.

"zstd_dict" records can only be read from a store that has the dictionary (dict.zr). The dictionary is loaded on open and shared, lock-free, across all reader threads.

Important: Do not use "zstd_dict" while data is still changing (through deletes or compaction below the Python layer). The dictionary captures a snapshot of the data; training it before the data settles wastes compression. Train the dictionary once preprocessing is complete and the content is final.

Integrating with JAX/Flax

For practical integration examples, please refer to the Data2Latent repository

Benchmarks

scripts/bench.py measures zrecord and loaderx against the alternatives in the index sampler, the record store, and the full data loader — each through the binding a client actually uses, so cffi, the GIL and the NumPy allocation are all inside the timings. Numbers are the median of 3 sampler runs, 5 store runs and 5 loader runs on the box below, warm page cache. Reproduce with python3 scripts/bench.py all — it prints this machine block first.

Machine — one box, inside a memory-limited container:

machine	value
CPU	AMD EPYC 7303 16-Core (Zen 3), 1 socket, 16 cores / 32 threads
frequency	1500–3437 MHz
caches	L1d 512 KiB, L1i 512 KiB, L2 8 MiB, L3 64 MiB
NUMA	4 nodes
memory	16 GiB (16 GiB cgroup limit)
HugePages	16384 × 2 MiB = 32 GiB configured, usable in this container
OS	Debian GNU/Linux forky/sid, kernel 6.12.95
python	3.14.6, numpy 2.5.1

The container sees all 32 threads and has no CPU quota. HugePages are configured by the host and usable here, but the benchmark runs on the ordinary page cache and does not depend on them.

1. Sampler — index generation on its own, IID (with replacement), 1M index space, against both NumPy random APIs.

sampler	batch	per batch	indices	vs default_rng
numpy randint	256	6.1 µs	42 M/s	1.02x
numpy default_rng	256	6.2 µs	41 M/s	1.00x
zsampler	256	3.1 µs	83 M/s	2.00x
numpy randint	1024	9.0 µs	114 M/s	0.91x
numpy default_rng	1024	8.2 µs	125 M/s	1.00x
zsampler	1024	4.4 µs	233 M/s	1.86x
numpy randint	8192	37.4 µs	219 M/s	0.83x
numpy default_rng	8192	31.1 µs	263 M/s	1.00x
zsampler	8192	17.2 µs	476 M/s	1.81x

2. Store — Zrecord against the alternatives: random batch gather, 20000 records, batch 256 × 200, structured (mildly compressible) data.

store	write	gather	per batch	on disk	ratio
zrecord-zstd	375 MiB/s	8938 MiB/s	0.34 ms	216 MiB	1.08x
zrecord-zstdict	10 MiB/s	3846 MiB/s	0.78 ms	182 MiB	1.29x
zrecord-raw	421 MiB/s	13565 MiB/s	0.22 ms	235 MiB	1.00x
npy-mmap	2465 MiB/s	12004 MiB/s	0.25 ms	234 MiB	1.00x
arrayrecord	218 MiB/s	641 MiB/s	4.68 ms	217 MiB	1.08x
hdf5	548 MiB/s	149 MiB/s	20.18 ms	236 MiB	0.99x
hdf5-gzip	46 MiB/s	87 MiB/s	34.33 ms	211 MiB	1.11x
blosc2	20 MiB/s	75 MiB/s	40.03 ms	216 MiB	1.09x
tensorstore	562 MiB/s	40 MiB/s	74.10 ms	214 MiB	1.09x

3. Loader — the full input pipeline end to end (sample, fetch, collate, hand over a batch), same workload, 4 workers.

loader	batches/s	per batch	samples/s	throughput
loaderx	3397.4	0.29 ms	869,734	10192 MiB/s
loaderx-raw	4104.9	0.24 ms	1,050,854	12315 MiB/s
grain	52.0	19.23 ms	13,312	156 MiB/s
torch	604.7	1.65 ms	154,803	1814 MiB/s

Conclusion — why the numbers look like this.

Every hot path is one native call. Zsampler draws a whole batch of indices, zrecord gathers a whole batch and decompresses it, in a single cffi call into Zig with the GIL released and the batch copied straight into its destination buffer. The contenders do the same work one record at a time from Python. That one fact runs through all three tables: the sampler's margin is widest at the smallest batches (a batch costs one call either way, so the per-index work is what divides them), and the store gather column is where one-record-per-call costs the most. The speedup is not bought with distribution shortcuts: the IID draw is unbiased like NumPy's (Lemire with rejection, so uniformity costs nothing over a real index space).

The layout matches what a training loader does. zrecord is built for random record access — the meta table is mmap'd, any record is found in O(1), a dense store gathers at a fixed stride. Array stores are built for contiguous scans, so a scattered batch — exactly what a loader reads — fights their layout. That is why uncompressed .npy, which does no per-record work at all, still loses to zrecord-raw on gather.

Compression is in the kernel, and there is one codec. zrecord-zstd is not "storage plus a codec": the layout, the multi-core decompress and the GIL-free copy are one path, so turning compression on costs part of a margin, not an order of magnitude — ~8.7 GiB/s here, roughly two orders of magnitude ahead of the other compressed stores (87–40 MiB/s) at the same ~1.1x ratio. The modest ratio is the data, not the format: these samples are barely compressible, and on a smooth image set plain zstd reaches 7.6x and the dictionary 16.5x.

The loader gap is architecture, not storage. loaderx uses threads and never ends an epoch, so a step pays no IPC and never waits on an epoch boundary; torch and Grain restart per epoch with worker processes. That is most of the 5–65x end to end. Storage also differs per loader — each reads from what it was built for — so the loader table is a different comparison than the store table, not a rerun of it.

What these numbers do not claim. Everything runs with a warm page cache: this measures the access path, not cold storage or disk. on disk is allocated blocks, so zrecord's sparse preallocation is not charged to it while TensorStore's sharding is; both are layout, not compression. Writing is a one-time cost zrecord does not chase. The box is shared — the µs-scale sampler timings and the loader batches/s fluctuate (with 16 physical cores the compressed loader trails the raw one by ~17%; on a 24-core box they measured within 2%), which is why the numbers above are medians of several runs.

Current Limitations

• Single-host only; multi-host training is not supported.
• A single sample must be at most 1 MiB, and a store holds at most 2^24 (16777216) records.
• Metadata is mapped and used in place, so a store carries the host's byte order and is not portable to a machine of the opposite endianness. Every published platform is little-endian, so this only matters if you build for one yourself.

Build

zig build                       # host shared objects, into loaderx/lib/
zig build test                  # Zrecord suite, in both Debug and ReleaseFast
python3 tests/test_loaderx.py   # Python layer, against whichever build is importable
python3 scripts/bench.py        # throughput, against other stores

The Zig side is tested for behaviour only; throughput is measured from Python, through the binding a client actually uses. scripts/bench.py runs in three layers (sampler, store, loader, or all) and skips contenders that are not installed. See Benchmarks.

Publishing

Zig cross-compiles every target from one machine, so releases need no CI matrix:

zig build dist                    # every platform, into zig-out/dist/<wheel tag>/
python3 scripts/build_wheels.py   # one wheel per platform, plus the sdist

The dist directories are named after their Python wheel platform tag, so the tag mapping lives in exactly one place (dist_targets in build.zig). glibc and macOS minimums are pinned in the target triple, which is what makes manylinux_2_17 and macosx_11_0 honest rather than aspirational. Each wheel is checked after packing: it must carry this platform's libraries and no others.

The sdist ships sources only. Installing from it runs zig build through setup.py, so it needs the Zig compiler; wheel users never hit that path.

---

Zsampler

Index Generator: a high-performance sampler implemented in Zig. Every mode is a pure function of (seed, step), so a run resumes exactly by seeking to a step — there is no epoch to track, in keeping with the endless step-based loader.

1. Sequential — traverse the index space in order through a fixed-size sliding window, treating the space as a circular queue so the tail never truncates.
2. IID — draw each index uniformly at random with replacement. Unbiased (Lemire with rejection), matching NumPy. Simplest, but coverage is uneven over any short run.
3. Cyclic — without replacement, round-robin. Each cycle traverses a fresh permutation of the whole index space, so within a cycle every record appears exactly once and no batch repeats an index — coverage is even by construction, which keeps how often each sample is seen uniform. The permutation is a stateless bijection (a small Feistel network over the index space, brought into range by cycle-walking), so a million-record shuffle materializes nothing the size of the dataset and reshuffling each cycle is free. Every batch is exactly batch_size: an endless step-based loader has no final partial batch to special-case, so when batch_size does not divide the length the cycle's remainder is dropped — a different remainder each cycle, since the permutation changes, so every record is still reached over time.

---

Zrecord

A record-based data runtime, focused on delivering extreme throughput and low latency.

1. Zrecord is an unordered physical store made of N records. Records are independent and carry no ordering, so every index and slice operation is equivalent to a gather.
2. Zrecord hands the client a dense index space: the live records are always exactly 0..N. Deletion preserves that by swapping the tail into the hole, which means an index is stable only until something is deleted. A multi-stream dataset is purely a client-side notion; the client keeps its own stream index table pointing at records.
3. Zrecord returns byte arrays. Type interpretation is the client's job.
4. The IO model (append | read | delete) is batch-oriented and shape-agnostic. A record is a packed byte range located by an offset; there is no notion of fixed vs variable length, and a single-record operation is just the batch_size == 1 case. Equal length and single records are special cases, not parallel code paths.
5. Zrecord owns its memory internally — allocation and release are explicit.
6. Each record is compressed and decompressed independently:

| tag   |  name     | algorithm                              |
|-------|-----------|----------------------------------------|
|   0   |  raw      | none                                   |
|   1   |  zstd     | zstd (plain, level 3)                  |
|   2   |  zstdict  | zstd with a trained dictionary (level 19) |

7. Compression is transparent to the client:
   • Compression runs concurrently across all cores. A record that fails to shrink is stored raw automatically, so reading never costs more than the data itself.
   • Decompression writes straight into the caller's destination memory (gather), with no intermediate buffer and no extra copy.
   • zstd is the one transparent codec — faster than Deflate at both ends and a better ratio, so there is no reason to carry a second. It is vendored C, built for every platform by Zig, so the one-wheel-per-platform story is unchanged.
   • zstd_dict additionally trains one dictionary on a sample of the data (stored as dict.zr) and compresses every record against it. Because each record is still independent, random access is unchanged — but the dictionary carries the structure shared across records, which per-record compression cannot see. On many small, similar records (image tiles, token sequences) this is a large win: a smooth-image set that plain zstd takes to 7.6x compresses 16.5x with a dictionary. The dictionary is loaded once on open and shared, lock-free, across all reader threads. The dictionary size adapts to the training data automatically.
   • Records stored with zstd_dict are explicitly tagged; they cannot be read by a store that lacks the dictionary. Plain zstd records never use a dictionary even if one is present.

Persistence format

Zrecord storage is metadata plus chunked data:

zrecord/
  ├── meta.zr
  ├── 0.zr
  └── 1.zr

Metadata (meta.zr)

meta.zr is used as the values it holds, not decoded into them. Both structs below are extern, 16 bytes, and naturally aligned; an mmap is page aligned, so the header is a *Header and the record table is a []RecordLoc. Looking up a record is table[idx], and there is no serializer anywhere in the code.

The fields are deliberately wider than the limits need — chunk_id is a u16 for 4096 chunks, the lengths are u32 for a 1 MiB cap. Bit-packing them would save 5 bytes per record and cost a hand-written codec on the hottest lookup path. The slack buys a machine-shaped array; it is worth it.

1. Header — 16 bytes of global state.

• magic (ZREC) and version are ordinary fields, so opening a directory that is not a Zrecord store fails immediately instead of decoding garbage.
• length is the total number of records | tail_chunk/tail_offset mark the last write position.
• There is no chunk count. Chunks are created in order and the frontier is always in the last one, so the store holds exactly chunks 0..=tail_chunk — a count would be a second copy of that fact to keep in sync.

const Header = extern struct {
    magic: [4]u8, version: u16,
    tail_chunk: u16, tail_offset: u32,
    length: u32,
};

2. Record table — 16 bytes per entry, indexed directly. Mapping an index to a physical address is what makes random access efficient.

• chunk_id is the containing chunk | offset is the position within it | phys_length/logic_length are the stored and original sizes | compress is the algorithm.

const Codec = enum(u8) { raw = 0, zstd = 1, zstdict = 2, _ };
const RecordLoc = extern struct {
    offset: u32, phys_length: u32, logic_length: u32,
    chunk_id: u16, compress: Codec, _reserved: u8,
};

There is no liveness flag. Every entry below length is live, because deletion swaps the tail into the hole rather than tombstoning.

3. Maximum length. Running out of space is effectively impossible, so the metadata is sized statically at the maximum.

• Current design: up to 2^12 (4096) chunks, 2^32 (4 GiB) per chunk, 2^20 (1 MiB) per record.
• Derivation: N = chunks × 2^32 (chunk size) / 2^20 (max record size) ⇒ a maximum length of 2^24 (16777216). Equality here is the point: the store can never run out of chunks before it runs out of record slots, which is what lets the table be sized statically. It is a compile-time assertion.
• meta.zr is a fixed 16 + 2^24 × 16 bytes (256 MiB). It is sparse — a store with one record allocates 4 KiB of it — and the mapping does not prefault.

Chunk data (x.zr)

Densely packed record data.

Executor

1. Write. Writes are append-only; everything else is offset redirection.

• Append: compress concurrently → assign physical locations serially → flush concurrently → commit metadata. Data is durable-ordered before the record table, and the table before length, so a crash truncates rather than corrupts. A record never straddles two chunks; one that would not fit rolls over to a fresh chunk.
• Delete: swap the last table entry into the deleted slot and drop the length by one. A batch is applied in descending index order, so each swap pulls from a slot no later target refers to. The index space stays dense — which is what the sampler needs, since it draws uniformly from 0..N and would otherwise keep hitting holes. The deleted record's bytes become garbage.

2. Read. Fill the destination memory concurrently, in place from the Python side (executed on async threads).

• Committed records are immutable, so the read path is lock free; length is published to readers through an atomic.

3. Concurrency model. Io.Group.async shards work by CPU count, and shards beyond the limit run inline on the calling thread.

• Shards receive contiguous blocks rather than a strided subset, keeping each worker's reads and writes sequential.
• Each shard creates one zstd context (ZSTD_CCtx to write, ZSTD_DCtx to read) and reuses it across every record it handles, rather than paying that setup per record. The dictionary (ZSTD_CDict/ZSTD_DDict) is immutable, so all shards share one, lock-free.
• Decompression writes straight into the caller's destination buffer, so there is no intermediate copy.

4. Garbage collection. Deletion leaves the record's bytes stranded, so space is reclaimed by an offline compact that rewrites the live records in place.

• Records are visited in physical order — which after swap-last deletes no longer matches table order — and repacked densely in that same order. A record therefore never moves to a higher address than it already had.
• Two things follow. Writing a record can never land on the bytes of a record not yet moved, so the rewrite is safe in place with no scratch copy of the store. And each table entry can be updated the instant its record lands, so the store is consistent at every point: an interrupted compaction leaves some records moved and the rest where they were, and re-running finishes the job.
• Records that are already in the right place are skipped, so a store with a small amount of garbage near the end is cheap to compact.
• Chunk files past the new frontier are closed and deleted.
• stats() reports live_bytes against chunk_bytes so callers can decide when it is worth running. Note the difference is an upper bound: a record never straddles a chunk boundary, so up to one record's worth per chunk is slack that compaction cannot remove.

5. File access.

• Metadata: a 256 MiB file created at init, accessed by mmap thereafter.
• Chunk data: 4 GiB files created at init, accessed concurrently through readPositionalAll/writePositionalAll.

Concurrency contract. gather and append are safe to call concurrently from many threads. delete and compact mutate the table in ways a lock-free reader would observe half-applied, so they require exclusive access to the store.

Distributed extension

Not implemented yet; architecture noted in advance.

1. A cluster layer is added; a node becomes a shard holding several chunks.
2. An indirection table maps a global path to a specific node, so the index path becomes idx → indirection[idx] = node, lidx → offset[lidx].