feat: Introduce a visit limit and optimize HNSW search by refining candidate heap management and conditional result set population.

Author: shisoft

src/apps/hnsw/partition/search.rs

@@ -1396,11 +1396,15 @@ impl HnswOnlinePartition {
         // Ensure ef is at least as large as k
         let ef = ef.max(k);
 
+        // Safety cap to prevent unbounded exploration in pathological cases
+        let max_visits = ef * 20;
+        let mut visit_count = 0;
+
         // Max-heap for results (to maintain closest ef elements, with worst at top for easy removal)
         let mut result_set: BinaryHeap<(OrderedFloat, Id, Id)> = BinaryHeap::with_capacity(ef + 1);
 
-        // Min-heap for candidates (sort by distance, closest first)
-        let mut candidates: BinaryHeap<ReverseOrd<(OrderedFloat, Id, Id)>> =
+        // Min-heap for candidates (sort by distance, closest first) - only stores (distance, vertex_id)
+        let mut candidates: BinaryHeap<ReverseOrd<(OrderedFloat, Id)>> =
             BinaryHeap::with_capacity(ef + 1);
 
         // Track visited vertices
@@ -1426,12 +1430,13 @@ impl HnswOnlinePartition {
                 logger,
                 job_id,
             ) {
-                let cell_id = self
-                    .get_vertex_cell_id(&start_id, &self.chunks, logger, metadata, job_id)
-                    .unwrap();
-                // Add to both candidates and result_set
-                candidates.push(ReverseOrd((OrderedFloat(distance), start_id, cell_id)));
-                result_set.push((OrderedFloat(distance), start_id, cell_id));
+                // Only add to result_set if we can get the cell_id
+                if let Some(cell_id) =
+                    self.get_vertex_cell_id(&start_id, &self.chunks, logger, metadata, job_id)
+                {
+                    candidates.push(ReverseOrd((OrderedFloat(distance), start_id)));
+                    result_set.push((OrderedFloat(distance), start_id, cell_id));
+                }
             }
         }
 
@@ -1441,19 +1446,25 @@ impl HnswOnlinePartition {
         }
 
         // Start with the worst distance in our result set
-        let mut dist_threshold = if result_set.len() > 0 {
+        let mut dist_threshold = if !result_set.is_empty() {
             result_set.peek().unwrap().0
         } else {
             OrderedFloat(Distance::INFINITY)
         };
 
         // Main search loop
-        while let Some(ReverseOrd((dist, current_vertex, _))) = candidates.pop() {
+        while let Some(ReverseOrd((dist, current_vertex))) = candidates.pop() {
             // Stop if minimum distance to candidates is worse than our worst result
             if dist > dist_threshold {
                 break;
             }
 
+            // Safety cap on visits
+            visit_count += 1;
+            if visit_count >= max_visits {
+                break;
+            }
+
             // Get neighbors of current vertex
             match self.visit_vertex_and_edge(&current_vertex, level, metadata, logger, job_id) {
                 Ok(neighbors) => {
@@ -1478,41 +1489,38 @@ impl HnswOnlinePartition {
                                     logger,
                                     job_id,
                                 ) {
-                                    // Only process if better than our worst result or if we need more results
-                                    if result_set.len() < ef
-                                        || OrderedFloat(distance) < dist_threshold
-                                    {
-                                        let cell_id = self
-                                            .get_vertex_cell_id(
-                                                &neighbor_id,
-                                                &self.chunks,
-                                                logger,
-                                                metadata,
-                                                job_id,
-                                            )
-                                            .unwrap();
-                                        // Add to candidates queue
-                                        candidates.push(ReverseOrd((
-                                            OrderedFloat(distance),
-                                            neighbor_id,
-                                            cell_id,
-                                        )));
-
-                                        // Add to results
-                                        result_set.push((
-                                            OrderedFloat(distance),
-                                            neighbor_id,
-                                            cell_id,
-                                        ));
-
-                                        // If we exceed capacity, remove the worst element
-                                        if result_set.len() > ef {
-                                            result_set.pop();
-                                        }
+                                    let ord_distance = OrderedFloat(distance);
 
-                                        // Update threshold to be the distance of the worst element
-                                        if let Some(worst) = result_set.peek() {
-                                            dist_threshold = worst.0;
+                                    // Only process if better than our worst result or if we need more results
+                                    if result_set.len() < ef || ord_distance < dist_threshold {
+                                        // Only add to results if we can get the cell_id
+                                        if let Some(cell_id) = self.get_vertex_cell_id(
+                                            &neighbor_id,
+                                            &self.chunks,
+                                            logger,
+                                            metadata,
+                                            job_id,
+                                        ) {
+                                            // Only add to candidates if strictly better (avoids redundant exploration)
+                                            if result_set.len() < ef
+                                                || ord_distance < dist_threshold
+                                            {
+                                                candidates
+                                                    .push(ReverseOrd((ord_distance, neighbor_id)));
+                                            }
+
+                                            // Add to results
+                                            result_set.push((ord_distance, neighbor_id, cell_id));
+
+                                            // If we exceed capacity, remove the worst element
+                                            if result_set.len() > ef {
+                                                result_set.pop();
+                                            }
+
+                                            // Update threshold to be the distance of the worst element
+                                            if let Some(worst) = result_set.peek() {
+                                                dist_threshold = worst.0;
+                                            }
                                         }
                                     }
                                 }