beam: HyperLogLog cardinality estimator with small/large-range corrections + tests

wizard-work/beam/hyperloglog.py

@@ -0,0 +1,102 @@
+"""HyperLogLog cardinality estimator.
+
+Standard HyperLogLog (Flajolet et al., 2007) with the small- and large-range
+corrections described in the original paper. Uses a 64-bit hash so the
+large-range correction path is reachable.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Iterable, Optional
+
+
+def _splitmix64(x: int) -> int:
+    x = (x + 0x9E3779B97F4A7C15) & ((1 << 64) - 1)
+    z = x
+    z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9 & ((1 << 64) - 1)
+    z = (z ^ (z >> 27)) * 0x94D049BB133111B9 & ((1 << 64) - 1)
+    z = z ^ (z >> 31)
+    return z & ((1 << 64) - 1)
+
+
+class HyperLogLog:
+    def __init__(self, p: int = 14):
+        if not (4 <= p <= 16):
+            raise ValueError("p must be in [4, 16]")
+        self.p = p
+        self.m = 1 << p
+        self.mask = self.m - 1
+        # registers store rho values (0 means empty)
+        self.registers = [0] * self.m
+        # alpha_m constant
+        if self.m == 16:
+            self.alpha = 0.673
+        elif self.m == 32:
+            self.alpha = 0.697
+        elif self.m == 64:
+            self.alpha = 0.709
+        else:
+            self.alpha = 0.7213 / (1.0 + 1.079 / self.m)
+
+    def _hash(self, value) -> int:
+        # deterministic 64-bit hash via splitmix on a stable integer key
+        if isinstance(value, int):
+            key = value
+        elif isinstance(value, str):
+            key = 0
+            for ch in value:
+                key = (key * 31 + ord(ch)) & ((1 << 64) - 1)
+        elif isinstance(value, bytes):
+            key = 0
+            for b in value:
+                key = (key * 31 + b) & ((1 << 64) - 1)
+        else:
+            key = hash(value) & ((1 << 64) - 1)
+        return _splitmix64(key)
+
+    def add(self, value) -> None:
+        h = self._hash(value)
+        idx = h & self.mask
+        w = h >> self.p
+        # rho: position of first 1-bit in w, 1-indexed; if w==0, rho = 64 - p + 1
+        if w == 0:
+            rho = (64 - self.p) + 1
+        else:
+            rho = (w & -w).bit_length()  # position of lowest set bit
+        if rho > self.registers[idx]:
+            self.registers[idx] = rho
+
+    def extend(self, values: Iterable) -> None:
+        for v in values:
+            self.add(v)
+
+    def count(self) -> float:
+        # raw estimate
+        inv = 0.0
+        for r in self.registers:
+            inv += 2.0 ** (-r)
+        E = self.alpha * self.m * self.m / inv
+
+        # small range correction
+        zeros = self.registers.count(0)
+        if E <= 2.5 * self.m and zeros > 0:
+            E = self.m * math.log(self.m / zeros)
+
+        # large range correction
+        two64 = 1 << 64
+        if E > (1.0 / 30.0) * two64:
+            E = -two64 * math.log(1.0 - E / two64)
+
+        return E
+
+    def merge(self, other: "HyperLogLog") -> "HyperLogLog":
+        if self.p != other.p:
+            raise ValueError("precision mismatch")
+        out = HyperLogLog(self.p)
+        for i in range(self.m):
+            out.registers[i] = max(self.registers[i], other.registers[i])
+        return out
+
+    def __len__(self) -> int:
+        return int(round(self.count()))

wizard-work/beam/test_hyperloglog.py

@@ -0,0 +1,99 @@
+import math
+import random
+from hyperloglog import HyperLogLog
+
+
+def test_empty():
+    h = HyperLogLog(p=10)
+    assert h.count() == 0.0
+
+
+def test_exact_small():
+    h = HyperLogLog(p=12)
+    n = 1000
+    for i in range(n):
+        h.add(i)
+    est = h.count()
+    assert abs(est - n) / n < 0.05, (est, n)
+
+
+def test_large():
+    h = HyperLogLog(p=14)
+    n = 100_000
+    for i in range(n):
+        h.add(i)
+    est = h.count()
+    rel = abs(est - n) / n
+    assert rel < 0.02, (est, n, rel)
+
+
+def test_duplicates():
+    h = HyperLogLog(p=12)
+    for i in range(5000):
+        h.add(i % 1000)
+    est = h.count()
+    assert abs(est - 1000) / 1000 < 0.05, est
+
+
+def test_strings():
+    h = HyperLogLog(p=12)
+    words = [f"item-{i}" for i in range(2000)]
+    for w in words:
+        h.add(w)
+    est = h.count()
+    assert abs(est - 2000) / 2000 < 0.05, est
+
+
+def test_merge():
+    a = HyperLogLog(p=12)
+    b = HyperLogLog(p=12)
+    for i in range(0, 1000, 2):
+        a.add(i)
+    for i in range(1, 1000, 2):
+        b.add(i)
+    m = a.merge(b)
+    est = m.count()
+    assert abs(est - 1000) / 1000 < 0.05, est
+
+
+def test_merge_precision_mismatch():
+    try:
+        HyperLogLog(p=10).merge(HyperLogLog(p=12))
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError")
+
+
+def test_extend():
+    h = HyperLogLog(p=10)
+    h.extend(range(500))
+    assert abs(h.count() - 500) / 500 < 0.06
+
+
+def test_len():
+    h = HyperLogLog(p=10)
+    for i in range(300):
+        h.add(i)
+    assert isinstance(len(h), int)
+
+
+def test_determinism():
+    a = HyperLogLog(p=12)
+    b = HyperLogLog(p=12)
+    for i in range(1000):
+        a.add(i); b.add(i)
+    assert a.count() == b.count()
+
+
+def test_invalid_p():
+    for p in (3, 17):
+        try:
+            HyperLogLog(p=p)
+        except ValueError:
+            continue
+        raise AssertionError("expected ValueError")
+
+
+if __name__ == "__main__":
+    import pytest
+    raise SystemExit(pytest.main([__file__, "-q"]))