loaderx 0.5.4__tar.gz → 0.7.3__tar.gz

{loaderx-0.5.4 → loaderx-0.7.3}/MANIFEST.in

@@ -7,5 +7,9 @@ recursive-include src *.zig
 recursive-include tests *.py
 recursive-include scripts *.py
 
+# Vendored zstd: the sdist rebuilds with Zig, so it needs the C sources.
+recursive-include vendor *.c *.h
+include vendor/zstd/LICENSE vendor/zstd/COPYING
+
 exclude loaderx/lib/lib_here
 recursive-exclude loaderx/lib *.so *.dylib *.dll *.pdb *.lib

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: loaderx
-Version: 0.5.4
+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
@@ -67,12 +67,17 @@ loaderx is built around several core principles:
 
 ## Quick Start
 ```python
-from loaderx import Dataset, DataLoader
+from loaderx import DataLoader, Dataset, write_dataset
 
-dataset = Dataset('train_data.zr')
-labelset = Dataset('train_label.zr')
+# 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'))
 
-loader = DataLoader(dataset, labelset)
+# 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:
@@ -82,34 +87,91 @@ 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
+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'))
 
-Dataset.write('train_data.zr', np.load('data.npy', mmap_mode='r'))
-Dataset.write('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.
+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')
 
-`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
+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.
 
-### Removing samples
+### 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
-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': ...}
+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
 ```
-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.
+
+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
 
@@ -117,7 +179,7 @@ For practical integration examples, please refer to the **[Data2Latent](https://
 
 ## Benchmarks
 
-`scripts/bench.py` measures loaderx against the alternatives in three layers —
+`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
@@ -141,40 +203,44 @@ space, against both NumPy random APIs.
 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.
+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** — random batch gather, 20000 records × 12 KiB `uint8(3,64,64)`,
+**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 |
-|---------------|------------|-------------|------------|---------|-------|
-| 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
+| 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;
-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,
+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
@@ -187,20 +253,21 @@ for: loaderx from Zrecord, Grain from ArrayRecord, torch from a memory-mapped
 
 | 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.
+| **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.
@@ -244,12 +311,26 @@ The sdist ships sources only. Installing from it runs `zig build` through
 ---
 
 # 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.
+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.
 
 ---
 
@@ -271,10 +352,11 @@ A record-based data runtime, focused on delivering extreme throughput and low la
 6. Each record is compressed and decompressed independently:
 
 ```
-| tag   |  name   | algorithm |
-|-------|---------|-----------|
-|   0   |   raw   | none      |
-|   1   |  flate  | Deflate   |
+| 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:
@@ -283,6 +365,22 @@ A record-based data runtime, focused on delivering extreme throughput and low la
       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:
@@ -329,7 +427,7 @@ a physical address is what makes random access efficient.
   the algorithm.
 
 ```zig
-const Codec = enum(u8) { raw = 0, flate = 1, _ };
+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,
@@ -376,10 +474,12 @@ side (executed on async threads).
 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.
+* 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.
@@ -400,7 +500,7 @@ is reclaimed by an offline `compact` that rewrites the live records in place.
   that compaction cannot remove.
 
 **5. File access.**
-* Metadata: a 192 MiB file created at init, accessed by mmap thereafter.
+* 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`.