loaderx 0.5.0__py3-none-win_arm64.whl

loaderx/__init__.py

@@ -0,0 +1,5 @@
+from ._zrecord import Zrecord, ZrecordError
+from .dataset import Dataset
+from .dataloader import DataLoader
+
+__version__ = "0.5.0"

loaderx/_lib.py

@@ -0,0 +1,45 @@
+"""Locate the bundled native library for the running platform.
+
+A wheel carries the shared objects for exactly one platform, under the name
+that platform uses (`lib*.so`, `lib*.dylib`, `*.dll`). Resolving by name here
+keeps that detail out of the binding modules.
+"""
+
+import sys
+from pathlib import Path
+
+LIB_DIR = Path(__file__).resolve().parent / "lib"
+
+
+def _filename(stem):
+    if sys.platform == "win32":
+        return f"{stem}.dll"
+    if sys.platform == "darwin":
+        return f"lib{stem}.dylib"
+    return f"lib{stem}.so"
+
+
+def path(stem):
+    """
+    Absolute path to the bundled `stem` library.
+
+    Args:
+        stem (str): Library name without prefix or extension, e.g. "zrecord".
+
+    Returns:
+        str: Path to pass to `ffi.dlopen`.
+
+    Raises:
+        ImportError: If the wheel carries no library for this platform.
+    """
+    found = LIB_DIR / _filename(stem)
+    if found.is_file():
+        return str(found)
+
+    present = sorted(p.name for p in LIB_DIR.glob("*") if p.suffix in (".so", ".dylib", ".dll"))
+    raise ImportError(
+        f"loaderx has no {stem} library for {sys.platform} ({found.name} not found in {LIB_DIR}).\n"
+        f"Present: {present or 'nothing'}.\n"
+        "From a source checkout, run `zig build` to produce it; otherwise this "
+        "wheel was built for a different platform."
+    )

loaderx/_sampler.py

@@ -0,0 +1,97 @@
+import time
+
+import numpy as np
+from cffi import FFI
+
+from . import _lib
+
+ffi = FFI()
+
+ffi.cdef("""
+typedef struct Sampler Sampler;
+
+Sampler* cinit(
+    uint64_t length,
+    uint64_t batch_size,
+    uint8_t mode,
+    uint64_t seed
+);
+
+void cdeinit(Sampler* s);
+
+void cnext(Sampler* s, uint64_t* indices_ptr, size_t indices_len);
+""")
+
+lib = ffi.dlopen(_lib.path("sampler"))
+
+class Sampler:
+    class Mode:
+        SEQUENTIAL = 0
+        IID = 1
+
+    def __init__(self, length: int, batch_size: int, mode: int, seed: int = 42):
+        # sampler
+        self.sampler = lib.cinit(length, batch_size, mode, seed)
+        if self.sampler == ffi.NULL:
+            raise MemoryError("Failed to create sampler")
+        # indices
+        self.indices = np.zeros(batch_size, dtype=np.uint64)
+
+    # step
+    def next(self):
+        ptr = ffi.cast("uint64_t *", self.indices.ctypes.data)
+        lib.cnext(self.sampler, ptr, self.indices.size)
+        # return copy to avoid loss
+        return self.indices.copy()
+
+    # iterator
+    def __iter__(self):
+        return self
+    def __next__(self):
+        return self.next()
+
+    # clean
+    def __del__(self):
+        if hasattr(self, "sampler") and self.sampler != ffi.NULL:
+            lib.cdeinit(self.sampler)
+            self.sampler = ffi.NULL
+
+def run():
+    from itertools import islice
+    length = 10
+    batch_size = 4
+    n_steps = 5
+
+    print("=== Sampler SEQUENTIAL ===")
+    sampler = Sampler(length, batch_size, Sampler.Mode.SEQUENTIAL)
+    for indices in islice(sampler, n_steps):
+        print(indices)
+
+    print("=== Sampler IID ===")
+    sampler = Sampler(length, batch_size, Sampler.Mode.IID)
+    for indices in islice(sampler, n_steps):
+        print(indices)
+
+# Benchmark, zig-sampler speeder 3.66x
+def bench():
+    length = 1_000_000
+    batch_size = 8192
+    n_steps = 10_000
+
+    print("=== NumPy IID ===")
+    t0 = time.time()
+    batch = np.zeros(batch_size, dtype=np.uint64)
+    for _ in range(n_steps):
+        batch[:] = np.random.randint(0, length, size=batch_size, dtype=np.uint64)
+    print(f"{(time.time() - t0)*1000:.2f} ms")
+
+    print("=== Sampler IID ===")
+    t0 = time.time()
+    sampler = Sampler(length, batch_size, Sampler.Mode.IID)
+    for _ in range(n_steps):
+        sampler.next()
+    print(f"{(time.time() - t0)*1000:.2f} ms")
+
+if __name__ == "__main__":
+    run()
+    bench()

loaderx/_zrecord.py

@@ -0,0 +1,283 @@
+"""Thin CFFI binding over the Zig zrecord runtime.
+
+zrecord stores opaque byte records; type interpretation is the client's job
+(see :mod:`loaderx.dataset`). Every call below releases the GIL, so several
+loader threads can gather concurrently.
+"""
+
+import numpy as np
+from cffi import FFI
+
+from . import _lib
+
+ffi = FFI()
+
+ffi.cdef("""
+typedef struct Zrecord Zrecord;
+
+Zrecord* zr_open(const char* path_ptr, size_t path_len, int32_t* code);
+void     zr_close(Zrecord* zr);
+
+uint32_t zr_len(Zrecord* zr);
+int32_t  zr_sync(Zrecord* zr);
+uint32_t zr_max_record_len(void);
+
+int32_t  zr_lengths(Zrecord* zr, const uint64_t* idxs, size_t n, uint32_t* out);
+
+int32_t  zr_gather(Zrecord* zr, const uint64_t* idxs, size_t n,
+                   uint8_t* dst, size_t dst_len, size_t stride);
+int32_t  zr_gather_var(Zrecord* zr, const uint64_t* idxs, size_t n,
+                       uint8_t* dst, size_t dst_len, const uint64_t* offsets);
+
+int64_t  zr_append(Zrecord* zr, const uint8_t* src, size_t src_len, size_t n,
+                   size_t stride, size_t item_len, uint8_t codec);
+int64_t  zr_append_var(Zrecord* zr, const uint8_t* src, size_t src_len, size_t n,
+                       const uint64_t* offsets, const uint32_t* lens, uint8_t codec);
+
+int32_t  zr_delete(Zrecord* zr, const uint64_t* idxs, size_t n);
+int32_t  zr_compact(Zrecord* zr);
+void     zr_stats(Zrecord* zr, uint32_t* records, uint64_t* live_bytes, uint64_t* chunk_bytes);
+""")
+
+lib = ffi.dlopen(_lib.path("zrecord"))
+
+# Compression tags. Must match `Codec` in src/zrecord.zig.
+CODECS = {"raw": 0, "flate": 1}
+
+MAX_RECORD_LEN = lib.zr_max_record_len()
+
+# Must match `Code` in src/zrecord.zig.
+_ERRORS = {
+    -1: "out of memory",
+    -2: "file I/O failed",
+    -3: "index out of range",
+    -4: f"record exceeds the {MAX_RECORD_LEN} byte limit",
+    -5: "store is full",
+    -6: "buffer too small",
+    -7: "stored record failed to decode",
+    -8: "not a zrecord store, or written by a newer version",
+}
+
+
+class ZrecordError(RuntimeError):
+    def __init__(self, code):
+        super().__init__(_ERRORS.get(code, f"unknown zrecord error {code}"))
+        self.code = code
+
+
+def _check(code):
+    if code < 0:
+        raise ZrecordError(code)
+    return code
+
+
+def _u64(array):
+    """View `array` as a C-contiguous uint64 buffer, copying only if needed."""
+    return np.ascontiguousarray(array, dtype=np.uint64)
+
+
+def _ptr(ctype, array):
+    return ffi.cast(ctype, array.ctypes.data)
+
+
+def _flat(array, name):
+    """Assert `array` is safe to hand over as a flat byte range."""
+    if not array.flags.c_contiguous:
+        raise ValueError(f"{name} must be C-contiguous")
+    return array
+
+
+class Zrecord:
+    """A zrecord store: an unordered, append-only collection of byte records.
+
+    Indices are stable — once assigned, a record's index never changes meaning.
+    """
+
+    def __init__(self, path, codec="flate"):
+        """
+        Open (or create) the store at `path`.
+
+        Args:
+            path (str | Path): Directory holding the store.
+            codec (str): Default compression for `append`; "raw" or "flate".
+        """
+        if codec not in CODECS:
+            raise ValueError(f"unknown codec {codec!r}, expected one of {sorted(CODECS)}")
+
+        # cffi converts `bytes` to `const char*` and keeps it alive for the call.
+        raw = str(path).encode()
+        code = ffi.new("int32_t*", -2)
+        self._zr = lib.zr_open(raw, len(raw), code)
+        if self._zr == ffi.NULL:
+            raise ZrecordError(code[0])
+        self.codec = codec
+
+    def __len__(self):
+        return lib.zr_len(self._zr)
+
+    def sync(self):
+        """Flush chunk data and the record table to disk."""
+        _check(lib.zr_sync(self._zr))
+
+    def lengths(self, idxs):
+        """
+        Return the logical (uncompressed) size of each requested record.
+
+        Args:
+            idxs (array-like of int): Record indices.
+
+        Returns:
+            np.ndarray: uint32 sizes in bytes.
+        """
+        idxs = _u64(idxs)
+        out = np.empty(idxs.size, dtype=np.uint32)
+        _check(lib.zr_lengths(
+            self._zr, _ptr("const uint64_t*", idxs), idxs.size, _ptr("uint32_t*", out),
+        ))
+        return out
+
+    def gather(self, idxs, dst, stride):
+        """
+        Decompress records `idxs` straight into `dst`, one every `stride` bytes.
+
+        `dst` is written in place — no intermediate buffer is allocated.
+
+        Args:
+            idxs (np.ndarray): uint64 record indices.
+            dst (np.ndarray): C-contiguous destination buffer.
+            stride (int): Byte distance between consecutive record slots.
+        """
+        idxs = _u64(idxs)
+        _flat(dst, "dst")
+        _check(lib.zr_gather(
+            self._zr, _ptr("const uint64_t*", idxs), idxs.size,
+            _ptr("uint8_t*", dst), dst.nbytes, stride,
+        ))
+
+    def gather_var(self, idxs, dst, offsets):
+        """
+        Decompress ragged records `idxs` into `dst` at explicit byte `offsets`.
+
+        Args:
+            idxs (np.ndarray): uint64 record indices.
+            dst (np.ndarray): C-contiguous destination buffer.
+            offsets (np.ndarray): uint64 byte offset of each record within `dst`.
+        """
+        idxs = _u64(idxs)
+        offsets = _u64(offsets)
+        _flat(dst, "dst")
+        _check(lib.zr_gather_var(
+            self._zr, _ptr("const uint64_t*", idxs), idxs.size,
+            _ptr("uint8_t*", dst), dst.nbytes, _ptr("const uint64_t*", offsets),
+        ))
+
+    def append(self, src, n, item_len, codec=None):
+        """
+        Append `n` equally sized records taken from `src`.
+
+        Args:
+            src (np.ndarray): C-contiguous buffer of `n * item_len` bytes.
+            n (int): Number of records.
+            item_len (int): Size of each record in bytes.
+            codec (str, optional): Override the store's default compression.
+
+        Returns:
+            int: Index assigned to the first appended record.
+        """
+        tag = CODECS[codec or self.codec]
+        _flat(src, "src")
+        return _check(lib.zr_append(
+            self._zr, _ptr("const uint8_t*", src), src.nbytes,
+            n, item_len, item_len, tag,
+        ))
+
+    def append_var(self, src, offsets, lens, codec=None):
+        """
+        Append ragged records; record `i` is `src[offsets[i]:][:lens[i]]`.
+
+        Args:
+            src (np.ndarray): C-contiguous source buffer.
+            offsets (np.ndarray): uint64 byte offset of each record.
+            lens (np.ndarray): uint32 byte length of each record.
+            codec (str, optional): Override the store's default compression.
+
+        Returns:
+            int: Index assigned to the first appended record.
+        """
+        tag = CODECS[codec or self.codec]
+        _flat(src, "src")
+        offsets = _u64(offsets)
+        lens = np.ascontiguousarray(lens, dtype=np.uint32)
+        return _check(lib.zr_append_var(
+            self._zr, _ptr("const uint8_t*", src), src.nbytes, lens.size,
+            _ptr("const uint64_t*", offsets), _ptr("const uint32_t*", lens), tag,
+        ))
+
+    def delete(self, idxs):
+        """
+        Remove records, keeping the index space dense.
+
+        The tail of the table is swapped into each hole, so afterwards the live
+        records are exactly `0..len(self)`. The trade-off is that deletion moves
+        records: an index naming a survivor may name a different survivor
+        afterwards. Indices are stable only while nothing is deleted.
+
+        Only the table changes. The removed bytes stay on disk until `compact`.
+
+        Must not run concurrently with reads.
+
+        Args:
+            idxs (array-like of int): Record indices to remove. Duplicates and
+                any order are fine.
+        """
+        idxs = _u64(idxs)
+        _check(lib.zr_delete(self._zr, _ptr("const uint64_t*", idxs), idxs.size))
+
+    def compact(self):
+        """
+        Reclaim the space left behind by `delete`, in place.
+
+        Offline operation: nothing else may touch the store while it runs. Safe
+        to interrupt — records are moved one at a time and the table is updated
+        as each lands, so re-running finishes the job.
+        """
+        _check(lib.zr_compact(self._zr))
+
+    def stats(self):
+        """
+        Report storage usage.
+
+        Returns:
+            dict: `records`, `live_bytes` (what the live records occupy),
+                `chunk_bytes` (how far the chunk files have grown) and
+                `reclaimable` (an upper bound on what `compact` would free —
+                it also counts the chunk-boundary slack compaction cannot
+                remove).
+        """
+        records = ffi.new("uint32_t*")
+        live = ffi.new("uint64_t*")
+        chunk = ffi.new("uint64_t*")
+        lib.zr_stats(self._zr, records, live, chunk)
+        return {
+            "records": records[0],
+            "live_bytes": live[0],
+            "chunk_bytes": chunk[0],
+            "reclaimable": chunk[0] - live[0],
+        }
+
+    # statement
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    # close
+    def close(self):
+        zr = getattr(self, "_zr", ffi.NULL)
+        if zr != ffi.NULL:
+            lib.zr_close(zr)
+            self._zr = ffi.NULL
+
+    def __del__(self):
+        self.close()

loaderx/dataloader.py

@@ -0,0 +1,98 @@
+import threading
+import numpy as np
+from queue import Queue
+
+from ._sampler import Sampler
+
+class DataLoader:
+    def __init__(self, dataset, labelset, num_workers=4, batch_size=256, prefetch_size=4, mode=1, seed=42, transform=(lambda x: x)):
+        """
+        Initialize a DataLoader.
+
+        Args:
+            dataset (BaseDataset): The dataset to load from.
+            labelset (BaseDataset): The labelset to load from.
+            batch_size (int, optional): The batch size to use. Defaults to 256.
+            prefetch_size (int, optional): The number of batches to prefetch. Defaults to 8.
+            mode (int, optional): 
+            seed (int, optional): The seed to use for shuffling. Defaults to 42.
+            transform (callable, optional): A function to apply to the data and label. Defaults to lambda x: x.
+
+        Raises:
+            ValueError: If the dataset and labelset have different lengths.
+        """
+        self.dataset = dataset
+        self.labelset = labelset
+        if len(dataset) != len(labelset):
+            raise ValueError("dataset and labelset must have the same length")
+
+        self.indices = Queue(maxsize=prefetch_size)
+        self.batches = Queue(maxsize=prefetch_size)
+
+        self.sampler = Sampler(len(dataset), batch_size, mode, seed)
+        
+        self.stop_signal = threading.Event()
+
+        self.threads = [
+            threading.Thread(target=self._sampler),
+            *[threading.Thread(target=self._prefetch, args=(transform, )) for _ in range(num_workers)]
+        ]
+
+        for thread in self.threads:
+            thread.daemon = True
+            thread.start()
+
+    def _sampler(self):
+        """
+        Sample indices from the dataset and put them into the index queue.
+        """
+        while not self.stop_signal.is_set():
+            self.indices.put(self.sampler.next())
+
+    def _prefetch(self, transform):
+        """
+        Fetch the data and label from the dataset and labelset based on the indices
+        in the index queue and put them into the batches queue.
+        """
+        while not self.stop_signal.is_set():
+            idxs = self.indices.get()
+            data, label = transform((self.dataset.__getitems__(idxs), self.labelset.__getitems__(idxs)))
+            self.batches.put({'data': data, 'label': label})
+
+    def __len__(self):
+        """
+        Raises a TypeError since an external loader has no length.
+        """
+        raise TypeError("Eternal loader has no length.")
+    
+    # iterator
+    def __iter__(self):
+        return self
+    def __next__(self):
+        # debug: monitor bottlenecks
+        # print(self.indices.qsize(), self.batches.qsize())
+        return self.batches.get()
+    
+    # statement
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    # close
+    def close(self):
+        self.stop_signal.set()
+        
+        for queue in [self.indices, self.batches]:
+            while not queue.empty():
+                try:
+                    queue.get_nowait()
+                except:
+                    break
+        
+        for thread in self.threads:
+            thread.join()
+
+    def __del__(self):
+        self.close()

loaderx/dataset.py

@@ -0,0 +1,222 @@
+"""NumPy-shaped view over a zrecord store.
+
+zrecord returns raw bytes, so the tensor type lives here: a small `schema.json`
+inside the store records the per-sample dtype and shape. Reads decompress
+directly into the freshly allocated output array — the bytes are never staged
+in a temporary buffer.
+"""
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+from ._zrecord import MAX_RECORD_LEN, CODECS, Zrecord
+
+SCHEMA = "schema.json"
+FORMAT = "loaderx-zrecord"
+
+# Bytes of source data handed to a single append call. Caps the compression
+# arena the Zig side allocates while still giving every worker plenty to chew on.
+WRITE_CHUNK_BYTES = 64 << 20
+
+
+class Dataset:
+    def __init__(self, path):
+        """
+        Open a dataset.
+
+        Args:
+            path (str | Path): Directory of a zrecord store written by
+                `Dataset.write`.
+
+        Raises:
+            FileNotFoundError: If `path` holds no loaderx schema.
+        """
+        self.path = Path(path)
+
+        schema_path = self.path / SCHEMA
+        if not schema_path.is_file():
+            raise FileNotFoundError(
+                f"{schema_path} not found; {self.path} is not a loaderx dataset"
+            )
+        schema = json.loads(schema_path.read_text())
+        if schema.get("format") != FORMAT:
+            raise ValueError(f"{schema_path} is not a {FORMAT} schema")
+
+        self.dtype = np.dtype(schema["dtype"])
+        self.shape = tuple(schema["shape"])
+        self.codec = schema["codec"]
+
+        self._itemsize = int(np.prod(self.shape, dtype=np.int64)) * self.dtype.itemsize
+        self._store = Zrecord(self.path, codec=self.codec)
+
+    def __len__(self):
+        """
+        Return the number of samples in the dataset.
+
+        Returns:
+            int: The number of samples in the dataset.
+        """
+        return len(self._store)
+
+    def __getitem__(self, idx):
+        """
+        Get the item at index idx from the dataset.
+
+        Args:
+            idx (int | slice | array-like): Index, slice or fancy index.
+
+        Returns:
+            np.ndarray: A single sample for an integer index, otherwise a
+                stacked batch.
+        """
+        if isinstance(idx, (int, np.integer)):
+            if idx < 0:
+                idx += len(self)
+            return self.__getitems__(np.array([idx], dtype=np.uint64))[0]
+        if isinstance(idx, slice):
+            idx = np.arange(*idx.indices(len(self)))
+        return self.__getitems__(idx)
+
+    def __getitems__(self, idxs):
+        """
+        Get the items at index idxs from the dataset.
+
+        Args:
+            idxs (array-like of int): indices of the items to retrieve.
+
+        Returns:
+            np.ndarray: The items at index idxs, shaped `(len(idxs), *self.shape)`.
+        """
+        idxs = np.asarray(idxs)
+        # A boolean mask would otherwise cast to a run of 0s and 1s.
+        if idxs.dtype == np.bool_:
+            idxs = np.flatnonzero(idxs)
+        idxs = np.ascontiguousarray(idxs, dtype=np.uint64)
+        out = np.empty((idxs.size, *self.shape), dtype=self.dtype)
+        self._store.gather(idxs, out, self._itemsize)
+        return out
+
+    # write
+
+    @classmethod
+    def write(cls, path, array, codec="flate", chunk_bytes=WRITE_CHUNK_BYTES):
+        """
+        Persist `array` as a zrecord dataset, one record per leading-axis slice.
+
+        Streams through `array` in bounded chunks, so a `np.load(...,
+        mmap_mode='r')` array converts without being read into memory at once.
+
+        Args:
+            path (str | Path): Destination directory. Appends if it already
+                holds a compatible dataset.
+            array (array-like): Samples stacked along axis 0. A 1-D array is a
+                dataset of scalars, which is what a label set usually is.
+            codec (str): "flate" for transparent compression, "raw" to store
+                verbatim. Records that fail to shrink are stored raw regardless.
+            chunk_bytes (int): Upper bound on the bytes submitted per call.
+
+        Returns:
+            Dataset: The dataset, opened for reading.
+
+        Raises:
+            ValueError: If a sample exceeds the zrecord record size limit, or
+                the array does not match an existing dataset at `path`.
+        """
+        if codec not in CODECS:
+            raise ValueError(f"unknown codec {codec!r}, expected one of {sorted(CODECS)}")
+
+        array = np.asanyarray(array)
+        if array.ndim < 1:
+            raise ValueError("array must have a sample axis")
+
+        dtype = np.dtype(array.dtype)
+        shape = tuple(int(d) for d in array.shape[1:])
+        itemsize = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
+        if itemsize > MAX_RECORD_LEN:
+            raise ValueError(
+                f"sample of {itemsize} bytes exceeds the {MAX_RECORD_LEN} byte record limit"
+            )
+
+        path = Path(path)
+        path.mkdir(parents=True, exist_ok=True)
+
+        schema = {
+            "format": FORMAT,
+            "version": 1,
+            "dtype": dtype.str,
+            "shape": list(shape),
+            "codec": codec,
+        }
+        schema_path = path / SCHEMA
+        if schema_path.is_file():
+            existing = json.loads(schema_path.read_text())
+            if (existing.get("dtype"), tuple(existing.get("shape", ()))) != (dtype.str, shape):
+                raise ValueError(
+                    f"{path} holds {existing.get('dtype')}{tuple(existing.get('shape', ()))} "
+                    f"records, cannot append {dtype.str}{shape}"
+                )
+            schema = existing
+        else:
+            schema_path.write_text(json.dumps(schema, indent=2) + "\n")
+
+        step = max(1, chunk_bytes // max(itemsize, 1))
+        with Zrecord(path, codec=schema["codec"]) as store:
+            for start in range(0, len(array), step):
+                block = np.ascontiguousarray(array[start:start + step], dtype=dtype)
+                store.append(block, len(block), itemsize)
+            store.sync()
+
+        return cls(path)
+
+    def delete(self, idxs):
+        """
+        Remove samples, keeping the index space dense.
+
+        `len(self)` shrinks and the survivors stay addressable as
+        `0..len(self)`, which is what the sampler assumes. The tail of the
+        dataset is swapped into each hole, so surviving samples may change
+        index — hold no indices across a delete.
+
+        The removed bytes stay on disk until `compact`. Must not run
+        concurrently with reads, so do not call it while a `DataLoader` is
+        running over this dataset.
+
+        Args:
+            idxs (array-like of int): Indices to remove.
+        """
+        self._store.delete(idxs)
+
+    def compact(self):
+        """
+        Reclaim the space left behind by `delete`.
+
+        Offline operation: nothing else may touch the dataset while it runs.
+        """
+        self._store.compact()
+
+    def stats(self):
+        """
+        Report storage usage; see `loaderx.Zrecord.stats`.
+
+        Returns:
+            dict: record count and byte counters, including `reclaimable`.
+        """
+        return self._store.stats()
+
+    # statement
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    # close
+
+    def close(self):
+        store = getattr(self, "_store", None)
+        if store is not None:
+            store.close()
+            self._store = None

loaderx-0.5.0.dist-info/METADATA

@@ -0,0 +1,334 @@
+Metadata-Version: 2.4
+Name: loaderx
+Version: 0.5.0
+Summary: A record-based data runtime, focused on delivering extreme throughput and low latency
+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 Dataset, DataLoader
+
+dataset = Dataset('train_data.zr')
+labelset = Dataset('train_label.zr')
+
+loader = DataLoader(dataset, labelset)
+
+for i, batch in enumerate(loader):
+    if i >= 256:
+        break
+
+print(batch['data'].shape)
+print(batch['label'].shape)
+```
+
+### Converting a NumPy tensor
+```python
+import numpy as np
+from loaderx import Dataset
+
+Dataset.write('train_data.zr', np.load('data.npy', mmap_mode='r'))
+Dataset.write('train_label.zr', np.load('label.npy', mmap_mode='r'))
+```
+One record per slice along axis 0; a 1-D array (the usual shape of a label set)
+becomes a dataset of scalars. The conversion streams in bounded chunks, so an
+mmapped array is never fully materialized.
+
+`Dataset` keeps a single `schema.json` inside the store recording the per-sample
+dtype and shape — that is the only loaderx-level metadata. Everything beneath it
+is plain Zrecord, reachable through `loaderx.Zrecord` when you want raw byte
+records (ragged samples, pre-encoded images) instead of tensors.
+
+### Removing samples
+```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': ...}
+```
+Deletion keeps the index space dense, which is what the sampler needs — but it
+does so by moving the tail into the hole, so surviving samples may change index.
+Do not hold indices across a delete. Both calls require exclusive access; do not
+run them while a `DataLoader` is iterating the dataset.
+
+### Integrating with JAX/Flax
+
+For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
+
+## 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
+zig build bench                 # append/gather throughput
+python3 tests/test_loaderx.py   # Python layer, against whichever build is importable
+```
+
+Two things about the shared objects are deliberate rather than incidental:
+
+* **libc is linked.** They are `dlopen`ed into CPython, and without libc they
+  export an unresolved `__tls_get_addr`, so Zig's thread-locals never bind and
+  `std.Thread` deadlocks on the first batch.
+* **Tests run in ReleaseFast too.** The shipped libraries are ReleaseFast, and a
+  codegen difference there has already produced a release-blocking bug that
+  Debug did not have. Test the configuration that ships.
+
+### Publishing
+Zig cross-compiles every target from one machine, so releases need no CI matrix:
+
+```
+zig build dist                    # every platform, into zig-out/dist/<wheel tag>/
+python3 scripts/build_wheels.py   # one wheel per platform, plus the sdist
+```
+
+The dist directories are named after their Python wheel platform tag, so the tag
+mapping lives in exactly one place (`dist_targets` in `build.zig`). glibc and
+macOS minimums are pinned in the target triple, which is what makes
+`manylinux_2_17` and `macosx_11_0` honest rather than aspirational. Each wheel is
+checked after packing: it must carry this platform's libraries and no others.
+
+The sdist ships sources only. Installing from it runs `zig build` through
+`setup.py`, so it needs the Zig compiler; wheel users never hit that path.
+
+---
+
+# Zsampler
+Index Generator: a high-performance sampler implemented in Zig
+
+1. Sequential generation: indices are produced by traversing the index space in order.
+    * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
+2. Random generation: indices are sampled randomly from the index space.
+    * Global random: a set of samples is drawn randomly from the entire index space.
+
+---
+
+# Zrecord
+A record-based data runtime, focused on delivering extreme throughput and low latency.
+
+1. Zrecord is an unordered physical store made of N records. Records are
+   independent and carry no ordering, so every index and slice operation is
+   equivalent to a gather.
+2. Zrecord hands the client a dense index space: the live records are always
+   exactly `0..N`. Deletion preserves that by swapping the tail into the hole,
+   which means an index is stable only until something is deleted. A
+   multi-stream dataset is purely a client-side notion; the client keeps its
+   own stream index table pointing at records.
+3. Zrecord returns byte arrays. Type interpretation is the client's job.
+4. For performance, the IO model (`append | read | delete`) is batch-oriented;
+   a single-record operation is just the `batch_size == 1` case.
+5. Zrecord owns its memory internally — allocation and release are explicit.
+6. Each record is compressed and decompressed independently:
+
+```
+| tag   |  name   | algorithm |
+|-------|---------|-----------|
+|   0   |   raw   | none      |
+|   1   |  flate  | Deflate   |
+```
+
+7. Compression is transparent to the client:
+    * Compression runs concurrently across all cores. A record that fails to
+      shrink is stored raw automatically, so reading never costs more than the
+      data itself.
+    * Decompression writes straight into the caller's destination memory
+      (`gather`), with no intermediate buffer and no extra copy.
+
+## Persistence format
+Zrecord storage is metadata plus chunked data:
+```
+zrecord/
+  ├── meta.zr
+  ├── 0.zr
+  └── 1.zr
+```
+
+### Metadata (meta.zr)
+
+meta.zr is **used as the values it holds, not decoded into them**. Both structs
+below are `extern`, 16 bytes, and naturally aligned; an mmap is page aligned, so
+the header is a `*Header` and the record table is a `[]RecordLoc`. Looking up a
+record is `table[idx]`, and there is no serializer anywhere in the code.
+
+The fields are deliberately wider than the limits need — `chunk_id` is a `u16`
+for 4096 chunks, the lengths are `u32` for a 1 MiB cap. Bit-packing them would
+save 5 bytes per record and cost a hand-written codec on the hottest lookup
+path. The slack buys a machine-shaped array; it is worth it.
+
+**1. Header** — 16 bytes of global state.
+* `magic` (`ZREC`) and `version` are ordinary fields, so opening a directory
+  that is not a Zrecord store fails immediately instead of decoding garbage.
+* `length` is the total number of records | `tail_chunk`/`tail_offset` mark the
+  last write position.
+* There is no chunk count. Chunks are created in order and the frontier is
+  always in the last one, so the store holds exactly chunks `0..=tail_chunk` —
+  a count would be a second copy of that fact to keep in sync.
+
+```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, flate = 1, _ };
+const RecordLoc = extern struct {
+    offset: u32, phys_length: u32, logic_length: u32,
+    chunk_id: u16, compress: Codec, _reserved: u8,
+};
+```
+
+There is no liveness flag. Every entry below `length` is live, because deletion
+swaps the tail into the hole rather than tombstoning.
+
+**3. Maximum length.** Running out of space is effectively impossible, so the
+metadata is sized statically at the maximum.
+* Current design: up to 2^12 (4096) chunks, 2^32 (4 GiB) per chunk, 2^20 (1 MiB)
+  per record.
+* Derivation: N = chunks × 2^32 (chunk size) / 2^20 (max record size) ⇒ a
+  maximum length of 2^24 (16777216). Equality here is the point: the store can
+  never run out of chunks before it runs out of record slots, which is what
+  lets the table be sized statically. It is a compile-time assertion.
+* meta.zr is a fixed 16 + 2^24 × 16 bytes (256 MiB). It is sparse — a store
+  with one record allocates 4 KiB of it — and the mapping does not prefault.
+
+### Chunk data (x.zr)
+Densely packed record data.
+
+## Executor
+
+**1. Write.** Writes are append-only; everything else is offset redirection.
+* Append: compress concurrently → assign physical locations serially → flush
+  concurrently → commit metadata. Data is durable-ordered before the record
+  table, and the table before `length`, so a crash truncates rather than
+  corrupts. A record never straddles two chunks; one that would not fit rolls
+  over to a fresh chunk.
+* Delete: swap the last table entry into the deleted slot and drop the length by
+  one. A batch is applied in descending index order, so each swap pulls from a
+  slot no later target refers to. The index space stays dense — which is what
+  the sampler needs, since it draws uniformly from `0..N` and would otherwise
+  keep hitting holes. The deleted record's bytes become garbage.
+
+**2. Read.** Fill the destination memory concurrently, in place from the Python
+side (executed on async threads).
+* Committed records are immutable, so the read path is lock free; `length` is
+  published to readers through an atomic.
+
+**3. Concurrency model.** `Io.Group.async` shards work by CPU count, and shards
+beyond the limit run inline on the calling thread.
+* Shards receive contiguous blocks rather than a strided subset, keeping each
+  worker's reads and writes sequential.
+* Each shard reuses one compressor scratch (`flate.Compress` is ~230 KiB — far
+  too large to rebuild per record, let alone to place on a task stack).
+* Decompression runs in direct mode: matches are resolved against the bytes
+  already written into the destination, so there is no window buffer.
+
+**4. Garbage collection.** Deletion leaves the record's bytes stranded, so space
+is reclaimed by an offline `compact` that rewrites the live records in place.
+* Records are visited in *physical* order — which after swap-last deletes no
+  longer matches table order — and repacked densely in that same order. A record
+  therefore never moves to a higher address than it already had.
+* Two things follow. Writing a record can never land on the bytes of a record
+  not yet moved, so the rewrite is safe in place with no scratch copy of the
+  store. And each table entry can be updated the instant its record lands, so
+  the store is consistent at every point: an interrupted compaction leaves some
+  records moved and the rest where they were, and re-running finishes the job.
+* Records that are already in the right place are skipped, so a store with a
+  small amount of garbage near the end is cheap to compact.
+* Chunk files past the new frontier are closed and deleted.
+* `stats()` reports `live_bytes` against `chunk_bytes` so callers can decide when
+  it is worth running. Note the difference is an upper bound: a record never
+  straddles a chunk boundary, so up to one record's worth per chunk is slack
+  that compaction cannot remove.
+
+**5. File access.**
+* Metadata: a 192 MiB file created at init, accessed by mmap thereafter.
+* Chunk data: 4 GiB files created at init, accessed concurrently through
+  `readPositionalAll`/`writePositionalAll`.
+
+**Concurrency contract.** `gather` and `append` are safe to call concurrently
+from many threads. `delete` and `compact` mutate the table in ways a lock-free
+reader would observe half-applied, so they require exclusive access to the
+store.
+
+## Distributed extension
+**Not implemented yet; architecture noted in advance.**
+
+1. A cluster layer is added; a node becomes a shard holding several chunks.
+2. An indirection table maps a global path to a specific node, so the index path
+   becomes `idx → indirection[idx] = node, lidx → offset[lidx]`.

loaderx-0.5.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+loaderx/__init__.py,sha256=g-tkCE23AVFH3E1-Wy3Gm4r6-m9GTIzy1UlN7pQQAoQ,131
+loaderx/_lib.py,sha256=kj4ZoHy_U6kU5KVjZQeyfPu8dmEl7diWLoldUhZZaoo,1343
+loaderx/_sampler.py,sha256=4ImeP_RBysfgsEmz118i0e84TJNisqoL28bKR74Tmfs,2370
+loaderx/_zrecord.py,sha256=PKW0IGab-s5EN1LJqIXV7RGBW_M47tDJXEBEQH1rJQQ,9334
+loaderx/dataloader.py,sha256=dmj-86YFxLualASdU28QPDSoidNer6PjVULSG72mPuM,3229
+loaderx/dataset.py,sha256=IyZXH4SLWh715vqDGwsXzyMgXNkmeJN6mOdgEZi-QgA,7442
+loaderx/lib/sampler.dll,sha256=KdCUyy4cQ8ReiPO_irLE1jYzMp372aHyNnBRiKU_5WU,210432
+loaderx/lib/zrecord.dll,sha256=t8lGwrDO8FnwEC7_rEt5yNltOzxWLbZBnIcXgb_BFFk,776704
+loaderx-0.5.0.dist-info/licenses/LICENSE,sha256=6TzYUKn7HWuu0ug9ZepYbKY7s1LOh43uBNWIN6i7nsI,1075
+loaderx-0.5.0.dist-info/METADATA,sha256=k4MF7ux4SVR_dPWhRpdZQDC5fJ4kWY-Y0pIrSQMx1Wg,15394
+loaderx-0.5.0.dist-info/WHEEL,sha256=84gBEa4fGSDR6CJj9KMeLIZ46AIp14gh61YGBn-qQNM,97
+loaderx-0.5.0.dist-info/top_level.txt,sha256=g8WeEbc7sZyfT8dsWS7dJy-Q-W1_PVlG4bNkPWNmQfM,8
+loaderx-0.5.0.dist-info/RECORD,,

loaderx-0.5.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-win_arm64
+

loaderx-0.5.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+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.

loaderx-0.5.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+loaderx