use std::hint::black_box;
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use lance_select::{
NullableIndexExprResult, NullableRowAddrMask, NullableRowAddrSet, RowAddrTreeMap,
};
const ROW_COUNTS: &[u64] = &[10_000, 100_000, 1_000_000, 10_000_000];
fn allow_run(n: u64) -> NullableRowAddrMask {
let mut tree = RowAddrTreeMap::new();
tree.insert_range(0..n);
NullableRowAddrMask::AllowList(NullableRowAddrSet::new(tree, RowAddrTreeMap::new()))
}
fn allow_middle(n: u64) -> NullableRowAddrMask {
let mut tree = RowAddrTreeMap::new();
tree.insert_range((n / 4)..(3 * n / 4));
NullableRowAddrMask::AllowList(NullableRowAddrSet::new(tree, RowAddrTreeMap::new()))
}
fn bench_not(c: &mut Criterion) {
let mut group = c.benchmark_group("not");
for &n in ROW_COUNTS {
group.throughput(Throughput::Elements(n));
for label in ["Exact", "AtMost", "AtLeast"] {
let id = BenchmarkId::new(label, n);
let make = move || match label {
"Exact" => NullableIndexExprResult::exact(allow_run(n)),
"AtMost" => NullableIndexExprResult::at_most(allow_run(n)),
"AtLeast" => NullableIndexExprResult::at_least(allow_run(n)),
_ => unreachable!(),
};
group.bench_function(id, |b| {
b.iter_batched(make, |r| black_box(!r), criterion::BatchSize::SmallInput);
});
}
}
group.finish();
}
type PairFn = Box<dyn Fn() -> (NullableIndexExprResult, NullableIndexExprResult)>;
fn pair_cases(n: u64) -> Vec<(&'static str, PairFn)> {
let lhs = move || allow_run(n);
let rhs = move || allow_middle(n);
let exact = |m| NullableIndexExprResult::exact(m);
let at_most = |m| NullableIndexExprResult::at_most(m);
let at_least = |m| NullableIndexExprResult::at_least(m);
vec![
(
"Exact_Exact",
Box::new(move || (exact(lhs()), exact(rhs()))),
),
(
"AtMost_AtMost",
Box::new(move || (at_most(lhs()), at_most(rhs()))),
),
(
"AtLeast_AtLeast",
Box::new(move || (at_least(lhs()), at_least(rhs()))),
),
(
"Exact_AtMost",
Box::new(move || (exact(lhs()), at_most(rhs()))),
),
(
"Exact_AtLeast",
Box::new(move || (exact(lhs()), at_least(rhs()))),
),
(
"AtMost_AtLeast",
Box::new(move || (at_most(lhs()), at_least(rhs()))),
),
]
}
fn bench_and(c: &mut Criterion) {
let mut group = c.benchmark_group("and");
for &n in ROW_COUNTS {
group.throughput(Throughput::Elements(n));
for (label, make) in pair_cases(n) {
let id = BenchmarkId::new(label, n);
group.bench_function(id, |b| {
b.iter_batched(
&make,
|(l, r)| black_box(l & r),
criterion::BatchSize::SmallInput,
);
});
}
}
group.finish();
}
fn bench_or(c: &mut Criterion) {
let mut group = c.benchmark_group("or");
for &n in ROW_COUNTS {
group.throughput(Throughput::Elements(n));
for (label, make) in pair_cases(n) {
let id = BenchmarkId::new(label, n);
group.bench_function(id, |b| {
b.iter_batched(
&make,
|(l, r)| black_box(l | r),
criterion::BatchSize::SmallInput,
);
});
}
}
group.finish();
}
criterion_group!(benches, bench_not, bench_and, bench_or);
criterion_main!(benches);