faiss 0.1.2 → 0.1.3

data/vendor/faiss/{IndexBinaryIVF.cpp → faiss/IndexBinaryIVF.cpp}

@@ -5,11 +5,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Copyright 2004-present Facebook. All Rights Reserved
 // -*- c++ -*-
 
 #include <faiss/IndexBinaryIVF.h>
 
+#include <cinttypes>
 #include <cstdio>
 #include <omp.h>
 
@@ -97,7 +97,7 @@ void IndexBinaryIVF::add_core(idx_t n, const uint8_t *x, const idx_t *xids,
     n_add++;
   }
   if (verbose) {
-    printf("IndexBinaryIVF::add_with_ids: added %ld / %ld vectors\n",
+    printf("IndexBinaryIVF::add_with_ids: added %ld / %" PRId64 " vectors\n",
            n_add, n);
   }
   ntotal += n_add;
@@ -222,7 +222,7 @@ void IndexBinaryIVF::train(idx_t n, const uint8_t *x) {
     }
   } else {
     if (verbose) {
-      printf("Training quantizer on %ld vectors in %dD\n", n, d);
+      printf("Training quantizer on %" PRId64 " vectors in %dD\n", n, d);
     }
 
     Clustering clus(d, nlist, cp);
@@ -333,7 +333,7 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
                            const uint8_t *codes,
                            const idx_t *ids,
                            int radius,
-                           RangeQueryResult &result) const
+                           RangeQueryResult &result) const override
     {
         size_t nup = 0;
         for (size_t j = 0; j < n; j++) {
@@ -351,12 +351,6 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
 };
 
 
-template <bool store_pairs>
-BinaryInvertedListScanner *select_IVFBinaryScannerL2 (size_t code_size) {
-
-}
-
-
 void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
                              size_t n,
                              const uint8_t *x,
@@ -383,7 +377,7 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
             (ivf.get_InvertedListScanner (store_pairs));
 
 #pragma omp for
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             const uint8_t *xi = x + i * ivf.code_size;
             scanner->set_query(xi);
 
@@ -407,7 +401,7 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
                 }
                 FAISS_THROW_IF_NOT_FMT
                     (key < (idx_t) ivf.nlist,
-                     "Invalid key=%ld  at ik=%ld nlist=%ld\n",
+                     "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
                      key, ik, ivf.nlist);
 
                 scanner->set_list (key, coarse_dis[i * nprobe + ik]);
@@ -481,7 +475,7 @@ void search_knn_hamming_count(const IndexBinaryIVF& ivf,
   size_t nlistv = 0, ndis = 0;
 
 #pragma omp parallel for reduction(+: nlistv, ndis)
-  for (size_t i = 0; i < nx; i++) {
+  for (int64_t i = 0; i < nx; i++) {
     const idx_t * keysi = keys + i * nprobe;
     HCounterState<HammingComputer>& csi = cs[i];
 
@@ -495,7 +489,7 @@ void search_knn_hamming_count(const IndexBinaryIVF& ivf,
       }
       FAISS_THROW_IF_NOT_FMT (
         key < (idx_t) ivf.nlist,
-        "Invalid key=%ld  at ik=%ld nlist=%ld\n",
+        "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
         key, ik, ivf.nlist);
 
       nlistv++;
@@ -669,7 +663,7 @@ void IndexBinaryIVF::range_search(
             if (key < 0) return;
             FAISS_THROW_IF_NOT_FMT (
                     key < (idx_t) nlist,
-                    "Invalid key=%ld  at ik=%ld nlist=%ld\n",
+                    "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
                     key, ik, nlist);
             const size_t list_size = invlists->list_size(key);
 
@@ -686,7 +680,7 @@ void IndexBinaryIVF::range_search(
         };
 
 #pragma omp for
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             scanner->set_query (x + i * code_size);
 
             RangeQueryResult & qres = pres.new_result (i);

data/vendor/faiss/{IndexBinaryIVF.h → faiss/IndexBinaryIVF.h}

@@ -33,7 +33,7 @@ struct BinaryInvertedListScanner;
  * Otherwise the object is similar to the IndexIVF
  */
 struct IndexBinaryIVF : IndexBinary {
-    /// Acess to the actual data
+    /// Access to the actual data
     InvertedLists *invlists;
     bool own_invlists;
 

data/vendor/faiss/{IndexHNSW.cpp → faiss/IndexHNSW.cpp}

@@ -14,6 +14,7 @@
 #include <cassert>
 #include <cstring>
 #include <cstdio>
+#include <cinttypes>
 #include <cmath>
 #include <omp.h>
 
@@ -22,7 +23,6 @@
 
 #include <sys/types.h>
 #include <sys/stat.h>
-#include <unistd.h>
 #include <stdint.h>
 
 #ifdef __SSE__
@@ -119,7 +119,7 @@ void hnsw_add_vertices(IndexHNSW &index_hnsw,
     size_t ntotal = n0 + n;
     double t0 = getmillisecs();
     if (verbose) {
-        printf("hnsw_add_vertices: adding %ld elements on top of %ld "
+        printf("hnsw_add_vertices: adding %zd elements on top of %zd "
                "(preset_levels=%d)\n",
                n, n0, int(preset_levels));
     }
@@ -289,7 +289,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
 {
     FAISS_THROW_IF_NOT_MSG(storage,
        "Please use IndexHSNWFlat (or variants) instead of IndexHNSW directly");
-    size_t nreorder = 0;
+    size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
 
     idx_t check_period = InterruptCallback::get_period_hint (
           hnsw.max_level * d * hnsw.efSearch);
@@ -297,22 +297,26 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
 
-#pragma omp parallel reduction(+ : nreorder)
+#pragma omp parallel
         {
             VisitedTable vt (ntotal);
 
             DistanceComputer *dis = storage_distance_computer(storage);
             ScopeDeleter1<DistanceComputer> del(dis);
 
-#pragma omp for
+#pragma omp for reduction (+ : n1, n2, n3, ndis, nreorder)
             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);
 
                 maxheap_heapify (k, simi, idxi);
-                hnsw.search(*dis, k, idxi, simi, vt);
-
+                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt);
+                n1 += stats.n1;
+                n2 += stats.n2;
+                n3 += stats.n3;
+                ndis += stats.ndis;
+                nreorder += stats.nreorder;
                 maxheap_reorder (k, simi, idxi);
 
                 if (reconstruct_from_neighbors &&
@@ -341,7 +345,7 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
         }
     }
 
-    hnsw_stats.nreorder += nreorder;
+    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
 }
 
 
@@ -416,6 +420,8 @@ void IndexHNSW::search_level_0(
 {
 
     storage_idx_t ntotal = hnsw.levels.size();
+    size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
+
 #pragma omp parallel
     {
         DistanceComputer *qdis = storage_distance_computer(storage);
@@ -423,7 +429,7 @@ void IndexHNSW::search_level_0(
 
         VisitedTable vt (ntotal);
 
-#pragma omp for
+#pragma omp for reduction (+ : n1, n2, n3, ndis, nreorder)
         for(idx_t i = 0; i < n; i++) {
             idx_t * idxi = labels + i * k;
             float * simi = distances + i * k;
@@ -447,10 +453,17 @@ void IndexHNSW::search_level_0(
 
                     candidates.push(cj, nearest_d[i * nprobe + j]);
 
+                    HNSWStats search_stats;
                     nres = hnsw.search_from_candidates(
                       *qdis, k, idxi, simi,
-                      candidates, vt, 0, nres
+                      candidates, vt, search_stats, 0, nres
                     );
+                    n1 += search_stats.n1;
+                    n2 += search_stats.n2;
+                    n3 += search_stats.n3;
+                    ndis += search_stats.ndis;
+                    nreorder += search_stats.nreorder;
+
                 }
             } else if (search_type == 2) {
 
@@ -464,10 +477,17 @@ void IndexHNSW::search_level_0(
                     if (cj < 0) break;
                     candidates.push(cj, nearest_d[i * nprobe + j]);
                 }
+
+                HNSWStats search_stats;
                 hnsw.search_from_candidates(
                   *qdis, k, idxi, simi,
-                  candidates, vt, 0
+                  candidates, vt, search_stats, 0
                 );
+                n1 += search_stats.n1;
+                n2 += search_stats.n2;
+                n3 += search_stats.n3;
+                ndis += search_stats.ndis;
+                nreorder += search_stats.nreorder;
 
             }
             vt.advance();
@@ -477,7 +497,7 @@ void IndexHNSW::search_level_0(
         }
     }
 
-
+    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
 }
 
 void IndexHNSW::init_level_0_from_knngraph(
@@ -488,9 +508,9 @@ void IndexHNSW::init_level_0_from_knngraph(
 #pragma omp parallel for
     for (idx_t i = 0; i < ntotal; i++) {
         DistanceComputer *qdis = storage_distance_computer(storage);
-        float vec[d];
-        storage->reconstruct(i, vec);
-        qdis->set_query(vec);
+        std::vector<float> vec(d);
+        storage->reconstruct(i, vec.data());
+        qdis->set_query(vec.data());
 
         std::priority_queue<NodeDistFarther> initial_list;
 
@@ -533,14 +553,14 @@ void IndexHNSW::init_level_0_from_entry_points(
 
         DistanceComputer *dis = storage_distance_computer(storage);
         ScopeDeleter1<DistanceComputer> del(dis);
-        float vec[storage->d];
+        std::vector<float> vec(storage->d);
 
 #pragma omp  for schedule(dynamic)
         for (int i = 0; i < n; i++) {
             storage_idx_t pt_id = points[i];
             storage_idx_t nearest = nearests[i];
-            storage->reconstruct (pt_id, vec);
-            dis->set_query (vec);
+            storage->reconstruct (pt_id, vec.data());
+            dis->set_query (vec.data());
 
             hnsw.add_links_starting_from(*dis, pt_id,
                                          nearest, (*dis)(nearest),
@@ -624,7 +644,7 @@ void IndexHNSW::link_singletons()
         }
     }
 
-    printf("  Found %d / %ld singletons (%d appear in a level above)\n",
+    printf("  Found %d / %" PRId64 " singletons (%d appear in a level above)\n",
            n_sing, ntotal, n_sing_l1);
 
     std::vector<float>recons(singletons.size() * d);
@@ -721,7 +741,7 @@ void ReconstructFromNeighbors::reconstruct(storage_idx_t i, float *x, float *tmp
                 x[l] += w * tmp[l];
         }
     } else {
-        const float *betas[nsq];
+        std::vector<const float *> betas(nsq);
         {
             const float *b = codebook.data();
             const uint8_t *c = &codes[i * code_size];
@@ -948,6 +968,7 @@ int search_from_candidates_2(const HNSW & hnsw,
                              idx_t *I, float * D,
                              MinimaxHeap &candidates,
                              VisitedTable &vt,
+                             HNSWStats &stats,
                              int level, int nres_in = 0)
 {
     int nres = nres_in;
@@ -996,15 +1017,9 @@ int search_from_candidates_2(const HNSW & hnsw,
         }
     }
 
-    if (level == 0) {
-#pragma omp critical
-        {
-            hnsw_stats.n1 ++;
-            if (candidates.size() == 0)
-                hnsw_stats.n2 ++;
-        }
-    }
-
+    stats.n1 ++;
+    if (candidates.size() == 0)
+        stats.n2 ++;
 
     return nres;
 }
@@ -1019,6 +1034,7 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
         IndexHNSW::search (n, x, k, distances, labels);
 
     } else { // "mixed" search
+        size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
 
         const IndexIVFPQ *index_ivfpq =
             dynamic_cast<const IndexIVFPQ*>(storage);
@@ -1044,7 +1060,7 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
             int candidates_size = hnsw.upper_beam;
             MinimaxHeap candidates(candidates_size);
 
-#pragma omp for
+#pragma omp for reduction (+ : n1, n2, n3, ndis, nreorder)
             for(idx_t i = 0; i < n; i++) {
                 idx_t * idxi = labels + i * k;
                 float * simi = distances + i * k;
@@ -1081,10 +1097,16 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
                     // reorder from sorted to heap
                     maxheap_heapify (k, simi, idxi, simi, idxi, k);
 
+                    HNSWStats search_stats;
                     hnsw.search_from_candidates(
                       *dis, k, idxi, simi,
-                      candidates, vt, 0, k
+                      candidates, vt, search_stats, 0, k
                     );
+                    n1 += search_stats.n1;
+                    n2 += search_stats.n2;
+                    n3 += search_stats.n3;
+                    ndis += search_stats.ndis;
+                    nreorder += search_stats.nreorder;
 
                     vt.advance();
 
@@ -1098,9 +1120,16 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
                     // 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, 0, k);
+                        candidates, vt, search_stats, 0, k);
+                    n1 += search_stats.n1;
+                    n2 += search_stats.n2;
+                    n3 += search_stats.n3;
+                    ndis += search_stats.ndis;
+                    nreorder += search_stats.nreorder;
+
                     vt.advance ();
                     vt.advance ();
 
@@ -1109,6 +1138,8 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
                 maxheap_reorder (k, simi, idxi);
             }
         }
+
+        hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
     }
 
 

data/vendor/faiss/{IndexIVF.cpp → faiss/IndexIVF.cpp}

@@ -11,7 +11,10 @@
 
 
 #include <omp.h>
+#include <mutex>
 
+#include <algorithm>
+#include <cinttypes>
 #include <cstdio>
 #include <memory>
 
@@ -72,7 +75,7 @@ void Level1Quantizer::train_q1 (size_t n, const float *x, bool verbose, MetricTy
                           "nlist not consistent with quantizer size");
     } else if (quantizer_trains_alone == 0) {
         if (verbose)
-            printf ("Training level-1 quantizer on %ld vectors in %ldD\n",
+            printf ("Training level-1 quantizer on %zd vectors in %zdD\n",
                     n, d);
 
         Clustering clus (d, nlist, cp);
@@ -87,7 +90,7 @@ void Level1Quantizer::train_q1 (size_t n, const float *x, bool verbose, MetricTy
     } else if (quantizer_trains_alone == 2) {
         if (verbose)
             printf (
-                "Training L2 quantizer on %ld vectors in %ldD%s\n",
+                "Training L2 quantizer on %zd vectors in %zdD%s\n",
                 n, d,
                 clustering_index ? "(user provided index)" : "");
         FAISS_THROW_IF_NOT (metric_type == METRIC_L2);
@@ -188,7 +191,7 @@ void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
         for (idx_t i0 = 0; i0 < n; i0 += bs) {
             idx_t i1 = std::min (n, i0 + bs);
             if (verbose) {
-                printf("   IndexIVF::add_with_ids %ld:%ld\n", i0, i1);
+                printf("   IndexIVF::add_with_ids %" PRId64 ":%" PRId64 "\n", i0, i1);
             }
             add_with_ids (i1 - i0, x + i0 * d,
                           xids ? xids + i0 : nullptr);
@@ -238,7 +241,7 @@ void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
 
 
     if (verbose) {
-        printf("    added %ld / %ld vectors (%ld -1s)\n", nadd, n, nminus1);
+        printf("    added %zd / %" PRId64 " vectors (%zd -1s)\n", nadd, n, nminus1);
     }
 
     ntotal += n;
@@ -253,6 +256,8 @@ void IndexIVF::make_direct_map (bool b)
     }
 }
 
+
+
 void IndexIVF::set_direct_map_type (DirectMap::Type type)
 {
     direct_map.set_type (type, invlists, ntotal);
@@ -278,7 +283,6 @@ void IndexIVF::search (idx_t n, const float *x, idx_t k,
 }
 
 
-
 void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
                                    const idx_t *keys,
                                    const float *coarse_dis ,
@@ -295,15 +299,18 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
     using HeapForL2 = CMax<float, idx_t>;
 
     bool interrupt = false;
+    std::mutex exception_mutex;
+    std::string exception_string;
 
     int pmode = this->parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT;
     bool do_heap_init = !(this->parallel_mode & PARALLEL_MODE_NO_HEAP_INIT);
 
     // don't start parallel section if single query
-    bool do_parallel =
-        pmode == 0 ? n > 1 :
-        pmode == 1 ? nprobe > 1 :
-        nprobe * n > 1;
+    bool do_parallel = omp_get_max_threads() >= 2 && (
+            pmode == 0 ? n > 1 :
+            pmode == 1 ? nprobe > 1 :
+            nprobe * n > 1);
+
 
 #pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis, nheap)
     {
@@ -327,6 +334,19 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
             }
         };
 
+        auto add_local_results = [&](
+                const float * local_dis, const idx_t * local_idx,
+                float *simi, idx_t *idxi)
+        {
+            if (metric_type == METRIC_INNER_PRODUCT) {
+                heap_addn<HeapForIP>
+                    (k, simi, idxi, local_dis, local_idx, k);
+            } else {
+                heap_addn<HeapForL2>
+                    (k, simi, idxi, local_dis, local_idx, k);
+            }
+        };
+
         auto reorder_result = [&] (float *simi, idx_t *idxi) {
             if (!do_heap_init) return;
             if (metric_type == METRIC_INNER_PRODUCT) {
@@ -346,7 +366,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
                 return (size_t)0;
             }
             FAISS_THROW_IF_NOT_FMT (key < (idx_t) nlist,
-                                    "Invalid key=%ld nlist=%ld\n",
+                                    "Invalid key=%" PRId64 " nlist=%zd\n",
                                     key, nlist);
 
             size_t list_size = invlists->list_size(key);
@@ -360,18 +380,27 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
 
             nlistv++;
 
-            InvertedLists::ScopedCodes scodes (invlists, key);
+            try {
+                InvertedLists::ScopedCodes scodes (invlists, key);
 
-            std::unique_ptr<InvertedLists::ScopedIds> sids;
-            const Index::idx_t * ids = nullptr;
+                std::unique_ptr<InvertedLists::ScopedIds> sids;
+                const Index::idx_t * ids = nullptr;
 
-            if (!store_pairs)  {
-                sids.reset (new InvertedLists::ScopedIds (invlists, key));
-                ids = sids->get();
-            }
+                if (!store_pairs)  {
+                    sids.reset (new InvertedLists::ScopedIds (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);
+            } catch(const std::exception & e) {
+                std::lock_guard<std::mutex> lock(exception_mutex);
+                exception_string =
+                    demangle_cpp_symbol(typeid(e).name()) + "  " + e.what();
+                interrupt = true;
+                return size_t(0);
+            }
 
             return list_size;
         };
@@ -383,7 +412,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
         if (pmode == 0) {
 
 #pragma omp for
-            for (size_t i = 0; i < n; i++) {
+            for (idx_t i = 0; i < n; i++) {
 
                 if (interrupt) {
                     continue;
@@ -429,7 +458,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
                 init_result (local_dis.data(), local_idx.data());
 
 #pragma omp for schedule(dynamic)
-                for (size_t ik = 0; ik < nprobe; ik++) {
+                for (long ik = 0; ik < nprobe; ik++) {
                     ndis += scan_one_list
                         (keys [i * nprobe + ik],
                          coarse_dis[i * nprobe + ik],
@@ -447,20 +476,42 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
 #pragma omp barrier
 #pragma omp critical
                 {
-                    if (metric_type == METRIC_INNER_PRODUCT) {
-                        heap_addn<HeapForIP>
-                            (k, simi, idxi,
-                             local_dis.data(), local_idx.data(), k);
-                    } else {
-                        heap_addn<HeapForL2>
-                            (k, simi, idxi,
-                             local_dis.data(), local_idx.data(), k);
-                    }
+                    add_local_results (local_dis.data(), local_idx.data(),
+                                       simi, idxi);
                 }
 #pragma omp barrier
 #pragma omp single
                 reorder_result (simi, idxi);
             }
+        } else if (pmode == 2) {
+            std::vector <idx_t> local_idx (k);
+            std::vector <float> local_dis (k);
+
+#pragma omp single
+            for (int64_t i = 0; i < n; i++) {
+                init_result (distances + i * k, labels + i * k);
+            }
+
+#pragma omp for schedule(dynamic)
+            for (int64_t ij = 0; ij < n * nprobe; ij++) {
+                size_t i = ij / nprobe;
+                size_t j = ij % nprobe;
+
+                scanner->set_query (x + i * d);
+                init_result (local_dis.data(), local_idx.data());
+                ndis += scan_one_list (
+                        keys [ij], coarse_dis[ij],
+                        local_dis.data(), local_idx.data());
+#pragma omp critical
+                {
+                    add_local_results (local_dis.data(), local_idx.data(),
+                                       distances + i * k, labels + i * k);
+                }
+            }
+#pragma omp single
+            for (int64_t i = 0; i < n; i++) {
+                reorder_result (distances + i * k, labels + i * k);
+            }
         } else {
             FAISS_THROW_FMT ("parallel_mode %d not supported\n",
                              pmode);
@@ -468,7 +519,12 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
     } // parallel section
 
     if (interrupt) {
-        FAISS_THROW_MSG ("computation interrupted");
+        if (!exception_string.empty()) {
+            FAISS_THROW_FMT ("search interrupted with: %s",
+                             exception_string.c_str());
+        } else {
+            FAISS_THROW_MSG ("computation interrupted");
+        }
     }
 
     indexIVF_stats.nq += n;
@@ -503,11 +559,18 @@ void IndexIVF::range_search (idx_t nx, const float *x, float radius,
 void IndexIVF::range_search_preassigned (
          idx_t nx, const float *x, float radius,
          const idx_t *keys, const float *coarse_dis,
-         RangeSearchResult *result) const
+         RangeSearchResult *result,
+         bool store_pairs,
+         const IVFSearchParameters *params) const
 {
+    long nprobe = params ? params->nprobe : this->nprobe;
+    long max_codes = params ? params->max_codes : this->max_codes;
 
     size_t nlistv = 0, ndis = 0;
-    bool store_pairs = false;
+
+    bool interrupt = false;
+    std::mutex exception_mutex;
+    std::string exception_string;
 
     std::vector<RangeSearchPartialResult *> all_pres (omp_get_max_threads());
 
@@ -527,26 +590,36 @@ void IndexIVF::range_search_preassigned (
             if (key < 0) return;
             FAISS_THROW_IF_NOT_FMT (
                   key < (idx_t) nlist,
-                  "Invalid key=%ld  at ik=%ld nlist=%ld\n",
+                  "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
                   key, ik, nlist);
             const size_t list_size = invlists->list_size(key);
 
             if (list_size == 0) return;
 
-            InvertedLists::ScopedCodes scodes (invlists, key);
-            InvertedLists::ScopedIds ids (invlists, key);
+            try {
+
+                InvertedLists::ScopedCodes scodes (invlists, key);
+                InvertedLists::ScopedIds ids (invlists, key);
+
+                scanner->set_list (key, coarse_dis[i * nprobe + ik]);
+                nlistv++;
+                ndis += list_size;
+                scanner->scan_codes_range (list_size, scodes.get(),
+                                        ids.get(), radius, qres);
+
+            } catch(const std::exception & e) {
+                std::lock_guard<std::mutex> lock(exception_mutex);
+                exception_string =
+                    demangle_cpp_symbol(typeid(e).name()) + "  " + e.what();
+                interrupt = true;
+            }
 
-            scanner->set_list (key, coarse_dis[i * nprobe + ik]);
-            nlistv++;
-            ndis += list_size;
-            scanner->scan_codes_range (list_size, scodes.get(),
-                                       ids.get(), radius, qres);
         };
 
         if (parallel_mode == 0) {
 
 #pragma omp for
-            for (size_t i = 0; i < nx; i++) {
+            for (idx_t i = 0; i < nx; i++) {
                 scanner->set_query (x + i * d);
 
                 RangeQueryResult & qres = pres.new_result (i);
@@ -565,7 +638,7 @@ void IndexIVF::range_search_preassigned (
                 RangeQueryResult & qres = pres.new_result (i);
 
 #pragma omp for schedule(dynamic)
-                for (size_t ik = 0; ik < nprobe; ik++) {
+                for (int64_t ik = 0; ik < nprobe; ik++) {
                     scan_list_func (i, ik, qres);
                 }
             }
@@ -574,9 +647,9 @@ void IndexIVF::range_search_preassigned (
             RangeQueryResult *qres = nullptr;
 
 #pragma omp for schedule(dynamic)
-            for (size_t iik = 0; iik < nx * nprobe; iik++) {
-                size_t i = iik / nprobe;
-                size_t ik = iik % nprobe;
+            for (idx_t iik = 0; iik < nx * (idx_t)nprobe; iik++) {
+                idx_t i = iik / (idx_t)nprobe;
+                idx_t ik = iik % (idx_t)nprobe;
                 if (qres == nullptr || qres->qno != i) {
                     FAISS_ASSERT (!qres || i > qres->qno);
                     qres = &pres.new_result (i);
@@ -597,6 +670,16 @@ void IndexIVF::range_search_preassigned (
 
         }
     }
+
+    if (interrupt) {
+        if (!exception_string.empty()) {
+            FAISS_THROW_FMT ("search interrupted with: %s",
+                             exception_string.c_str());
+        } else {
+            FAISS_THROW_MSG ("computation interrupted");
+        }
+    }
+
     indexIVF_stats.nq += nx;
     indexIVF_stats.nlist += nlistv;
     indexIVF_stats.ndis += ndis;