faiss 0.6.0 → 0.6.2

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,14 +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;
 
@@ -45,12 +47,17 @@ HNSWStats hnsw_stats;
 
 namespace {
 
+// Returns the storage's native distance computer. For similarity metrics
+// (e.g. METRIC_INNER_PRODUCT), distance values are real similarity scores
+// (larger = better); HNSW handles the ordering via `hnsw.is_similarity`.
+//
+// NOTE: callers that drive the legacy max-heap-only code paths (notably
+// `search_from_candidates_2` in the IndexHNSW2Level mixed search) cannot
+// consume similarity scores directly; they assume smaller-is-better.
+// Those paths only fire for the (default) L2 IndexHNSW2Level + Index2Layer
+// configuration today, so passing the native DC is safe in practice.
 DistanceComputer* storage_distance_computer(const Index* storage) {
-    if (is_similarity_metric(storage->metric_type)) {
-        return new NegativeDistanceComputer(storage->get_distance_computer());
-    } else {
-        return storage->get_distance_computer();
-    }
+    return storage->get_distance_computer();
 }
 
 void hnsw_add_vertices(
@@ -82,10 +89,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;
@@ -94,10 +99,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]++;
@@ -105,13 +110,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;
         }
@@ -123,39 +128,42 @@ 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, hnsw.use_visited_hashset);
+                std::unique_ptr<VisitedTable> vt =
+                        VisitedTable::create(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);
 
@@ -169,12 +177,12 @@ void hnsw_add_vertices(
                             pt_level,
                             pt_id,
                             locks,
-                            vt,
+                            *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) {
@@ -199,9 +207,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();
     }
 }
 
@@ -211,12 +218,17 @@ 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) {
+    hnsw.is_similarity = is_similarity_metric(metric);
+}
 
-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;
+    hnsw.is_similarity = is_similarity_metric(metric_type);
 }
 
 IndexHNSW::~IndexHNSW() {
@@ -263,28 +275,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, hnsw.use_visited_hashset);
-            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 = VisitedTable::create(
+                        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();
     }
 
@@ -302,16 +334,14 @@ void IndexHNSW::search(
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
 
-    using RH = HeapBlockResultHandler<HNSW::C>;
-    RH bres(n, distances, labels, k);
-
-    hnsw_search(this, n, x, bres, params);
-
     if (is_similarity_metric(this->metric_type)) {
-        // we need to revert the negated distances
-        for (size_t i = 0; i < k * n; i++) {
-            distances[i] = -distances[i];
-        }
+        using RH = HeapBlockResultHandler<HNSW::C_similarity>;
+        RH bres(n, distances, labels, k);
+        hnsw_search(this, n, x, bres, params);
+    } else {
+        using RH = HeapBlockResultHandler<HNSW::C_distance>;
+        RH bres(n, distances, labels, k);
+        hnsw_search(this, n, x, bres, params);
     }
 }
 
@@ -321,16 +351,14 @@ void IndexHNSW::range_search(
         float radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
-    using RH = RangeSearchBlockResultHandler<HNSW::C>;
-    RH bres(result, is_similarity_metric(metric_type) ? -radius : radius);
-
-    hnsw_search(this, n, x, bres, params);
-
-    if (is_similarity_metric(this->metric_type)) {
-        // we need to revert the negated distances
-        for (size_t i = 0; i < result->lims[result->nq]; i++) {
-            result->distances[i] = -result->distances[i];
-        }
+    if (is_similarity_metric(metric_type)) {
+        using RH = RangeSearchBlockResultHandler<HNSW::C_similarity>;
+        RH bres(result, radius);
+        hnsw_search(this, n, x, bres, params);
+    } else {
+        using RH = RangeSearchBlockResultHandler<HNSW::C_distance>;
+        RH bres(result, radius);
+        hnsw_search(this, n, x, bres, params);
     }
 }
 
@@ -338,8 +366,13 @@ void IndexHNSW::search1(
         const float* x,
         ResultHandler& handler,
         SearchParameters* params) const {
-    SingleQueryBlockResultHandler<HNSW::C, false> bres(handler);
-    hnsw_search(this, 1, x, bres, params);
+    if (is_similarity_metric(metric_type)) {
+        SingleQueryBlockResultHandler<HNSW::C_similarity, false> bres(handler);
+        hnsw_search(this, 1, x, bres, params);
+    } else {
+        SingleQueryBlockResultHandler<HNSW::C_distance, false> bres(handler);
+        hnsw_search(this, 1, x, bres, params);
+    }
 }
 
 void IndexHNSW::add(idx_t n, const float* x) {
@@ -347,15 +380,22 @@ void IndexHNSW::add(idx_t n, const float* x) {
             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;
 }
@@ -427,48 +467,66 @@ 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);
+    auto run = [&]<class C>() {
+        using RH = HeapBlockResultHandler<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));
-        HNSWStats search_stats;
-        VisitedTable vt(ntotal, hnsw.use_visited_hashset);
-        RH::SingleResultHandler res(bres);
+        {
+            std::unique_ptr<DistanceComputer> qdis;
+            HNSWStats search_stats;
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<typename RH::SingleResultHandler> res;
+            try {
+                qdis.reset(storage_distance_computer(storage));
+                vt = VisitedTable::create(
+                        hnsw_ntotal, hnsw.use_visited_hashset);
+                res = std::make_unique<typename 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();
-        }
+            for (idx_t i = 0; i < n; i++) {
+                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);
+            {
+                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
-        for (int64_t i = 0; i < k * n; i++) {
-            distances[i] = -distances[i];
-        }
+        run.template operator()<HNSW::C_similarity>();
+    } else {
+        run.template operator()<HNSW::C_distance>();
     }
 }
 
@@ -487,8 +545,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;
             }
@@ -518,14 +576,12 @@ 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, hnsw.use_visited_hashset);
+        std::unique_ptr<VisitedTable> vt =
+                VisitedTable::create(ntotal, hnsw.use_visited_hashset);
 
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
@@ -539,7 +595,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);
@@ -551,8 +607,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();
     }
 }
 
@@ -595,7 +651,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++) {
@@ -624,7 +680,7 @@ void IndexHNSW::link_singletons() {
            n_sing_l1);
 
     std::vector<float> recons(singletons.size() * d);
-    for (int i = 0; i < singletons.size(); i++) {
+    for (size_t i = 0; i < singletons.size(); i++) {
         FAISS_ASSERT(false); // not implemented
     }
 }
@@ -649,10 +705,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;
@@ -663,17 +719,20 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
  **************************************************************/
 
 IndexHNSWFlatPanorama::IndexHNSWFlatPanorama()
-        : IndexHNSWFlat(), cum_sums(), pano(0, 1, 1), num_panorama_levels(0) {}
+        : IndexHNSWFlat(),
+          cum_sums(),
+          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(),
-          pano(d * sizeof(float), num_panorama_levels, 1),
-          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.
@@ -718,12 +777,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;
 }
@@ -738,11 +797,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;
 }
@@ -777,7 +836,7 @@ int search_from_candidates_2(
         idx_t* I,
         float* D,
         MinimaxHeap& candidates,
-        VisitedTable& vt,
+        VisitedTableVector& vt,
         HNSWStats& stats,
         int level,
         int nres_in = 0) {
@@ -855,8 +914,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]);
@@ -874,72 +936,87 @@ void IndexHNSW2Level::search(
                 labels,
                 false);
 
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel
         {
             // visited table (not hash set) for tri-state flags.
-            VisitedTable vt(ntotal, /*use_hashset=*/false);
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
-
+            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 = VisitedTable::create(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
+                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 j = 0; j < nprobe; j++) {
-                    idx_t key = coarse_assign[j + i * nprobe];
-                    if (key < 0) {
-                        break;
+                        for (size_t jj = 0; jj < list_length; jj++) {
+                            vt->set(ids[jj]);
+                        }
                     }
-                    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]);
+                    candidates->clear();
+
+                    for (int j = 0; j < k; j++) {
+                        if (idxi[j] < 0) {
+                            break;
+                        }
+                        candidates->push(
+                                static_cast<storage_idx_t>(idxi[j]), simi[j]);
                     }
-                }
 
-                candidates.clear();
+                    // reorder from sorted to heap
+                    maxheap_heapify(k, simi, idxi, simi, idxi, k);
 
-                for (int j = 0; j < k; j++) {
-                    if (idxi[j] < 0) {
-                        break;
-                    }
-                    candidates.push(idxi[j], simi[j]);
+                    HNSWStats search_stats;
+                    search_from_candidates_2(
+                            hnsw,
+                            *dis,
+                            k,
+                            idxi,
+                            simi,
+                            *candidates,
+                            static_cast<VisitedTableVector&>(*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; });
                 }
-
-                // 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();
-
-                maxheap_reorder(k, simi, idxi);
             }
         }
+        omp_rethrow_if_exception(ex);
 
         hnsw_stats.combine({n1, n2, ndis, nhops});
     }
@@ -976,11 +1053,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");
@@ -1024,31 +1101,51 @@ 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
-        for (idx_t i = 0; i < n; i++) {
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(this->storage));
-            dis->set_query(x + i * d);
-            nearest[i] = -1;
-            nearest_d[i] = std::numeric_limits<float>::max();
-
-            std::random_device rd;
-            std::mt19937 gen(rd());
-            std::uniform_int_distribution<idx_t> distrib(0, this->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[i]) {
-                    nearest[i] = idx;
-                    nearest_d[i] = distance;
+        auto pick_entrypoints = [&]<class C>() {
+#pragma omp parallel for
+            for (idx_t i = 0; i < n; i++) {
+                std::unique_ptr<DistanceComputer> dis(
+                        storage_distance_computer(this->storage));
+                dis->set_query(x + i * d);
+                nearest[i] = -1;
+                // C::neutral() is the "worst possible" value: +inf for
+                // CMax (distance) and -inf for CMin (similarity). The
+                // first real candidate will always be strictly better.
+                nearest_d[i] = C::neutral();
+
+                std::random_device rd;
+                std::mt19937 gen(rd());
+                std::uniform_int_distribution<idx_t> distrib(
+                        0, this->ntotal - 1);
+
+                for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
+                    auto idx = distrib(gen);
+                    auto distance = (*dis)(idx);
+                    if (C::cmp(nearest_d[i], distance)) {
+                        nearest[i] = static_cast<storage_idx_t>(idx);
+                        nearest_d[i] = distance;
+                    }
                 }
+                FAISS_THROW_IF_NOT_MSG(
+                        nearest[i] >= 0, "Could not find a valid entrypoint.");
             }
-            FAISS_THROW_IF_NOT_MSG(
-                    nearest[i] >= 0, "Could not find a valid entrypoint.");
+        };
+
+        if (is_similarity_metric(metric_type)) {
+            pick_entrypoints.template operator()<HNSW::C_similarity>();
+        } else {
+            pick_entrypoints.template operator()<HNSW::C_distance>();
         }
 
         search_level_0(
@@ -1065,6 +1162,77 @@ 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;
+    }
+
+    auto run = [&]<class C>() {
+        const HNSW& hnsw = this->hnsw;
+        size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
+        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;
+            // C::neutral() is the "worst possible" value under C: +inf for
+            // CMax (distance) and -inf for CMin (similarity). The first
+            // real candidate will always be strictly better.
+            float nearest_d = C::neutral();
+
+            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);
+                // C::cmp(nearest_d, distance) is true iff distance is
+                // strictly better than the current nearest_d.
+                if (C::cmp(nearest_d, distance)) {
+                    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<C> res(&qres, radius);
+            std::unique_ptr<VisitedTable> vt =
+                    VisitedTable::create(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)) {
+        run.template operator()<HNSW::C_similarity>();
+    } else {
+        run.template operator()<HNSW::C_distance>();
+    }
+}
+
 faiss::NumericType IndexHNSWCagra::get_numeric_type() const {
     return numeric_type_;
 }