perf_test/
perf_test.rs

1use diskann_rs::{DiskAnnError, DistanceMetric, SingleFileDiskANN};
2use rand::prelude::*;
3use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
4use std::sync::Arc;
5use std::time::Instant;
6
7fn main() -> Result<(), DiskAnnError> {
8    const NUM_VECTORS: usize = 1_000_000;
9    const DIM: usize = 1536;
10    const MAX_DEGREE: usize = 32;
11    const FRACTION_TOP: f64 = 0.01;
12    const FRACTION_MID: f64 = 0.1;
13    let distance_metric = DistanceMetric::Cosine;
14
15    let singlefile_path = "diskann_parallel.db";
16
17    // Build if missing
18    if !std::path::Path::new(singlefile_path).exists() {
19        println!(
20            "Building single-file index with parallel adjacency + distance={:?}",
21            distance_metric
22        );
23        let start = Instant::now();
24        let _index = SingleFileDiskANN::build_index_singlefile(
25            NUM_VECTORS,
26            DIM,
27            MAX_DEGREE,
28            FRACTION_TOP,
29            FRACTION_MID,
30            distance_metric,
31            singlefile_path,
32        )?;
33        let elapsed = start.elapsed().as_secs_f32();
34        println!("Done building index in {:.2} s", elapsed);
35    } else {
36        println!(
37            "Index file {} already exists, skipping build.",
38            singlefile_path
39        );
40    }
41
42    // open
43    let open_start = Instant::now();
44    let index = Arc::new(SingleFileDiskANN::open_index_singlefile(singlefile_path)?);
45    let open_time = open_start.elapsed().as_secs_f32();
46    println!(
47        "Opened index with {} vectors, dim={}, metric={:?} in {:.2} s",
48        index.num_vectors, index.dim, index.distance_metric, open_time
49    );
50
51    // Create queries
52    let queries = 5;
53    let k = 10;
54    let beam_width = 64;
55
56    // Generate all queries in a batch
57    let mut rng = rand::thread_rng();
58    let mut query_batch: Vec<Vec<f32>> = Vec::with_capacity(queries);
59    for _ in 0..queries {
60        let q: Vec<f32> = (0..index.dim).map(|_| rng.gen()).collect();
61        query_batch.push(q);
62    }
63
64    // Now run queries in parallel
65    let search_start = Instant::now();
66    query_batch.par_iter().enumerate().for_each(|(i, query)| {
67        let neighbors = index.search(query, k, beam_width);
68        println!("Query {i} => top-{k} neighbors = {:?}", neighbors);
69    });
70    let search_time = search_start.elapsed().as_secs_f32();
71    println!("Performed {queries} queries in {:.2} s", search_time);
72
73    Ok(())
74}