# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "pyroaring"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Roaring against a plain bitmap against a sorted array.

Sets of row ids over a universe of 10^8 rows at densities from 0.001% to
50%: serialized size of a plain bitmap, a sorted int32 array and a Roaring
bitmap. Then the intersection of two 1% sets with each.
"""
import argparse, json, time
import numpy as np
from pyroaring import BitMap

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--universe", type=int, default=100_000_000, help="Rows in the universe.")
p.add_argument("--densities", type=lambda s: [float(v) for v in s.split(",")], default=[0.00001, 0.001, 0.01, 0.1, 0.5], help="Densities.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)
U = args.universe

def ids(d): return np.flatnonzero(rng.random(U) < d).astype(np.uint32)

rows = []
for d in args.densities:
    s = ids(d)
    rows += [{"density_pct": d * 100, "format": "plain bitmap", "mb": U / 8 / 1e6},
             {"density_pct": d * 100, "format": "sorted int32 array", "mb": len(s) * 4 / 1e6},
             {"density_pct": d * 100, "format": "Roaring", "mb": len(BitMap(s).serialize()) / 1e6}]
    print(f"density {d:8.3%}: plain {U / 8 / 1e6:6.2f} MB   sorted array {len(s) * 4 / 1e6:8.3f} MB   Roaring {rows[-1]['mb']:8.3f} MB")

a, b = ids(0.01), ids(0.01)
ra, rb = BitMap(a), BitMap(b)
pa, pb = np.zeros(U, dtype=bool), np.zeros(U, dtype=bool); pa[a] = True; pb[b] = True
pa8, pb8 = np.packbits(pa), np.packbits(pb)
times = []
for name, fn in (("sorted arrays (merge)", lambda: np.intersect1d(a, b, assume_unique=True)),
                 ("plain bitmaps (AND)", lambda: np.bitwise_and(pa8, pb8)),
                 ("Roaring (AND)", lambda: ra & rb)):
    fn(); t = time.perf_counter(); [fn() for _ in range(5)]; dt = (time.perf_counter() - t) / 5
    times.append({"representation": name, "ms": dt * 1000}); print(f"intersect two 1% sets, {name:<22} {dt * 1000:8.2f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"sizes": rows, "intersections": times}}, f)