1use diskann_rs::{DiskANN, DiskAnnError, DistanceMetric};
3use rand::prelude::*;
4use std::sync::Arc;
5
6fn main() -> Result<(), DiskAnnError> {
7 let singlefile_path = "diskann.db";
8 let num_vectors = 100_000;
9 let dim = 128;
10 let max_degree = 32;
11 let build_beam_width = 128;
12 let alpha = 1.2;
13 let distance_metric = DistanceMetric::Cosine;
14
15 if !std::path::Path::new(singlefile_path).exists() {
17 println!("Building DiskANN index at {singlefile_path}...");
18
19 println!("Generating {} sample vectors of dimension {}...", num_vectors, dim);
21 let mut rng = thread_rng();
22 let mut vectors = Vec::new();
23 for _ in 0..num_vectors {
24 let v: Vec<f32> = (0..dim).map(|_| rng.gen()).collect();
25 vectors.push(v);
26 }
27
28 let index = DiskANN::build_index(
29 &vectors,
30 max_degree,
31 build_beam_width,
32 alpha,
33 distance_metric,
34 singlefile_path,
35 )?;
36 println!("Build done. Index contains {} vectors", index.num_vectors);
37 } else {
38 println!("Index file {singlefile_path} already exists, skipping build.");
39 }
40
41 let index = Arc::new(DiskANN::open_index(singlefile_path)?);
43 println!(
44 "Opened index: {} vectors, dimension={}, max_degree={}",
45 index.num_vectors, index.dim, index.max_degree
46 );
47
48 let mut rng = thread_rng();
50 let query: Vec<f32> = (0..index.dim).map(|_| rng.gen()).collect();
51 let k = 10;
52 let search_beam_width = 64;
53
54 println!("\nSearching for {} nearest neighbors with beam_width={}...", k, search_beam_width);
55 let start = std::time::Instant::now();
56 let neighbors = index.search(&query, k, search_beam_width);
57 let elapsed = start.elapsed();
58
59 println!("Search completed in {:?}", elapsed);
60 println!("Found {} neighbors:", neighbors.len());
61 for (i, &id) in neighbors.iter().enumerate() {
62 println!(" {}: node {}", i + 1, id);
63 }
64
65 Ok(())
66}