correct DiskANN algorithm (silly bug with greedy search)

diskann/src/lib.rs

@@ -146,7 +146,14 @@ impl NeighbourBuffer {
         self.scores.truncate(self.size);
         self.visited.truncate(self.size);
 
-        self.next_unvisited = Some(loc as u32);
+        match self.next_unvisited {
+            Some(ref mut next_unvisited) => {
+                *next_unvisited = (loc as u32).min(*next_unvisited);
+            },
+            None => {
+                self.next_unvisited = Some(loc as u32);
+            }
+        }
     }
 
     pub fn clear(&mut self) {
@@ -194,7 +201,6 @@ pub fn greedy_search(scratch: &mut Scratch, start: u32, query: VectorRef, vecs:
     let mut counters = GreedySearchCounters { distances: 0 };
 
     while let Some(pt) = scratch.neighbour_buffer.next_unvisited() {
-        //println!("pt {} {:?}", pt, graph.out_neighbours(pt));
         scratch.neighbour_pre_buffer.clear();
         for &neighbour in graph.out_neighbours(pt).iter() {
             if scratch.visited.insert(neighbour) {
@@ -296,14 +302,12 @@ pub fn build_graph(rng: &mut Rng, graph: &mut IndexGraph, medioid: u32, vecs: &V
         let neighbours = graph.out_neighbours(sigma_i).to_owned();
         for neighbour in neighbours {
             let mut neighbour_neighbours = graph.out_neighbours_mut(neighbour);
-            // To cut down pruning time slightly, allow accumulating more neighbours than usual limit
-            if neighbour_neighbours.len() == config.r_cap {
-                let mut n = neighbour_neighbours.to_vec();
+            if neighbour_neighbours.len() == config.r {
                 scratch.visited_list.clear();
                 merge_existing_neighbours(&mut scratch.visited_list, neighbour, &neighbour_neighbours, vecs, config);
                 merge_existing_neighbours(&mut scratch.visited_list, neighbour, &vec![sigma_i], vecs, config);
-                robust_prune(scratch, neighbour, &mut n, vecs, config);
-            } else if !neighbour_neighbours.contains(&sigma_i) && neighbour_neighbours.len() < config.r_cap {
+                robust_prune(scratch, neighbour, &mut neighbour_neighbours, vecs, config);
+            } else if !neighbour_neighbours.contains(&sigma_i) && neighbour_neighbours.len() < config.r {
                 neighbour_neighbours.push(sigma_i);
             }
         }
@@ -387,3 +391,18 @@ impl Drop for Timer {
         println!("{}: {:.2}s", self.0, self.1.elapsed().as_secs_f32());
     }
 }
+
+pub fn report_degrees(graph: &IndexGraph) {
+    let mut total_degree = 0;
+    let mut degrees = Vec::with_capacity(graph.graph.len());
+    for out_neighbours in graph.graph.iter() {
+        let deg = out_neighbours.read().unwrap().len();
+        total_degree += deg;
+        degrees.push(deg);
+    }
+    degrees.sort_unstable();
+    println!("average degree {}", (total_degree as f64) / (graph.graph.len() as f64));
+    println!("median degree {}", degrees[degrees.len() / 2]);
+    println!("min degree {}", degrees[0]);
+    println!("max degree {}", degrees[degrees.len() - 1]);
+}

diskann/src/main.rs

@@ -7,7 +7,7 @@ use std::{io::Read, time::Instant};
 use anyhow::Result;
 use half::f16;
 
-use diskann::{build_graph, IndexBuildConfig, medioid, IndexGraph, greedy_search, Scratch, vector::{fast_dot, SCALE, dot, VectorList, self}, Timer};
+use diskann::{build_graph, IndexBuildConfig, medioid, IndexGraph, greedy_search, Scratch, vector::{fast_dot, SCALE, dot, VectorList, self}, Timer, report_degrees, random_fill_graph};
 use simsimd::SpatialSimilarity;
 
 const D_EMB: usize = 1152;
@@ -26,12 +26,13 @@ const PQ_TEST_SIZE: usize = 1000;
 fn main() -> Result<()> {
     tracing_subscriber::fmt::init();
 
+    /*/
     {
         let file = std::fs::File::open("opq.msgpack")?;
         let codec: vector::ProductQuantizer = rmp_serde::from_read(file)?;
         let input = load_file("embeddings.bin", Some(D_EMB * PQ_TEST_SIZE))?.data.into_iter().map(|a| a.to_f32()).collect::<Vec<_>>();
         let codes = codec.quantize_batch(&input);
-        println!("{:?}", codes);
+        //println!("{:?}", codes);
         let raw_query = load_file("query.bin", Some(D_EMB))?.data.into_iter().map(|a| a.to_f32()).collect::<Vec<_>>();
         let query = codec.preprocess_query(&raw_query);
         let mut real_scores = vec![];
@@ -41,17 +42,17 @@ fn main() -> Result<()> {
         let pq_scores = codec.asymmetric_dot_product(&query, &codes);
         for (x, y) in real_scores.iter().zip(pq_scores.iter()) {
             let y = (*y as f32) / SCALE;
-            println!("{} {} {} {}", x, y, x - y, (x - y) / x);
+            //println!("{} {} {} {}", x, y, x - y, (x - y) / x);
         }
-    }
+    }*/
 
     let mut rng = fastrand::Rng::with_seed(1);
 
-    let n = 100000;
+    let n = 100_000;
     let vecs = {
         let _timer = Timer::new("loaded vectors");
 
-        &load_file("embeddings.bin", Some(D_EMB * n))?
+        &load_file("query.bin", Some(D_EMB * n))?
     };
 
     let (graph, medioid) = {
@@ -59,10 +60,10 @@ fn main() -> Result<()> {
 
         let mut config = IndexBuildConfig {
             r: 64,
-            r_cap: 80,
-            l: 128,
+            r_cap: 64,
+            l: 192,
             maxc: 750,
-            alpha: 65536,
+            alpha: 65200,
         };
 
         let mut graph = IndexGraph::random_r_regular(&mut rng, vecs.len(), config.r, config.r_cap);
@@ -70,8 +71,11 @@ fn main() -> Result<()> {
         let medioid = medioid(&vecs);
 
         build_graph(&mut rng, &mut graph, medioid, &vecs, config);
-        config.alpha = 58000;
-        build_graph(&mut rng, &mut graph, medioid, &vecs, config);
+        report_degrees(&graph);
+        //random_fill_graph(&mut rng, &mut graph, config.r);
+        //config.alpha = 65536;
+        //build_graph(&mut rng, &mut graph, medioid, &vecs, config);
+        report_degrees(&graph);
 
         (graph, medioid)
     };
@@ -82,8 +86,6 @@ fn main() -> Result<()> {
         edge_ctr += adjlist.read().unwrap().len();
     }
 
-    println!("average degree: {}", edge_ctr as f32 / graph.graph.len() as f32);
-
     let time = Instant::now();
     let mut recall = 0;
     let mut cmps_ctr = 0;