Metadata-Version: 2.4
Name: loaderx
Version: 0.7.3
Summary: 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
Author-email: Ben0i0d <ben0i0d@foxmail.com>
License: MIT License
        
        Copyright (c) 2025 EOELAB AI Research
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://codeberg.org/eoelab/loaderx
Project-URL: Documentation, https://codeberg.org/eoelab/loaderx
Project-URL: Source, https://codeberg.org/eoelab/loaderx
Project-URL: Bug Tracker, https://codeberg.org/eoelab/loaderx
Keywords: flax,python,dataloader
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: POSIX :: Linux
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: Microsoft :: Windows
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: cffi
Dynamic: license-file

# 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
```python
from loaderx import DataLoader, Dataset, write_dataset

# Write (zrecord layer — one‑time conversion from NumPy)
write_dataset('train_data.zr', np.load('data.npy', mmap_mode='r'))
write_dataset('train_label.zr', np.load('label.npy', mmap_mode='r'))

# Read (loader layer — pure loader, no mutations)
loader = DataLoader({
    'data': Dataset('train_data.zr'),
    'label': Dataset('train_label.zr'),
})

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

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

A batch is a dict with one entry per stream. Every stream is gathered at the same
indices, so record `i` lines up across them — a multi-field sample is just more
streams sharing a length: `DataLoader({'image': ..., 'label': ..., 'meta': ...})`.

### Converting a NumPy tensor
```python
import numpy as np
from loaderx import Dataset, write_dataset

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

# many small, similar records compress far better against a trained dictionary:
write_dataset('train_data.zr', 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. `codec` selects compression —
`"zstd"` (default), `"raw"`, or `"zstd_dict"` (see [Zrecord](#zrecord)). Open
the result with :class:`Dataset`::

    ds = Dataset('train_data.zr')

The store carries a single ``schema.json`` 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.

### Variable-length records
When records share a dtype but not a length — token sequences, waveforms, point
sets — use `RaggedDataset`. It writes an iterable of variable-length arrays and
reads a batch back in the packed layout Zrecord stores natively: one flat buffer
of every row plus `row_splits`, gathered in place across cores with nothing
padded. The store stays a plain blob of bytes; how a batch becomes dense is a
separate, visible step.

```python
import numpy as np
from loaderx import RaggedDataset, pad, write_ragged

seqs = [np.arange(L, dtype=np.int32) for L in (3, 1, 4, 1, 5)]
write_ragged('tokens.zr', seqs)               # dtype and per-row shape inferred
rds = RaggedDataset('tokens.zr')              # open for reading

values, row_splits = rds[[0, 2, 4]]            # packed: values is flat, row_splits marks edges
first = values[row_splits[0]:row_splits[1]]    # record 0's rows — the client slices

padded, lengths = pad(values, row_splits)      # dense (B, max_len, *tail) + row counts
padded, lengths = rds.getitems_padded([0, 2, 4])  # or fuse the gather and the pad in one pass
```

Records may carry a fixed trailing shape (`item_shape=(3,)` for sequences of
3-vectors); a plain 1-D array is a sequence of scalar rows. In a `DataLoader`,
padding is just the collate `transform` (a batch dict in, a batch dict out), so
ragged streams and fixed-shape streams flow through the same loader:

```python
def collate(batch):
    values, row_splits = batch['tokens']
    padded, lengths = pad(values, row_splits)
    return {'input_ids': padded, 'lengths': lengths, 'label': batch['label']}

loader = DataLoader({'tokens': rds, '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](https://codeberg.org/eoelab/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. 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 IID draw is unbiased like
NumPy's: it uses Lemire's method with rejection, which over any real index space
effectively never redraws, so uniformity costs nothing here.

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

| store             | write      | gather      | per batch  | on disk | ratio |
|-------------------|------------|-------------|------------|---------|-------|
| zrecord-zstd      |  531 MiB/s | 10157 MiB/s |   0.30 ms  | 216 MiB | 1.08x |
| zrecord-zstdict   |   40 MiB/s |  4387 MiB/s |   0.68 ms  | 201 MiB | 1.17x |
| zrecord-raw       |  794 MiB/s | 21317 MiB/s |   0.14 ms  | 235 MiB | 1.00x |
| npy-mmap          | 2633 MiB/s | 10057 MiB/s |   0.30 ms  | 234 MiB | 1.00x |
| arrayrecord       |  329 MiB/s |   744 MiB/s |   4.03 ms  | 217 MiB | 1.08x |
| hdf5              | 1105 MiB/s |   486 MiB/s |   6.18 ms  | 236 MiB | 0.99x |
| hdf5-gzip         |   78 MiB/s |   193 MiB/s |  15.51 ms  | 211 MiB | 1.11x |
| lmdb              |  666 MiB/s |  4655 MiB/s |   0.64 ms  | 313 MiB | 0.75x |
| lmdb-zlib         |   70 MiB/s |   306 MiB/s |   9.81 ms  | 235 MiB | 1.00x |
| blosc2            |   30 MiB/s |   217 MiB/s |  13.81 ms  | 216 MiB | 1.09x |
| zarr              |   78 MiB/s |    23 MiB/s | 131.13 ms  | 217 MiB | 1.08x |
| tensorstore       |  545 MiB/s |    89 MiB/s |  33.76 ms  | 214 MiB | 1.09x |

The gather column is what the store is built for. zrecord-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; and
zrecord-zstd keeps most of that speed *while* decompressing — 10 GiB/s, several
times faster than any other compressed store (hdf5-gzip, lmdb-zlib, blosc2, zarr,
tensorstore) and faster even than the uncompressed `.npy`. That is the point of
one fast codec: turning compression on barely costs a read. The array stores lag
on gather because a scattered single-record read fights a layout meant for
contiguous slabs — exactly the access pattern a training loader has. Where zrecord
does not lead is write: it pays compression and its own layout up front, a
one-time cost it does not chase (`zrecord-zstdict` most of all, at level 19).
Compression is modest here (1.08–1.17x) only because this data is barely
compressible; on the smooth-image set of the [Zrecord](#zrecord) notes plain zstd
reaches 7.6x and the dictionary 16.5x.

The reads are not all single-threaded: zrecord 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**    |    4113.9 |  0.24 ms  | 1,053,168 | 12342 MiB/s |
| loaderx-raw    |    4838.9 |  0.21 ms  | 1,238,752 | 14517 MiB/s |
| grain          |      62.9 | 15.90 ms  |    16,097 |   189 MiB/s |
| torch          |     961.0 |  1.04 ms  |   246,005 |  2883 MiB/s |
| jax-dataloader |    3176.8 |  0.31 ms  |   813,270 |  9531 MiB/s |

End to end, loaderx serves batches several times faster than torch or Grain and
ahead of jax-dataloader — and it does so from disk-backed *compressed* storage,
where jax-dataloader's number is an in-memory ceiling with no storage path at all.
zstd decompresses fast enough that the compressed loader nearly matches the raw
one (4114 vs 4839 batches/s), so on-disk compression is close to free here.
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](#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. 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   |  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.

```zig
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.

```zig
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]`.
