faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -8,6 +8,7 @@
 #include <faiss/IndexHNSW.h>
 
 #include <omp.h>
+#include <atomic>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
@@ -26,13 +27,15 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/FaissException.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/VisitedTable.h>
+#include <faiss/impl/hnsw/MinimaxHeap.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
 
 namespace faiss {
 
-using MinimaxHeap = HNSW::MinimaxHeap;
 using storage_idx_t = HNSW::storage_idx_t;
 using NodeDistFarther = HNSW::NodeDistFarther;
 
@@ -81,10 +84,8 @@ void hnsw_add_vertices(
         printf("  max_level = %d\n", max_level);
     }
 
-    std::vector<omp_lock_t> locks(ntotal);
-    for (int i = 0; i < ntotal; i++) {
-        omp_init_lock(&locks[i]);
-    }
+    auto& locks = index_hnsw.locks;
+    locks.prepare(ntotal);
 
     // add vectors from highest to lowest level
     std::vector<int> hist;
@@ -93,10 +94,10 @@ void hnsw_add_vertices(
     { // make buckets with vectors of the same level
 
         // build histogram
-        for (int i = 0; i < n; i++) {
-            storage_idx_t pt_id = i + n0;
+        for (size_t i = 0; i < n; i++) {
+            storage_idx_t pt_id = static_cast<storage_idx_t>(i + n0);
             int pt_level = hnsw.levels[pt_id] - 1;
-            while (pt_level >= hist.size()) {
+            while (pt_level >= static_cast<int>(hist.size())) {
                 hist.push_back(0);
             }
             hist[pt_level]++;
@@ -104,13 +105,13 @@ void hnsw_add_vertices(
 
         // accumulate
         std::vector<int> offsets(hist.size() + 1, 0);
-        for (int i = 0; i < hist.size() - 1; i++) {
+        for (size_t i = 0; i < hist.size() - 1; i++) {
             offsets[i + 1] = offsets[i] + hist[i];
         }
 
         // bucket sort
-        for (int i = 0; i < n; i++) {
-            storage_idx_t pt_id = i + n0;
+        for (size_t i = 0; i < n; i++) {
+            storage_idx_t pt_id = static_cast<storage_idx_t>(i + n0);
             int pt_level = hnsw.levels[pt_id] - 1;
             order[offsets[pt_level]++] = pt_id;
         }
@@ -122,39 +123,41 @@ void hnsw_add_vertices(
     { // perform add
         RandomGenerator rng2(789);
 
-        int i1 = n;
+        size_t i1 = static_cast<int>(n);
 
-        for (int pt_level = hist.size() - 1;
+        for (int pt_level = static_cast<int>(hist.size()) - 1;
              pt_level >= int(!index_hnsw.init_level0);
              pt_level--) {
-            int i0 = i1 - hist[pt_level];
+            size_t i0 = i1 - hist[pt_level];
 
             if (verbose) {
-                printf("Adding %d elements at level %d\n", i1 - i0, pt_level);
+                printf("Adding %zu elements at level %d\n", i1 - i0, pt_level);
             }
 
             // random permutation to get rid of dataset order bias
-            for (int j = i0; j < i1; j++) {
-                std::swap(order[j], order[j + rng2.rand_int(i1 - j)]);
+            for (size_t j = i0; j < i1; j++) {
+                std::swap(
+                        order[j],
+                        order[j + rng2.rand_int(static_cast<int>(i1 - j))]);
             }
 
             bool interrupt = false;
 
 #pragma omp parallel if (i1 > i0 + 100)
             {
-                VisitedTable vt(ntotal);
+                VisitedTable vt(ntotal, hnsw.use_visited_hashset);
 
                 std::unique_ptr<DistanceComputer> dis(
                         storage_distance_computer(index_hnsw.storage));
-                int prev_display =
-                        verbose && omp_get_thread_num() == 0 ? 0 : -1;
+                bool do_display = verbose && omp_get_thread_num() == 0;
+                size_t prev_display = 0;
                 size_t counter = 0;
 
                 // here we should do schedule(dynamic) but this segfaults for
                 // some versions of LLVM. The performance impact should not be
                 // too large when (i1 - i0) / num_threads >> 1
 #pragma omp for schedule(static)
-                for (int i = i0; i < i1; i++) {
+                for (int64_t i = i0; i < i1; i++) {
                     storage_idx_t pt_id = order[i];
                     dis->set_query(x + (pt_id - n0) * d);
 
@@ -171,9 +174,9 @@ void hnsw_add_vertices(
                             vt,
                             index_hnsw.keep_max_size_level0 && (pt_level == 0));
 
-                    if (prev_display >= 0 && i - i0 > prev_display + 10000) {
+                    if (do_display && i - i0 > prev_display + 10000) {
                         prev_display = i - i0;
-                        printf("  %d / %d\r", i - i0, i1 - i0);
+                        printf("  %zu / %zu\r", i - i0, i1 - i0);
                         fflush(stdout);
                     }
                     if (counter % check_period == 0) {
@@ -198,9 +201,8 @@ void hnsw_add_vertices(
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
     }
-
-    for (int i = 0; i < ntotal; i++) {
-        omp_destroy_lock(&locks[i]);
+    if (!index_hnsw.retain_locks) {
+        locks.clear();
     }
 }
 
@@ -210,11 +212,13 @@ void hnsw_add_vertices(
  * IndexHNSW implementation
  **************************************************************/
 
-IndexHNSW::IndexHNSW(int d, int M, MetricType metric)
-        : Index(d, metric), hnsw(M) {}
+IndexHNSW::IndexHNSW(int d_in, int M, MetricType metric)
+        : Index(d_in, metric), hnsw(M) {}
 
-IndexHNSW::IndexHNSW(Index* storage, int M)
-        : Index(storage->d, storage->metric_type), hnsw(M), storage(storage) {
+IndexHNSW::IndexHNSW(Index* storage_in, int M)
+        : Index(storage_in->d, storage_in->metric_type),
+          hnsw(M),
+          storage(storage_in) {
     metric_arg = storage->metric_arg;
 }
 
@@ -262,28 +266,48 @@ void hnsw_search(
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 
 #pragma omp parallel if (i1 - i0 > 1)
         {
-            VisitedTable vt(index->ntotal);
-            typename BlockResultHandler::SingleResultHandler res(bres);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(index->storage));
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<typename BlockResultHandler::SingleResultHandler>
+                    res;
+            std::unique_ptr<DistanceComputer> dis;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        index->ntotal, hnsw.use_visited_hashset);
+                res = std::make_unique<
+                        typename BlockResultHandler::SingleResultHandler>(bres);
+                dis.reset(storage_distance_computer(index->storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
             for (idx_t i = i0; i < i1; i++) {
-                res.begin(i);
-                dis->set_query(x + i * index->d);
-
-                HNSWStats stats = hnsw.search(*dis, index, res, vt, params);
-                n1 += stats.n1;
-                n2 += stats.n2;
-                ndis += stats.ndis;
-                nhops += stats.nhops;
-                res.end();
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    res->begin(i);
+                    dis->set_query(x + i * index->d);
+
+                    HNSWStats stats =
+                            hnsw.search(*dis, index, *res, *vt, params);
+                    n1 += stats.n1;
+                    n2 += stats.n2;
+                    ndis += stats.ndis;
+                    nhops += stats.nhops;
+                    res->end();
+                    vt->advance();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 
@@ -308,7 +332,7 @@ void IndexHNSW::search(
 
     if (is_similarity_metric(this->metric_type)) {
         // we need to revert the negated distances
-        for (size_t i = 0; i < k * n; i++) {
+        for (idx_t i = 0; i < k * n; i++) {
             distances[i] = -distances[i];
         }
     }
@@ -333,20 +357,35 @@ void IndexHNSW::range_search(
     }
 }
 
+void IndexHNSW::search1(
+        const float* x,
+        ResultHandler& handler,
+        SearchParameters* params) const {
+    SingleQueryBlockResultHandler<HNSW::C, false> bres(handler);
+    hnsw_search(this, 1, x, bres, params);
+}
+
 void IndexHNSW::add(idx_t n, const float* x) {
     FAISS_THROW_IF_NOT_MSG(
             storage,
             "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
     FAISS_THROW_IF_NOT(is_trained);
-    int n0 = ntotal;
+    size_t n0 = ntotal;
     storage->add(n, x);
     ntotal = storage->ntotal;
 
-    hnsw_add_vertices(*this, n0, n, x, verbose, hnsw.levels.size() == ntotal);
+    hnsw_add_vertices(
+            *this,
+            n0,
+            n,
+            x,
+            verbose,
+            hnsw.levels.size() == static_cast<size_t>(ntotal));
 }
 
 void IndexHNSW::reset() {
     hnsw.reset();
+    locks.clear();
     storage->reset();
     ntotal = 0;
 }
@@ -418,42 +457,59 @@ void IndexHNSW::search_level_0(
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
-    storage_idx_t ntotal = hnsw.levels.size();
+    size_t hnsw_ntotal = hnsw.levels.size();
 
     using RH = HeapBlockResultHandler<HNSW::C>;
     RH bres(n, distances, labels, k);
 
+    std::exception_ptr ex;
+    std::atomic<bool> interrupt{false};
 #pragma omp parallel
     {
-        std::unique_ptr<DistanceComputer> qdis(
-                storage_distance_computer(storage));
+        std::unique_ptr<DistanceComputer> qdis;
         HNSWStats search_stats;
-        VisitedTable vt(ntotal);
-        RH::SingleResultHandler res(bres);
+        std::unique_ptr<VisitedTable> vt;
+        std::unique_ptr<RH::SingleResultHandler> res;
+        try {
+            qdis.reset(storage_distance_computer(storage));
+            vt = std::make_unique<VisitedTable>(
+                    hnsw_ntotal, hnsw.use_visited_hashset);
+            res = std::make_unique<RH::SingleResultHandler>(bres);
+        } catch (...) {
+            omp_capture_exception(ex, [&] { interrupt = true; });
+        }
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            res.begin(i);
-            qdis->set_query(x + i * d);
-
-            hnsw.search_level_0(
-                    *qdis.get(),
-                    res,
-                    nprobe,
-                    nearest + i * nprobe,
-                    nearest_d + i * nprobe,
-                    search_type,
-                    search_stats,
-                    vt,
-                    params);
-            res.end();
-            vt.advance();
+            if (interrupt.load(std::memory_order_relaxed)) {
+                continue;
+            }
+            try {
+                res->begin(i);
+                qdis->set_query(x + i * d);
+
+                hnsw.search_level_0(
+                        *qdis.get(),
+                        *res,
+                        nprobe,
+                        nearest + i * nprobe,
+                        nearest_d + i * nprobe,
+                        search_type,
+                        search_stats,
+                        *vt,
+                        params);
+                res->end();
+                vt->advance();
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
         }
 #pragma omp critical
         {
             hnsw_stats.combine(search_stats);
         }
     }
+    omp_rethrow_if_exception(ex);
     if (is_similarity_metric(this->metric_type)) {
 // we need to revert the negated distances
 #pragma omp parallel for
@@ -478,8 +534,8 @@ void IndexHNSW::init_level_0_from_knngraph(
 
         std::priority_queue<NodeDistFarther> initial_list;
 
-        for (size_t j = 0; j < k; j++) {
-            int v1 = I[i * k + j];
+        for (int j = 0; j < k; j++) {
+            int v1 = static_cast<int>(I[i * k + j]);
             if (v1 == i) {
                 continue;
             }
@@ -509,14 +565,11 @@ void IndexHNSW::init_level_0_from_entry_points(
         int n,
         const storage_idx_t* points,
         const storage_idx_t* nearests) {
-    std::vector<omp_lock_t> locks(ntotal);
-    for (int i = 0; i < ntotal; i++) {
-        omp_init_lock(&locks[i]);
-    }
+    locks.prepare(ntotal);
 
 #pragma omp parallel
     {
-        VisitedTable vt(ntotal);
+        VisitedTable vt(ntotal, hnsw.use_visited_hashset);
 
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
@@ -530,7 +583,7 @@ void IndexHNSW::init_level_0_from_entry_points(
             dis->set_query(vec.data());
 
             hnsw.add_links_starting_from(
-                    *dis, pt_id, nearest, (*dis)(nearest), 0, locks.data(), vt);
+                    *dis, pt_id, nearest, (*dis)(nearest), 0, locks, vt);
 
             if (verbose && i % 10000 == 0) {
                 printf("  %d / %d\r", i, n);
@@ -542,8 +595,8 @@ void IndexHNSW::init_level_0_from_entry_points(
         printf("\n");
     }
 
-    for (int i = 0; i < ntotal; i++) {
-        omp_destroy_lock(&locks[i]);
+    if (!retain_locks) {
+        locks.clear();
     }
 }
 
@@ -586,7 +639,7 @@ void IndexHNSW::link_singletons() {
 
     std::vector<bool> seen(ntotal);
 
-    for (size_t i = 0; i < ntotal; i++) {
+    for (idx_t i = 0; i < ntotal; i++) {
         size_t begin, end;
         hnsw.neighbor_range(i, 0, &begin, &end);
         for (size_t j = begin; j < end; j++) {
@@ -615,8 +668,8 @@ void IndexHNSW::link_singletons() {
            n_sing_l1);
 
     std::vector<float> recons(singletons.size() * d);
-    for (int i = 0; i < singletons.size(); i++) {
-        FAISS_ASSERT(!"not implemented");
+    for (size_t i = 0; i < singletons.size(); i++) {
+        FAISS_ASSERT(false); // not implemented
     }
 }
 
@@ -640,10 +693,10 @@ IndexHNSWFlat::IndexHNSWFlat() {
     is_trained = true;
 }
 
-IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
+IndexHNSWFlat::IndexHNSWFlat(int d_in, int M, MetricType metric)
         : IndexHNSW(
-                  (metric == METRIC_L2) ? new IndexFlatL2(d)
-                                        : new IndexFlat(d, metric),
+                  (metric == METRIC_L2) ? new IndexFlatL2(d_in)
+                                        : new IndexFlat(d_in, metric),
                   M) {
     own_fields = true;
     is_trained = true;
@@ -653,44 +706,21 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
  * IndexHNSWFlatPanorama implementation
  **************************************************************/
 
-void IndexHNSWFlatPanorama::compute_cum_sums(
-        const float* x,
-        float* dst_cum_sums,
-        int d,
-        int num_panorama_levels,
-        int panorama_level_width) {
-    // Iterate backwards through levels, accumulating sum as we go.
-    // This avoids computing the suffix sum for each vector, which takes
-    // extra memory.
-
-    float sum = 0.0f;
-    dst_cum_sums[num_panorama_levels] = 0.0f;
-    for (int level = num_panorama_levels - 1; level >= 0; level--) {
-        int start_idx = level * panorama_level_width;
-        int end_idx = std::min(start_idx + panorama_level_width, d);
-        for (int j = start_idx; j < end_idx; j++) {
-            sum += x[j] * x[j];
-        }
-        dst_cum_sums[level] = std::sqrt(sum);
-    }
-}
-
 IndexHNSWFlatPanorama::IndexHNSWFlatPanorama()
         : IndexHNSWFlat(),
           cum_sums(),
-          panorama_level_width(0),
+          pano(sizeof(float), 1, 1),
           num_panorama_levels(0) {}
 
 IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
-        int d,
+        int d_in,
         int M,
-        int num_panorama_levels,
+        int num_panorama_levels_in,
         MetricType metric)
-        : IndexHNSWFlat(d, M, metric),
+        : IndexHNSWFlat(d_in, M, metric),
           cum_sums(),
-          panorama_level_width(
-                  (d + num_panorama_levels - 1) / num_panorama_levels),
-          num_panorama_levels(num_panorama_levels) {
+          pano(d_in * sizeof(float), num_panorama_levels_in, 1),
+          num_panorama_levels(num_panorama_levels_in) {
     // For now, we only support L2 distance.
     // Supporting dot product and cosine distance is a trivial addition
     // left for future work.
@@ -704,18 +734,8 @@ IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
 
 void IndexHNSWFlatPanorama::add(idx_t n, const float* x) {
     idx_t n0 = ntotal;
-    cum_sums.resize((ntotal + n) * (num_panorama_levels + 1));
-
-    for (size_t idx = 0; idx < n; idx++) {
-        const float* vector = x + idx * d;
-        compute_cum_sums(
-                vector,
-                &cum_sums[(n0 + idx) * (num_panorama_levels + 1)],
-                d,
-                num_panorama_levels,
-                panorama_level_width);
-    }
-
+    cum_sums.resize((ntotal + n) * (pano.n_levels + 1));
+    pano.compute_cumulative_sums(cum_sums.data(), n0, n, x);
     IndexHNSWFlat::add(n, x);
 }
 
@@ -725,13 +745,13 @@ void IndexHNSWFlatPanorama::reset() {
 }
 
 void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) {
-    std::vector<float> new_cum_sums(ntotal * (num_panorama_levels + 1));
+    std::vector<float> new_cum_sums(ntotal * (pano.n_levels + 1));
 
     for (idx_t i = 0; i < ntotal; i++) {
         idx_t src = perm[i];
-        memcpy(&new_cum_sums[i * (num_panorama_levels + 1)],
-               &cum_sums[src * (num_panorama_levels + 1)],
-               (num_panorama_levels + 1) * sizeof(float));
+        memcpy(&new_cum_sums[i * (pano.n_levels + 1)],
+               &cum_sums[src * (pano.n_levels + 1)],
+               (pano.n_levels + 1) * sizeof(float));
     }
 
     std::swap(cum_sums, new_cum_sums);
@@ -745,12 +765,12 @@ void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) {
 IndexHNSWPQ::IndexHNSWPQ() = default;
 
 IndexHNSWPQ::IndexHNSWPQ(
-        int d,
+        int d_in,
         int pq_m,
         int M,
         int pq_nbits,
         MetricType metric)
-        : IndexHNSW(new IndexPQ(d, pq_m, pq_nbits, metric), M) {
+        : IndexHNSW(new IndexPQ(d_in, pq_m, pq_nbits, metric), M) {
     own_fields = true;
     is_trained = false;
 }
@@ -765,11 +785,11 @@ void IndexHNSWPQ::train(idx_t n, const float* x) {
  **************************************************************/
 
 IndexHNSWSQ::IndexHNSWSQ(
-        int d,
+        int d_in,
         ScalarQuantizer::QuantizerType qtype,
         int M,
         MetricType metric)
-        : IndexHNSW(new IndexScalarQuantizer(d, qtype, metric), M) {
+        : IndexHNSW(new IndexScalarQuantizer(d_in, qtype, metric), M) {
     is_trained = this->storage->is_trained;
     own_fields = true;
 }
@@ -882,8 +902,11 @@ void IndexHNSW2Level::search(
 
         const IndexIVFPQ* index_ivfpq =
                 dynamic_cast<const IndexIVFPQ*>(storage);
+        FAISS_THROW_IF_NOT_MSG(
+                index_ivfpq,
+                "IndexHNSW2Level mixed search requires IndexIVFPQ storage");
 
-        int nprobe = index_ivfpq->nprobe;
+        size_t nprobe = index_ivfpq->nprobe;
 
         std::unique_ptr<idx_t[]> coarse_assign(new idx_t[n * nprobe]);
         std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
@@ -901,71 +924,88 @@ void IndexHNSW2Level::search(
                 labels,
                 false);
 
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel
         {
-            VisitedTable vt(ntotal);
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
-
+            // visited table (not hash set) for tri-state flags.
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<DistanceComputer> dis;
             constexpr int candidates_size = 1;
-            MinimaxHeap candidates(candidates_size);
+            std::unique_ptr<MinimaxHeap> candidates;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        ntotal, /*use_hashset=*/false);
+                dis.reset(storage_distance_computer(storage));
+                candidates = std::make_unique<MinimaxHeap>(candidates_size);
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for reduction(+ : n1, n2, ndis, nhops)
             for (idx_t i = 0; i < n; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
-
-                // mark all inverted list elements as visited
-
-                for (int j = 0; j < nprobe; j++) {
-                    idx_t key = coarse_assign[j + i * nprobe];
-                    if (key < 0) {
-                        break;
-                    }
-                    size_t list_length = index_ivfpq->get_list_size(key);
-                    const idx_t* ids = index_ivfpq->invlists->get_ids(key);
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    idx_t* idxi = labels + i * k;
+                    float* simi = distances + i * k;
+                    dis->set_query(x + i * d);
+
+                    // mark all inverted list elements as visited
+                    for (size_t j = 0; j < nprobe; j++) {
+                        idx_t key = coarse_assign[j + i * nprobe];
+                        if (key < 0) {
+                            break;
+                        }
+                        size_t list_length = index_ivfpq->get_list_size(key);
+                        const idx_t* ids = index_ivfpq->invlists->get_ids(key);
 
-                    for (int jj = 0; jj < list_length; jj++) {
-                        vt.set(ids[jj]);
+                        for (size_t jj = 0; jj < list_length; jj++) {
+                            vt->set(ids[jj]);
+                        }
                     }
-                }
 
-                candidates.clear();
+                    candidates->clear();
 
-                for (int j = 0; j < k; j++) {
-                    if (idxi[j] < 0) {
-                        break;
+                    for (int j = 0; j < k; j++) {
+                        if (idxi[j] < 0) {
+                            break;
+                        }
+                        candidates->push(
+                                static_cast<storage_idx_t>(idxi[j]), simi[j]);
                     }
-                    candidates.push(idxi[j], simi[j]);
-                }
-
-                // reorder from sorted to heap
-                maxheap_heapify(k, simi, idxi, simi, idxi, k);
-
-                HNSWStats search_stats;
-                search_from_candidates_2(
-                        hnsw,
-                        *dis,
-                        k,
-                        idxi,
-                        simi,
-                        candidates,
-                        vt,
-                        search_stats,
-                        0,
-                        k);
-                n1 += search_stats.n1;
-                n2 += search_stats.n2;
-                ndis += search_stats.ndis;
-                nhops += search_stats.nhops;
 
-                vt.advance();
-                vt.advance();
+                    // reorder from sorted to heap
+                    maxheap_heapify(k, simi, idxi, simi, idxi, k);
 
-                maxheap_reorder(k, simi, idxi);
+                    HNSWStats search_stats;
+                    search_from_candidates_2(
+                            hnsw,
+                            *dis,
+                            k,
+                            idxi,
+                            simi,
+                            *candidates,
+                            *vt,
+                            search_stats,
+                            0,
+                            k);
+                    n1 += search_stats.n1;
+                    n2 += search_stats.n2;
+                    ndis += search_stats.ndis;
+                    nhops += search_stats.nhops;
+
+                    vt->advance();
+                    vt->advance();
+
+                    maxheap_reorder(k, simi, idxi);
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
 
         hnsw_stats.combine({n1, n2, ndis, nhops});
     }
@@ -1002,11 +1042,11 @@ IndexHNSWCagra::IndexHNSWCagra() {
 }
 
 IndexHNSWCagra::IndexHNSWCagra(
-        int d,
+        int d_in,
         int M,
         MetricType metric,
         NumericType numeric_type)
-        : IndexHNSW(d, M, metric) {
+        : IndexHNSW(d_in, M, metric) {
     FAISS_THROW_IF_NOT_MSG(
             ((metric == METRIC_L2) || (metric == METRIC_INNER_PRODUCT)),
             "unsupported metric type for IndexHNSWCagra");
@@ -1050,10 +1090,18 @@ void IndexHNSWCagra::search(
     if (!base_level_only) {
         IndexHNSW::search(n, x, k, distances, labels, params);
     } else {
+        if (ntotal == 0) {
+            std::fill(
+                    distances,
+                    distances + n * k,
+                    std::numeric_limits<float>::max());
+            std::fill(labels, labels + n * k, -1);
+            return;
+        }
         std::vector<storage_idx_t> nearest(n);
         std::vector<float> nearest_d(n);
 
-#pragma omp for
+#pragma omp parallel for
         for (idx_t i = 0; i < n; i++) {
             std::unique_ptr<DistanceComputer> dis(
                     storage_distance_computer(this->storage));
@@ -1069,7 +1117,7 @@ void IndexHNSWCagra::search(
                 auto idx = distrib(gen);
                 auto distance = (*dis)(idx);
                 if (distance < nearest_d[i]) {
-                    nearest[i] = idx;
+                    nearest[i] = static_cast<storage_idx_t>(idx);
                     nearest_d[i] = distance;
                 }
             }
@@ -1091,6 +1139,70 @@ void IndexHNSWCagra::search(
     }
 }
 
+void IndexHNSWCagra::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    if (!base_level_only) {
+        IndexHNSW::range_search(n, x, radius, result, params);
+        return;
+    }
+
+    const HNSW& hnsw = this->hnsw;
+    size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
+    float threshold = is_similarity_metric(metric_type) ? -radius : radius;
+    RangeSearchPartialResult pres(result);
+
+    for (idx_t i = 0; i < n; i++) {
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
+        dis->set_query(x + i * d);
+
+        storage_idx_t nearest = -1;
+        float nearest_d = std::numeric_limits<float>::max();
+
+        std::random_device rd;
+        std::mt19937 gen(rd());
+        std::uniform_int_distribution<idx_t> distrib(0, ntotal - 1);
+
+        for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
+            auto idx = distrib(gen);
+            auto distance = (*dis)(idx);
+            if (distance < nearest_d) {
+                nearest = idx;
+                nearest_d = distance;
+            }
+        }
+        FAISS_THROW_IF_NOT_MSG(
+                nearest >= 0, "Could not find a valid entrypoint.");
+
+        RangeQueryResult& qres = pres.new_result(i);
+        RangeResultHandler<HNSW::C> res(&qres, threshold);
+        VisitedTable vt(ntotal, hnsw.use_visited_hashset);
+        HNSWStats stats;
+        hnsw.search_level_0(
+                *dis, res, 1, &nearest, &nearest_d, 1, stats, vt, params);
+        n1 += stats.n1;
+        n2 += stats.n2;
+        ndis += stats.ndis;
+        nhops += stats.nhops;
+    }
+
+    pres.set_lims();
+    result->do_allocation();
+    pres.copy_result();
+
+    hnsw_stats.combine({n1, n2, ndis, nhops});
+
+    if (is_similarity_metric(metric_type)) {
+        for (size_t i = 0; i < result->lims[result->nq]; i++) {
+            result->distances[i] = -result->distances[i];
+        }
+    }
+}
+
 faiss::NumericType IndexHNSWCagra::get_numeric_type() const {
     return numeric_type_;
 }