use apache_datasketches::tuple::{
array_of_doubles_jaccard_similarity, ArrayOfDoublesIntersection, ArrayOfDoublesSketch,
ArrayOfDoublesSketchBuilder, ArrayOfDoublesUnionBuilder, CompactArrayOfDoublesSketch,
};
use std::hint::black_box;
use std::time::{Duration, Instant};
const LG_K: u8 = 12;
const SER_CALLS: u64 = 20_000;
const DESER_CALLS: u64 = 5_000;
const OP_CALLS: u64 = 5_000;
const NUM_VALUES: u8 = 3;
const HOT_KEY_SPACE: u64 = 1 << 10;
const VALUES: [f64; 3] = [1.0, 2.0, 3.0];
const STR_KEY_SPACE: u64 = 1 << 16;
fn string_keys() -> Vec<String> {
(0..STR_KEY_SPACE).map(|i| format!("key_{i:010}")).collect()
}
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;
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);
}
}
}
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)
}
fn report(label: &str, items: u64, passes: &[Pass]) {
report_line(label, items, items, passes, String::new());
}
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}"
);
}
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}"
);
}
struct Pass {
elapsed: Duration,
estimate: f64,
}
struct JaccardPass {
elapsed: Duration,
lower_bound: f64,
estimate: f64,
upper_bound: f64,
}
fn build() -> ArrayOfDoublesSketch {
ArrayOfDoublesSketchBuilder::new()
.lg_k(LG_K)
.num_values(NUM_VALUES)
.build()
.expect("builder rejected fixed valid parameters")
}
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, &VALUES)
.expect("update rejected a correctly-sized value slice");
}
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, &VALUES)
.expect("update rejected a correctly-sized value slice");
}
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], &VALUES)
.expect("update rejected a correctly-sized value slice");
}
let elapsed = start.elapsed();
passes.push(Pass {
elapsed,
estimate: sketch.get_estimate(),
});
}
report("str", items, &passes);
}
fn bench_serde(items: u64, reps: usize) {
let mut update_sketch = build();
for key in 0..items {
update_sketch
.update_u64(key, &VALUES)
.expect("update rejected a correctly-sized value slice");
}
let sketch = update_sketch.compact(true);
let reference = sketch.serialize();
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()).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 = || {
CompactArrayOfDoublesSketch::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());
}
fn build_operands(items: u64) -> (CompactArrayOfDoublesSketch, CompactArrayOfDoublesSketch) {
let mut a = build();
for key in 0..items {
a.update_u64(key, &VALUES)
.expect("update rejected a correctly-sized value slice");
}
let mut b = build();
for key in (items / 2)..(items + items / 2) {
b.update_u64(key, &VALUES)
.expect("update rejected a correctly-sized value slice");
}
(a.compact(true), b.compact(true))
}
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 = ArrayOfDoublesUnionBuilder::new()
.lg_k(LG_K)
.num_values(NUM_VALUES)
.build()
.expect("fixed valid parameters were rejected");
union
.update(&a)
.expect("operands match the union's num_values");
union
.update(&b)
.expect("operands match the union's num_values");
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());
}
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 = ArrayOfDoublesIntersection::new(NUM_VALUES)
.expect("fixed valid num_values was rejected");
intersection
.update(&a)
.expect("operands match the intersection's num_values");
intersection
.update(&b)
.expect("operands match the intersection's num_values");
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());
}
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 =
array_of_doubles_jaccard_similarity(&a, &b).expect("operands agree on num_values");
for _ in 0..OP_CALLS {
bounds =
array_of_doubles_jaccard_similarity(&a, &b).expect("operands agree on num_values");
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} num_values={NUM_VALUES} 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);
}
}