"""Lightweight approximate data structures for population statistics. All numpy-based, mergeable (so they can be maintained per partition and reduced), and each clearly reports whether an answer is exact or approximate. - QuantileHistogram: fixed-bin histogram over a known range; percentile queries by linear interpolation inside the winning bin. Stand-in for a t-digest. - HyperLogLog: classic HLL with 2^p registers and numpy vectorized updates. - Reservoir: uniform reservoir sample of exemplar ids (Algorithm R, vectorized-ish). """ from __future__ import annotations from typing import Dict, List, Optional import numpy as np # ---------------------------------------------------------------- hashing _SPLITMIX_1 = np.uint64(0x9E3779B97F4A7C15) _SPLITMIX_2 = np.uint64(0xBF58476D1CE4E5B9) _SPLITMIX_3 = np.uint64(0x94D049BB133111EB) def _hash64(values: np.ndarray) -> np.ndarray: """Deterministic 64-bit mix (splitmix64) of integer identifiers.""" with np.errstate(over="ignore"): z = values.astype(np.uint64) + _SPLITMIX_1 z = (z ^ (z >> np.uint64(30))) * _SPLITMIX_2 z = (z ^ (z >> np.uint64(27))) * _SPLITMIX_3 return z ^ (z >> np.uint64(31)) # ---------------------------------------------------------------- quantile histogram class QuantileHistogram: """Fixed-bin histogram over [lo, hi]; approximate percentiles via interpolation. Mergeable: histograms with identical bin config add counts element-wise. """ def __init__(self, lo: float = 0.0, hi: float = 1.0, bins: int = 64, counts: Optional[np.ndarray] = None): self.lo, self.hi, self.bins = float(lo), float(hi), int(bins) self.counts = counts if counts is not None else np.zeros(bins, dtype=np.int64) def add(self, values: np.ndarray) -> None: if values.size == 0: return clipped = np.clip(values, self.lo, self.hi) idx = np.minimum( ((clipped - self.lo) / (self.hi - self.lo) * self.bins).astype(np.int64), self.bins - 1) self.counts += np.bincount(idx, minlength=self.bins) def merge(self, other: "QuantileHistogram") -> "QuantileHistogram": assert (self.lo, self.hi, self.bins) == (other.lo, other.hi, other.bins) return QuantileHistogram(self.lo, self.hi, self.bins, self.counts + other.counts) @property def total(self) -> int: return int(self.counts.sum()) def percentile(self, q: float) -> Optional[float]: """Approximate q-th percentile (q in [0, 100]).""" total = self.total if total == 0: return None target = total * (q / 100.0) cum = np.cumsum(self.counts) bin_idx = int(np.searchsorted(cum, target, side="left")) bin_idx = min(bin_idx, self.bins - 1) prev_cum = cum[bin_idx - 1] if bin_idx > 0 else 0 in_bin = self.counts[bin_idx] frac = 0.5 if in_bin == 0 else float(np.clip((target - prev_cum) / in_bin, 0.0, 1.0)) width = (self.hi - self.lo) / self.bins return self.lo + (bin_idx + frac) * width def to_dict(self) -> Dict: return {"lo": self.lo, "hi": self.hi, "bins": self.bins, "counts": self.counts.tolist(), "exact": False} @staticmethod def from_arrays(lo: float, hi: float, counts: np.ndarray) -> "QuantileHistogram": return QuantileHistogram(lo, hi, len(counts), counts.astype(np.int64)) # ---------------------------------------------------------------- HyperLogLog class HyperLogLog: """HLL cardinality estimator with 2^p registers (approximate). Mergeable: element-wise max of registers. """ def __init__(self, p: int = 12, registers: Optional[np.ndarray] = None): self.p = int(p) self.m = 1 << self.p self.registers = registers if registers is not None else np.zeros(self.m, dtype=np.uint8) def add_ids(self, ids: np.ndarray) -> None: if ids.size == 0: return h = _hash64(np.asarray(ids)) reg_idx = (h >> np.uint64(64 - self.p)).astype(np.int64) rest = (h << np.uint64(self.p)) | np.uint64((1 << self.p) - 1) # rank = leading zeros of the remaining bits + 1 ranks = np.zeros(ids.shape, dtype=np.uint8) remaining = rest.copy() # 64-bit leading-zero count via float log2 is unsafe; loop over 64 bits max via vector ops lz = np.full(ids.shape, 64, dtype=np.int64) nonzero = remaining != 0 # bit_length via log2 on uint64 split into hi/lo 32-bit halves (exact for ints < 2^53 per half) hi = (remaining >> np.uint64(32)).astype(np.float64) lo = (remaining & np.uint64(0xFFFFFFFF)).astype(np.float64) bl = np.zeros(ids.shape, dtype=np.int64) hi_nz = hi > 0 bl[hi_nz] = 32 + np.floor(np.log2(hi[hi_nz])).astype(np.int64) + 1 lo_only = (~hi_nz) & (lo > 0) bl[lo_only] = np.floor(np.log2(lo[lo_only])).astype(np.int64) + 1 lz[nonzero] = 64 - bl[nonzero] ranks = np.minimum(lz + 1, 64 - self.p + 1).astype(np.uint8) np.maximum.at(self.registers, reg_idx, ranks) def merge(self, other: "HyperLogLog") -> "HyperLogLog": assert self.p == other.p return HyperLogLog(self.p, np.maximum(self.registers, other.registers)) def estimate(self) -> float: m = float(self.m) alpha = 0.7213 / (1.0 + 1.079 / m) inv_sum = np.sum(np.exp2(-self.registers.astype(np.float64))) raw = alpha * m * m / inv_sum zeros = int(np.sum(self.registers == 0)) if raw <= 2.5 * m and zeros > 0: # small-range correction return m * np.log(m / zeros) return float(raw) # ---------------------------------------------------------------- reservoir sample class Reservoir: """Uniform reservoir sample of ids (exemplars for drill-down).""" def __init__(self, k: int = 200, seed: int = 0): self.k = int(k) self.rng = np.random.default_rng(seed) self.items: List[int] = [] self.seen = 0 def add_ids(self, ids: np.ndarray) -> None: for v in np.asarray(ids).tolist(): self.seen += 1 if len(self.items) < self.k: self.items.append(int(v)) else: j = int(self.rng.integers(0, self.seen)) if j < self.k: self.items[j] = int(v) def merge(self, other: "Reservoir") -> "Reservoir": merged = Reservoir(self.k, seed=self.seen + other.seen + 1) pool = self.items + other.items merged.seen = self.seen + other.seen if len(pool) <= self.k: merged.items = pool else: idx = merged.rng.choice(len(pool), size=self.k, replace=False) merged.items = [pool[i] for i in idx] return merged