fix: improve HNSW search quality and tests

Author: shisoft

src/apps/embedding/mod.rs

@@ -627,7 +627,10 @@ impl EmbeddingIndexerCore for EmbeddingIndexer {
 
             // Search HNSW (on the embedding schema's vector field)
             let query_vector = OwnedPrimArray::F32(result.vector);
-            let ef = (limit * 2) as u64;
+            let ef_min = 32usize;
+            let ef_max = 200usize;
+            let ef_scaled = (limit * 4).clamp(ef_min, ef_max);
+            let ef = std::cmp::max(limit, ef_scaled) as u64;
 
             let hits = self
                 .hnsw_coordinator

src/apps/hnsw/coordinator/tests.rs

@@ -46,6 +46,59 @@ impl TestVector {
     }
 }
 
+fn compute_l2_distance(query: &[f32], vector: &OwnedPrimArray) -> Distance {
+    match vector {
+        OwnedPrimArray::F32(values) => values
+            .iter()
+            .zip(query.iter())
+            .map(|(a, b)| {
+                let diff = a - b;
+                diff * diff
+            })
+            .sum(),
+        _ => Distance::INFINITY,
+    }
+}
+
+fn exact_top_k(query: &[f32], vectors: &[TestVector], k: usize) -> Vec<(Id, Distance)> {
+    let mut scored: Vec<(Id, Distance)> = vectors
+        .iter()
+        .map(|v| (v.cell_id, compute_l2_distance(query, &v.vector)))
+        .collect();
+    scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
+    scored.truncate(k);
+    scored
+}
+
+fn recall_at_k(results: &[(Id, Distance)], exact: &[(Id, Distance)]) -> f64 {
+    if exact.is_empty() {
+        return 0.0;
+    }
+    let exact_ids: std::collections::HashSet<Id> = exact.iter().map(|(id, _)| *id).collect();
+    let hit_count = results
+        .iter()
+        .filter(|(id, _)| exact_ids.contains(id))
+        .count();
+    hit_count as f64 / exact.len() as f64
+}
+
+fn top1_distance_ratio(results: &[(Id, Distance)], exact: &[(Id, Distance)]) -> Distance {
+    if results.is_empty() || exact.is_empty() {
+        return Distance::INFINITY;
+    }
+    let exact_top = exact[0].1.max(1e-9);
+    results[0].1 / exact_top
+}
+
+fn overlap_at_k(a: &[(Id, Distance)], b: &[(Id, Distance)]) -> f64 {
+    if b.is_empty() {
+        return 0.0;
+    }
+    let ids_a: std::collections::HashSet<Id> = a.iter().map(|(id, _)| *id).collect();
+    let hits = b.iter().filter(|(id, _)| ids_a.contains(id)).count();
+    hits as f64 / b.len() as f64
+}
+
 /// Test environment structure to hold all components needed for HNSW coordinator tests
 pub struct TestEnvironment {
     pub job_counter: RefCell<u64>,
@@ -348,10 +401,24 @@ mod tests {
         env.index_cells().await.unwrap();
 
         let query = vec![1.0, 2.0, 3.0];
-        let results = env.top_k(query, 1).await.unwrap();
+        let results = env.top_k(query.clone(), 1).await.unwrap();
         assert_eq!(results.len(), 1);
-        assert_eq!(results[0].0.lower, 1000); // First cell
-        assert_eq!(results[0].1, 0.0); // Exact match
+
+        let exact = exact_top_k(&query, &env.test_vectors, 1);
+        let recall = recall_at_k(&results, &exact);
+        assert!(recall >= 1.0, "Recall@1 should be 1.0, got {:.2}", recall);
+        let ratio = top1_distance_ratio(&results, &exact);
+        assert!(
+            ratio <= 1.0 + 1e-6,
+            "Top-1 distance ratio should be 1.0, got {:.6}",
+            ratio
+        );
+        let overlap = overlap_at_k(&results, &exact);
+        assert!(
+            overlap >= 1.0,
+            "Overlap@1 should be 1.0, got {:.2}",
+            overlap
+        );
         env.cleanup().await;
     }
 
@@ -382,11 +449,24 @@ mod tests {
         env.index_cells().await.unwrap();
 
         let query = vec![4.0, 5.0, 6.0];
-        let results = env.top_k(query, 2).await.unwrap();
+        let results = env.top_k(query.clone(), 2).await.unwrap();
         assert_eq!(results.len(), 2);
-        assert_eq!(results[0].0.lower, 1001); // Second cell
-        assert_eq!(results[0].1, 0.0); // Exact match
-        assert!(results[1].0.lower == 1000 || results[1].0.lower == 1002); // First or third cell
+
+        let exact = exact_top_k(&query, &env.test_vectors, 2);
+        let recall = recall_at_k(&results, &exact);
+        assert!(recall >= 1.0, "Recall@2 should be 1.0, got {:.2}", recall);
+        let ratio = top1_distance_ratio(&results, &exact);
+        assert!(
+            ratio <= 1.0 + 1e-6,
+            "Top-1 distance ratio should be 1.0, got {:.6}",
+            ratio
+        );
+        let overlap = overlap_at_k(&results, &exact);
+        assert!(
+            overlap >= 1.0,
+            "Overlap@2 should be 1.0, got {:.2}",
+            overlap
+        );
         env.cleanup().await;
     }
 
@@ -420,11 +500,28 @@ mod tests {
 
         let query = vec![10.0, 11.0, 12.0];
         let k = 5;
-        let results = env.top_k(query, k).await.unwrap();
+        let results = env.top_k(query.clone(), k).await.unwrap();
         assert_eq!(results.len(), k as usize);
-        assert_eq!(results[0].0.lower, 1010); // Tenth cell
-        assert_eq!(results[0].1, 0.0); // Exact match
-        assert!(results[1].0.lower == 1009 || results[1].0.lower == 1011); // Neighboring cells
+
+        let exact = exact_top_k(&query, &env.test_vectors, k as usize);
+        let recall = recall_at_k(&results, &exact);
+        assert!(
+            recall >= 0.9,
+            "Recall@5 should be >= 0.9, got {:.2}",
+            recall
+        );
+        let ratio = top1_distance_ratio(&results, &exact);
+        assert!(
+            ratio <= 1.1,
+            "Top-1 distance ratio should be <= 1.1, got {:.3}",
+            ratio
+        );
+        let overlap = overlap_at_k(&results, &exact);
+        assert!(
+            overlap >= 0.8,
+            "Overlap@5 should be >= 0.8, got {:.2}",
+            overlap
+        );
         env.cleanup().await;
     }
 

src/apps/hnsw/mod.rs

@@ -121,8 +121,9 @@ impl VectorIndexerCore for VectorIndexer {
     ) -> BoxFuture<Result<Vec<VectorHit>, IndexError>> {
         // Convert the query vector slice to OwnedPrimArray
         let query = OwnedPrimArray::F32(query_vector.to_vec());
-        // Use L2 metric as default, ef = limit * 2 for better recall
-        let ef = (limit * 2) as u64;
+        let ef_min = 32usize;
+        let ef_max = 200usize;
+        let ef = (limit * 4).clamp(ef_min, ef_max) as u64;
         self.coordinator
             .query_top_k(
                 schema_id,

src/apps/hnsw/partition/search.rs

@@ -11,7 +11,7 @@ use bifrost::{conshash::ConsistentHashing, raft::client::RaftClient};
 use dovahkiin::types::Map as DovMap;
 use dovahkiin::{
     data_map,
-    types::{Id, OwnedPrimArray, SharedValue},
+    types::{Id, OwnedPrimArray, SharedPrimArray, SharedValue},
 };
 use futures::stream::FuturesUnordered;
 use lightning::map::Map;
@@ -820,13 +820,30 @@ impl HnswOnlinePartition {
         _logger: &Arc<JobLogger>,
         _job_id: JobId,
     ) -> Result<Vec<Id>, NeighbourhoodError> {
-        // Use edge cache for unified read path (includes pending writes)
         if let Some(neighbors) = self.edge_cache.get_neighbors(*vertex_id, level) {
-            Ok(neighbors)
-        } else {
-            // Vertex not found in cache - return empty neighbors (new vertex or not yet indexed)
-            Ok(vec![])
+            return Ok(neighbors);
+        }
+
+        if level >= MAX_LEVELS {
+            return Ok(vec![]);
         }
+
+        let cell_field_id = super::schema::CELL_FIELD_ID;
+        let mut field_ids = vec![cell_field_id];
+        field_ids.extend(NEIGHBORS_FIELD_IDS.iter().copied());
+
+        let cell = match self.chunks.read_selected(vertex_id, &field_ids, false) {
+            Ok(cell) => cell,
+            Err(_) => return Ok(vec![]),
+        };
+
+        let neighbors = match &cell.data[level + 1] {
+            SharedValue::PrimArray(SharedPrimArray::Id(list)) => list.to_vec(),
+            SharedValue::Null | SharedValue::NA => vec![],
+            _ => vec![],
+        };
+
+        Ok(neighbors)
     }
 
     fn bfs_top_k_pruned(
@@ -907,12 +924,7 @@ impl HnswOnlinePartition {
                 }
             };
 
-            // Deterministic neighbor sampling: limit neighbors explored per vertex
-            // HNSW neighbor lists are often ordered by distance, so taking the first N
-            // prioritizes closer (more relevant) neighbors while reducing overhead
-            const MAX_NEIGHBORS_TO_EXPLORE: usize = 64;
-
-            let sample_limit = MAX_NEIGHBORS_TO_EXPLORE.min(neighbors.len());
+            let sample_limit = neighbors.len();
             for nid in neighbors.iter().take(sample_limit) {
                 let nid_index = metadata.registry.get_or_insert(*nid);
                 if !visited.contains(nid_index) {
@@ -1977,8 +1989,8 @@ impl HnswOnlinePartition {
             }
             last_best_distance = current_best;
 
-            // Early termination: if best distance is stable for 3 iterations and we have enough candidates
-            if stable_count >= 3 && search.history.len() >= ef && priority_frontier.len() >= ef {
+            if stable_count >= 5 && search.history.len() >= ef * 2 && priority_frontier.len() >= ef
+            {
                 break;
             }