Loaderx

Zrecord is a rebuildable, typed, ordered record sequence built from authoritative source data and scripts. A creator consumes records in append order, close publishes one immutable container, and readers support both sequential slicing and indexed gather without changing record identity. To change content or order, rebuild it at a new path.

Zrecord is the typed on-disk container; Loaderx is the sampler and prefetch loader that consumes Zrecord streams. They currently ship together while both layers mature, but their public responsibilities remain separate.

pip install loaderx

Wheels are published for CPython 3.10+ on glibc Linux (x86-64 and ARM64), Apple Silicon macOS, and Windows AMD64. Free-threaded CPython 3.14 is also supported.

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. NumPy-native typed records with explicit schemas.
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. Dense batches are arrays; ragged batches are lists of arrays. Loaderx never pads either representation.
6. Logical IDs are stable sequence positions. Append order defines 0..N-1; published containers are immutable.

Usage

Quick Start

import numpy as np

from loaderx.zrecord import Dense
from loaderx.dataloader import DataLoader

data = np.load('data.npy', mmap_mode='r')
label = np.load('label.npy', mmap_mode='r')
with Dense.create('train_data', data.dtype, data.shape[1:], codec='zstd') as ds:
    ds.append(data)
with Dense.create('train_label', label.dtype, label.shape[1:], codec='zstd') as ds:
    ds.append(label)

data_store = Dense.open('train_data')
label_store = Dense.open('train_label')
loader = DataLoader({'data': data_store, 'label': label_store},
                    transform=lambda batch: batch)
for i, batch in enumerate(loader):
    if i >= 256:
        break

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

loader.close()
data_store.close()
label_store.close()

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.

Records

Dense stores one fixed-shape array per record and returns a stacked array for multi-record reads:

import numpy as np
from loaderx.zrecord import Dense

data = np.arange(64, dtype=np.float32).reshape(8, 2, 4)
with Dense.create('data', data.dtype, data.shape[1:], codec='zstd') as ds:
    ds.append(data)
ds = Dense.open('data')
ds[0, 5, 2]                          # (3, 2, 4) — shape from the persisted schema
ds.close()

A store is an ordered sequence of records, not an ndarray, so ds[0, 5, 2] selects sequence positions 0, 5 and 2 — never ds[0][5][2]. A scalar selects one record; indices must be in 0..len(ds)-1. The ragged example below reads the same way.

Ragged stores variable-shape arrays with one shared dtype and ndim. Reads return a list and restore every record's exact shape. Every axis may vary; scalar records use Dense with item_shape=(). Padding remains an application policy:

from loaderx.zrecord import Ragged

seqs = [np.arange(L, dtype=np.int32) for L in (3, 1, 4, 1, 5)]
with Ragged.create('tokens', np.int32, ndim=1, codec='zstd') as rs:
    rs.append(seqs)                       # dtype/rank fixed; lengths remain per-record
rs = Ragged.open('tokens')

records = rs[0, 2, 4]                # list of ndarray — one per record, exact shapes
rs.close()
lengths = np.array([len(r) for r in records])
padded = np.zeros((len(records), lengths.max()), dtype=records[0].dtype)
for i, r in enumerate(records):
    padded[i, :len(r)] = r             # (B, max_len) — your policy, your loop

Logical ID is the stable sequence position. Each append preserves input order, and successive calls extend the sequence. After close, the sequence is immutable: there is no delete, update, compaction, or reopen-for-append.

Large sequences can be consumed in bounded ordered batches with ordinary slices:

with Dense.open("events") as events:
    for start in range(0, len(events), 1024):
        batch = events[start:start + 1024]
        consume(batch)

Creating containers

Dense.create and Ragged.create return append-only builders. The required codec is "raw", "zstd", or "zstd_dict". Schema, shape, and dtype are explicit and are never inferred from the input. Dense accepts the complete NumPy array or memory map; Ragged accepts the complete finite iterable of NumPy arrays.

A creator consumes a finite build stream and publishes an immutable sequence; readers do not tail an active writer. Append is synchronous and returns after the complete input has been persisted.

from loaderx.zrecord import Dense, Ragged

with Dense.create('mnist/x', dtype=np.uint8, item_shape=(28, 28),
                  codec='zstd', data_shards=4) as ds:
    ds.append(images)

with Ragged.create('tokens', dtype=np.int32, ndim=1, codec='zstd') as tok:
    tok.append(sequences)

with Dense.open('mnist/x') as ds:
    first_four = ds[:4]               # opened containers are read-only

data_shards controls write parallelism, accepts 1..255, and defaults to four. Readers discover it automatically.

Append errors are reported by the current call. A writer cannot be read and a reader cannot be appended to. Writer close() publishes the container; reader close() releases it. Dense, Ragged, and DataLoader support with.

Structured, subarray, object, metadata-bearing, and zero-itemsize dtypes are not supported. Append does not coerce Python containers or bytes; convert them with np.asarray or np.frombuffer first. Raw bytes and encoded files can be stored as np.uint8 Ragged records.

Codec notes

"zstd" compresses each record independently with plain zstd (level 3). Use it for general-purpose compression; it is fast and must be selected explicitly.

"zstd_dict" uses a shared dictionary trained from representative settled data. It is useful for large corpora of small, similar records such as token sequences and image tiles. Plain "zstd" or "raw" is usually better for large records and small corpora.

Train the dictionary manually with train_dict on settled data, then pass its bytes to dict_bytes. The source may be a NumPy array, an iterable of records, or an existing typed container. zstd_dict never trains automatically:

from loaderx.utils import train_dict
from loaderx.zrecord import Dense, Ragged

# train once on the settled data — a standalone, reusable artifact
d = train_dict(settled_array)

# then any new store can install it and append explicitly
with Ragged.create('tokens', np.int32, ndim=1,
                   codec='zstd_dict', dict_bytes=d) as ds:
    ds.append(token_generator)
with Dense.create('data', data.dtype, data.shape[1:],
                       codec='zstd_dict', dict_bytes=d) as ds:
    ds.append(data)

Changing an existing store's codec requires writing a new store:

from loaderx.utils import train_dict
from loaderx.zrecord import Dense

with Dense.open("src") as s, \
     Dense.create("dst", dtype=s.dtype, item_shape=s.item_shape,
                       codec="zstd_dict", dict_bytes=train_dict(s)) as d:
    d.append(s[:])                             # ndarray -> dense append

For Ragged, s[:] already returns the list accepted by append; create the destination with s.dtype and s.ndim. The destination path must be new.

Important: Train the dictionary from settled authoritative input before building the container. Training it before preprocessing is complete wastes compression and does not describe the final records.

Offline Hugging Face conversion

The optional converter uses Hugging Face Datasets for remote discovery, download, caching, revision handling and source-format decoding. It then writes explicit typed Zrecord streams, so training needs neither datasets nor Arrow:

pip install 'loaderx[converter]'

from loaderx.converter import convert, huggingface

dataset = huggingface("ylecun/mnist", revision="main", token=None)
convert(
    dataset,
    "mnist",
    codec="zstd",
)

Mainland China or a private Hub can select an endpoint directly without setting process environment variables:

dataset = huggingface(
    "ylecun/mnist",
    endpoint="https://hf-mirror.com",
)

huggingface resolves the requested revision to an immutable snapshot and returns its DatasetDict. Pass config=... when a repository has no default configuration, and token="hf_..." for private or gated repositories.

convert preserves numeric dtypes and record shapes. Fixed-shape columns become Dense stores and variable-shape columns become Ragged stores. Unsupported features fail explicitly; conversion does not filter, pad, tokenize, reshape, or apply a transform.

The result groups aligned streams under one published root:

mnist/
  train/
    image/
    label/
  test/
    image/
    label/

The output is published only after all columns have been written and their record counts agree. Each column remains an ordinary Zrecord store.

DataLoader and multi-stream stores

A zrecord store is one stream; a training sample is usually several named streams (skeleton + label + id, tokens + label, ...). Composition is a plain Python dict passed to DataLoader. There is no persistent wrapper, manifest, directory convention or bundle mutation API. DataLoader verifies that all streams have the same length, then gathers every stream with the same indices.

from loaderx.zrecord import Dense, Ragged
from loaderx.dataloader import DataLoader

root = "xsub/train"
with Dense.create(root + "/joint", joint.dtype, joint.shape[1:], codec="zstd") as s:
    s.append(joint)
with Dense.create(root + "/label", label.dtype, label.shape[1:], codec="zstd") as s:
    s.append(label)
with Ragged.create(root + "/token", np.int32, ndim=1, codec="zstd") as s:
    s.append(seqs)

streams = {
    "joint": Dense.open(root + "/joint"),
    "label": Dense.open(root + "/label"),
    "token": Ragged.open(root + "/token"),
}
streams["joint"][0, 5, 2]            # each stream keeps its own index API
loader = DataLoader(streams, batch_size=256)
batch = next(loader)                  # {name: values}, index-aligned

loader.close()
for stream in streams.values():
    stream.close()

Equal-length and variable-length are both just stores — Dense (one fixed-shape record per sample) and Ragged (variable row count per record). The loader does not interpret either: it fetches by index and packs a dict, so a dense stream's batch value is the stacked (B, *item_shape) array and a ragged stream's is a list of per-record arrays. No padding is imposed — densify to a fixed shape however the model needs (a plain numpy loop), or reshape/stack in the transform collate.

Collation is the transform: a batch dict in, an arbitrary result out.

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

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

The transform runs once per gathered batch. Calls may run concurrently and complete out of order, so the callback must be thread-safe. Its return value is passed to the consumer unchanged, and exceptions propagate to the consumer.

Numba can optionally accelerate a CPU-heavy Dense transform while releasing the GIL. Compile it before timing, then call it from the ordinary transform:

import numba
import numpy as np

@numba.njit(nogil=True, parallel=False)
def normalize_u8(x):
    out = np.empty(x.shape, dtype=np.float32)
    for i in range(x.size):
        out.flat[i] = x.flat[i] / 255.0
    return out

normalize_u8(np.zeros((1, 3, 224, 224), dtype=np.uint8))  # compile warmup

def transform(batch):
    batch["image"] = normalize_u8(batch["image"])
    return batch

Numba is optional. Compile it before measuring loader throughput.

CPU → GPU transfer

Loaderx produces CPU batches. Device transfer belongs in transform, where the training framework is already available:

import torch

device = "cuda:0"
def to_device(batch):
    return {k: torch.from_numpy(v).to(device, non_blocking=True)
            for k, v in batch.items()}

streams = {
    "joint": Dense.open(root + "/joint"),
    "label": Dense.open(root + "/label"),
}
loader = DataLoader(streams, transform=to_device)
for batch in loader:
    model(batch)                        # already on device

Call loader.close() when the training loop exits, then close its caller-owned streams. Do not close either from transform.

A non_blocking=True copy is asynchronous only from pinned memory. Loaderx does not pin memory; use the framework's API in the transform when needed:

def to_device(batch):
    return {k: torch.from_numpy(v).pin_memory().to(device, non_blocking=True)
            for k, v in batch.items()}

For JAX, use jax.device_put:

import jax

def to_device(batch):
    return {k: jax.device_put(v) for k, v in batch.items()}

For a practical JAX/Flax integration, see the MNIST layer-representation example.

Zsampler

Zsampler is the batch index generator used by DataLoader. A run resumes exactly with seek(step), without tracking an epoch.

from loaderx.zsampler import Sampler

sampler = Sampler(1_000_000, 256, Sampler.Mode.IID, seed=42)
indices = sampler.next()       # borrowed until this sampler's next draw
saved = indices.copy()         # retain across draws only when needed

next() and iteration return a view of one reusable uint64 batch buffer. The contents stay unchanged until the next explicit draw from that Sampler; copy only plans that must outlive it. DataLoader consumes each view synchronously before drawing again.

1. Sequential traverses records in order and wraps at the end.
2. IID samples uniformly with replacement.
3. Cyclic traverses a fresh permutation without replacement. It does not allocate a dataset-sized permutation. Full batches are returned; a different remainder is omitted on each cycle when the length is not divisible by the batch size.

Benchmarks

Dense and Ragged are measured separately because they expose different contracts, but every store table uses the same columns. scripts/bench_dense.py measures fixed-shape random gather, scripts/bench_ragged.py measures variable-shape records, and scripts/bench.py covers machine, sampler, and the end-to-end loader comparison. Every path runs through the public Python binding, including output allocation and Ragged reconstruction.

Methodology

The tables are one complete pass on a warm page cache, refreshed on 2.5.1 with the default data_shards=4. They are not three-run medians. Every Store backend receives the same prepared real records, then a full sweep and independent IID batches verify exact dtype, shape, order, and values before timing. Loader backends use the same source and seed but their own shipped samplers, so exact permutations may differ.

logical write and logical gather report uncompressed NumPy payload bytes per elapsed second, not physical storage bandwidth. Write timing includes public-API adaptation and logical finalization, but excludes source preparation, dictionary training, writer setup, cleanup, and stable-media sync. Zrecord writers receive the complete prepared source in one append and bound execution internally. Gather timing includes allocation, reads, decompression, and reconstruction; it excludes open, close, sampler time, and destruction after return. Fresh uniform IID gathers accumulate at least two timed seconds. krecords/s is record throughput, p95 is one random batch's 95th-percentile latency, and disk is allocated blocks. Compare results within a workload table, not across different record geometries. Every listed backend and codec is required for the published matrix.

The vision source is the Oxford-IIIT Pet train split prepared by scripts/prepare_vision.py. Dense uses RGB photographs resized on the short side to 256 and center-cropped to (3, 224, 224); Ragged uses the same ordered source at native RGB resolution. Both are mmap-loaded uint8 CHW records.

Machine — one local workstation (AMD Ryzen AI 9 HX PRO 370, 12 cores / 24 threads):

The process sees all 24 threads, is not memory-limited by cgroup, and uses the ordinary page cache. The 16 GiB shared-memory mount accommodates the Torch run.

Large Vision Records

Fixed-Shape Dense

Zrecord against array-store alternatives: random batch gather over the first 2,500 Oxford-IIIT Pet train images, batch 256. Every prepared image is uint8[3,224,224]; metadata.json records the source revision, transform, decoder versions and logical SHA-256.

Fixed-resolution vision records — 147 KiB per record, 36.8 MiB per batch:

At 147 KiB per record, Zrecord-raw reaches 15.1 GiB/s and is 3.1x npy-mmap-raw. Decoded photographs have little remaining redundancy: plain zstd reduces them only 1.15x, while dictionary mode falls to 1.12x. This is why large real images should use plain zstd or raw. LMDB and Arrow IPC are competitive raw record stores, while codecs tied to whole IPC batches or Parquet row groups pay read amplification on random gathers. Dense demonstrates that typed record ownership and per-record compression do not turn fixed tensors into an object-store slow path.

Native-Resolution Ragged

This workload contains the first 2,500 native-resolution Oxford-IIIT Pet CHW RGB images. Height is 108..2606 (median 375) and width is 117..3264 (median 500), totaling 1252.3 MiB of logical uint8 payload. Each of 50 random batches contains 256 records. Every backend persists payload plus exact shape and must return an ordered list[np.ndarray] of (3, H, W) arrays; a flat byte list or a one-dimensional variable-length abstraction is not enough.

Zrecord-raw is 1.5x LMDB and 3.0x Arrow IPC in logical gather. Plain and dictionary zstd both reach 1.21x, confirming that dictionary mode is not useful for these large photographs. The Ragged matrix is intentionally asymmetric: HDF5 gzip, Arrow IPC zstd and Parquet zstd adapters are not implemented because the real 256-record batch took 1.1–1.7 seconds; TileDB variable queries took 8.6–9.6 seconds and the backend was removed entirely. ArrayRecord zstd remains because its batch p95 is 97 ms.

Small Token Records

Both token workloads come from the same real corpus: WikiText-103 raw train, tokenized with GPT-2 and stored as int32 IDs. Preparation is outside every measurement. scripts/prepare_tokens.py preserves nonempty text boundaries, combines fragments shorter than 16 tokens, and splits records at 512 tokens into tokens.npy plus offsets.npy; both benchmark scripts mmap those files.

Fixed Token Blocks

The Dense workload ignores text boundaries and packs the stream into 200,000 fixed int32[512] records: 2 KiB per record and 390.6 MiB logical payload. Each IID batch gathers 256 records; fresh draws continue until timed gathers accumulate at least two seconds.

The contiguous NumPy baseline is strongest for gather when the whole corpus is one fixed typed matrix. Zrecord-raw reaches 3.87 Mrecords/s while retaining independent record semantics and writes 2.3x faster than npy-mmap-raw; dictionary zstd writes 19% faster than plain zstd, uses 13% less disk and gathers 14% faster in this pass. LMDB's B-tree/page overhead is visible in both throughput and disk.

Variable Token Sequences

The Ragged workload keeps 200,000 real text records of 16..512 tokens: p10 28, median 129, mean 138.8, p90 254, totaling 105.9 MiB. It uses the same 256-record, random plan and every backend must return ordered list[np.ndarray] with exact int32 values and original one-dimensional shapes.

Here the record contract, not bulk byte bandwidth, is the useful scale. Zrecord's three codecs return 1.17–1.74 Mrecords/s with 0.19–0.29 ms p95. Dictionary zstd writes 76% faster than plain zstd, uses 21% less disk and gathers 29% faster in this pass. Raw reaches 1902 MiB/s write.

Sampler

Index generation on its own, IID (with replacement), 1M index space, against NumPy's modern API. The µs-scale figures fluctuate with box load; this pass shows a 2.25–10.64x margin across batch sizes.

End-to-End DataLoader

The full pipeline comparison (sample, fetch, collate, handoff) uses the Dense vision source above: 2,500 (3,224,224) uint8 records, 4 workers, batch 256 (36.8 MiB), and one 200-batch pass after warmup. peak PSS apportions shared and copy-on-write mappings across the process tree; peak RSS sums each process's resident set and therefore double-counts shared pages on Linux. These are process-memory diagnostics, not total pipeline physical memory.

This is an end-to-end systems comparison with each loader's normal storage: Torch uses read-only NumPy mmap, Loaderx uses Zrecord, and Grain uses ArrayRecord. Fork represents normal Linux Torch operation; spawn is the no-fork control. Grain is selected explicitly with --only grain.

At 36.8 MiB per batch, gather dominates sampler cost and overlaps with transform work. With source geometry and entropy fixed, raw Loaderx is 1.53x compressed Loaderx. Compressed Loaderx is 1.47x Torch fork, 1.13x Torch spawn, and 2.52x Grain; raw Loaderx is 2.25x, 1.72x, and 3.85x faster, respectively. RSS must not be read as a physical-memory ratio because Linux counts shared mappings repeatedly and page-cache accounting differs from mmap.

The Torch-only worker sweep runs each count once. Worker 0 is the in-process baseline: 67.7 and 68.2 batches/s with 1411/1413 MiB peak PSS in two equivalent runs. The process comparison starts at one worker:

Spawn peak PSS grows from 1708 to 4038 MiB as workers rise from one to eight, while fork grows from 1429 to 1799 MiB because it retains COW sharing. At eight workers spawn uses 2.24x fork's peak PSS and throughput has already flattened. This demonstrates no-fork memory pressure, not OOM: the 31 GiB machine completed the run. Aggregate RSS remains diagnostic rather than physical memory.

CPU-heavy transform solutions. The same Python per-sample transform exposes the GIL bottleneck on standard CPython. Numba is an explicit solution, not the default: --transform numba-nogil compiles the equivalent batch transform with nogil=True. The other explicit solution runs the unchanged Python transform on free-threaded CPython. Numba compilation is warmed before timing; both interpreters use the same prepared real corpus.

Peak PSS for compressed/raw was 860/871 MiB with GIL Python, 933/957 MiB with Numba, and 879/887 MiB with free-threaded Python. These are two deployment solutions to the transform bottleneck, not claims that Store itself was optimized for either runtime.

Conclusion. Loaderx batches sampling, gather, decompression, and Ragged reconstruction while preserving exact record semantics. IID sampling remains uniform with replacement. Dense serves fixed-shape arrays directly; Ragged restores each record's original shape. This favors random training gathers, where formats organized around contiguous scans or larger storage groups can pay read amplification.

Compression results depend on the data. In the Dense photograph workload, plain zstd reaches only 1.15x and dictionary zstd 1.12x; dictionary mode is most useful for repeated small records, not as a universal default.

The loader table combines execution model and storage, rather than comparing schedulers over one shared backend. Compressed Loaderx is 1.47x Torch fork, 1.13x Torch spawn, and 2.52x Grain; raw Loaderx is 2.25x, 1.72x, and 3.85x faster, respectively. CPU-heavy Python transforms remain limited by the GIL; the measured alternatives are Numba nogil=True and free-threaded Python.

These numbers measure warm-cache access, not cold disk or durability. Disk means allocated blocks; write timing excludes stable-media sync and dictionary training. Arrow IPC and Parquet use 256-record groups, so their random gathers include group-level read amplification. This is one complete benchmark pass; Store reads accumulate at least two timed seconds, while sampler and loader figures can fluctuate with machine load. Treat absolute values as ballpark and cross-backend margins as the main signal.

Real-data verification

Production dataset-specific preprocessing and verification remain in the DataPipe repository. The built-in converter covers standardized Hugging Face datasets; DataPipe handles sources without a common remote protocol and implements derived modalities as loader transforms without NumPy dump intermediates. Oxford-IIIT Pet is only the shared benchmark fixture.

Current Limitations

• Single-host only; multi-host training is not supported.
• A single sample must be at most 2 GiB. Store size is practically bounded by disk capacity and platform file limits.
• Stores are not portable between machines with different byte orders. All published platforms are little-endian.

设计文档

• 性能优化与工程决策记录
• 分布式架构草案

Build

python3 setup.py build_ext --inplace
zig build test
python3 scripts/test_loaderx.py

Build the release wheel matrix with:

python3 scripts/build_release.py

Source distributions are intentionally not published. The wheel matrix covers the supported platforms. Source builds require Zig.