beam: add jump consistent hash primitive with tests

wizard-work/beam/jump_consistent_hash.py

@@ -0,0 +1,74 @@
+"""Jump consistent hash (Lamping & Veach, 2014).
+
+Maps a 64-bit key to a bucket in [0, num_buckets) with minimal
+rebalancing when the bucket count changes: only ~1/(k+1) of keys move
+when going from k to k+1 buckets.
+
+The reference algorithm uses 64-bit wrapping arithmetic and a double
+for the jump computation. This implementation reproduces it exactly
+using Python ints masked to 64 bits, so results match the C reference.
+"""
+from __future__ import annotations
+
+import hashlib
+import struct
+from typing import Any, Hashable
+
+
+_MASK64 = (1 << 64) - 1
+_MULT = 2862933555777941757
+_SCALE = 1 << 31  # 2^31
+
+
+def _to_u64(key: Any) -> int:
+    if isinstance(key, int):
+        return key & _MASK64
+    if isinstance(key, bytes):
+        raw = key
+    else:
+        raw = str(key).encode("utf-8")
+    # md5 -> 128 bits -> fold to 64
+    digest = hashlib.md5(raw).digest()
+    (hi, lo) = struct.unpack(">QQ", digest)
+    return (hi ^ lo) & _MASK64
+
+
+def jump_hash(key: Any, num_buckets: int) -> int:
+    """Return bucket index in [0, num_buckets) for key."""
+    if not isinstance(num_buckets, int) or isinstance(num_buckets, bool):
+        raise TypeError("num_buckets must be an int")
+    if num_buckets < 0:
+        raise ValueError("num_buckets must be >= 0")
+    if num_buckets == 0:
+        raise ValueError("num_buckets must be >= 1")
+    if num_buckets == 1:
+        return 0
+
+    k = _to_u64(key)
+    b = -1
+    j = 0
+    while j < num_buckets:
+        b = j
+        k = (k * _MULT + 1) & _MASK64
+        # j = (b+1) * 2^31 / ((key >> 33) + 1)
+        denom = (k >> 33) + 1
+        j = (b + 1) * _SCALE // denom
+    return b
+
+
+class JumpConsistentHash:
+    """Convenience wrapper keeping a bucket count."""
+
+    def __init__(self, num_buckets: int):
+        if num_buckets < 1:
+            raise ValueError("num_buckets must be >= 1")
+        self.num_buckets = num_buckets
+
+    def bucket(self, key: Any) -> int:
+        return jump_hash(key, self.num_buckets)
+
+    def resize(self, new_num_buckets: int) -> "JumpConsistentHash":
+        if new_num_buckets < 1:
+            raise ValueError("num_buckets must be >= 1")
+        self.num_buckets = new_num_buckets
+        return self

wizard-work/beam/test_jump_consistent_hash.py

@@ -0,0 +1,60 @@
+import random
+import unittest
+
+from jump_consistent_hash import jump_hash, JumpConsistentHash
+
+
+class TestJumpHash(unittest.TestCase):
+    def test_range(self):
+        for n in (1, 2, 5, 10, 100, 1000):
+            for i in range(200):
+                b = jump_hash(f"k{i}", n)
+                self.assertGreaterEqual(b, 0)
+                self.assertLess(b, n)
+
+    def test_single_bucket(self):
+        for i in range(50):
+            self.assertEqual(jump_hash(i, 1), 0)
+
+    def test_deterministic(self):
+        self.assertEqual(jump_hash("abc", 7), jump_hash("abc", 7))
+
+    def test_invalid(self):
+        with self.assertRaises(ValueError):
+            jump_hash("x", 0)
+        with self.assertRaises(ValueError):
+            jump_hash("x", -1)
+        with self.assertRaises(TypeError):
+            jump_hash("x", 1.5)
+
+    def test_minimal_movement(self):
+        # Moving from k to k+1 buckets: only ~1/(k+1) keys should move.
+        keys = [f"key-{i}" for i in range(10000)]
+        for k in (4, 8, 16, 32):
+            before = {key: jump_hash(key, k) for key in keys}
+            after = {key: jump_hash(key, k + 1) for key in keys}
+            moved = sum(1 for key in keys if before[key] != after[key])
+            # expected ~ 10000/(k+1); allow generous slack
+            expected = 10000 / (k + 1)
+            self.assertLess(moved, expected * 2.5, f"k={k} moved={moved}")
+
+    def test_uniformity(self):
+        n = 10
+        counts = [0] * n
+        for i in range(100000):
+            counts[jump_hash(i, n)] += 1
+        # each bucket within 5% of mean
+        mean = 100000 / n
+        for c in counts:
+            self.assertGreater(c, mean * 0.9)
+            self.assertLess(c, mean * 1.1)
+
+    def test_wrapper(self):
+        jch = JumpConsistentHash(5)
+        self.assertEqual(jch.bucket("a"), jump_hash("a", 5))
+        jch.resize(7)
+        self.assertEqual(jch.num_buckets, 7)
+
+
+if __name__ == "__main__":
+    unittest.main()