apache-datasketches 0.2.3

//! Throughput harness for `ThetaSketch::update_u64`.
//!
//! Run with (release matters -- a debug build measures nothing useful):
//!   cargo run --release --example bench_theta_update --features theta
//!   cargo run --release --example bench_theta_update --features theta -- 100000000
//!   cargo run --release --example bench_theta_update --features theta -- --ladder
//!
//! Accepts `[ITEMS] [--reps N] [--ladder]`. Every figure printed is the lower
//! median of `--reps` passes (default 3), with the spread alongside it, so a
//! single noisy pass cannot become a published number. `--ladder` sweeps a
//! range of item counts instead of one, because a family's per-update cost is
//! not constant as the sketch fills.
//!
//! Pair with the native C++ counterpart on the same item count to get the
//! binding's overhead, which is the number worth tracking:
//!
//!   ./benches/cpp_reference/run.sh 10000000
//!
//! That harness takes the same arguments and prints the same format. Keep the
//! two in step -- in particular the output format, the median definition and
//! the string-key format; the mechanism lives in its `bench_common.h`.
//!
//! Fixed parameters so runs are comparable: `lg_k = 12`, defaults elsewhere.
//! Keep them in sync with `benches/cpp_reference/theta_update.cc`; both print
//! the estimate, and the estimates must match.
//!
//! Theta *does* have a theta screen, so as with the Tuple harnesses `distinct`
//! and `hot` exercise different halves of the update: once theta drops, most
//! distinct keys are rejected early, while hot keys always reach the table.
//!
//! `str` covers the string update path, which crosses the boundary as a
//! borrowed `(pointer, length)` pair rather than an integer. Its C++
//! counterpart calls the same `(data, length)` overload the shim does, so the
//! difference between them is binding overhead and not a choice of overload.

use apache_datasketches::theta::{
    jaccard_similarity, CompactThetaSketch, ThetaIntersection, ThetaSketch, ThetaSketchBuilder,
    ThetaUnionBuilder,
};
use std::hint::black_box;
use std::time::{Duration, Instant};

const LG_K: u8 = 12;
const HOT_KEY_SPACE: u64 = 1 << 10;

/// The `ser` and `deser` scenarios are the one place where the item count is
/// not the divisor: serialization cost tracks the serialized *size*, and at
/// this harness's `lg_k` the sketch saturates well below the ladder's bottom
/// rung, so the same buffer is produced at 1M items as at 100M. The printed
/// `ns/op` is therefore per serialize call, over a call count fixed here
/// rather than taken from the command line -- otherwise the number would
/// silently mean something different at each rung.
///
/// Keep in step with `bench_common.h`.
const SER_CALLS: u64 = 20_000;
const DESER_CALLS: u64 = 5_000;

/// Union, intersection and Jaccard cost tracks the retained-entries table (at
/// most `2^lg_k` entries), not the item count -- like `ser`/`deser`, the same
/// operand sketches produce the same cost at every ladder rung, so this is a
/// fixed call count rather than one taken from the command line.
///
/// Keep in step with `bench_common.h`.
const OP_CALLS: u64 = 5_000;

/// Size of the pre-built string-key pool. See `string_keys`.
const STR_KEY_SPACE: u64 = 1 << 16;

/// Built once, outside every timed region: formatting a key costs more than
/// the update does, and it costs a different amount in each language, so
/// including it would swamp the per-call delta this harness exists to show.
/// Keep the format identical to the C++ counterpart or the estimates diverge.
fn string_keys() -> Vec<String> {
    (0..STR_KEY_SPACE).map(|i| format!("key_{i:010}")).collect()
}

/// Item counts for `--ladder`, which exists because a single item count hides
/// the shape: a family's per-update cost is not constant as the sketch fills.
///
/// Starts at 1M rather than lower. Below that the cheap families are still in
/// a warm-up regime -- HLL's coupon list, CPC's flavour transitions -- so the
/// printed ns/op would be an average taken across a regime change rather than
/// a steady-state cost, which is precisely the kind of number the ladder
/// exists to stop people quoting.
const LADDER: [u64; 3] = [1_000_000, 10_000_000, 100_000_000];
const DEFAULT_ITEMS: u64 = 10_000_000;
const DEFAULT_REPS: usize = 3;

/// Parses `[ITEMS] [--reps N] [--ladder]`. Hand-rolled: three flags do not
/// justify pulling an argument crate into a bench example.
fn parse_args() -> (Vec<u64>, usize) {
    let mut items = None;
    let mut reps = DEFAULT_REPS;
    let mut ladder = false;
    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--ladder" => ladder = true,
            "--reps" => {
                reps = args
                    .next()
                    .and_then(|v| v.parse().ok())
                    .filter(|&n| n > 0)
                    .expect("--reps needs a positive integer")
            }
            other => {
                let n = other
                    .parse()
                    .expect("item count must be a positive integer");
                assert!(n > 0, "item count must be a positive integer");
                items = Some(n);
            }
        }
    }
    // Rejected rather than resolved by precedence: silently ignoring an
    // explicit item count would make a mis-typed invocation look like it
    // measured what was asked for.
    assert!(
        !(ladder && items.is_some()),
        "pass an item count or --ladder, not both"
    );
    let counts = if ladder {
        LADDER.to_vec()
    } else {
        vec![items.unwrap_or(DEFAULT_ITEMS)]
    };
    (counts, reps)
}

/// Prints the lower median of the passes plus the spread, so a published
/// figure is never a single noisy point -- the AGENTS.md rule that a
/// performance claim rest on a median of at least three runs is enforced here
/// rather than left to whoever happens to be running it.
///
/// Lower median (`sorted[(n - 1) / 2]`), not the average of the two middle
/// values: every number printed is then one that an actual pass produced. At
/// the default `reps = 3` the two definitions agree; this only matters for an
/// even `--reps`.
///
/// The estimates are asserted equal across reps rather than merely reported.
/// These workloads are deterministic, so a disagreement means the reps are not
/// running the same thing -- most likely a sketch reused across reps instead
/// of rebuilt, which would quietly lower the ns/op of every rep after the
/// first.
///
/// `ns/op`, `reps` and `estimate` are printed as labelled values rather than
/// as bare numbers in fixed columns, so that reading them back does not mean
/// counting awk fields that shift whenever a column is added.
fn report(label: &str, items: u64, passes: &[Pass]) {
    report_line(label, items, items, passes, String::new());
}

/// As [`report`], plus the serialized size, and dividing by an explicit `ops`
/// rather than by the item count.
///
/// The size is worth printing for its own sake -- it is the quantity a `ser`
/// or `deser` `ns/op` is proportional to, so without it the timing cannot be
/// interpreted -- but it is also a check the estimate cannot make. Two sides
/// can agree exactly on the estimate while one compacts ordered and the other
/// does not, or serializes a different format; the byte count differs the
/// moment they do.
fn report_bytes(label: &str, items: u64, ops: u64, passes: &[Pass], bytes: usize) {
    report_line(label, items, ops, passes, format!(" bytes={bytes}"));
}

fn report_line(label: &str, items: u64, ops: u64, passes: &[Pass], suffix: String) {
    for (i, pass) in passes.iter().enumerate() {
        assert_eq!(
            pass.estimate, passes[0].estimate,
            "rep {i} estimated {} but rep 0 estimated {}: the reps are not running \
             the same workload",
            pass.estimate, passes[0].estimate
        );
    }
    let mut ns_per_op: Vec<f64> = passes
        .iter()
        .map(|p| p.elapsed.as_secs_f64() * 1e9 / ops as f64)
        .collect();
    ns_per_op.sort_by(f64::total_cmp);
    let median = ns_per_op[(ns_per_op.len() - 1) / 2];
    let (min, max) = (ns_per_op[0], ns_per_op[ns_per_op.len() - 1]);
    let rate = 1000.0 / median;
    let (reps, estimate) = (passes.len(), passes[0].estimate);
    println!(
        "{label:9} {items:>12} items  {median:>7.2} ns/op  min {min:>7.2}  max {max:>7.2}  \
         {rate:>8.1} M/s  reps={reps} estimate={estimate:.0}{suffix}"
    );
}

/// As [`report`], but for Jaccard: the result is a confidence interval in
/// `[0.0, 1.0]`, not a scale-free count, so `report`'s fixed `{:.0}` precision
/// would round every printed value to 0 and make the parity check vacuous.
/// Nine decimal digits instead -- both sides call into the exact same
/// vendored jaccard implementation, so the bits already agree and only the
/// print format needs to.
fn report_jaccard(label: &str, items: u64, ops: u64, passes: &[JaccardPass]) {
    for (i, pass) in passes.iter().enumerate() {
        assert_eq!(
            (pass.lower_bound, pass.estimate, pass.upper_bound),
            (
                passes[0].lower_bound,
                passes[0].estimate,
                passes[0].upper_bound
            ),
            "rep {i} did not reproduce rep 0's bounds: the reps are not running the same workload"
        );
    }
    let mut ns_per_op: Vec<f64> = passes
        .iter()
        .map(|p| p.elapsed.as_secs_f64() * 1e9 / ops as f64)
        .collect();
    ns_per_op.sort_by(f64::total_cmp);
    let median = ns_per_op[(ns_per_op.len() - 1) / 2];
    let (min, max) = (ns_per_op[0], ns_per_op[ns_per_op.len() - 1]);
    let rate = 1000.0 / median;
    let reps = passes.len();
    let (lower_bound, estimate, upper_bound) = (
        passes[0].lower_bound,
        passes[0].estimate,
        passes[0].upper_bound,
    );
    println!(
        "{label:9} {items:>12} items  {median:>7.2} ns/op  min {min:>7.2}  max {max:>7.2}  \
         {rate:>8.1} M/s  reps={reps} lower={lower_bound:.9} estimate={estimate:.9} \
         upper={upper_bound:.9}"
    );
}

/// One timed pass over `items` updates, and the estimate the sketch held
/// afterwards. Reading the estimate also keeps the update loop from being
/// optimised out.
struct Pass {
    elapsed: Duration,
    estimate: f64,
}

/// As [`Pass`], but for the three-field Jaccard result.
struct JaccardPass {
    elapsed: Duration,
    lower_bound: f64,
    estimate: f64,
    upper_bound: f64,
}

fn build() -> ThetaSketch {
    ThetaSketchBuilder::new()
        .lg_k(LG_K)
        .build()
        .expect("fixed valid parameters were rejected")
}

/// Each rep rebuilds the sketch: a reused one would already be full, so every
/// rep after the first would measure a different workload. `report` asserts
/// the per-rep estimates agree, which is what catches that if it regresses.
fn bench_distinct(items: u64, reps: usize) {
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let mut sketch = build();
        let start = Instant::now();
        for key in 0..items {
            sketch.update_u64(key);
        }
        let elapsed = start.elapsed();
        passes.push(Pass {
            elapsed,
            estimate: sketch.get_estimate(),
        });
    }
    report("distinct", items, &passes);
}

fn bench_hot(items: u64, reps: usize) {
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let mut sketch = build();
        let start = Instant::now();
        for i in 0..items {
            sketch.update_u64(i % HOT_KEY_SPACE);
        }
        let elapsed = start.elapsed();
        passes.push(Pass {
            elapsed,
            estimate: sketch.get_estimate(),
        });
    }
    report("hot", items, &passes);
}

fn bench_str(items: u64, reps: usize) {
    let keys = string_keys();
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let mut sketch = build();
        let start = Instant::now();
        for i in 0..items {
            sketch.update_str(&keys[(i % STR_KEY_SPACE) as usize]);
        }
        let elapsed = start.elapsed();
        passes.push(Pass {
            elapsed,
            estimate: sketch.get_estimate(),
        });
    }
    report("str", items, &passes);
}

/// Serialization, measured per call rather than per item: its cost tracks the
/// serialized size, which at `lg_k = 12` is the same at every ladder rung.
///
/// The sketch is built once and shared by both directions and every rep.
/// Serializing does not mutate it, so unlike the update scenarios there is no
/// state that a second rep would find already dirtied -- and rebuilding at the
/// 100M rung would cost more than the measurement itself.
fn bench_serde(items: u64, reps: usize) {
    let mut update_sketch = build();
    for key in 0..items {
        update_sketch.update_u64(key);
    }
    let sketch = update_sketch.compact(true);
    let reference = sketch.serialize_compact();

    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let start = Instant::now();
        let mut total = 0usize;
        for _ in 0..SER_CALLS {
            total += black_box(sketch.serialize_compact()).len();
        }
        let elapsed = start.elapsed();
        black_box(total);
        passes.push(Pass {
            elapsed,
            estimate: sketch.get_estimate(),
        });
    }
    report_bytes("ser", items, SER_CALLS, &passes, reference.len());

    let deserialize =
        || CompactThetaSketch::deserialize(&reference).expect("the bytes came from serialize");
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let start = Instant::now();
        let mut total = 0.0;
        for _ in 0..DESER_CALLS {
            total += deserialize().get_estimate();
        }
        let elapsed = start.elapsed();
        black_box(total);
        passes.push(Pass {
            elapsed,
            estimate: deserialize().get_estimate(),
        });
    }
    report_bytes("deser", items, DESER_CALLS, &passes, reference.len());
}

/// Two operands with 50% overlap, built once outside every timed region: the
/// operand-construction cost belongs to the setup, not to the union/
/// intersection/jaccard call being measured.
fn build_operands(items: u64) -> (CompactThetaSketch, CompactThetaSketch) {
    let mut a = build();
    for key in 0..items {
        a.update_u64(key);
    }
    let mut b = build();
    for key in (items / 2)..(items + items / 2) {
        b.update_u64(key);
    }
    (a.compact(true), b.compact(true))
}

/// A fresh union is built inside the timed loop, so the figure is
/// construct + two updates + get_result, not the merge alone -- reusing one
/// accumulator across `OP_CALLS` iterations would have each iteration merge
/// into an ever-growing result, measuring a different workload every time.
fn bench_union(items: u64, reps: usize) {
    let (a, b) = build_operands(items);
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let start = Instant::now();
        let mut total = 0.0;
        let mut estimate = 0.0;
        for _ in 0..OP_CALLS {
            let mut union = ThetaUnionBuilder::new()
                .lg_k(LG_K)
                .build()
                .expect("fixed valid parameters were rejected");
            union.update(&a);
            union.update(&b);
            estimate = union.get_result(true).get_estimate();
            total += estimate;
        }
        let elapsed = start.elapsed();
        black_box(total);
        passes.push(Pass { elapsed, estimate });
    }
    report_line("union", items, OP_CALLS, &passes, String::new());
}

/// As [`bench_union`]: a fresh intersection per iteration, so the figure is
/// construct + two updates + get_result.
fn bench_intersect(items: u64, reps: usize) {
    let (a, b) = build_operands(items);
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let start = Instant::now();
        let mut total = 0.0;
        let mut estimate = 0.0;
        for _ in 0..OP_CALLS {
            let mut intersection = ThetaIntersection::new();
            intersection.update(&a);
            intersection.update(&b);
            estimate = intersection
                .get_result(true)
                .expect("both operands were non-empty")
                .get_estimate();
            total += estimate;
        }
        let elapsed = start.elapsed();
        black_box(total);
        passes.push(Pass { elapsed, estimate });
    }
    report_line("intersect", items, OP_CALLS, &passes, String::new());
}

/// `jaccard_similarity` is a pure function of its two operands -- no
/// accumulator to rebuild, unlike union and intersection.
fn bench_jaccard(items: u64, reps: usize) {
    let (a, b) = build_operands(items);
    let mut passes = Vec::with_capacity(reps);
    for _ in 0..reps {
        let start = Instant::now();
        let mut total = 0.0;
        let mut bounds = jaccard_similarity(&a, &b);
        for _ in 0..OP_CALLS {
            bounds = jaccard_similarity(&a, &b);
            total += bounds.estimate;
        }
        let elapsed = start.elapsed();
        black_box(total);
        passes.push(JaccardPass {
            elapsed,
            lower_bound: bounds.lower_bound,
            estimate: bounds.estimate,
            upper_bound: bounds.upper_bound,
        });
    }
    report_jaccard("jaccard", items, OP_CALLS, &passes);
}

fn main() {
    let (counts, reps) = parse_args();
    for items in counts {
        println!("lg_k={LG_K} items={items} reps={reps}");
        bench_distinct(items, reps);
        bench_hot(items, reps);
        bench_str(items, reps);
        bench_serde(items, reps);
        bench_union(items, reps);
        bench_intersect(items, reps);
        bench_jaccard(items, reps);
    }
}