103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
"""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()))
|