fix entire index algorithm (very silly bug)

src/dump_processor.rs

@@ -14,6 +14,7 @@ use itertools::Itertools;
 use simsimd::SpatialSimilarity;
 use std::hash::Hasher;
 use foldhash::{HashSet, HashSetExt};
+use std::os::unix::prelude::FileExt;
 
 use diskann::vector::{scale_dot_result_f64, ProductQuantizer};
 
@@ -161,15 +162,29 @@ fn main() -> Result<()> {
     let (mut queries_index, max_query_id) = if let Some(queries_file) = args.queries {
         println!("constructing index");
         // not memory-efficient but this is small
-        let data = fs::read(queries_file).context("read queries file")?;
+        let mut file = fs::File::open(queries_file).context("read queries file")?;
+        let mut size = file.metadata()?.len();
         //let mut index = faiss::index_factory(D_EMB, "HNSW32,SQfp16", faiss::MetricType::InnerProduct)?;
-        let mut index = faiss::index_factory(D_EMB, "HNSW32,SQfp16", faiss::MetricType::InnerProduct)?;
+        let mut index = faiss::index_factory(D_EMB, "HNSW64,SQ8", faiss::MetricType::InnerProduct)?;
         //let mut index = faiss::index_factory(D_EMB, "IVF4096,SQfp16", faiss::MetricType::InnerProduct)?;
-        let unpacked = common::decode_fp16_buffer(&data);
-        index.train(&unpacked)?;
-        index.add(&unpacked)?;
+        let mut buf = vec![0; (D_EMB as usize) * (1<<18)];
+        loop {
+            if size == 0 {
+                break;
+            }
+            if size < (buf.len() as u64) {
+                buf.resize(size as usize, 0);
+            }
+            file.read_exact(&mut buf)?;
+            size -= buf.len() as u64;
+            let unpacked = common::decode_fp16_buffer(&buf);
+            if !index.is_trained() { index.train(&unpacked)?; print!("train"); }
+            index.add(&unpacked)?;
+            print!(".");
+        }
         println!("done");
-        (Some(index), unpacked.len() / D_EMB as usize)
+        let ntotal = index.ntotal();
+        (Some(index), ntotal as usize)
     } else {
         (None, 0)
     };
@@ -267,14 +282,15 @@ fn main() -> Result<()> {
                 let shard = shard as usize;
                 // this random access is almost certainly rather slow
                 // parallelize?
-                files[shard].1.seek(SeekFrom::Start(offset))?;
                 let mut buf = vec![0; len as usize];
-                files[shard].1.read_exact(&mut buf)?;
-                let s: &mut [u32] = bytemuck::cast_slice_mut(&mut *buf);
-                for within_shard_id in s.iter_mut() {
-                    *within_shard_id = shard_id_mappings[shard].1[*within_shard_id as usize];
+                files[shard].1.read_exact_at(&mut buf, offset)?;
+                let s: &[u32] = bytemuck::cast_slice(&mut *buf);
+                for within_shard_id in s.iter() {
+                    let global_id = shard_id_mappings[shard].1[*within_shard_id as usize];
+                    if !out_vertices.contains(&global_id) {
+                        out_vertices.push(global_id);
+                    }
                 }
-                out_vertices.extend(s.iter().unique());
             }
 
             Ok((out_vertices, shards))
@@ -422,7 +438,7 @@ fn main() -> Result<()> {
             let codes = quantizer.quantize_batch(&batch_embeddings);
 
             for (i, (x, _embedding)) in batch.into_iter().enumerate() {
-                let (vertices, shards) = read_out_vertices(count)?; // TODO: could parallelize this given the batching
+                let (vertices, shards) = read_out_vertices(count + i as u32)?; // TODO: could parallelize this given the batching
                 let mut entry = PackedIndexEntry {
                     id: count + i as u32,
                     vertices,

src/generate_index_shard.rs

@@ -3,7 +3,7 @@ use itertools::Itertools;
 use std::io::{BufReader, BufWriter, Write};
 use rmp_serde::decode::Error as DecodeError;
 use std::fs;
-use diskann::{augment_bipartite, build_graph, project_bipartite, random_fill_graph, vector::{dot, VectorList}, IndexBuildConfig, IndexGraph, Timer, report_degrees};
+use diskann::{augment_bipartite, build_graph, random_fill_graph, vector::{dot, VectorList}, IndexBuildConfig, IndexGraph, Timer, report_degrees, medioid};
 use half::f16;
 
 mod common;
@@ -41,10 +41,11 @@ fn main() -> Result<()> {
     }
 
     let mut config = IndexBuildConfig {
-        r: 40,
-        l: 200,
+        r: 64,
+        l: 192,
         maxc: 750,
-        alpha: 65300
+        alpha: 65200,
+        saturate_graph: false
     };
 
     let vecs = VectorList {
@@ -67,9 +68,7 @@ fn main() -> Result<()> {
 
     report_degrees(&graph);
 
-    let medioid = vecs.iter().position_max_by_key(|&v| {
-        dot(v, &centroid_fp16)
-    }).unwrap() as u32;
+    let medioid = medioid(&vecs);
 
     {
         let _timer = Timer::new("first pass");
@@ -101,7 +100,8 @@ fn main() -> Result<()> {
 
     {
         let _timer = Timer::new("augment bipartite");
-        augment_bipartite(&mut rng, &mut graph, query_knns, query_knns_bwd, config);
+        //augment_bipartite(&mut rng, &mut graph, query_knns, query_knns_bwd, config, 50);
+        //random_fill_graph(&mut rng, &mut graph, config.r);
     }
 
     let len = original_ids.len();

src/query_disk_index.rs

@@ -2,6 +2,7 @@ use anyhow::{bail, Context, Result};
 use diskann::vector::scale_dot_result_f64;
 use serde::{Serialize, Deserialize};
 use std::io::{BufReader, Read, Seek, SeekFrom, Write};
+use std::os::unix::prelude::FileExt;
 use std::path::PathBuf;
 use std::fs;
 use base64::Engine;
@@ -37,9 +38,8 @@ struct CLIArguments {
 
 fn read_node(id: u32, data_file: &mut fs::File, header: &IndexHeader) -> Result<PackedIndexEntry> {
     let offset = id as usize * header.record_pad_size;
-    data_file.seek(SeekFrom::Start(offset as u64))?;
     let mut buf = vec![0; header.record_pad_size as usize];
-    data_file.read_exact(&mut buf)?;
+    data_file.read_exact_at(&mut buf, offset as u64)?;
     let len = u16::from_le_bytes(buf[0..2].try_into().unwrap()) as usize;
     Ok(bitcode::decode(&buf[2..len+2])?)
 }
@@ -117,9 +117,11 @@ fn summary_stats(ranks: &mut [usize]) {
     ranks.sort_unstable();
     let median = ranks[ranks.len() / 2] + 1;
     let harmonic_mean = ranks.iter().map(|x| 1.0 / ((x+1) as f64)).sum::<f64>() / ranks.len() as f64;
-    println!("median {} mean {} max {} min {} harmonic mean {}", median, mean, ranks[ranks.len() - 1] + 1, ranks[0] + 1, 1.0 / harmonic_mean);
+    println!("median {} mean {:.2} max {} min {} harmonic mean {:.2}", median, mean, ranks[ranks.len() - 1] + 1, ranks[0] + 1, 1.0 / harmonic_mean);
 }
 
+const K: usize = 20;
+
 fn main() -> Result<()> {
     let args: CLIArguments = argh::from_env();
 
@@ -150,8 +152,11 @@ fn main() -> Result<()> {
 
     println!("{} items {} dead {} shards", header.count, header.dead_count, header.shards.len());
 
-    let mut top_20_ranks_best_shard = vec![];
+    let mut top_k_ranks_best_shard = vec![];
     let mut top_rank_best_shard = vec![];
+    let mut pq_cmps = vec![];
+    let mut cmps = vec![];
+    let mut recall_total = 0;
 
     for query_vector in queries.iter() {
         let query_vector_fp32 = query_vector.iter().map(|x| x.to_f32()).collect::<Vec<f32>>();
@@ -183,26 +188,30 @@ fn main() -> Result<()> {
             println!("brute force: {} {} {} {:?}", id, distance, url, shards);
         }*/
 
-        let mut top_ranks = vec![usize::MAX; 20];
+        let mut top_ranks = vec![usize::MAX; K];
 
         for shard in 0..header.shards.len() {
             let selected_start = header.shards[shard].1;
 
             let mut scratch = Scratch {
                 visited: HashSet::new(),
-                neighbour_buffer: NeighbourBuffer::new(5000),
+                neighbour_buffer: NeighbourBuffer::new(1000),
                 neighbour_pre_buffer: Vec::new(),
                 visited_list: Vec::new()
             };
 
             //let query_vector = diskann::vector::quantize(&query_vector, &header.quantizer, &mut rng);
-            let cmps = greedy_search(&mut scratch, selected_start, &query_vector, &query_preprocessed, IndexRef {
+            let cmps_result = greedy_search(&mut scratch, selected_start, &query_vector, &query_preprocessed, IndexRef {
                 data_file: &mut data_file,
                 header: &header,
                 pq_codes: &pq_codes,
                 pq_code_size: header.quantizer.n_dims / header.quantizer.n_dims_per_code,
             }, args.disable_pq)?;
 
+            // slightly dubious because this is across shards
+            pq_cmps.push(cmps_result.1);
+            cmps.push(cmps_result.0);
+
             if args.verbose {
                 println!("index scan {}: {:?} cmps", shard, cmps);
             }
@@ -221,14 +230,26 @@ fn main() -> Result<()> {
             if args.verbose { println!("") }
         }
 
+        // results list is always correctly sorted
+        for &rank in top_ranks.iter() {
+            if rank < K {
+                recall_total += 1;
+            }
+        }
+
         top_rank_best_shard.push(top_ranks[0]);
-        top_20_ranks_best_shard.extend(top_ranks);
+        top_k_ranks_best_shard.extend(top_ranks);
     }
 
     println!("ranks of top 20:");
-    summary_stats(&mut top_20_ranks_best_shard);
+    summary_stats(&mut top_k_ranks_best_shard);
     println!("ranks of top 1:");
     summary_stats(&mut top_rank_best_shard);
+    println!("pq comparisons:");
+    summary_stats(&mut pq_cmps);
+    println!("comparisons:");
+    summary_stats(&mut cmps);
+    println!("recall@{}: {}", K, recall_total as f64 / (K * queries.len()) as f64);
 
     Ok(())
 }