Implement kh57 algorithm and tests

Author: arkkln

src/kh57/__init__.py

@@ -1 +1,9 @@
 """kh57 - efficient range reservoir sampling from massive sorted KV datasets."""
+
+from kh57.backends import Backend, MemBackend
+from kh57.encoding import kh57, recover, uniform_hash
+from kh57.sampling import sample
+
+
+__all__ = ["Backend", "MemBackend", "kh57", "recover", "sample", "uniform_hash"]
+__version__ = "0.1.0"

src/kh57/backends/__init__.py

@@ -1 +1,6 @@
 """kh57 backends - storage adapters for sampling sources."""
+
+from kh57.backends.mem import Backend, MemBackend
+
+
+__all__ = ["Backend", "MemBackend"]

src/kh57/backends/mem.py

@@ -1,4 +1,87 @@
-"""In-memory sorted-dict backend for kh57.
+"""In-memory backend and the Backend protocol.
 
-Stub. Implementation coming later.
+the Backend protocol is the minimum sorted-KV surface `kh57.sample()` needs.
+keys are raw bytes (callers big-endian encode their ints), so the protocol
+stays agnostic of key semantics and disk-backed adapters (rocksdb, lmdb)
+just need to match the same shape.
 """
+
+from __future__ import annotations
+
+from bisect import bisect_left
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+__all__ = ["Backend", "MemBackend"]
+
+
+@runtime_checkable
+class Backend(Protocol):
+    """Minimum sorted-KV interface required by `kh57.sample()`."""
+
+    def get(self, key: bytes) -> bytes | None:
+        """Return the value for `key`, or None if absent."""
+        ...
+
+    def put(self, key: bytes, value: bytes) -> None:
+        """Insert or overwrite `key` with `value`."""
+        ...
+
+    def delete(self, key: bytes) -> None:
+        """Remove `key` if present."""
+        ...
+
+    def range_scan(self, lo: bytes, hi: bytes) -> Iterator[tuple[bytes, bytes]]:
+        """Yield (key, value) pairs with lo <= key < hi, ascending by key."""
+        ...
+
+
+class MemBackend:
+    """Reference in-memory backend: a dict plus a lazily sorted key list.
+
+    a plain dict holds the data; the sorted key list is rebuilt on the first
+    `range_scan` after a mutation. this avoids a third-party sorted-container
+    dependency and is fast for the dominant pattern (bulk load, then scan):
+    writes are O(1), and sorting is paid once per write burst instead of per
+    write. not safe to mutate while a `range_scan` iterator is live.
+    """
+
+    def __init__(self) -> None:
+        """Create an empty backend."""
+        self._data: dict[bytes, bytes] = {}
+        self._sorted_keys: list[bytes] = []
+        self._dirty = False
+
+    def get(self, key: bytes) -> bytes | None:
+        """Return the value for `key`, or None if absent."""
+        return self._data.get(key)
+
+    def put(self, key: bytes, value: bytes) -> None:
+        """Insert or overwrite `key` with `value`."""
+        if key not in self._data:
+            self._dirty = True
+        self._data[key] = value
+
+    def delete(self, key: bytes) -> None:
+        """Remove `key` if present."""
+        if self._data.pop(key, None) is not None and not self._dirty:
+            # cheap incremental removal keeps the sorted list valid
+            idx = bisect_left(self._sorted_keys, key)
+            del self._sorted_keys[idx]
+
+    def range_scan(self, lo: bytes, hi: bytes) -> Iterator[tuple[bytes, bytes]]:
+        """Yield (key, value) pairs with lo <= key < hi, ascending by key."""
+        if self._dirty:
+            self._sorted_keys = sorted(self._data)
+            self._dirty = False
+        start = bisect_left(self._sorted_keys, lo)
+        stop = bisect_left(self._sorted_keys, hi, lo=start)
+        for key in self._sorted_keys[start:stop]:
+            yield key, self._data[key]
+
+    def __len__(self) -> int:
+        return len(self._data)

src/kh57/encoding.pyx

@@ -1,5 +1,108 @@
 # cython: language_level=3
 """kh57 encoding and recovery.
 
-Stub. Algorithm implementation coming later.
+encodes a 57-bit non-negative integer key into a 64-bit compound sort key:
+the top 7 bits carry a level id derived from `bit_length(uniform_hash(key))`,
+the bottom 57 bits carry the original key. the first bit is zero by
+construction, keeping the sign bit free so big-endian encoding preserves
+order even where keys are treated as signed. within a level the original
+key order is preserved; levels are deterministic random subsets of the key
+space, with each higher level holding roughly twice as many keys.
 """
+
+from libc.stdint cimport int64_t, uint8_t
+
+from kh57.siphash cimport isiphash64
+
+
+SALT = b"0123456789abcdef"
+
+cdef unsigned char* _SALT = SALT
+
+
+cpdef int64_t uniform_hash(int64_t key):
+    """Uniform, not necessarily cryptographic, hash of a signed 64-bit int.
+
+    uses siphash-2-4 underneath: 64-bit message, 16-byte key (the module
+    SALT), 64-bit hash out.
+    """
+    return isiphash64(key, _SALT)
+
+
+def _uniform_hash_with_salt(key: int, salt: bytes) -> int:
+    """Like `uniform_hash` but with an explicit 16-byte salt.
+
+    the top-level `uniform_hash` is intentionally not parametrized: the salt
+    determines the level structure, so it must stay constant for a dataset.
+    use this only if you need a different (but again constant) salt.
+    """
+    if len(salt) != 16:
+        raise ValueError("salt must be exactly 16 bytes")
+    cdef const uint8_t[:] s = salt
+    return isiphash64(key, <uint8_t*>&s[0])
+
+
+def kh57(key: int) -> int:
+    """Encode a 57-bit non-negative integer key into a 64-bit compound key.
+
+    Args:
+        key: non-negative integer in the range [0, 2^57 - 1].
+
+    Returns:
+        the 64-bit encoded key: 7 level bits then 57 key bits.
+
+    Raises:
+        ValueError: if the key is negative or does not fit in 57 bits.
+    """
+    # supports range 0, 1, 2, ..., (2^57 - 2), (2^57 - 1)
+    if key < 0:
+        raise ValueError("key must be non-negative")
+    if key >> 57:
+        raise ValueError("key must fit in 57 bits")
+
+    some_deterministic_hash = uniform_hash(key)
+
+    # Cast to unsigned int
+    buffer = some_deterministic_hash.to_bytes(8, "big", signed=True)
+    some_deterministic_hash = int.from_bytes(buffer, "big", signed=False)
+
+    # Take the most significant bit as the level id.
+    # Higher levels hold most of the records; the highest available level
+    # holds almost half of them.
+    level_idx = some_deterministic_hash.bit_length()
+
+    if level_idx:
+        # Levels with id N are stored as (N-1) to avoid wasting a bit:
+        # 2^6 = 64 needs 7 bits, but level 0 is a near-impossible
+        # special case (hash == 0), so the ids can be shifted down safely.
+        level_idx -= 1
+
+    # Format: 8 bytes in total
+    # +-------+---------------------------------------------------------------+
+    # | 7 bit |         .        .        .        .        .         57 bits |
+    # | level |         .        .        .        .        .      big-endian |
+    # |    id |         .        .        .        .        .unsigned integer |
+    # +-------+---------------------------------------------------------------+
+    # |      0 .      1 .      2 .      3 .      4 .      5 .      6 .      7 |
+    # +-----------------------------------------------------------------------+
+    # Note: the first bit is zero by design, so storage layers that treat
+    #       the first bit as a sign bit still sort correctly: big-endian
+    #       encoding preserves the order.
+    h = (level_idx << 57) | key
+
+    return h
+
+
+def recover(h: int) -> tuple[int, int]:
+    """Recover the level id and the original key from a `kh57` encoded key.
+
+    Args:
+        h: the 64-bit encoded key returned by `kh57`.
+
+    Returns:
+        (level_id, key) tuple.
+    """
+    level_idx = h >> 57
+    mask = (1 << 57) - 1
+    key = h & mask
+    return level_idx, key

src/kh57/sampling.pyx

@@ -1,5 +1,108 @@
 # cython: language_level=3
-"""Range reservoir sampling over sorted KV datasets.
+# cython: annotation_typing=False
+# cython: embedsignature=True
+"""Range reservoir sampling over kh57-encoded sorted KV backends.
 
-Stub. Algorithm implementation coming later.
+backends store items under 8-byte big-endian `kh57(key)` encoded keys, so
+physical order is: level 0 slice, level 1 slice, ..., each slice internally
+in original key order. `sample()` walks that physical order, restricted to
+the queried key range on every level, and stops as soon as it has enough.
 """
+
+import random
+
+
+_MASK57 = (1 << 57) - 1
+_KEYSPACE = 1 << 57
+# 7 level bits give 128 possible levels; the current encoding only ever
+# produces 0..63, the rest scan as empty slices.
+_LEVELS = 128
+
+_default_rng = random.Random()
+
+
+def sample(
+    backend,
+    n: int,
+    begin: int | None = None,
+    end: int | None = None,
+    *,
+    rng: random.Random | None = None,
+) -> list:
+    """Pull up to `n` uniformly-sampled (key, value) pairs from a sub-range.
+
+    walks the levels of a kh57-encoded backend in physical (ascending) order:
+    the sparsest levels come first, each next level holds roughly twice as
+    many in-range items. every level whose slice fits in the remaining quota
+    is taken whole; the first level that overflows it gets reservoir-sampled
+    for the remainder, then iteration stops. each level is a deterministic
+    uniform subset of the range, so the union is a uniform sample, total
+    reads stay within ~2x of `n`, and appends outside `[begin, end)` never
+    touch the scanned slices (stability).
+
+    Args:
+        backend: object implementing the `kh57.backends.Backend` protocol,
+            holding values under 8-byte big-endian `kh57(key)` encoded keys.
+        n: number of items to sample. fewer are returned if the range holds
+            fewer than `n` items.
+        begin: inclusive range start as an original (not encoded) int key.
+            None means unbounded from below.
+        end: exclusive range end as an original int key. None means
+            unbounded from above.
+        rng: pass a seeded `random.Random` for deterministic sampling
+            (default: a module-level instance).
+
+    Returns:
+        list of (original_key, value) pairs. Order is unspecified.
+    """
+    if n < 0:
+        raise ValueError("n must be non-negative")
+    if n == 0:
+        return []
+
+    lo_key = 0 if begin is None else begin
+    hi_key = _KEYSPACE if end is None else end
+    if lo_key < 0:
+        raise ValueError("begin must be non-negative")
+    if hi_key > _KEYSPACE:
+        raise ValueError("end must be <= 2**57")
+    if lo_key >= hi_key:
+        return []
+
+    if rng is None:
+        rng = _default_rng
+
+    out = []
+    remaining = n
+    for level in range(_LEVELS):
+        base = level << 57
+        lo = (base + lo_key).to_bytes(8, "big")
+        hi_int = base + hi_key
+        if hi_int < (1 << 64):
+            hi = hi_int.to_bytes(8, "big")
+        else:
+            # level 127 with an unbounded end: 2^64 does not fit in 8 bytes,
+            # any 9-byte string with the max 8-byte prefix sorts above all keys
+            hi = b"\xff" * 8 + b"\x00"
+
+        # reservoir sampling, algorithm R: the slices this loop sees are at
+        # most ~2x the quota by construction, so algorithm L's skip-ahead
+        # would buy nothing over R's simplicity.
+        reservoir = []
+        seen = 0
+        for encoded, value in backend.range_scan(lo, hi):
+            item = (int.from_bytes(encoded, "big") & _MASK57, value)
+            if seen < remaining:
+                reservoir.append(item)
+            else:
+                j = rng.randrange(seen + 1)
+                if j < remaining:
+                    reservoir[j] = item
+            seen += 1
+
+        out.extend(reservoir)
+        remaining -= len(reservoir)
+        if remaining == 0:
+            break
+
+    return out

src/kh57/siphash.pxd

@@ -1,2 +1,21 @@
 # cython: language_level=3
-# siphash declarations - stub.
+# siphash-2-4 low-level declarations, cimport-able from other kh57 modules.
+
+from libc.stdint cimport int64_t, uint8_t, uint64_t
+
+
+cdef uint64_t usiphash64(
+    uint64_t data,
+    uint8_t* key
+) noexcept nogil
+
+cdef int64_t isiphash64(
+    int64_t data,
+    uint8_t* key
+) noexcept nogil
+
+cdef uint64_t low_level_siphash(
+    uint8_t* data,
+    size_t datalen,
+    uint8_t* key
+) noexcept nogil