loaderx 0.5.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.

Quick Start

from loaderx import Dataset, DataLoader

dataset = Dataset('train_data.zr')
labelset = Dataset('train_label.zr')

loader = DataLoader(dataset, labelset)

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

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

Converting a NumPy tensor

import numpy as np
from loaderx import Dataset

Dataset.write('train_data.zr', np.load('data.npy', mmap_mode='r'))
Dataset.write('train_label.zr', np.load('label.npy', mmap_mode='r'))

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.

Dataset keeps a single schema.json inside the store recording the per-sample dtype and shape — that is the only loaderx-level metadata. Everything beneath it is plain Zrecord, reachable through loaderx.Zrecord when you want raw byte records (ragged samples, pre-encoded images) instead of tensors.

Removing samples

ds = Dataset('train_data.zr')
ds.delete([7, 12, 40])   # len shrinks; survivors stay addressable as 0..len
ds.compact()             # reclaim the freed space on disk
print(ds.stats())        # {'records': ..., 'live_bytes': ..., 'reclaimable': ...}

Deletion keeps the index space dense, which is what the sampler needs — but it does so by moving the tail into the hole, so surviving samples may change index. Do not hold indices across a delete. Both calls require exclusive access; do not run them while a DataLoader is iterating the dataset.

Integrating with JAX/Flax

For practical integration examples, please refer to the Data2Latent repository

Benchmarks

scripts/bench.py measures loaderx against the alternatives in three layers — 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. The numbers below are one run on a 24-core machine, Python 3.14, warm page cache. Reproduce with python3 scripts/bench.py all.

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	5.9 µs	44 M/s	0.64x
numpy default_rng	256	3.7 µs	69 M/s	1.00x
zsampler	256	1.9 µs	138 M/s	2.01x
numpy randint	1024	7.2 µs	142 M/s	0.83x
numpy default_rng	1024	6.0 µs	171 M/s	1.00x
zsampler	1024	3.0 µs	340 M/s	1.99x
numpy randint	8192	25.6 µs	320 M/s	0.65x
numpy default_rng	8192	16.6 µs	493 M/s	1.00x
zsampler	8192	7.8 µs	1048 M/s	2.13x

Zsampler draws indices about 2x faster than NumPy's faster API — more against the legacy randint most code still calls — and the lead is widest at the small batches a loader actually draws, where a batch is one cffi call either way and NumPy's vectorised fill has nothing to amortise. The one concession: zsampler's IID draw is modulo-biased where NumPy's is not, ~1e-13 over a u64 draw into a million-record space, which is below anything training will notice.

2. Store — random batch gather, 20000 records × 12 KiB uint8(3,64,64), batch 256 × 200, structured (mildly compressible) data.

store	write	gather	per batch	on disk	ratio
loaderx-flate	362 MiB/s	1671 MiB/s	1.80 ms	209 MiB	1.12x
loaderx-raw	784 MiB/s	21997 MiB/s	0.14 ms	235 MiB	1.00x
npy-mmap	3064 MiB/s	13692 MiB/s	0.22 ms	234 MiB	1.00x
arrayrecord	374 MiB/s	847 MiB/s	3.54 ms	217 MiB	1.08x
hdf5	1327 MiB/s	511 MiB/s	5.87 ms	236 MiB	0.99x
hdf5-gzip	83 MiB/s	206 MiB/s	14.57 ms	211 MiB	1.11x
lmdb	677 MiB/s	5337 MiB/s	0.56 ms	313 MiB	0.75x
lmdb-zlib	71 MiB/s	336 MiB/s	8.93 ms	235 MiB	1.00x
blosc2	32 MiB/s	226 MiB/s	13.27 ms	216 MiB	1.09x
zarr	83 MiB/s	28 MiB/s	107.36 ms	217 MiB	1.08x
tensorstore	531 MiB/s	98 MiB/s	30.69 ms	214 MiB	1.09x

The gather column is what the store is built for. loaderx-raw returns a scattered batch faster than everything here, a memory-mapped .npy included, because it copies straight into the destination across all cores with the GIL released; loaderx-flate keeps most of that speed while decompressing, and still gathers several times faster than any other compressed store (hdf5-gzip, lmdb-zlib, blosc2, zarr, tensorstore). The array stores lag on gather because a scattered single-record read fights a layout meant for contiguous slabs — which is exactly the access pattern a training loader has. Where loaderx does not lead is write: it pays compression and its own layout up front, a one-time cost it does not chase. Compression is modest (1.12x) only because this data is barely compressible; the row's point is that turning it on costs almost nothing on read.

The reads are not all single-threaded: loaderx and blosc2 decompress across cores, the rest one record at a time. Zarr and TensorStore are sharded so on disk reflects compression rather than per-file block rounding, and TensorStore is given gzip (its zarr3 default is uncompressed). on disk is allocated blocks, which is why LMDB's preallocated map lands below 1.00x.

3. Loader — the full input pipeline end to end (sample, fetch, collate, hand over a batch), same workload, 4 workers. Each loader reads from what it is built for: loaderx from Zrecord, Grain from ArrayRecord, torch from a memory-mapped .npy, jax-dataloader from an in-memory array.

loader	batches/s	per batch	samples/s	throughput
loaderx	724.2	1.38 ms	185,406	2173 MiB/s
loaderx-raw	5271.0	0.19 ms	1,349,378	15813 MiB/s
grain	67.3	14.85 ms	17,240	202 MiB/s
torch	832.2	1.20 ms	213,038	2497 MiB/s
jax-dataloader	3284.6	0.30 ms	840,864	9854 MiB/s

End to end, loaderx-raw serves batches several times faster than torch or Grain and ahead of jax-dataloader — and it does so from disk-backed storage, where jax-dataloader's number is an in-memory ceiling with no storage path at all. The default loaderx trades some of that for on-disk compression and still lands in torch's range while reading compressed records. loaderx and jax use threads; torch and Grain use worker processes that escape the GIL but pay IPC per batch. loaderx never ends an epoch, so no step waits on a boundary; the others restart per epoch, a cost that is in their numbers as it would be in training.

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

1. Sequential generation: indices are produced by traversing the index space in order.
   • Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
2. Random generation: indices are sampled randomly from the index space.
   • Global random: a set of samples is drawn randomly from the entire index space.

---

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. For performance, the IO model (append | read | delete) is batch-oriented; a single-record operation is just the batch_size == 1 case.
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   |  flate  | Deflate   |

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.

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, flate = 1, _ };
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 reuses one compressor scratch (flate.Compress is ~230 KiB — far too large to rebuild per record, let alone to place on a task stack).
• Decompression runs in direct mode: matches are resolved against the bytes already written into the destination, so there is no window buffer.

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 192 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].