faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #include <faiss/IndexHNSW.h>
 
 #include <omp.h>
@@ -17,43 +15,26 @@
 #include <cstdlib>
 #include <cstring>
 
+#include <limits>
+#include <memory>
 #include <queue>
+#include <random>
 #include <unordered_set>
 
-#include <stdint.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <cstdint>
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/Heap.h>
+#include <faiss/impl/ResultHandler.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
 
-extern "C" {
-
-/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
-
-int sgemm_(
-        const char* transa,
-        const char* transb,
-        FINTEGER* m,
-        FINTEGER* n,
-        FINTEGER* k,
-        const float* alpha,
-        const float* a,
-        FINTEGER* lda,
-        const float* b,
-        FINTEGER* ldb,
-        float* beta,
-        float* c,
-        FINTEGER* ldc);
-}
-
 namespace faiss {
 
 using MinimaxHeap = HNSW::MinimaxHeap;
@@ -68,35 +49,6 @@ HNSWStats hnsw_stats;
 
 namespace {
 
-/* Wrap the distance computer into one that negates the
-   distances. This makes supporting INNER_PRODUCE search easier */
-
-struct NegativeDistanceComputer : DistanceComputer {
-    /// owned by this
-    DistanceComputer* basedis;
-
-    explicit NegativeDistanceComputer(DistanceComputer* basedis)
-            : basedis(basedis) {}
-
-    void set_query(const float* x) override {
-        basedis->set_query(x);
-    }
-
-    /// compute distance of vector i to current query
-    float operator()(idx_t i) override {
-        return -(*basedis)(i);
-    }
-
-    /// compute distance between two stored vectors
-    float symmetric_dis(idx_t i, idx_t j) override {
-        return -basedis->symmetric_dis(i, j);
-    }
-
-    virtual ~NegativeDistanceComputer() {
-        delete basedis;
-    }
-};
-
 DistanceComputer* storage_distance_computer(const Index* storage) {
     if (is_similarity_metric(storage->metric_type)) {
         return new NegativeDistanceComputer(storage->get_distance_computer());
@@ -175,7 +127,9 @@ void hnsw_add_vertices(
 
         int i1 = n;
 
-        for (int pt_level = hist.size() - 1; pt_level >= 0; pt_level--) {
+        for (int pt_level = hist.size() - 1;
+             pt_level >= !index_hnsw.init_level0;
+             pt_level--) {
             int i0 = i1 - hist[pt_level];
 
             if (verbose) {
@@ -192,9 +146,8 @@ void hnsw_add_vertices(
             {
                 VisitedTable vt(ntotal);
 
-                DistanceComputer* dis =
-                        storage_distance_computer(index_hnsw.storage);
-                ScopeDeleter1<DistanceComputer> del(dis);
+                std::unique_ptr<DistanceComputer> dis(
+                        storage_distance_computer(index_hnsw.storage));
                 int prev_display =
                         verbose && omp_get_thread_num() == 0 ? 0 : -1;
                 size_t counter = 0;
@@ -212,7 +165,13 @@ void hnsw_add_vertices(
                         continue;
                     }
 
-                    hnsw.add_with_locks(*dis, pt_level, pt_id, locks, vt);
+                    hnsw.add_with_locks(
+                            *dis,
+                            pt_level,
+                            pt_id,
+                            locks,
+                            vt,
+                            index_hnsw.keep_max_size_level0 && (pt_level == 0));
 
                     if (prev_display >= 0 && i - i0 > prev_display + 10000) {
                         prev_display = i - i0;
@@ -232,7 +191,11 @@ void hnsw_add_vertices(
             }
             i1 = i0;
         }
-        FAISS_ASSERT(i1 == 0);
+        if (index_hnsw.init_level0) {
+            FAISS_ASSERT(i1 == 0);
+        } else {
+            FAISS_ASSERT((i1 - hist[0]) == 0);
+        }
     }
     if (verbose) {
         printf("Done in %.3f ms\n", getmillisecs() - t0);
@@ -250,18 +213,10 @@ void hnsw_add_vertices(
  **************************************************************/
 
 IndexHNSW::IndexHNSW(int d, int M, MetricType metric)
-        : Index(d, metric),
-          hnsw(M),
-          own_fields(false),
-          storage(nullptr),
-          reconstruct_from_neighbors(nullptr) {}
+        : Index(d, metric), hnsw(M) {}
 
 IndexHNSW::IndexHNSW(Index* storage, int M)
-        : Index(storage->d, storage->metric_type),
-          hnsw(M),
-          own_fields(false),
-          storage(storage),
-          reconstruct_from_neighbors(nullptr) {}
+        : Index(storage->d, storage->metric_type), hnsw(M), storage(storage) {}
 
 IndexHNSW::~IndexHNSW() {
     if (own_fields) {
@@ -278,18 +233,21 @@ void IndexHNSW::train(idx_t n, const float* x) {
     is_trained = true;
 }
 
-void IndexHNSW::search(
+namespace {
+
+template <class BlockResultHandler>
+void hnsw_search(
+        const IndexHNSW* index,
         idx_t n,
         const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels,
-        const SearchParameters* params_in) const {
-    FAISS_THROW_IF_NOT(k > 0);
+        BlockResultHandler& bres,
+        const SearchParameters* params_in) {
     FAISS_THROW_IF_NOT_MSG(
-            storage,
-            "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
+            index->storage,
+            "No storage index, please use IndexHNSWFlat (or variants) "
+            "instead of IndexHNSW directly");
     const SearchParametersHNSW* params = nullptr;
+    const HNSW& hnsw = index->hnsw;
 
     int efSearch = hnsw.efSearch;
     if (params_in) {
@@ -297,63 +255,82 @@ void IndexHNSW::search(
         FAISS_THROW_IF_NOT_MSG(params, "params type invalid");
         efSearch = params->efSearch;
     }
-    size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
+    size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
 
-    idx_t check_period =
-            InterruptCallback::get_period_hint(hnsw.max_level * d * efSearch);
+    idx_t check_period = InterruptCallback::get_period_hint(
+            hnsw.max_level * index->d * efSearch);
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
 
-#pragma omp parallel
+#pragma omp parallel if (i1 - i0 > 1)
         {
-            VisitedTable vt(ntotal);
+            VisitedTable vt(index->ntotal);
+            typename BlockResultHandler::SingleResultHandler res(bres);
 
-            DistanceComputer* dis = storage_distance_computer(storage);
-            ScopeDeleter1<DistanceComputer> del(dis);
+            std::unique_ptr<DistanceComputer> dis(
+                    storage_distance_computer(index->storage));
 
-#pragma omp for reduction(+ : n1, n2, n3, ndis, nreorder)
+#pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
             for (idx_t i = i0; i < i1; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
+                res.begin(i);
+                dis->set_query(x + i * index->d);
 
-                maxheap_heapify(k, simi, idxi);
-                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt, params);
+                HNSWStats stats = hnsw.search(*dis, res, vt, params);
                 n1 += stats.n1;
                 n2 += stats.n2;
-                n3 += stats.n3;
                 ndis += stats.ndis;
-                nreorder += stats.nreorder;
-                maxheap_reorder(k, simi, idxi);
-
-                if (reconstruct_from_neighbors &&
-                    reconstruct_from_neighbors->k_reorder != 0) {
-                    int k_reorder = reconstruct_from_neighbors->k_reorder;
-                    if (k_reorder == -1 || k_reorder > k)
-                        k_reorder = k;
-
-                    nreorder += reconstruct_from_neighbors->compute_distances(
-                            k_reorder, idxi, x + i * d, simi);
-
-                    // sort top k_reorder
-                    maxheap_heapify(
-                            k_reorder, simi, idxi, simi, idxi, k_reorder);
-                    maxheap_reorder(k_reorder, simi, idxi);
-                }
+                nhops += stats.nhops;
+                res.end();
             }
         }
         InterruptCallback::check();
     }
 
-    if (is_similarity_metric(metric_type)) {
+    hnsw_stats.combine({n1, n2, ndis, nhops});
+}
+
+} // anonymous namespace
+
+void IndexHNSW::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params_in) 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_in);
+
+    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];
         }
     }
+}
+
+void IndexHNSW::range_search(
+        idx_t n,
+        const float* x,
+        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);
 
-    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
+    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];
+        }
+    }
 }
 
 void IndexHNSW::add(idx_t n, const float* x) {
@@ -381,8 +358,8 @@ void IndexHNSW::reconstruct(idx_t key, float* recons) const {
 void IndexHNSW::shrink_level_0_neighbors(int new_size) {
 #pragma omp parallel
     {
-        DistanceComputer* dis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(dis);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
 
 #pragma omp for
         for (idx_t i = 0; i < ntotal; i++) {
@@ -423,45 +400,59 @@ void IndexHNSW::search_level_0(
         float* distances,
         idx_t* labels,
         int nprobe,
-        int search_type) const {
+        int search_type,
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
+    const SearchParametersHNSW* params = nullptr;
+
+    if (params_in) {
+        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "params type invalid");
+    }
+
     storage_idx_t ntotal = hnsw.levels.size();
 
+    using RH = HeapBlockResultHandler<HNSW::C>;
+    RH bres(n, distances, labels, k);
+
 #pragma omp parallel
     {
         std::unique_ptr<DistanceComputer> qdis(
                 storage_distance_computer(storage));
         HNSWStats search_stats;
         VisitedTable vt(ntotal);
+        RH::SingleResultHandler res(bres);
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            idx_t* idxi = labels + i * k;
-            float* simi = distances + i * k;
-
+            res.begin(i);
             qdis->set_query(x + i * d);
-            maxheap_heapify(k, simi, idxi);
 
             hnsw.search_level_0(
                     *qdis.get(),
-                    k,
-                    idxi,
-                    simi,
+                    res,
                     nprobe,
                     nearest + i * nprobe,
                     nearest_d + i * nprobe,
                     search_type,
                     search_stats,
-                    vt);
-
+                    vt,
+                    params);
+            res.end();
             vt.advance();
-            maxheap_reorder(k, simi, idxi);
         }
 #pragma omp critical
         { hnsw_stats.combine(search_stats); }
     }
+    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];
+        }
+    }
 }
 
 void IndexHNSW::init_level_0_from_knngraph(
@@ -515,8 +506,8 @@ void IndexHNSW::init_level_0_from_entry_points(
     {
         VisitedTable vt(ntotal);
 
-        DistanceComputer* dis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(dis);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
         std::vector<float> vec(storage->d);
 
 #pragma omp for schedule(dynamic)
@@ -551,8 +542,8 @@ void IndexHNSW::reorder_links() {
         std::vector<float> distances(M);
         std::vector<size_t> order(M);
         std::vector<storage_idx_t> tmp(M);
-        DistanceComputer* dis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(dis);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
 
 #pragma omp for
         for (storage_idx_t i = 0; i < ntotal; i++) {
@@ -614,245 +605,16 @@ void IndexHNSW::link_singletons() {
     }
 }
 
-/**************************************************************
- * ReconstructFromNeighbors implementation
- **************************************************************/
-
-ReconstructFromNeighbors::ReconstructFromNeighbors(
-        const IndexHNSW& index,
-        size_t k,
-        size_t nsq)
-        : index(index), k(k), nsq(nsq) {
-    M = index.hnsw.nb_neighbors(0);
-    FAISS_ASSERT(k <= 256);
-    code_size = k == 1 ? 0 : nsq;
-    ntotal = 0;
-    d = index.d;
-    FAISS_ASSERT(d % nsq == 0);
-    dsub = d / nsq;
-    k_reorder = -1;
-}
-
-void ReconstructFromNeighbors::reconstruct(
-        storage_idx_t i,
-        float* x,
-        float* tmp) const {
-    const HNSW& hnsw = index.hnsw;
-    size_t begin, end;
-    hnsw.neighbor_range(i, 0, &begin, &end);
-
-    if (k == 1 || nsq == 1) {
-        const float* beta;
-        if (k == 1) {
-            beta = codebook.data();
-        } else {
-            int idx = codes[i];
-            beta = codebook.data() + idx * (M + 1);
-        }
-
-        float w0 = beta[0]; // weight of image itself
-        index.storage->reconstruct(i, tmp);
-
-        for (int l = 0; l < d; l++)
-            x[l] = w0 * tmp[l];
-
-        for (size_t j = begin; j < end; j++) {
-            storage_idx_t ji = hnsw.neighbors[j];
-            if (ji < 0)
-                ji = i;
-            float w = beta[j - begin + 1];
-            index.storage->reconstruct(ji, tmp);
-            for (int l = 0; l < d; l++)
-                x[l] += w * tmp[l];
-        }
-    } else if (nsq == 2) {
-        int idx0 = codes[2 * i];
-        int idx1 = codes[2 * i + 1];
-
-        const float* beta0 = codebook.data() + idx0 * (M + 1);
-        const float* beta1 = codebook.data() + (idx1 + k) * (M + 1);
-
-        index.storage->reconstruct(i, tmp);
-
-        float w0;
-
-        w0 = beta0[0];
-        for (int l = 0; l < dsub; l++)
-            x[l] = w0 * tmp[l];
-
-        w0 = beta1[0];
-        for (int l = dsub; l < d; l++)
-            x[l] = w0 * tmp[l];
-
-        for (size_t j = begin; j < end; j++) {
-            storage_idx_t ji = hnsw.neighbors[j];
-            if (ji < 0)
-                ji = i;
-            index.storage->reconstruct(ji, tmp);
-            float w;
-            w = beta0[j - begin + 1];
-            for (int l = 0; l < dsub; l++)
-                x[l] += w * tmp[l];
-
-            w = beta1[j - begin + 1];
-            for (int l = dsub; l < d; l++)
-                x[l] += w * tmp[l];
-        }
-    } else {
-        std::vector<const float*> betas(nsq);
-        {
-            const float* b = codebook.data();
-            const uint8_t* c = &codes[i * code_size];
-            for (int sq = 0; sq < nsq; sq++) {
-                betas[sq] = b + (*c++) * (M + 1);
-                b += (M + 1) * k;
-            }
-        }
-
-        index.storage->reconstruct(i, tmp);
-        {
-            int d0 = 0;
-            for (int sq = 0; sq < nsq; sq++) {
-                float w = *(betas[sq]++);
-                int d1 = d0 + dsub;
-                for (int l = d0; l < d1; l++) {
-                    x[l] = w * tmp[l];
-                }
-                d0 = d1;
-            }
-        }
-
-        for (size_t j = begin; j < end; j++) {
-            storage_idx_t ji = hnsw.neighbors[j];
-            if (ji < 0)
-                ji = i;
-
-            index.storage->reconstruct(ji, tmp);
-            int d0 = 0;
-            for (int sq = 0; sq < nsq; sq++) {
-                float w = *(betas[sq]++);
-                int d1 = d0 + dsub;
-                for (int l = d0; l < d1; l++) {
-                    x[l] += w * tmp[l];
-                }
-                d0 = d1;
-            }
-        }
-    }
-}
-
-void ReconstructFromNeighbors::reconstruct_n(
-        storage_idx_t n0,
-        storage_idx_t ni,
-        float* x) const {
-#pragma omp parallel
-    {
-        std::vector<float> tmp(index.d);
-#pragma omp for
-        for (storage_idx_t i = 0; i < ni; i++) {
-            reconstruct(n0 + i, x + i * index.d, tmp.data());
-        }
-    }
-}
-
-size_t ReconstructFromNeighbors::compute_distances(
-        size_t n,
-        const idx_t* shortlist,
-        const float* query,
-        float* distances) const {
-    std::vector<float> tmp(2 * index.d);
-    size_t ncomp = 0;
-    for (int i = 0; i < n; i++) {
-        if (shortlist[i] < 0)
-            break;
-        reconstruct(shortlist[i], tmp.data(), tmp.data() + index.d);
-        distances[i] = fvec_L2sqr(query, tmp.data(), index.d);
-        ncomp++;
-    }
-    return ncomp;
-}
-
-void ReconstructFromNeighbors::get_neighbor_table(storage_idx_t i, float* tmp1)
-        const {
-    const HNSW& hnsw = index.hnsw;
-    size_t begin, end;
-    hnsw.neighbor_range(i, 0, &begin, &end);
-    size_t d = index.d;
-
-    index.storage->reconstruct(i, tmp1);
-
-    for (size_t j = begin; j < end; j++) {
-        storage_idx_t ji = hnsw.neighbors[j];
-        if (ji < 0)
-            ji = i;
-        index.storage->reconstruct(ji, tmp1 + (j - begin + 1) * d);
-    }
-}
-
-/// called by add_codes
-void ReconstructFromNeighbors::estimate_code(
-        const float* x,
-        storage_idx_t i,
-        uint8_t* code) const {
-    // fill in tmp table with the neighbor values
-    float* tmp1 = new float[d * (M + 1) + (d * k)];
-    float* tmp2 = tmp1 + d * (M + 1);
-    ScopeDeleter<float> del(tmp1);
-
-    // collect coordinates of base
-    get_neighbor_table(i, tmp1);
-
-    for (size_t sq = 0; sq < nsq; sq++) {
-        int d0 = sq * dsub;
-
-        {
-            FINTEGER ki = k, di = d, m1 = M + 1;
-            FINTEGER dsubi = dsub;
-            float zero = 0, one = 1;
-
-            sgemm_("N",
-                   "N",
-                   &dsubi,
-                   &ki,
-                   &m1,
-                   &one,
-                   tmp1 + d0,
-                   &di,
-                   codebook.data() + sq * (m1 * k),
-                   &m1,
-                   &zero,
-                   tmp2,
-                   &dsubi);
-        }
-
-        float min = HUGE_VAL;
-        int argmin = -1;
-        for (size_t j = 0; j < k; j++) {
-            float dis = fvec_L2sqr(x + d0, tmp2 + j * dsub, dsub);
-            if (dis < min) {
-                min = dis;
-                argmin = j;
-            }
-        }
-        code[sq] = argmin;
-    }
+void IndexHNSW::permute_entries(const idx_t* perm) {
+    auto flat_storage = dynamic_cast<IndexFlatCodes*>(storage);
+    FAISS_THROW_IF_NOT_MSG(
+            flat_storage, "don't know how to permute this index");
+    flat_storage->permute_entries(perm);
+    hnsw.permute_entries(perm);
 }
 
-void ReconstructFromNeighbors::add_codes(size_t n, const float* x) {
-    if (k == 1) { // nothing to encode
-        ntotal += n;
-        return;
-    }
-    codes.resize(codes.size() + code_size * n);
-#pragma omp parallel for
-    for (int i = 0; i < n; i++) {
-        estimate_code(
-                x + i * index.d,
-                ntotal + i,
-                codes.data() + (ntotal + i) * code_size);
-    }
-    ntotal += n;
-    FAISS_ASSERT(codes.size() == ntotal * code_size);
+DistanceComputer* IndexHNSW::get_distance_computer() const {
+    return storage->get_distance_computer();
 }
 
 /**************************************************************
@@ -864,7 +626,10 @@ IndexHNSWFlat::IndexHNSWFlat() {
 }
 
 IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
-        : IndexHNSW(new IndexFlat(d, metric), M) {
+        : IndexHNSW(
+                  (metric == METRIC_L2) ? new IndexFlatL2(d)
+                                        : new IndexFlat(d, metric),
+                  M) {
     own_fields = true;
     is_trained = true;
 }
@@ -873,10 +638,15 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
  * IndexHNSWPQ implementation
  **************************************************************/
 
-IndexHNSWPQ::IndexHNSWPQ() {}
+IndexHNSWPQ::IndexHNSWPQ() = default;
 
-IndexHNSWPQ::IndexHNSWPQ(int d, int pq_m, int M)
-        : IndexHNSW(new IndexPQ(d, pq_m, 8), M) {
+IndexHNSWPQ::IndexHNSWPQ(
+        int d,
+        int pq_m,
+        int M,
+        int pq_nbits,
+        MetricType metric)
+        : IndexHNSW(new IndexPQ(d, pq_m, pq_nbits, metric), M) {
     own_fields = true;
     is_trained = false;
 }
@@ -896,11 +666,11 @@ IndexHNSWSQ::IndexHNSWSQ(
         int M,
         MetricType metric)
         : IndexHNSW(new IndexScalarQuantizer(d, qtype, metric), M) {
-    is_trained = false;
+    is_trained = this->storage->is_trained;
     own_fields = true;
 }
 
-IndexHNSWSQ::IndexHNSWSQ() {}
+IndexHNSWSQ::IndexHNSWSQ() = default;
 
 /**************************************************************
  * IndexHNSW2Level implementation
@@ -916,7 +686,7 @@ IndexHNSW2Level::IndexHNSW2Level(
     is_trained = false;
 }
 
-IndexHNSW2Level::IndexHNSW2Level() {}
+IndexHNSW2Level::IndexHNSW2Level() = default;
 
 namespace {
 
@@ -935,7 +705,6 @@ int search_from_candidates_2(
         int level,
         int nres_in = 0) {
     int nres = nres_in;
-    int ndis = 0;
     for (int i = 0; i < candidates.size(); i++) {
         idx_t v1 = candidates.ids[i];
         FAISS_ASSERT(v1 >= 0);
@@ -958,7 +727,6 @@ int search_from_candidates_2(
             if (vt.visited[v1] == vt.visno + 1) {
                 // nothing to do
             } else {
-                ndis++;
                 float d = qdis(v1);
                 candidates.push(v1, d);
 
@@ -1004,7 +772,7 @@ void IndexHNSW2Level::search(
         IndexHNSW::search(n, x, k, distances, labels);
 
     } else { // "mixed" search
-        size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
+        size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
 
         const IndexIVFPQ* index_ivfpq =
                 dynamic_cast<const IndexIVFPQ*>(storage);
@@ -1030,13 +798,13 @@ void IndexHNSW2Level::search(
 #pragma omp parallel
         {
             VisitedTable vt(ntotal);
-            DistanceComputer* dis = storage_distance_computer(storage);
-            ScopeDeleter1<DistanceComputer> del(dis);
+            std::unique_ptr<DistanceComputer> dis(
+                    storage_distance_computer(storage));
 
-            int candidates_size = hnsw.upper_beam;
+            constexpr int candidates_size = 1;
             MinimaxHeap candidates(candidates_size);
 
-#pragma omp for reduction(+ : n1, n2, n3, ndis, nreorder)
+#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;
@@ -1058,7 +826,7 @@ void IndexHNSW2Level::search(
 
                 candidates.clear();
 
-                for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
+                for (int j = 0; j < k; j++) {
                     if (idxi[j] < 0)
                         break;
                     candidates.push(idxi[j], simi[j]);
@@ -1081,9 +849,8 @@ void IndexHNSW2Level::search(
                         k);
                 n1 += search_stats.n1;
                 n2 += search_stats.n2;
-                n3 += search_stats.n3;
                 ndis += search_stats.ndis;
-                nreorder += search_stats.nreorder;
+                nhops += search_stats.nhops;
 
                 vt.advance();
                 vt.advance();
@@ -1092,7 +859,7 @@ void IndexHNSW2Level::search(
             }
         }
 
-        hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
+        hnsw_stats.combine({n1, n2, ndis, nhops});
     }
 }
 
@@ -1118,4 +885,86 @@ void IndexHNSW2Level::flip_to_ivf() {
     delete storage2l;
 }
 
+/**************************************************************
+ * IndexHNSWCagra implementation
+ **************************************************************/
+
+IndexHNSWCagra::IndexHNSWCagra() {
+    is_trained = true;
+}
+
+IndexHNSWCagra::IndexHNSWCagra(int d, int M, MetricType metric)
+        : IndexHNSW(
+                  (metric == METRIC_L2)
+                          ? static_cast<IndexFlat*>(new IndexFlatL2(d))
+                          : static_cast<IndexFlat*>(new IndexFlatIP(d)),
+                  M) {
+    FAISS_THROW_IF_NOT_MSG(
+            ((metric == METRIC_L2) || (metric == METRIC_INNER_PRODUCT)),
+            "unsupported metric type for IndexHNSWCagra");
+    own_fields = true;
+    is_trained = true;
+    init_level0 = true;
+    keep_max_size_level0 = true;
+}
+
+void IndexHNSWCagra::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(
+            !base_level_only,
+            "Cannot add vectors when base_level_only is set to True");
+
+    IndexHNSW::add(n, x);
+}
+
+void IndexHNSWCagra::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    if (!base_level_only) {
+        IndexHNSW::search(n, x, k, distances, labels, params);
+    } else {
+        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;
+                }
+            }
+            FAISS_THROW_IF_NOT_MSG(
+                    nearest[i] >= 0, "Could not find a valid entrypoint.");
+        }
+
+        search_level_0(
+                n,
+                x,
+                k,
+                nearest.data(),
+                nearest_d.data(),
+                distances,
+                labels,
+                1, // n_probes
+                1, // search_type
+                params);
+    }
+}
+
 } // namespace faiss