demo/
demo.rs

1// examples/demo.rs
2use 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    // Build if missing
16    if !std::path::Path::new(singlefile_path).exists() {
17        println!("Building DiskANN index at {singlefile_path}...");
18        
19        // Generate sample vectors (in real usage, you'd load your own data)
20        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    // Open the index
42    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    // Perform a sample query
49    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}