Spaces:
Paused
Paused
| """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) | |
| 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} | |
| 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 | |