fix: preserve HNSW search quality after recovery

Track top candidates independently of the history buffer and flush edge cache on shutdown so recovered indices keep high-recall neighbor data.

Author: shisoft

src/apps/hnsw/partition/search.rs

@@ -357,6 +357,9 @@ impl HnswOnlinePartition {
         logger: &Arc<JobLogger>,
         job_id: JobId,
     ) -> Option<Arc<OwnedPrimArray>> {
+        static LOG_REMAINING: std::sync::atomic::AtomicUsize =
+            std::sync::atomic::AtomicUsize::new(20);
+
         // Check if this is the query vector (special case)
         if vertex_id.is_unit_id() {
             return metadata.get_query_vector(field_id);
@@ -369,14 +372,24 @@ impl HnswOnlinePartition {
 
         // Cache miss - try to get cell_id from edge cache first (avoids storage I/O)
         // This optimization reduces vector lookup from 2 I/Os to 1 I/O when edge cache is warm
-        let cell_id = if let Some(cached_cell_id) = self.edge_cache.get_cached_cell_id(*vertex_id) {
-            cached_cell_id
+        let (cell_id, source) = if let Some(cached_cell_id) =
+            self.edge_cache.get_cached_cell_id(*vertex_id)
+        {
+            (cached_cell_id, "cache")
         } else {
             // Edge cache miss - fall back to storage read for cell_id
             let (cid, _) = self.get_vertex(vertex_id, chunks, logger, metadata, job_id)?;
-            cid
+            (cid, "storage")
         };
 
+        if LOG_REMAINING.load(std::sync::atomic::Ordering::Relaxed) > 0 {
+            LOG_REMAINING.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+            info!(
+                "HNSW vertex mapping: vertex={:?} -> cell={:?} (source={})",
+                vertex_id, cell_id, source
+            );
+        }
+
         // Read vector from data cell
         let cell = match chunks.read_selected(&cell_id, &[field_id], false) {
             Ok(cell) => cell,
@@ -398,6 +411,26 @@ impl HnswOnlinePartition {
         match field {
             SharedValue::PrimArray(vector) => {
                 let vector = vector.copy_into_owned();
+                if LOG_REMAINING.load(std::sync::atomic::Ordering::Relaxed) > 0 {
+                    LOG_REMAINING.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+                    match &vector {
+                        OwnedPrimArray::F32(values) => {
+                            let preview: Vec<f32> = values.iter().take(3).copied().collect();
+                            info!(
+                                "HNSW vertex vector: vertex={:?}, len={}, preview={:?}",
+                                vertex_id,
+                                values.len(),
+                                preview
+                            );
+                        }
+                        _ => {
+                            info!(
+                                "HNSW vertex vector: vertex={:?}, non-f32 array",
+                                vertex_id
+                            );
+                        }
+                    }
+                }
                 let arc_vector = Arc::new(vector);
                 // Cache in shared cache only (Moka handles eviction automatically)
                 self.shared_vertex_vector_cache
@@ -437,11 +470,38 @@ impl HnswOnlinePartition {
         }
 
         // Cache miss - compute distance
-        let va_array =
-            self.get_vertex_vector(vertex_a, field_id, metadata, chunks, logger, job_id)?;
-        let vb_array =
-            self.get_vertex_vector(vertex_b, field_id, metadata, chunks, logger, job_id)?;
-        let distance = metric_distance(metric, &va_array, &vb_array)?;
+        let va_array = match self.get_vertex_vector(vertex_a, field_id, metadata, chunks, logger, job_id)
+        {
+            Some(array) => array,
+            None => {
+                info!(
+                    "HNSW distance: missing vector for vertex={:?}, field={}",
+                    vertex_a, field_id
+                );
+                return None;
+            }
+        };
+        let vb_array = match self.get_vertex_vector(vertex_b, field_id, metadata, chunks, logger, job_id)
+        {
+            Some(array) => array,
+            None => {
+                info!(
+                    "HNSW distance: missing vector for vertex={:?}, field={}",
+                    vertex_b, field_id
+                );
+                return None;
+            }
+        };
+        let distance = match metric_distance(metric, &va_array, &vb_array) {
+            Some(dist) => dist,
+            None => {
+                info!(
+                    "HNSW distance: metric failed for vertices {:?} -> {:?}, field={}",
+                    vertex_a, vertex_b, field_id
+                );
+                return None;
+            }
+        };
 
         // Store in per-search cache
         metadata
@@ -609,12 +669,12 @@ impl HnswOnlinePartition {
             };
             temp_metadata.set_query_vector(query.clone(), field_id);
 
-            // For large indices, sample up to 50 entry points to balance quality vs performance
-            let vertices_to_check: Vec<Id> = if starting_vertices.len() > 50 {
-                starting_vertices.iter().take(50).cloned().collect()
-            } else {
-                starting_vertices.clone()
-            };
+        // For large indices, sample up to 50 entry points to balance quality vs performance
+        let vertices_to_check: Vec<Id> = if starting_vertices.len() > 50 {
+            starting_vertices.iter().take(50).cloned().collect()
+        } else {
+            starting_vertices.clone()
+        };
 
             for vertex_id in &vertices_to_check {
                 let distance_opt = self.get_vertex_distance(
@@ -630,6 +690,7 @@ impl HnswOnlinePartition {
                 if let Some(distance) = distance_opt {
                     entry_distances.push((*vertex_id, distance));
                 }
+
             }
 
             // Sort by distance and take the closest entry point(s)
@@ -641,6 +702,7 @@ impl HnswOnlinePartition {
                 .take(std::cmp::min(5, entry_distances.len()))
                 .map(|(id, _)| *id)
                 .collect();
+
         }
 
         //println!("Starting search with vertex: {:?}", starting_vertices);
@@ -1967,6 +2029,10 @@ impl HnswOnlinePartition {
         // Early termination tracking: stop when top-k results are stable
         let mut stable_count = 0;
         let mut last_best_distance = Distance::INFINITY;
+        // Track best candidates separately to avoid history eviction.
+        let best_candidates_capacity = ef;
+        let mut best_candidates: BinaryHeap<(OrderedFloat, Id)> =
+            BinaryHeap::with_capacity(best_candidates_capacity + 1);
 
         // Main search loop with priority-ordered multi-vertex exploration
         while !priority_frontier.is_empty() && iterations < max_iterations {
@@ -2004,6 +2070,11 @@ impl HnswOnlinePartition {
                 // Add to history
                 search.history.push((vertex_id, vertex_distance));
 
+                best_candidates.push((OrderedFloat(vertex_distance), vertex_id));
+                if best_candidates.len() > best_candidates_capacity {
+                    best_candidates.pop();
+                }
+
                 // Get neighbors and add unvisited ones to priority frontier
                 if let Ok(neighbors) = self.get_neighbors_at_level(vertex_id, 0) {
                     for neighbor_id in neighbors {
@@ -2048,13 +2119,12 @@ impl HnswOnlinePartition {
             }
         }
 
-        // Search phase complete
-
-        // LOCAL EXPANSION PHASE: Explore neighbors to get k results
-        let mut all_candidates: Vec<(Id, Distance)> = search.history.to_vec();
-
-        // Sort initial candidates by distance to prioritize exploring closest vertices first
-        all_candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
+        // LOCAL EXPANSION PHASE: Use best_candidates from heap (already sorted by distance)
+        let mut all_candidates: Vec<(Id, Distance)> = best_candidates
+            .into_sorted_vec()
+            .into_iter()
+            .map(|(dist, id)| (id, dist.0))
+            .collect();
 
         // Use aggressive expansion when we have fewer candidates than needed
         let needs_aggressive_expansion = all_candidates.len() < k || all_candidates.len() < ef;
@@ -2172,6 +2242,23 @@ impl HnswOnlinePartition {
             None => std::cmp::Ordering::Equal,
         });
 
+        for (vertex_id, distance) in final_candidates.iter().take(10) {
+            if let Some((cell_id, _)) = self.get_vertex(
+                vertex_id,
+                &self.chunks,
+                &*EMPTY_LOGGER,
+                &mut search.metadata,
+                dummy_job_id,
+            ) {
+                info!(
+                    "HNSW final candidate: vertex={:?}, cell={:?}, dist={}",
+                    vertex_id,
+                    cell_id,
+                    distance
+                );
+            }
+        }
+
         // Convert vertex IDs to cell IDs and take top-k
         let mut results = Vec::new();
         for (vertex_id, distance) in final_candidates.into_iter() {

src/apps/hnsw/recovery_tests.rs

@@ -22,7 +22,9 @@ use neb::ram::schema::{Field, Schema as NebSchema};
 
 use crate::apps::hnsw::coordinator::AsyncServiceClient;
 use crate::apps::hnsw::measurements::MetricEncoding;
-use crate::apps::hnsw::partition::schema::{initialize_schemas, TOP_LEVEL_VERTICES};
+use crate::apps::hnsw::partition::schema::{
+    initialize_schemas, NEIGHBORS_FIELDS, TOP_LEVEL_VERTICES,
+};
 use crate::apps::hnsw::HNSWPartitionService;
 use crate::server::MorpheusServer;
 use crate::tests::start_server_with_storage_and_recovery;
@@ -1246,6 +1248,19 @@ mod tests {
         let mut total_cells_before = 0usize;
         let mut top_level_vertices_before: Vec<Id> = Vec::new();
 
+        let test_vectors: Vec<TestVector> = data_cell_ids
+            .iter()
+            .enumerate()
+            .map(|(i, cell_id)| {
+                let vector = vec![
+                    i as f32 * 10.0,
+                    (i as f32 + 1.0) * 10.0,
+                    (i as f32 + 2.0) * 10.0,
+                ];
+                TestVector::new(cell_id.higher, cell_id.lower, vector)
+            })
+            .collect();
+
         // Phase 1: Create index and insert vectors
         {
             log::info!(
@@ -1254,20 +1269,6 @@ mod tests {
             );
             let (server, client) = setup_server_with_index(port, group_name).await.unwrap();
 
-            // Create test vectors
-            let test_vectors: Vec<TestVector> = data_cell_ids
-                .iter()
-                .enumerate()
-                .map(|(i, cell_id)| {
-                    let vector = vec![
-                        i as f32 * 10.0,
-                        (i as f32 + 1.0) * 10.0,
-                        (i as f32 + 2.0) * 10.0,
-                    ];
-                    TestVector::new(cell_id.higher, cell_id.lower, vector)
-                })
-                .collect();
-
             // Insert vectors
             insert_vectors(&server, &client, &test_vectors)
                 .await
@@ -1310,6 +1311,7 @@ mod tests {
                 );
             }
 
+
             // Verify search works before shutdown
             let query = vec![0.0, 10.0, 20.0];
             let results = client
@@ -1366,6 +1368,8 @@ mod tests {
 
             let neb = &server.neb_client;
 
+            let _partition_service = server.init_hnsw_index_partition_service().await.unwrap();
+
             // Assert: All data cells recovered
             let mut data_cells_recovered = 0;
             for cell_id in &data_cell_ids {
@@ -1401,8 +1405,7 @@ mod tests {
                     verts.len()
                 );
 
-                // Verify each top-level vertex cell exists, points to valid data, and has edges
-                let mut total_edges_found = 0;
+                // Verify each top-level vertex cell exists and points to valid data
                 for v in verts {
                     let vertex_cell = neb
                         .read_cell(*v)
@@ -1421,38 +1424,7 @@ mod tests {
                             data_cell_id
                         );
                     }
-
-                    // Verify vertex has edges (graph connectivity)
-                    use crate::apps::hnsw::partition::schema::HNSW_VERTEX_SCHEMA_ID;
-                    use crate::graph::fields::UNDIRECTED_KEY_ID;
-                    use crate::graph::id_list::AtomicIdList;
-
-                    let mut edge_list = AtomicIdList::new(
-                        server.neb_client.clone(),
-                        *v,
-                        *UNDIRECTED_KEY_ID,
-                        HNSW_VERTEX_SCHEMA_ID,
-                    );
-
-                    match edge_list.all().await {
-                        Ok(edges) => {
-                            total_edges_found += edges.len();
-                            // Each top-level vertex should have at least some edges
-                            // (in a well-connected HNSW graph)
-                        }
-                        Err(e) => {
-                            // Edge list might not exist for some vertices (e.g., isolated at top level)
-                            log::debug!("Vertex {:?} has no edges: {:?}", v, e);
-                        }
-                    }
                 }
-
-                // With 100 vectors, we expect a reasonable number of edges in the graph
-                // Each vector connects to several neighbors (M parameter typically 16-64)
-                log::info!(
-                    "Total edges found from top-level vertices: {}",
-                    total_edges_found
-                );
             } else {
                 panic!("top_level_vertices should be an Id array");
             }
@@ -1549,7 +1521,6 @@ mod tests {
                 unique_result_ids.len()
             );
 
-            log::info!("Phase 2 complete: All assertions passed");
             server.shutdown().await;
         }
 

src/server/mod.rs

@@ -341,8 +341,11 @@ impl MorpheusServer {
     pub async fn shutdown(&self) {
         log::info!("MorpheusServer: starting graceful shutdown");
 
-        // TODO: Add explicit HNSW index saving here if needed
-        // Currently, HNSW indices are saved after each insertion in the partition service
+        // Ensure HNSW edge cache is flushed before storage/raft shutdown.
+        if let Some(service) = self.services.hnsw_partition.get() {
+            log::info!("MorpheusServer: flushing HNSW edge cache before shutdown");
+            service.partition.flush_edge_cache();
+        }
 
         // Delegate to NebServer's shutdown method which handles:
         // 1. LSM tree flushing (if RangedIndexer is enabled)