search_bench/
search_bench.rs

1//! Search latency on a saved index, without Python.
2//!
3//! cargo run --release --example search_bench -- <index.rvec> <queries.f32> <truth.u32> [ef,...]
4//!
5//! `queries.f32` holds row-major float32 query vectors; `truth.u32` the ids
6//! of the 10 true nearest neighbors of each query (record ids are their row
7//! numbers as strings, as written by bench/ann_bench.py).
8
9use std::time::Instant;
10
11use recern_vector::{Database, SearchOptions};
12
13fn read<T: Copy>(path: &str, from: fn([u8; 4]) -> T) -> Vec<T> {
14    std::fs::read(path)
15        .expect(path)
16        .as_chunks::<4>()
17        .0
18        .iter()
19        .map(|b| from(*b))
20        .collect()
21}
22
23fn main() {
24    let args: Vec<String> = std::env::args().skip(1).collect();
25    let db = Database::open(&args[0]).expect("index");
26    let c = db.collections().next().expect("a collection");
27    let dim = c.config().dim;
28    let queries = read(&args[1], f32::from_le_bytes);
29    let truth = read(&args[2], u32::from_le_bytes);
30    let efs: Vec<usize> = args
31        .get(3)
32        .map_or("10,20,40,80,160", String::as_str)
33        .split(',')
34        .map(|e| e.parse().unwrap())
35        .collect();
36    let n = queries.len() / dim;
37
38    println!(
39        "{:>5}  {:>8}  {:>9}  {:>9}  {:>10}",
40        "ef", "recall", "mean µs", "p50 µs", "distances"
41    );
42    for ef in efs {
43        let options = SearchOptions::default().ef(ef);
44        let (mut times, mut found, mut distances) = (Vec::with_capacity(n), 0, 0);
45        for (q, t) in queries
46            .chunks_exact(dim)
47            .zip(truth.as_chunks::<10>().0.iter())
48        {
49            let start = Instant::now();
50            let report = c.explain(q, 10, &options).unwrap();
51            times.push(start.elapsed().as_secs_f64() * 1e6);
52            distances += report.distance_computations;
53            found += report
54                .hits
55                .iter()
56                .filter(|h| t.contains(&h.id.parse::<u32>().unwrap()))
57                .count();
58        }
59        let mean = times.iter().sum::<f64>() / n as f64;
60        times.sort_by(f64::total_cmp);
61        println!(
62            "{ef:>5}  {:>8.4}  {mean:>9.1}  {:>9.1}  {:>10}",
63            found as f64 / (n * 10) as f64,
64            times[n / 2],
65            distances / n
66        );
67    }
68}
search_bench/search_bench.rs

search_bench/
search_bench.rs
