faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/IndexBinaryIVF.cpp

@@ -9,318 +9,344 @@
 
 #include <faiss/IndexBinaryIVF.h>
 
+#include <omp.h>
 #include <cinttypes>
 #include <cstdio>
-#include <omp.h>
 
+#include <algorithm>
 #include <memory>
 
-
-#include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
-#include <faiss/impl/AuxIndexStructures.h>
-#include <faiss/impl/FaissAssert.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexLSH.h>
-
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
-IndexBinaryIVF::IndexBinaryIVF(IndexBinary *quantizer, size_t d, size_t nlist)
-    : IndexBinary(d),
-      invlists(new ArrayInvertedLists(nlist, code_size)),
-      own_invlists(true),
-      nprobe(1),
-      max_codes(0),
-      quantizer(quantizer),
-      nlist(nlist),
-      own_fields(false),
-      clustering_index(nullptr)
-{
-  FAISS_THROW_IF_NOT (d == quantizer->d);
-  is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);
-
-  cp.niter = 10;
+IndexBinaryIVF::IndexBinaryIVF(IndexBinary* quantizer, size_t d, size_t nlist)
+        : IndexBinary(d),
+          invlists(new ArrayInvertedLists(nlist, code_size)),
+          own_invlists(true),
+          nprobe(1),
+          max_codes(0),
+          quantizer(quantizer),
+          nlist(nlist),
+          own_fields(false),
+          clustering_index(nullptr) {
+    FAISS_THROW_IF_NOT(d == quantizer->d);
+    is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);
+
+    cp.niter = 10;
 }
 
 IndexBinaryIVF::IndexBinaryIVF()
-    : invlists(nullptr),
-      own_invlists(false),
-      nprobe(1),
-      max_codes(0),
-      quantizer(nullptr),
-      nlist(0),
-      own_fields(false),
-      clustering_index(nullptr)
-{}
-
-void IndexBinaryIVF::add(idx_t n, const uint8_t *x) {
-  add_with_ids(n, x, nullptr);
+        : invlists(nullptr),
+          own_invlists(false),
+          nprobe(1),
+          max_codes(0),
+          quantizer(nullptr),
+          nlist(0),
+          own_fields(false),
+          clustering_index(nullptr) {}
+
+void IndexBinaryIVF::add(idx_t n, const uint8_t* x) {
+    add_with_ids(n, x, nullptr);
 }
 
-void IndexBinaryIVF::add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids) {
-  add_core(n, x, xids, nullptr);
+void IndexBinaryIVF::add_with_ids(
+        idx_t n,
+        const uint8_t* x,
+        const idx_t* xids) {
+    add_core(n, x, xids, nullptr);
 }
 
-void IndexBinaryIVF::add_core(idx_t n, const uint8_t *x, const idx_t *xids,
-                              const idx_t *precomputed_idx) {
-  FAISS_THROW_IF_NOT(is_trained);
-  assert(invlists);
-  direct_map.check_can_add (xids);
+void IndexBinaryIVF::add_core(
+        idx_t n,
+        const uint8_t* x,
+        const idx_t* xids,
+        const idx_t* precomputed_idx) {
+    FAISS_THROW_IF_NOT(is_trained);
+    assert(invlists);
+    direct_map.check_can_add(xids);
 
-  const idx_t * idx;
+    const idx_t* idx;
 
-  std::unique_ptr<idx_t[]> scoped_idx;
+    std::unique_ptr<idx_t[]> scoped_idx;
 
-  if (precomputed_idx) {
-    idx = precomputed_idx;
-  } else {
-    scoped_idx.reset(new idx_t[n]);
-    quantizer->assign(n, x, scoped_idx.get());
-    idx = scoped_idx.get();
-  }
+    if (precomputed_idx) {
+        idx = precomputed_idx;
+    } else {
+        scoped_idx.reset(new idx_t[n]);
+        quantizer->assign(n, x, scoped_idx.get());
+        idx = scoped_idx.get();
+    }
 
-  long n_add = 0;
-  for (size_t i = 0; i < n; i++) {
-    idx_t id = xids ? xids[i] : ntotal + i;
-    idx_t list_no = idx[i];
+    idx_t n_add = 0;
+    for (size_t i = 0; i < n; i++) {
+        idx_t id = xids ? xids[i] : ntotal + i;
+        idx_t list_no = idx[i];
 
-    if (list_no < 0) {
-        direct_map.add_single_id (id, -1, 0);
-    } else {
-        const uint8_t *xi = x + i * code_size;
-        size_t offset = invlists->add_entry(list_no, id, xi);
+        if (list_no < 0) {
+            direct_map.add_single_id(id, -1, 0);
+        } else {
+            const uint8_t* xi = x + i * code_size;
+            size_t offset = invlists->add_entry(list_no, id, xi);
 
-        direct_map.add_single_id (id, list_no, offset);
-    }
+            direct_map.add_single_id(id, list_no, offset);
+        }
 
-    n_add++;
-  }
-  if (verbose) {
-    printf("IndexBinaryIVF::add_with_ids: added %ld / %" PRId64 " vectors\n",
-           n_add, n);
-  }
-  ntotal += n_add;
+        n_add++;
+    }
+    if (verbose) {
+        printf("IndexBinaryIVF::add_with_ids: added "
+               "%" PRId64 " / %" PRId64 " vectors\n",
+               n_add,
+               n);
+    }
+    ntotal += n_add;
 }
 
-void IndexBinaryIVF::make_direct_map (bool b)
-{
+void IndexBinaryIVF::make_direct_map(bool b) {
     if (b) {
-        direct_map.set_type (DirectMap::Array, invlists, ntotal);
+        direct_map.set_type(DirectMap::Array, invlists, ntotal);
     } else {
-        direct_map.set_type (DirectMap::NoMap, invlists, ntotal);
+        direct_map.set_type(DirectMap::NoMap, invlists, ntotal);
     }
 }
 
-void IndexBinaryIVF::set_direct_map_type (DirectMap::Type type)
-{
-    direct_map.set_type (type, invlists, ntotal);
+void IndexBinaryIVF::set_direct_map_type(DirectMap::Type type) {
+    direct_map.set_type(type, invlists, ntotal);
 }
 
+void IndexBinaryIVF::search(
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        int32_t* distances,
+        idx_t* labels) const {
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(nprobe > 0);
 
-void IndexBinaryIVF::search(idx_t n, const uint8_t *x, idx_t k,
-                            int32_t *distances, idx_t *labels) const {
-  std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
-  std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
+    const size_t nprobe = std::min(nlist, this->nprobe);
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
 
-  double t0 = getmillisecs();
-  quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
-  indexIVF_stats.quantization_time += getmillisecs() - t0;
+    double t0 = getmillisecs();
+    quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
+    indexIVF_stats.quantization_time += getmillisecs() - t0;
 
-  t0 = getmillisecs();
-  invlists->prefetch_lists(idx.get(), n * nprobe);
+    t0 = getmillisecs();
+    invlists->prefetch_lists(idx.get(), n * nprobe);
 
-  search_preassigned(n, x, k, idx.get(), coarse_dis.get(),
-                     distances, labels, false);
-  indexIVF_stats.search_time += getmillisecs() - t0;
+    search_preassigned(
+            n, x, k, idx.get(), coarse_dis.get(), distances, labels, false);
+    indexIVF_stats.search_time += getmillisecs() - t0;
 }
 
-void IndexBinaryIVF::reconstruct(idx_t key, uint8_t *recons) const {
-    idx_t lo = direct_map.get (key);
-    reconstruct_from_offset (lo_listno(lo), lo_offset(lo), recons);
+void IndexBinaryIVF::reconstruct(idx_t key, uint8_t* recons) const {
+    idx_t lo = direct_map.get(key);
+    reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);
 }
 
-void IndexBinaryIVF::reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const {
-  FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
+void IndexBinaryIVF::reconstruct_n(idx_t i0, idx_t ni, uint8_t* recons) const {
+    FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
 
-  for (idx_t list_no = 0; list_no < nlist; list_no++) {
-    size_t list_size = invlists->list_size(list_no);
-    const Index::idx_t *idlist = invlists->get_ids(list_no);
+    for (idx_t list_no = 0; list_no < nlist; list_no++) {
+        size_t list_size = invlists->list_size(list_no);
+        const Index::idx_t* idlist = invlists->get_ids(list_no);
 
-    for (idx_t offset = 0; offset < list_size; offset++) {
-      idx_t id = idlist[offset];
-      if (!(id >= i0 && id < i0 + ni)) {
-        continue;
-      }
+        for (idx_t offset = 0; offset < list_size; offset++) {
+            idx_t id = idlist[offset];
+            if (!(id >= i0 && id < i0 + ni)) {
+                continue;
+            }
 
-      uint8_t *reconstructed = recons + (id - i0) * d;
-      reconstruct_from_offset(list_no, offset, reconstructed);
+            uint8_t* reconstructed = recons + (id - i0) * d;
+            reconstruct_from_offset(list_no, offset, reconstructed);
+        }
     }
-  }
 }
 
-void IndexBinaryIVF::search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k,
-                                            int32_t *distances, idx_t *labels,
-                                            uint8_t *recons) const {
-  std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
-  std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
-
-  quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
-
-  invlists->prefetch_lists(idx.get(), n * nprobe);
-
-  // search_preassigned() with `store_pairs` enabled to obtain the list_no
-  // and offset into `codes` for reconstruction
-  search_preassigned(n, x, k, idx.get(), coarse_dis.get(),
-                     distances, labels, /* store_pairs */true);
-  for (idx_t i = 0; i < n; ++i) {
-    for (idx_t j = 0; j < k; ++j) {
-      idx_t ij = i * k + j;
-      idx_t key = labels[ij];
-      uint8_t *reconstructed = recons + ij * d;
-      if (key < 0) {
-        // Fill with NaNs
-        memset(reconstructed, -1, sizeof(*reconstructed) * d);
-      } else {
-        int list_no = key >> 32;
-        int offset = key & 0xffffffff;
-
-        // Update label to the actual id
-        labels[ij] = invlists->get_single_id(list_no, offset);
-
-        reconstruct_from_offset(list_no, offset, reconstructed);
-      }
+void IndexBinaryIVF::search_and_reconstruct(
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        int32_t* distances,
+        idx_t* labels,
+        uint8_t* recons) const {
+    const size_t nprobe = std::min(nlist, this->nprobe);
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(nprobe > 0);
+
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
+
+    quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
+
+    invlists->prefetch_lists(idx.get(), n * nprobe);
+
+    // search_preassigned() with `store_pairs` enabled to obtain the list_no
+    // and offset into `codes` for reconstruction
+    search_preassigned(
+            n,
+            x,
+            k,
+            idx.get(),
+            coarse_dis.get(),
+            distances,
+            labels,
+            /* store_pairs */ true);
+    for (idx_t i = 0; i < n; ++i) {
+        for (idx_t j = 0; j < k; ++j) {
+            idx_t ij = i * k + j;
+            idx_t key = labels[ij];
+            uint8_t* reconstructed = recons + ij * d;
+            if (key < 0) {
+                // Fill with NaNs
+                memset(reconstructed, -1, sizeof(*reconstructed) * d);
+            } else {
+                int list_no = key >> 32;
+                int offset = key & 0xffffffff;
+
+                // Update label to the actual id
+                labels[ij] = invlists->get_single_id(list_no, offset);
+
+                reconstruct_from_offset(list_no, offset, reconstructed);
+            }
+        }
     }
-  }
 }
 
-void IndexBinaryIVF::reconstruct_from_offset(idx_t list_no, idx_t offset,
-                                             uint8_t *recons) const {
-  memcpy(recons, invlists->get_single_code(list_no, offset), code_size);
+void IndexBinaryIVF::reconstruct_from_offset(
+        idx_t list_no,
+        idx_t offset,
+        uint8_t* recons) const {
+    memcpy(recons, invlists->get_single_code(list_no, offset), code_size);
 }
 
 void IndexBinaryIVF::reset() {
-  direct_map.clear();
-  invlists->reset();
-  ntotal = 0;
+    direct_map.clear();
+    invlists->reset();
+    ntotal = 0;
 }
 
 size_t IndexBinaryIVF::remove_ids(const IDSelector& sel) {
-    size_t nremove = direct_map.remove_ids (sel, invlists);
+    size_t nremove = direct_map.remove_ids(sel, invlists);
     ntotal -= nremove;
     return nremove;
 }
 
-void IndexBinaryIVF::train(idx_t n, const uint8_t *x) {
-  if (verbose) {
-    printf("Training quantizer\n");
-  }
-
-  if (quantizer->is_trained && (quantizer->ntotal == nlist)) {
-    if (verbose) {
-      printf("IVF quantizer does not need training.\n");
-    }
-  } else {
+void IndexBinaryIVF::train(idx_t n, const uint8_t* x) {
     if (verbose) {
-      printf("Training quantizer on %" PRId64 " vectors in %dD\n", n, d);
+        printf("Training quantizer\n");
     }
 
-    Clustering clus(d, nlist, cp);
-    quantizer->reset();
-
-    IndexFlatL2 index_tmp(d);
-
-    if (clustering_index && verbose) {
-      printf("using clustering_index of dimension %d to do the clustering\n",
-             clustering_index->d);
-    }
+    if (quantizer->is_trained && (quantizer->ntotal == nlist)) {
+        if (verbose) {
+            printf("IVF quantizer does not need training.\n");
+        }
+    } else {
+        if (verbose) {
+            printf("Training quantizer on %" PRId64 " vectors in %dD\n", n, d);
+        }
 
-    // LSH codec that is able to convert the binary vectors to floats.
-    IndexLSH codec(d, d, false, false);
+        Clustering clus(d, nlist, cp);
+        quantizer->reset();
 
-    clus.train_encoded (n, x, &codec, clustering_index ? *clustering_index : index_tmp);
+        IndexFlatL2 index_tmp(d);
 
-    // convert clusters to binary
-    std::unique_ptr<uint8_t[]> x_b(new uint8_t[clus.k * code_size]);
-    real_to_binary(d * clus.k, clus.centroids.data(), x_b.get());
+        if (clustering_index && verbose) {
+            printf("using clustering_index of dimension %d to do the clustering\n",
+                   clustering_index->d);
+        }
 
-    quantizer->add(clus.k, x_b.get());
-    quantizer->is_trained = true;
-  }
+        // LSH codec that is able to convert the binary vectors to floats.
+        IndexLSH codec(d, d, false, false);
 
-  is_trained = true;
-}
+        clus.train_encoded(
+                n, x, &codec, clustering_index ? *clustering_index : index_tmp);
 
-void IndexBinaryIVF::merge_from(IndexBinaryIVF &other, idx_t add_id) {
-  // minimal sanity checks
-  FAISS_THROW_IF_NOT(other.d == d);
-  FAISS_THROW_IF_NOT(other.nlist == nlist);
-  FAISS_THROW_IF_NOT(other.code_size == code_size);
-  FAISS_THROW_IF_NOT_MSG(direct_map.no() && other.direct_map.no(),
-                         "direct map copy not implemented");
-  FAISS_THROW_IF_NOT_MSG(typeid (*this) == typeid (other),
-                         "can only merge indexes of the same type");
+        // convert clusters to binary
+        std::unique_ptr<uint8_t[]> x_b(new uint8_t[clus.k * code_size]);
+        real_to_binary(d * clus.k, clus.centroids.data(), x_b.get());
 
-  invlists->merge_from (other.invlists, add_id);
+        quantizer->add(clus.k, x_b.get());
+        quantizer->is_trained = true;
+    }
 
-  ntotal += other.ntotal;
-  other.ntotal = 0;
+    is_trained = true;
 }
 
-void IndexBinaryIVF::replace_invlists(InvertedLists *il, bool own) {
-  FAISS_THROW_IF_NOT(il->nlist == nlist &&
-                     il->code_size == code_size);
-  if (own_invlists) {
-    delete invlists;
-  }
-  invlists = il;
-  own_invlists = own;
+void IndexBinaryIVF::merge_from(IndexBinaryIVF& other, idx_t add_id) {
+    // minimal sanity checks
+    FAISS_THROW_IF_NOT(other.d == d);
+    FAISS_THROW_IF_NOT(other.nlist == nlist);
+    FAISS_THROW_IF_NOT(other.code_size == code_size);
+    FAISS_THROW_IF_NOT_MSG(
+            direct_map.no() && other.direct_map.no(),
+            "direct map copy not implemented");
+    FAISS_THROW_IF_NOT_MSG(
+            typeid(*this) == typeid(other),
+            "can only merge indexes of the same type");
+
+    invlists->merge_from(other.invlists, add_id);
+
+    ntotal += other.ntotal;
+    other.ntotal = 0;
 }
 
+void IndexBinaryIVF::replace_invlists(InvertedLists* il, bool own) {
+    FAISS_THROW_IF_NOT(il->nlist == nlist && il->code_size == code_size);
+    if (own_invlists) {
+        delete invlists;
+    }
+    invlists = il;
+    own_invlists = own;
+}
 
 namespace {
 
 using idx_t = Index::idx_t;
 
-
-template<class HammingComputer>
-struct IVFBinaryScannerL2: BinaryInvertedListScanner {
-
+template <class HammingComputer>
+struct IVFBinaryScannerL2 : BinaryInvertedListScanner {
     HammingComputer hc;
     size_t code_size;
     bool store_pairs;
 
-    IVFBinaryScannerL2 (size_t code_size, bool store_pairs):
-        code_size (code_size), store_pairs(store_pairs)
-    {}
+    IVFBinaryScannerL2(size_t code_size, bool store_pairs)
+            : code_size(code_size), store_pairs(store_pairs) {}
 
-    void set_query (const uint8_t *query_vector) override {
-        hc.set (query_vector, code_size);
+    void set_query(const uint8_t* query_vector) override {
+        hc.set(query_vector, code_size);
     }
 
     idx_t list_no;
-    void set_list (idx_t list_no, uint8_t /* coarse_dis */) override {
+    void set_list(idx_t list_no, uint8_t /* coarse_dis */) override {
         this->list_no = list_no;
     }
 
-    uint32_t distance_to_code (const uint8_t *code) const override {
-        return hc.hamming (code);
+    uint32_t distance_to_code(const uint8_t* code) const override {
+        return hc.hamming(code);
     }
 
-    size_t scan_codes (size_t n,
-                       const uint8_t *codes,
-                       const idx_t *ids,
-                       int32_t *simi, idx_t *idxi,
-                       size_t k) const override
-    {
+    size_t scan_codes(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            int32_t* simi,
+            idx_t* idxi,
+            size_t k) const override {
         using C = CMax<int32_t, idx_t>;
 
         size_t nup = 0;
         for (size_t j = 0; j < n; j++) {
-            uint32_t dis = hc.hamming (codes);
+            uint32_t dis = hc.hamming(codes);
             if (dis < simi[0]) {
                 idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                heap_replace_top<C> (k, simi, idxi, dis, id);
+                heap_replace_top<C>(k, simi, idxi, dis, id);
                 nup++;
             }
             codes += code_size;
@@ -328,40 +354,38 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
         return nup;
     }
 
-    void scan_codes_range (size_t n,
-                           const uint8_t *codes,
-                           const idx_t *ids,
-                           int radius,
-                           RangeQueryResult &result) const override
-    {
+    void scan_codes_range(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            int radius,
+            RangeQueryResult& result) const override {
         size_t nup = 0;
         for (size_t j = 0; j < n; j++) {
-            uint32_t dis = hc.hamming (codes);
+            uint32_t dis = hc.hamming(codes);
             if (dis < radius) {
-                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
-                result.add (dis, id);
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                result.add(dis, id);
             }
             codes += code_size;
         }
-
     }
-
-
 };
 
-
-void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
-                             size_t n,
-                             const uint8_t *x,
-                             idx_t k,
-                             const idx_t *keys,
-                             const int32_t * coarse_dis,
-                             int32_t *distances, idx_t *labels,
-                             bool store_pairs,
-                             const IVFSearchParameters *params)
-{
-    long nprobe = params ? params->nprobe : ivf.nprobe;
-    long max_codes = params ? params->max_codes : ivf.max_codes;
+void search_knn_hamming_heap(
+        const IndexBinaryIVF& ivf,
+        size_t n,
+        const uint8_t* x,
+        idx_t k,
+        const idx_t* keys,
+        const int32_t* coarse_dis,
+        int32_t* distances,
+        idx_t* labels,
+        bool store_pairs,
+        const IVFSearchParameters* params) {
+    idx_t nprobe = params ? params->nprobe : ivf.nprobe;
+    nprobe = std::min((idx_t)ivf.nlist, nprobe);
+    idx_t max_codes = params ? params->max_codes : ivf.max_codes;
     MetricType metric_type = ivf.metric_type;
 
     // almost verbatim copy from IndexIVF::search_preassigned
@@ -370,57 +394,57 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
     using HeapForIP = CMin<int32_t, idx_t>;
     using HeapForL2 = CMax<int32_t, idx_t>;
 
-#pragma omp parallel if(n > 1) reduction(+: nlistv, ndis, nheap)
+#pragma omp parallel if (n > 1) reduction(+ : nlistv, ndis, nheap)
     {
-        std::unique_ptr<BinaryInvertedListScanner> scanner
-            (ivf.get_InvertedListScanner (store_pairs));
+        std::unique_ptr<BinaryInvertedListScanner> scanner(
+                ivf.get_InvertedListScanner(store_pairs));
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            const uint8_t *xi = x + i * ivf.code_size;
+            const uint8_t* xi = x + i * ivf.code_size;
             scanner->set_query(xi);
 
-            const idx_t * keysi = keys + i * nprobe;
-            int32_t * simi = distances + k * i;
-            idx_t * idxi = labels + k * i;
+            const idx_t* keysi = keys + i * nprobe;
+            int32_t* simi = distances + k * i;
+            idx_t* idxi = labels + k * i;
 
             if (metric_type == METRIC_INNER_PRODUCT) {
-                heap_heapify<HeapForIP> (k, simi, idxi);
+                heap_heapify<HeapForIP>(k, simi, idxi);
             } else {
-                heap_heapify<HeapForL2> (k, simi, idxi);
+                heap_heapify<HeapForL2>(k, simi, idxi);
             }
 
             size_t nscan = 0;
 
             for (size_t ik = 0; ik < nprobe; ik++) {
-                idx_t key = keysi[ik];  /* select the list  */
+                idx_t key = keysi[ik]; /* select the list  */
                 if (key < 0) {
                     // not enough centroids for multiprobe
                     continue;
                 }
-                FAISS_THROW_IF_NOT_FMT
-                    (key < (idx_t) ivf.nlist,
-                     "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
-                     key, ik, ivf.nlist);
+                FAISS_THROW_IF_NOT_FMT(
+                        key < (idx_t)ivf.nlist,
+                        "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
+                        key,
+                        ik,
+                        ivf.nlist);
 
-                scanner->set_list (key, coarse_dis[i * nprobe + ik]);
+                scanner->set_list(key, coarse_dis[i * nprobe + ik]);
 
                 nlistv++;
 
                 size_t list_size = ivf.invlists->list_size(key);
-                InvertedLists::ScopedCodes scodes (ivf.invlists, key);
+                InvertedLists::ScopedCodes scodes(ivf.invlists, key);
                 std::unique_ptr<InvertedLists::ScopedIds> sids;
-                const Index::idx_t * ids = nullptr;
+                const Index::idx_t* ids = nullptr;
 
                 if (!store_pairs) {
-                    sids.reset (new InvertedLists::ScopedIds (ivf.invlists, key));
+                    sids.reset(new InvertedLists::ScopedIds(ivf.invlists, key));
                     ids = sids->get();
                 }
 
-                nheap += scanner->scan_codes (
-                        list_size, scodes.get(),
-                        ids, simi, idxi, k
-                );
+                nheap += scanner->scan_codes(
+                        list_size, scodes.get(), ids, simi, idxi, k);
 
                 nscan += list_size;
                 if (max_codes && nscan >= max_codes)
@@ -429,208 +453,205 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
 
             ndis += nscan;
             if (metric_type == METRIC_INNER_PRODUCT) {
-                heap_reorder<HeapForIP> (k, simi, idxi);
+                heap_reorder<HeapForIP>(k, simi, idxi);
             } else {
-                heap_reorder<HeapForL2> (k, simi, idxi);
+                heap_reorder<HeapForL2>(k, simi, idxi);
             }
 
         } // parallel for
-    } // parallel
+    }     // parallel
 
     indexIVF_stats.nq += n;
     indexIVF_stats.nlist += nlistv;
     indexIVF_stats.ndis += ndis;
     indexIVF_stats.nheap_updates += nheap;
-
 }
 
-template<class HammingComputer, bool store_pairs>
-void search_knn_hamming_count(const IndexBinaryIVF& ivf,
-                              size_t nx,
-                              const uint8_t *x,
-                              const idx_t *keys,
-                              int k,
-                              int32_t *distances,
-                              idx_t *labels,
-                              const IVFSearchParameters *params) {
-  const int nBuckets = ivf.d + 1;
-  std::vector<int> all_counters(nx * nBuckets, 0);
-  std::unique_ptr<idx_t[]> all_ids_per_dis(new idx_t[nx * nBuckets * k]);
-
-  long nprobe = params ? params->nprobe : ivf.nprobe;
-  long max_codes = params ? params->max_codes : ivf.max_codes;
-
-  std::vector<HCounterState<HammingComputer>> cs;
-  for (size_t i = 0; i < nx; ++i) {
-    cs.push_back(HCounterState<HammingComputer>(
-                   all_counters.data() + i * nBuckets,
-                   all_ids_per_dis.get() + i * nBuckets * k,
-                   x + i * ivf.code_size,
-                   ivf.d,
-                   k
-                 ));
-  }
-
-  size_t nlistv = 0, ndis = 0;
-
-#pragma omp parallel for reduction(+: nlistv, ndis)
-  for (int64_t i = 0; i < nx; i++) {
-    const idx_t * keysi = keys + i * nprobe;
-    HCounterState<HammingComputer>& csi = cs[i];
-
-    size_t nscan = 0;
-
-    for (size_t ik = 0; ik < nprobe; ik++) {
-      idx_t key = keysi[ik];  /* select the list  */
-      if (key < 0) {
-        // not enough centroids for multiprobe
-        continue;
-      }
-      FAISS_THROW_IF_NOT_FMT (
-        key < (idx_t) ivf.nlist,
-        "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
-        key, ik, ivf.nlist);
-
-      nlistv++;
-      size_t list_size = ivf.invlists->list_size(key);
-      InvertedLists::ScopedCodes scodes (ivf.invlists, key);
-      const uint8_t *list_vecs = scodes.get();
-      const Index::idx_t *ids = store_pairs
-        ? nullptr
-        : ivf.invlists->get_ids(key);
-
-      for (size_t j = 0; j < list_size; j++) {
-        const uint8_t * yj = list_vecs + ivf.code_size * j;
-
-        idx_t id = store_pairs ? (key << 32 | j) : ids[j];
-        csi.update_counter(yj, id);
-      }
-      if (ids)
-          ivf.invlists->release_ids (key, ids);
-
-      nscan += list_size;
-      if (max_codes && nscan >= max_codes)
-        break;
-    }
-    ndis += nscan;
-
-    int nres = 0;
-    for (int b = 0; b < nBuckets && nres < k; b++) {
-      for (int l = 0; l < csi.counters[b] && nres < k; l++) {
-        labels[i * k + nres] = csi.ids_per_dis[b * k + l];
-        distances[i * k + nres] = b;
-        nres++;
-      }
-    }
-    while (nres < k) {
-      labels[i * k + nres] = -1;
-      distances[i * k + nres] = std::numeric_limits<int32_t>::max();
-      ++nres;
+template <class HammingComputer, bool store_pairs>
+void search_knn_hamming_count(
+        const IndexBinaryIVF& ivf,
+        size_t nx,
+        const uint8_t* x,
+        const idx_t* keys,
+        int k,
+        int32_t* distances,
+        idx_t* labels,
+        const IVFSearchParameters* params) {
+    const int nBuckets = ivf.d + 1;
+    std::vector<int> all_counters(nx * nBuckets, 0);
+    std::unique_ptr<idx_t[]> all_ids_per_dis(new idx_t[nx * nBuckets * k]);
+
+    idx_t nprobe = params ? params->nprobe : ivf.nprobe;
+    nprobe = std::min((idx_t)ivf.nlist, nprobe);
+    idx_t max_codes = params ? params->max_codes : ivf.max_codes;
+
+    std::vector<HCounterState<HammingComputer>> cs;
+    for (size_t i = 0; i < nx; ++i) {
+        cs.push_back(HCounterState<HammingComputer>(
+                all_counters.data() + i * nBuckets,
+                all_ids_per_dis.get() + i * nBuckets * k,
+                x + i * ivf.code_size,
+                ivf.d,
+                k));
     }
-  }
 
-  indexIVF_stats.nq += nx;
-  indexIVF_stats.nlist += nlistv;
-  indexIVF_stats.ndis += ndis;
-}
+    size_t nlistv = 0, ndis = 0;
 
+#pragma omp parallel for reduction(+ : nlistv, ndis)
+    for (int64_t i = 0; i < nx; i++) {
+        const idx_t* keysi = keys + i * nprobe;
+        HCounterState<HammingComputer>& csi = cs[i];
 
+        size_t nscan = 0;
 
-template<bool store_pairs>
-void search_knn_hamming_count_1 (
-                        const IndexBinaryIVF& ivf,
-                        size_t nx,
-                        const uint8_t *x,
-                        const idx_t *keys,
-                        int k,
-                        int32_t *distances,
-                        idx_t *labels,
-                        const IVFSearchParameters *params) {
-    switch (ivf.code_size) {
-#define HANDLE_CS(cs)                                                  \
-    case cs:                                                            \
-       search_knn_hamming_count<HammingComputer ## cs, store_pairs>(    \
-           ivf, nx, x, keys, k, distances, labels, params);             \
-      break;
-      HANDLE_CS(4);
-      HANDLE_CS(8);
-      HANDLE_CS(16);
-      HANDLE_CS(20);
-      HANDLE_CS(32);
-      HANDLE_CS(64);
-#undef HANDLE_CS
-    default:
-        if (ivf.code_size % 8 == 0) {
-            search_knn_hamming_count<HammingComputerM8, store_pairs>
-                (ivf, nx, x, keys, k, distances, labels, params);
-        } else if (ivf.code_size % 4 == 0) {
-            search_knn_hamming_count<HammingComputerM4, store_pairs>
-                (ivf, nx, x, keys, k, distances, labels, params);
-        } else {
-            search_knn_hamming_count<HammingComputerDefault, store_pairs>
-                (ivf, nx, x, keys, k, distances, labels, params);
+        for (size_t ik = 0; ik < nprobe; ik++) {
+            idx_t key = keysi[ik]; /* select the list  */
+            if (key < 0) {
+                // not enough centroids for multiprobe
+                continue;
+            }
+            FAISS_THROW_IF_NOT_FMT(
+                    key < (idx_t)ivf.nlist,
+                    "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
+                    key,
+                    ik,
+                    ivf.nlist);
+
+            nlistv++;
+            size_t list_size = ivf.invlists->list_size(key);
+            InvertedLists::ScopedCodes scodes(ivf.invlists, key);
+            const uint8_t* list_vecs = scodes.get();
+            const Index::idx_t* ids =
+                    store_pairs ? nullptr : ivf.invlists->get_ids(key);
+
+            for (size_t j = 0; j < list_size; j++) {
+                const uint8_t* yj = list_vecs + ivf.code_size * j;
+
+                idx_t id = store_pairs ? (key << 32 | j) : ids[j];
+                csi.update_counter(yj, id);
+            }
+            if (ids)
+                ivf.invlists->release_ids(key, ids);
+
+            nscan += list_size;
+            if (max_codes && nscan >= max_codes)
+                break;
+        }
+        ndis += nscan;
+
+        int nres = 0;
+        for (int b = 0; b < nBuckets && nres < k; b++) {
+            for (int l = 0; l < csi.counters[b] && nres < k; l++) {
+                labels[i * k + nres] = csi.ids_per_dis[b * k + l];
+                distances[i * k + nres] = b;
+                nres++;
+            }
+        }
+        while (nres < k) {
+            labels[i * k + nres] = -1;
+            distances[i * k + nres] = std::numeric_limits<int32_t>::max();
+            ++nres;
         }
-        break;
     }
 
+    indexIVF_stats.nq += nx;
+    indexIVF_stats.nlist += nlistv;
+    indexIVF_stats.ndis += ndis;
 }
 
-}  // namespace
+template <bool store_pairs>
+void search_knn_hamming_count_1(
+        const IndexBinaryIVF& ivf,
+        size_t nx,
+        const uint8_t* x,
+        const idx_t* keys,
+        int k,
+        int32_t* distances,
+        idx_t* labels,
+        const IVFSearchParameters* params) {
+    switch (ivf.code_size) {
+#define HANDLE_CS(cs)                                               \
+    case cs:                                                        \
+        search_knn_hamming_count<HammingComputer##cs, store_pairs>( \
+                ivf, nx, x, keys, k, distances, labels, params);    \
+        break;
+        HANDLE_CS(4);
+        HANDLE_CS(8);
+        HANDLE_CS(16);
+        HANDLE_CS(20);
+        HANDLE_CS(32);
+        HANDLE_CS(64);
+#undef HANDLE_CS
+        default:
+            search_knn_hamming_count<HammingComputerDefault, store_pairs>(
+                    ivf, nx, x, keys, k, distances, labels, params);
+            break;
+    }
+}
 
-BinaryInvertedListScanner *IndexBinaryIVF::get_InvertedListScanner
-      (bool store_pairs) const
-{
+} // namespace
 
-#define HC(name) return new IVFBinaryScannerL2<name> (code_size, store_pairs)
+BinaryInvertedListScanner* IndexBinaryIVF::get_InvertedListScanner(
+        bool store_pairs) const {
+#define HC(name) return new IVFBinaryScannerL2<name>(code_size, store_pairs)
     switch (code_size) {
-    case 4: HC(HammingComputer4);
-    case 8: HC(HammingComputer8);
-    case 16: HC(HammingComputer16);
-    case 20: HC(HammingComputer20);
-    case 32: HC(HammingComputer32);
-    case 64: HC(HammingComputer64);
-    default:
-        if (code_size % 8 == 0) {
-            HC(HammingComputerM8);
-        } else if (code_size % 4 == 0) {
-            HC(HammingComputerM4);
-        } else {
+        case 4:
+            HC(HammingComputer4);
+        case 8:
+            HC(HammingComputer8);
+        case 16:
+            HC(HammingComputer16);
+        case 20:
+            HC(HammingComputer20);
+        case 32:
+            HC(HammingComputer32);
+        case 64:
+            HC(HammingComputer64);
+        default:
             HC(HammingComputerDefault);
-        }
     }
 #undef HC
-
 }
 
-void IndexBinaryIVF::search_preassigned(idx_t n, const uint8_t *x, idx_t k,
-                                        const idx_t *idx,
-                                        const int32_t * coarse_dis,
-                                        int32_t *distances, idx_t *labels,
-                                        bool store_pairs,
-                                        const IVFSearchParameters *params
-                                        ) const {
-
+void IndexBinaryIVF::search_preassigned(
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        const idx_t* idx,
+        const int32_t* coarse_dis,
+        int32_t* distances,
+        idx_t* labels,
+        bool store_pairs,
+        const IVFSearchParameters* params) const {
     if (use_heap) {
-        search_knn_hamming_heap (*this, n, x, k, idx, coarse_dis,
-                                 distances, labels, store_pairs,
-                                 params);
+        search_knn_hamming_heap(
+                *this,
+                n,
+                x,
+                k,
+                idx,
+                coarse_dis,
+                distances,
+                labels,
+                store_pairs,
+                params);
     } else {
         if (store_pairs) {
-            search_knn_hamming_count_1<true>
-                (*this, n, x, idx, k, distances, labels, params);
+            search_knn_hamming_count_1<true>(
+                    *this, n, x, idx, k, distances, labels, params);
         } else {
-            search_knn_hamming_count_1<false>
-                (*this, n, x, idx, k, distances, labels, params);
+            search_knn_hamming_count_1<false>(
+                    *this, n, x, idx, k, distances, labels, params);
         }
     }
 }
 
-
 void IndexBinaryIVF::range_search(
-        idx_t n, const uint8_t *x, int radius,
-        RangeSearchResult *res) const
-{
-
+        idx_t n,
+        const uint8_t* x,
+        int radius,
+        RangeSearchResult* res) const {
+    const size_t nprobe = std::min(nlist, this->nprobe);
     std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
     std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
 
@@ -641,77 +662,84 @@ void IndexBinaryIVF::range_search(
     t0 = getmillisecs();
     invlists->prefetch_lists(idx.get(), n * nprobe);
 
+    range_search_preassigned(n, x, radius, idx.get(), coarse_dis.get(), res);
+
+    indexIVF_stats.search_time += getmillisecs() - t0;
+}
+
+void IndexBinaryIVF::range_search_preassigned(
+        idx_t n,
+        const uint8_t* x,
+        int radius,
+        const idx_t* assign,
+        const int32_t* centroid_dis,
+        RangeSearchResult* res) const {
+    const size_t nprobe = std::min(nlist, this->nprobe);
     bool store_pairs = false;
     size_t nlistv = 0, ndis = 0;
 
-    std::vector<RangeSearchPartialResult *> all_pres (omp_get_max_threads());
+    std::vector<RangeSearchPartialResult*> all_pres(omp_get_max_threads());
 
-#pragma omp parallel reduction(+: nlistv, ndis)
+#pragma omp parallel reduction(+ : nlistv, ndis)
     {
         RangeSearchPartialResult pres(res);
-        std::unique_ptr<BinaryInvertedListScanner> scanner
-            (get_InvertedListScanner(store_pairs));
-        FAISS_THROW_IF_NOT (scanner.get ());
+        std::unique_ptr<BinaryInvertedListScanner> scanner(
+                get_InvertedListScanner(store_pairs));
+        FAISS_THROW_IF_NOT(scanner.get());
 
         all_pres[omp_get_thread_num()] = &pres;
 
-        auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult &qres)
-        {
-
-            idx_t key = idx[i * nprobe + ik];  /* select the list  */
-            if (key < 0) return;
-            FAISS_THROW_IF_NOT_FMT (
-                    key < (idx_t) nlist,
+        auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) {
+            idx_t key = assign[i * nprobe + ik]; /* select the list  */
+            if (key < 0)
+                return;
+            FAISS_THROW_IF_NOT_FMT(
+                    key < (idx_t)nlist,
                     "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
-                    key, ik, nlist);
+                    key,
+                    ik,
+                    nlist);
             const size_t list_size = invlists->list_size(key);
 
-            if (list_size == 0) return;
+            if (list_size == 0)
+                return;
 
-            InvertedLists::ScopedCodes scodes (invlists, key);
-            InvertedLists::ScopedIds ids (invlists, key);
+            InvertedLists::ScopedCodes scodes(invlists, key);
+            InvertedLists::ScopedIds ids(invlists, key);
 
-            scanner->set_list (key, coarse_dis[i * nprobe + ik]);
+            scanner->set_list(key, assign[i * nprobe + ik]);
             nlistv++;
             ndis += list_size;
-            scanner->scan_codes_range (list_size, scodes.get(),
-                                       ids.get(), radius, qres);
+            scanner->scan_codes_range(
+                    list_size, scodes.get(), ids.get(), radius, qres);
         };
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            scanner->set_query (x + i * code_size);
+            scanner->set_query(x + i * code_size);
 
-            RangeQueryResult & qres = pres.new_result (i);
+            RangeQueryResult& qres = pres.new_result(i);
 
             for (size_t ik = 0; ik < nprobe; ik++) {
-                scan_list_func (i, ik, qres);
+                scan_list_func(i, ik, qres);
             }
-
         }
 
         pres.finalize();
-
     }
     indexIVF_stats.nq += n;
     indexIVF_stats.nlist += nlistv;
     indexIVF_stats.ndis += ndis;
-    indexIVF_stats.search_time += getmillisecs() - t0;
-
 }
 
-
-
-
 IndexBinaryIVF::~IndexBinaryIVF() {
-  if (own_invlists) {
-    delete invlists;
-  }
+    if (own_invlists) {
+        delete invlists;
+    }
 
-  if (own_fields) {
-    delete quantizer;
-  }
+    if (own_fields) {
+        delete quantizer;
+    }
 }
 
-
-}  // namespace faiss
+} // namespace faiss