9
9
10
10
#include < faiss/impl/HNSW.h>
11
11
12
- #include < cmath>
13
12
#include < string>
14
13
15
14
#include < faiss/impl/AuxIndexStructures.h>
@@ -543,11 +542,12 @@ int search_from_candidates(
543
542
for (int i = 0 ; i < candidates.size (); i++) {
544
543
idx_t v1 = candidates.ids [i];
545
544
float d = candidates.dis [i];
546
- assert (v1 >= 0 );
545
+ FAISS_ASSERT (v1 >= 0 );
547
546
if (!sel || sel->is_member (v1)) {
548
- if (d < D[0 ]) {
549
- faiss::maxheap_replace_top (k, D, I, d, v1);
550
- nres++;
547
+ if (nres < k) {
548
+ faiss::maxheap_push (++nres, D, I, d, v1);
549
+ } else if (d < D[0 ]) {
550
+ faiss::maxheap_replace_top (nres, D, I, d, v1);
551
551
}
552
552
}
553
553
vt.set (v1);
@@ -612,9 +612,10 @@ int search_from_candidates(
612
612
613
613
auto add_to_heap = [&](const size_t idx, const float dis) {
614
614
if (!sel || sel->is_member (idx)) {
615
- if (dis < D[0 ]) {
616
- faiss::maxheap_replace_top (k, D, I, dis, idx);
617
- nres++;
615
+ if (nres < k) {
616
+ faiss::maxheap_push (++nres, D, I, dis, idx);
617
+ } else if (dis < D[0 ]) {
618
+ faiss::maxheap_replace_top (nres, D, I, dis, idx);
618
619
}
619
620
}
620
621
candidates.push (idx, dis);
@@ -667,7 +668,7 @@ int search_from_candidates(
667
668
stats.n3 += ndis;
668
669
}
669
670
670
- return std::min ( nres, k) ;
671
+ return nres;
671
672
}
672
673
673
674
std::priority_queue<HNSW::Node> search_from_candidate_unbounded (
@@ -815,11 +816,6 @@ HNSWStats HNSW::search(
815
816
// greedy search on upper levels
816
817
storage_idx_t nearest = entry_point;
817
818
float d_nearest = qdis (nearest);
818
- if (!std::isfinite (d_nearest)) {
819
- // means either the query or the entry point are NaN: in
820
- // both cases we can only return -1 as a result
821
- return stats;
822
- }
823
819
824
820
for (int level = max_level; level >= 1 ; level--) {
825
821
greedy_update_nearest (*this , qdis, level, nearest, d_nearest);
@@ -830,6 +826,7 @@ HNSWStats HNSW::search(
830
826
MinimaxHeap candidates (ef);
831
827
832
828
candidates.push (nearest, d_nearest);
829
+
833
830
search_from_candidates (
834
831
*this , qdis, k, I, D, candidates, vt, stats, 0 , 0 , params);
835
832
} else {
0 commit comments