loaderx 0.7.3__tar.gz → 0.9.4__tar.gz

{loaderx-0.7.3/loaderx.egg-info → loaderx-0.9.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: loaderx
-Version: 0.7.3
+Version: 0.9.4
 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
@@ -42,6 +42,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: cffi
+Requires-Dist: msgpack
 Dynamic: license-file
 
 # Loaderx
@@ -64,20 +65,27 @@ loaderx is built around several core principles:
 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
 ```python
-from loaderx import DataLoader, Dataset, write_dataset
+from loaderx.converter import from_numpy
+from loaderx.dataset import DenseDataset
+from loaderx.dataloader import DataLoader
 
-# 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'))
+from_numpy('train_data', np.load('data.npy', mmap_mode='r'))
+from_numpy('train_label', 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'),
-})
+loader = DataLoader({'data': DenseDataset('train_data'), 'label': DenseDataset('train_label')},
+                    transform=lambda batch: batch)
 
 for i, batch in enumerate(loader):
     if i >= 256:
@@ -87,69 +95,79 @@ 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': ...})`.
+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
 ```python
 import numpy as np
-from loaderx import Dataset, write_dataset
+from loaderx.converter import from_numpy
+from loaderx.dataset import DenseDataset
 
-write_dataset('train_data.zr', np.load('data.npy', mmap_mode='r'))
-write_dataset('train_label.zr', np.load('label.npy', mmap_mode='r'))
+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:
-write_dataset('train_data.zr', images, codec='zstd_dict')
+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. `codec` selects compression —
-`"zstd"` (default), `"raw"`, or `"zstd_dict"` (see [Zrecord](#zrecord)). Open
-the result with :class:`Dataset`::
+mmapped array is never fully materialized. Open the result with
+:class:`DenseDataset`:
 
-    ds = Dataset('train_data.zr')
+    ds = DenseDataset('train_data')
 
-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.
+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.
 
-### 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.
+### 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:
 
 ```python
 import numpy as np
-from loaderx import RaggedDataset, pad, write_ragged
+from loaderx.converter import from_numpy
+from loaderx.dataset import DenseDataset
 
-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
+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:
 
-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
+```python
+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')
 
-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 = rs[[0, 2, 4]]              # list of ndarray — one per record
+padded, lengths = pad(records)       # dense (B, max_len) + row counts
 ```
 
-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:
+A ``DataLoader`` works with dense streams, so collation is just the
+``transform`` — a batch dict in, a batch dict out:
 
 ```python
 def collate(batch):
-    values, row_splits = batch['tokens']
-    padded, lengths = pad(values, row_splits)
-    return {'input_ids': padded, 'lengths': lengths, 'label': batch['label']}
+    return {'input_ids': batch['tokens'], 'label': batch['label']}
 
-loader = DataLoader({'tokens': rds, 'label': labelset}, batch_size=32, transform=collate)
+loader = DataLoader({'tokens': dense_tokens, 'label': labelset}, batch_size=32,
+                    transform=collate)
 ```
 
 ### Codec notes
@@ -182,92 +200,110 @@ For practical integration examples, please refer to the **[Data2Latent](https://
 `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`.
+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 |   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.
+| 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
-batch 256 × 200, structured (mildly compressible) data.
+records, 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.
+| 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. 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.
+over a batch), same workload, 4 workers.
 
 | 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.
+| **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.
@@ -346,8 +382,11 @@ A record-based data runtime, focused on delivering extreme throughput and low la
    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.
+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: