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 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 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 let queries = 5;
53 let k = 10;
54 let beam_width = 64;
55
56 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 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}