faiss 0.1.1 → 0.1.2

data/vendor/faiss/IndexBinaryIVF.h

@@ -46,8 +46,7 @@ struct IndexBinaryIVF : IndexBinary {
     bool use_heap = true;
 
     /// map for direct access to the elements. Enables reconstruct().
-    bool maintain_direct_map;
-    std::vector<idx_t> direct_map;
+    DirectMap direct_map;
 
     IndexBinary *quantizer;   ///< quantizer that maps vectors to inverted lists
     size_t nlist;             ///< number of possible key values
@@ -110,8 +109,11 @@ struct IndexBinaryIVF : IndexBinary {
                                          bool store_pairs=false) const;
 
     /** assign the vectors, then call search_preassign */
-    virtual void search(idx_t n, const uint8_t *x, idx_t k,
-                        int32_t *distances, idx_t *labels) const override;
+    void search(idx_t n, const uint8_t *x, idx_t k,
+                int32_t *distances, idx_t *labels) const override;
+
+    void range_search(idx_t n, const uint8_t *x, int radius,
+                      RangeSearchResult *result) const override;
 
     void reconstruct(idx_t key, uint8_t *recons) const override;
 
@@ -168,6 +170,8 @@ struct IndexBinaryIVF : IndexBinary {
      */
     void make_direct_map(bool new_maintain_direct_map=true);
 
+    void set_direct_map_type (DirectMap::Type type);
+
     void replace_invlists(InvertedLists *il, bool own=false);
 };
 
@@ -201,6 +205,12 @@ struct BinaryInvertedListScanner {
                                int32_t *distances, idx_t *labels,
                                size_t k) const = 0;
 
+    virtual void scan_codes_range (size_t n,
+                                   const uint8_t *codes,
+                                   const idx_t *ids,
+                                   int radius,
+                                   RangeQueryResult &result) const = 0;
+
     virtual ~BinaryInvertedListScanner () {}
 
 };

data/vendor/faiss/IndexFlat.h

@@ -19,6 +19,7 @@ namespace faiss {
 
 /** Index that stores the full vectors and performs exhaustive search */
 struct IndexFlat: Index {
+
     /// database vectors, size ntotal * d
     std::vector<float> xb;
 
@@ -144,7 +145,7 @@ struct IndexRefineFlat: Index {
 };
 
 
-/// optimized version for 1D "vectors"
+/// optimized version for 1D "vectors".
 struct IndexFlat1D:IndexFlatL2 {
     bool continuous_update; ///< is the permutation updated continuously?
 

data/vendor/faiss/IndexHNSW.cpp

@@ -26,7 +26,6 @@
 #include <stdint.h>
 
 #ifdef __SSE__
-#include <immintrin.h>
 #endif
 
 #include <faiss/utils/distances.h>
@@ -55,7 +54,6 @@ namespace faiss {
 using idx_t = Index::idx_t;
 using MinimaxHeap = HNSW::MinimaxHeap;
 using storage_idx_t = HNSW::storage_idx_t;
-using NodeDistCloser = HNSW::NodeDistCloser;
 using NodeDistFarther = HNSW::NodeDistFarther;
 
 HNSWStats hnsw_stats;
@@ -67,6 +65,50 @@ 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 (storage->metric_type == METRIC_INNER_PRODUCT) {
+        return new NegativeDistanceComputer(storage->get_distance_computer());
+    } else {
+        return storage->get_distance_computer();
+    }
+}
+
+
+
 void hnsw_add_vertices(IndexHNSW &index_hnsw,
                        size_t n0,
                        size_t n, const float *x,
@@ -152,7 +194,7 @@ void hnsw_add_vertices(IndexHNSW &index_hnsw,
                 VisitedTable vt (ntotal);
 
                 DistanceComputer *dis =
-                    index_hnsw.storage->get_distance_computer();
+                    storage_distance_computer (index_hnsw.storage);
                 ScopeDeleter1<DistanceComputer> del(dis);
                 int prev_display = verbose && omp_get_thread_num() == 0 ? 0 : -1;
                 size_t counter = 0;
@@ -210,8 +252,8 @@ void hnsw_add_vertices(IndexHNSW &index_hnsw,
  * IndexHNSW implementation
  **************************************************************/
 
-IndexHNSW::IndexHNSW(int d, int M):
-    Index(d, METRIC_L2),
+IndexHNSW::IndexHNSW(int d, int M, MetricType metric):
+    Index(d, metric),
     hnsw(M),
     own_fields(false),
     storage(nullptr),
@@ -258,7 +300,8 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
 #pragma omp parallel reduction(+ : nreorder)
         {
             VisitedTable vt (ntotal);
-            DistanceComputer *dis = storage->get_distance_computer();
+
+            DistanceComputer *dis = storage_distance_computer(storage);
             ScopeDeleter1<DistanceComputer> del(dis);
 
 #pragma omp for
@@ -290,6 +333,14 @@ void IndexHNSW::search (idx_t n, const float *x, idx_t k,
         }
         InterruptCallback::check ();
     }
+
+    if (metric_type == METRIC_INNER_PRODUCT) {
+        // we need to revert the negated distances
+        for (size_t i = 0; i < k * n; i++) {
+            distances[i] = -distances[i];
+        }
+    }
+
     hnsw_stats.nreorder += nreorder;
 }
 
@@ -323,7 +374,7 @@ void IndexHNSW::shrink_level_0_neighbors(int new_size)
 {
 #pragma omp parallel
     {
-        DistanceComputer *dis = storage->get_distance_computer();
+        DistanceComputer *dis = storage_distance_computer(storage);
         ScopeDeleter1<DistanceComputer> del(dis);
 
 #pragma omp for
@@ -367,7 +418,7 @@ void IndexHNSW::search_level_0(
     storage_idx_t ntotal = hnsw.levels.size();
 #pragma omp parallel
     {
-        DistanceComputer *qdis = storage->get_distance_computer();
+        DistanceComputer *qdis = storage_distance_computer(storage);
         ScopeDeleter1<DistanceComputer> del(qdis);
 
         VisitedTable vt (ntotal);
@@ -436,7 +487,7 @@ void IndexHNSW::init_level_0_from_knngraph(
 
 #pragma omp parallel for
     for (idx_t i = 0; i < ntotal; i++) {
-        DistanceComputer *qdis = storage->get_distance_computer();
+        DistanceComputer *qdis = storage_distance_computer(storage);
         float vec[d];
         storage->reconstruct(i, vec);
         qdis->set_query(vec);
@@ -480,7 +531,7 @@ void IndexHNSW::init_level_0_from_entry_points(
     {
         VisitedTable vt (ntotal);
 
-        DistanceComputer *dis = storage->get_distance_computer();
+        DistanceComputer *dis = storage_distance_computer(storage);
         ScopeDeleter1<DistanceComputer> del(dis);
         float vec[storage->d];
 
@@ -518,7 +569,7 @@ 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->get_distance_computer();
+        DistanceComputer *dis = storage_distance_computer(storage);
         ScopeDeleter1<DistanceComputer> del(dis);
 
 #pragma omp for
@@ -826,8 +877,8 @@ IndexHNSWFlat::IndexHNSWFlat()
     is_trained = true;
 }
 
-IndexHNSWFlat::IndexHNSWFlat(int d, int M):
-    IndexHNSW(new IndexFlatL2(d), M)
+IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric):
+    IndexHNSW(new IndexFlat(d, metric), M)
 {
     own_fields = true;
     is_trained = true;
@@ -860,8 +911,9 @@ void IndexHNSWPQ::train(idx_t n, const float* x)
  **************************************************************/
 
 
-IndexHNSWSQ::IndexHNSWSQ(int d, ScalarQuantizer::QuantizerType qtype, int M):
-    IndexHNSW (new IndexScalarQuantizer (d, qtype), M)
+IndexHNSWSQ::IndexHNSWSQ(int d, ScalarQuantizer::QuantizerType qtype, int M,
+                         MetricType metric):
+    IndexHNSW (new IndexScalarQuantizer (d, qtype, metric), M)
 {
     is_trained = false;
     own_fields = true;
@@ -986,7 +1038,7 @@ void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
 #pragma omp parallel
         {
             VisitedTable vt (ntotal);
-            DistanceComputer *dis = storage->get_distance_computer();
+            DistanceComputer *dis = storage_distance_computer(storage);
             ScopeDeleter1<DistanceComputer> del(dis);
 
             int candidates_size = hnsw.upper_beam;

data/vendor/faiss/IndexHNSW.h

@@ -79,7 +79,7 @@ struct IndexHNSW : Index {
 
     ReconstructFromNeighbors *reconstruct_from_neighbors;
 
-    explicit IndexHNSW (int d = 0, int M = 32);
+    explicit IndexHNSW (int d = 0, int M = 32, MetricType metric = METRIC_L2);
     explicit IndexHNSW (Index *storage, int M = 32);
 
     ~IndexHNSW() override;
@@ -132,7 +132,7 @@ struct IndexHNSW : Index {
 
 struct IndexHNSWFlat : IndexHNSW {
     IndexHNSWFlat();
-    IndexHNSWFlat(int d, int M);
+    IndexHNSWFlat(int d, int M, MetricType metric = METRIC_L2);
 };
 
 /** PQ index topped with with a HNSW structure to access elements
@@ -149,7 +149,7 @@ struct IndexHNSWPQ : IndexHNSW {
  */
 struct IndexHNSWSQ : IndexHNSW {
     IndexHNSWSQ();
-    IndexHNSWSQ(int d, ScalarQuantizer::QuantizerType qtype, int M);
+    IndexHNSWSQ(int d, ScalarQuantizer::QuantizerType qtype, int M, MetricType metric = METRIC_L2);
 };
 
 /** 2-level code structure with fast random access

data/vendor/faiss/IndexIVF.cpp

@@ -157,8 +157,7 @@ IndexIVF::IndexIVF (Index * quantizer, size_t d,
     code_size (code_size),
     nprobe (1),
     max_codes (0),
-    parallel_mode (0),
-    maintain_direct_map (false)
+    parallel_mode (0)
 {
     FAISS_THROW_IF_NOT (d == quantizer->d);
     is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);
@@ -172,8 +171,7 @@ IndexIVF::IndexIVF (Index * quantizer, size_t d,
 IndexIVF::IndexIVF ():
     invlists (nullptr), own_invlists (false),
     code_size (0),
-    nprobe (1), max_codes (0), parallel_mode (0),
-    maintain_direct_map (false)
+    nprobe (1), max_codes (0), parallel_mode (0)
 {}
 
 void IndexIVF::add (idx_t n, const float * x)
@@ -199,6 +197,8 @@ void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
     }
 
     FAISS_THROW_IF_NOT (is_trained);
+    direct_map.check_can_add (xids);
+
     std::unique_ptr<idx_t []> idx(new idx_t[n]);
     quantizer->assign (n, x, idx.get());
     size_t nadd = 0, nminus1 = 0;
@@ -210,6 +210,8 @@ void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
     std::unique_ptr<uint8_t []> flat_codes(new uint8_t [n * code_size]);
     encode_vectors (n, x, idx.get(), flat_codes.get());
 
+    DirectMapAdd dm_adder(direct_map, n, xids);
+
 #pragma omp parallel reduction(+: nadd)
     {
         int nt = omp_get_num_threads();
@@ -220,13 +222,21 @@ void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
             idx_t list_no = idx [i];
             if (list_no >= 0 && list_no % nt == rank) {
                 idx_t id = xids ? xids[i] : ntotal + i;
-                invlists->add_entry (list_no, id,
-                                     flat_codes.get() + i * code_size);
+                size_t ofs = invlists->add_entry (
+                     list_no, id,
+                     flat_codes.get() + i * code_size
+                );
+
+                dm_adder.add (i, list_no, ofs);
+
                 nadd++;
+            } else if (rank == 0 && list_no == -1) {
+                dm_adder.add (i, -1, 0);
             }
         }
     }
 
+
     if (verbose) {
         printf("    added %ld / %ld vectors (%ld -1s)\n", nadd, n, nminus1);
     }
@@ -234,30 +244,18 @@ void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
     ntotal += n;
 }
 
-
-void IndexIVF::make_direct_map (bool new_maintain_direct_map)
+void IndexIVF::make_direct_map (bool b)
 {
-    // nothing to do
-    if (new_maintain_direct_map == maintain_direct_map)
-        return;
-
-    if (new_maintain_direct_map) {
-        direct_map.resize (ntotal, -1);
-        for (size_t key = 0; key < nlist; key++) {
-            size_t list_size = invlists->list_size (key);
-            ScopedIds idlist (invlists, key);
-
-            for (long ofs = 0; ofs < list_size; ofs++) {
-                FAISS_THROW_IF_NOT_MSG (
-                       0 <= idlist [ofs] && idlist[ofs] < ntotal,
-                       "direct map supported only for seuquential ids");
-                direct_map [idlist [ofs]] = key << 32 | ofs;
-            }
-        }
+    if (b) {
+        direct_map.set_type (DirectMap::Array, invlists, ntotal);
     } else {
-        direct_map.clear ();
+        direct_map.set_type (DirectMap::NoMap, invlists, ntotal);
     }
-    maintain_direct_map = new_maintain_direct_map;
+}
+
+void IndexIVF::set_direct_map_type (DirectMap::Type type)
+{
+    direct_map.set_type (type, invlists, ntotal);
 }
 
 
@@ -298,10 +296,13 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
 
     bool interrupt = false;
 
+    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 =
-        parallel_mode == 0 ? n > 1 :
-        parallel_mode == 1 ? nprobe > 1 :
+        pmode == 0 ? n > 1 :
+        pmode == 1 ? nprobe > 1 :
         nprobe * n > 1;
 
 #pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis, nheap)
@@ -318,6 +319,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
         // intialize + reorder a result heap
 
         auto init_result = [&](float *simi, idx_t *idxi) {
+            if (!do_heap_init) return;
             if (metric_type == METRIC_INNER_PRODUCT) {
                 heap_heapify<HeapForIP> (k, simi, idxi);
             } else {
@@ -326,6 +328,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
         };
 
         auto reorder_result = [&] (float *simi, idx_t *idxi) {
+            if (!do_heap_init) return;
             if (metric_type == METRIC_INNER_PRODUCT) {
                 heap_reorder<HeapForIP> (k, simi, idxi);
             } else {
@@ -377,7 +380,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
          * Actual loops, depending on parallel_mode
          ****************************************************/
 
-        if (parallel_mode == 0) {
+        if (pmode == 0) {
 
 #pragma omp for
             for (size_t i = 0; i < n; i++) {
@@ -417,7 +420,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
                 }
 
             } // parallel for
-        } else if (parallel_mode == 1) {
+        } else if (pmode == 1) {
             std::vector <idx_t> local_idx (k);
             std::vector <float> local_dis (k);
 
@@ -460,7 +463,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
             }
         } else {
             FAISS_THROW_FMT ("parallel_mode %d not supported\n",
-                             parallel_mode);
+                             pmode);
         }
     } // parallel section
 
@@ -608,13 +611,8 @@ InvertedListScanner *IndexIVF::get_InvertedListScanner (
 
 void IndexIVF::reconstruct (idx_t key, float* recons) const
 {
-    FAISS_THROW_IF_NOT_MSG (direct_map.size() == ntotal,
-                            "direct map is not initialized");
-    FAISS_THROW_IF_NOT_MSG (key >= 0 && key < direct_map.size(),
-                            "invalid key");
-    idx_t list_no = direct_map[key] >> 32;
-    idx_t offset = direct_map[key] & 0xffffffff;
-    reconstruct_from_offset (list_no, offset, recons);
+    idx_t lo = direct_map.get (key);
+    reconstruct_from_offset (lo_listno(lo), lo_offset(lo), recons);
 }
 
 
@@ -682,8 +680,8 @@ void IndexIVF::search_and_reconstruct (idx_t n, const float *x, idx_t k,
                 // Fill with NaNs
                 memset(reconstructed, -1, sizeof(*reconstructed) * d);
             } else {
-                int list_no = key >> 32;
-                int offset = key & 0xffffffff;
+                int list_no = lo_listno (key);
+                int offset = lo_offset (key);
 
                 // Update label to the actual id
                 labels[ij] = invlists->get_single_id (list_no, offset);
@@ -711,42 +709,41 @@ void IndexIVF::reset ()
 
 size_t IndexIVF::remove_ids (const IDSelector & sel)
 {
-    FAISS_THROW_IF_NOT_MSG (!maintain_direct_map,
-                    "direct map remove not implemented");
-
-    std::vector<idx_t> toremove(nlist);
-
-#pragma omp parallel for
-    for (idx_t i = 0; i < nlist; i++) {
-        idx_t l0 = invlists->list_size (i), l = l0, j = 0;
-        ScopedIds idsi (invlists, i);
-        while (j < l) {
-            if (sel.is_member (idsi[j])) {
-                l--;
-                invlists->update_entry (
-                     i, j,
-                     invlists->get_single_id (i, l),
-                     ScopedCodes (invlists, i, l).get());
-            } else {
-                j++;
-            }
-        }
-        toremove[i] = l0 - l;
-    }
-    // this will not run well in parallel on ondisk because of possible shrinks
-    size_t nremove = 0;
-    for (idx_t i = 0; i < nlist; i++) {
-        if (toremove[i] > 0) {
-            nremove += toremove[i];
-            invlists->resize(
-                i, invlists->list_size(i) - toremove[i]);
-        }
-    }
+    size_t nremove = direct_map.remove_ids (sel, invlists);
     ntotal -= nremove;
     return nremove;
 }
 
 
+void IndexIVF::update_vectors (int n, const idx_t *new_ids, const float *x)
+{
+
+    if (direct_map.type == DirectMap::Hashtable) {
+        // just remove then add
+        IDSelectorArray sel(n, new_ids);
+        size_t nremove = remove_ids (sel);
+        FAISS_THROW_IF_NOT_MSG (nremove == n,
+                                "did not find all entries to remove");
+        add_with_ids (n, x, new_ids);
+        return;
+    }
+
+    FAISS_THROW_IF_NOT (direct_map.type == DirectMap::Array);
+    // here it is more tricky because we don't want to introduce holes
+    // in continuous range of ids
+
+    FAISS_THROW_IF_NOT (is_trained);
+    std::vector<idx_t> assign (n);
+    quantizer->assign (n, x, assign.data());
+
+    std::vector<uint8_t> flat_codes (n * code_size);
+    encode_vectors (n, x, assign.data(), flat_codes.data());
+
+    direct_map.update_codes (invlists, n, new_ids, assign.data(), flat_codes.data());
+
+}
+
+
 
 
 void IndexIVF::train (idx_t n, const float *x)
@@ -779,15 +776,14 @@ void IndexIVF::check_compatible_for_merge (const IndexIVF &other) const
     FAISS_THROW_IF_NOT (other.code_size == code_size);
     FAISS_THROW_IF_NOT_MSG (typeid (*this) == typeid (other),
                   "can only merge indexes of the same type");
+    FAISS_THROW_IF_NOT_MSG (this->direct_map.no() && other.direct_map.no(),
+                            "merge direct_map not implemented");
 }
 
 
 void IndexIVF::merge_from (IndexIVF &other, idx_t add_id)
 {
     check_compatible_for_merge (other);
-    FAISS_THROW_IF_NOT_MSG ((!maintain_direct_map &&
-                             !other.maintain_direct_map),
-                  "direct map copy not implemented");
 
     invlists->merge_from (other.invlists, add_id);
 
@@ -817,7 +813,7 @@ void IndexIVF::copy_subset_to (IndexIVF & other, int subset_type,
 
     FAISS_THROW_IF_NOT (nlist == other.nlist);
     FAISS_THROW_IF_NOT (code_size == other.code_size);
-    FAISS_THROW_IF_NOT (!other.maintain_direct_map);
+    FAISS_THROW_IF_NOT (other.direct_map.no());
     FAISS_THROW_IF_NOT_FMT (
           subset_type == 0 || subset_type == 1 || subset_type == 2,
           "subset type %d not implemented", subset_type);