faiss 0.1.1 → 0.1.2

data/vendor/faiss/IndexIVF.h

@@ -12,10 +12,12 @@
 
 
 #include <vector>
+#include <unordered_map>
 #include <stdint.h>
 
 #include <faiss/Index.h>
 #include <faiss/InvertedLists.h>
+#include <faiss/DirectMap.h>
 #include <faiss/Clustering.h>
 #include <faiss/utils/Heap.h>
 
@@ -32,7 +34,6 @@ struct Level1Quantizer {
     Index * quantizer;        ///< quantizer that maps vectors to inverted lists
     size_t nlist;             ///< number of possible key values
 
-
     /**
      * = 0: use the quantizer as index in a kmeans training
      * = 1: just pass on the training set to the train() of the quantizer
@@ -107,14 +108,18 @@ struct IndexIVF: Index, Level1Quantizer {
     /** Parallel mode determines how queries are parallelized with OpenMP
      *
      * 0 (default): parallelize over queries
-     * 1: parallelize over over inverted lists
+     * 1: parallelize over inverted lists
      * 2: parallelize over both
+     *
+     * PARALLEL_MODE_NO_HEAP_INIT: binary or with the previous to
+     * prevent the heap to be initialized and finalized
      */
     int parallel_mode;
+    const int PARALLEL_MODE_NO_HEAP_INIT = 1024;
 
-    /// map for direct access to the elements. Enables reconstruct().
-    bool maintain_direct_map;
-    std::vector <idx_t> direct_map;
+    /** optional map that maps back ids to invlist entries. This
+     *  enables reconstruct() */
+    DirectMap direct_map;
 
     /** The Inverted file takes a quantizer (an Index) on input,
      * which implements the function mapping a vector to a list
@@ -195,8 +200,19 @@ struct IndexIVF: Index, Level1Quantizer {
     virtual InvertedListScanner *get_InvertedListScanner (
         bool store_pairs=false) const;
 
+    /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2 */
     void reconstruct (idx_t key, float* recons) const override;
 
+    /** Update a subset of vectors.
+     *
+     * The index must have a direct_map
+     *
+     * @param nv     nb of vectors to update
+     * @param idx    vector indices to update, size nv
+     * @param v      vectors of new values, size nv*d
+     */
+    virtual void update_vectors (int nv, const idx_t *idx, const float *v);
+
     /** Reconstruct a subset of the indexed vectors.
      *
      * Overrides default implementation to bypass reconstruct() which requires
@@ -268,6 +284,9 @@ struct IndexIVF: Index, Level1Quantizer {
      */
     void make_direct_map (bool new_maintain_direct_map=true);
 
+    void set_direct_map_type (DirectMap::Type type);
+
+
     /// replace the inverted lists, old one is deallocated if own_invlists
     void replace_invlists (InvertedLists *il, bool own=false);
 

data/vendor/faiss/IndexIVFFlat.cpp

@@ -45,8 +45,7 @@ void IndexIVFFlat::add_core (idx_t n, const float * x, const int64_t *xids,
 {
     FAISS_THROW_IF_NOT (is_trained);
     assert (invlists);
-    FAISS_THROW_IF_NOT_MSG (!(maintain_direct_map && xids),
-                            "cannot have direct map and add with ids");
+    direct_map.check_can_add (xids);
     const int64_t * idx;
     ScopeDeleter<int64_t> del;
 
@@ -60,19 +59,21 @@ void IndexIVFFlat::add_core (idx_t n, const float * x, const int64_t *xids,
     }
     int64_t n_add = 0;
     for (size_t i = 0; i < n; i++) {
-        int64_t id = xids ? xids[i] : ntotal + i;
-        int64_t list_no = idx [i];
-
-        if (list_no < 0)
-            continue;
-        const float *xi = x + i * d;
-        size_t offset = invlists->add_entry (
-              list_no, id, (const uint8_t*) xi);
+        idx_t id = xids ? xids[i] : ntotal + i;
+        idx_t list_no = idx [i];
+        size_t offset;
 
-        if (maintain_direct_map)
-            direct_map.push_back (list_no << 32 | offset);
-        n_add++;
+        if (list_no >= 0) {
+            const float *xi = x + i * d;
+            offset = invlists->add_entry (
+                     list_no, id, (const uint8_t*) xi);
+            n_add++;
+        } else {
+            offset = 0;
+        }
+        direct_map.add_single_id (id, list_no, offset);
     }
+
     if (verbose) {
         printf("IndexIVFFlat::add_core: added %ld / %ld vectors\n",
                n_add, n);
@@ -158,7 +159,7 @@ struct IVFFlatScanner: InvertedListScanner {
                 fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
             if (C::cmp (simi[0], dis)) {
                 heap_pop<C> (k, simi, idxi);
-                int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
+                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
                 heap_push<C> (k, simi, idxi, dis, id);
                 nup++;
             }
@@ -178,7 +179,7 @@ struct IVFFlatScanner: InvertedListScanner {
             float dis = metric == METRIC_INNER_PRODUCT ?
                 fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
             if (C::cmp (radius, dis)) {
-                int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
+                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
                 res.add (dis, id);
             }
         }
@@ -209,41 +210,6 @@ InvertedListScanner* IndexIVFFlat::get_InvertedListScanner
 
 
 
-void IndexIVFFlat::update_vectors (int n, idx_t *new_ids, const float *x)
-{
-
-    FAISS_THROW_IF_NOT (maintain_direct_map);
-    FAISS_THROW_IF_NOT (is_trained);
-    std::vector<idx_t> assign (n);
-    quantizer->assign (n, x, assign.data());
-
-    for (size_t i = 0; i < n; i++) {
-        idx_t id = new_ids[i];
-        FAISS_THROW_IF_NOT_MSG (0 <= id && id < ntotal,
-                                "id to update out of range");
-        { // remove old one
-            int64_t dm = direct_map[id];
-            int64_t ofs = dm & 0xffffffff;
-            int64_t il = dm >> 32;
-            size_t l = invlists->list_size (il);
-            if (ofs != l - 1) { // move l - 1 to ofs
-                int64_t id2 = invlists->get_single_id (il, l - 1);
-                direct_map[id2] = (il << 32) | ofs;
-                invlists->update_entry (il, ofs, id2,
-                                        invlists->get_single_code (il, l - 1));
-            }
-            invlists->resize (il, l - 1);
-        }
-        { // insert new one
-            int64_t il = assign[i];
-            size_t l = invlists->list_size (il);
-            int64_t dm = (il << 32) | l;
-            direct_map[id] = dm;
-            invlists->add_entry (il, id, (const uint8_t*)(x + i * d));
-        }
-    }
-
-}
 
 void IndexIVFFlat::reconstruct_from_offset (int64_t list_no, int64_t offset,
                                             float* recons) const
@@ -295,8 +261,7 @@ void IndexIVFFlatDedup::add_with_ids(
 
     FAISS_THROW_IF_NOT (is_trained);
     assert (invlists);
-    FAISS_THROW_IF_NOT_MSG (
-           !maintain_direct_map,
+    FAISS_THROW_IF_NOT_MSG (direct_map.no(),
            "IVFFlatDedup not implemented with direct_map");
     int64_t * idx = new int64_t [na];
     ScopeDeleter<int64_t> del (idx);
@@ -431,7 +396,7 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
 
     // mostly copied from IndexIVF.cpp
 
-    FAISS_THROW_IF_NOT_MSG (!maintain_direct_map,
+    FAISS_THROW_IF_NOT_MSG (direct_map.no(),
                     "direct map remove not implemented");
 
     std::vector<int64_t> toremove(nlist);
@@ -484,7 +449,7 @@ void IndexIVFFlatDedup::range_search(
     FAISS_THROW_MSG ("not implemented");
 }
 
-void IndexIVFFlatDedup::update_vectors (int , idx_t *, const float *)
+void IndexIVFFlatDedup::update_vectors (int , const idx_t *, const float *)
 {
     FAISS_THROW_MSG ("not implemented");
 }

data/vendor/faiss/IndexIVFFlat.h

@@ -44,15 +44,6 @@ struct IndexIVFFlat: IndexIVF {
     InvertedListScanner *get_InvertedListScanner (bool store_pairs)
         const override;
 
-    /** Update a subset of vectors.
-     *
-     * The index must have a direct_map
-     *
-     * @param nv     nb of vectors to update
-     * @param idx    vector indices to update, size nv
-     * @param v      vectors of new values, size nv*d
-     */
-    virtual void update_vectors (int nv, idx_t *idx, const float *v);
 
     void reconstruct_from_offset (int64_t list_no, int64_t offset,
                                   float* recons) const override;
@@ -99,8 +90,7 @@ struct IndexIVFFlatDedup: IndexIVFFlat {
         RangeSearchResult* result) const override;
 
     /// not implemented
-    void update_vectors (int nv, idx_t *idx, const float *v) override;
-
+    void update_vectors (int nv, const idx_t *idx, const float *v) override;
 
     /// not implemented
     void reconstruct_from_offset (int64_t list_no, int64_t offset,

data/vendor/faiss/IndexIVFPQ.cpp

@@ -36,8 +36,8 @@ namespace faiss {
  ******************************************/
 
 IndexIVFPQ::IndexIVFPQ (Index * quantizer, size_t d, size_t nlist,
-                        size_t M, size_t nbits_per_idx):
-    IndexIVF (quantizer, d, nlist, 0, METRIC_L2),
+                        size_t M, size_t nbits_per_idx, MetricType metric):
+    IndexIVF (quantizer, d, nlist, 0, metric),
     pq (d, M, nbits_per_idx)
 {
     FAISS_THROW_IF_NOT (nbits_per_idx <= 8);
@@ -278,6 +278,8 @@ void IndexIVFPQ::add_core_o (idx_t n, const float * x, const idx_t *xids,
 
     InterruptCallback::check();
 
+    direct_map.check_can_add (xids);
+
     FAISS_THROW_IF_NOT (is_trained);
     double t0 = getmillisecs ();
     const idx_t * idx;
@@ -312,13 +314,14 @@ void IndexIVFPQ::add_core_o (idx_t n, const float * x, const idx_t *xids,
     size_t n_ignore = 0;
     for (size_t i = 0; i < n; i++) {
         idx_t key = idx[i];
+        idx_t id = xids ? xids[i] : ntotal + i;
         if (key < 0) {
+            direct_map.add_single_id (id, -1, 0);
             n_ignore ++;
             if (residuals_2)
                 memset (residuals_2, 0, sizeof(*residuals_2) * d);
             continue;
         }
-        idx_t id = xids ? xids[i] : ntotal + i;
 
         uint8_t *code = xcodes + i * code_size;
         size_t offset = invlists->add_entry (key, id, code);
@@ -331,11 +334,9 @@ void IndexIVFPQ::add_core_o (idx_t n, const float * x, const idx_t *xids,
                 res2[j] = xi[j] - res2[j];
         }
 
-        if (maintain_direct_map)
-            direct_map.push_back (key << 32 | offset);
+        direct_map.add_single_id (id, key, offset);
     }
 
-
     double t3 = getmillisecs ();
     if(verbose) {
         char comment[100] = {0};
@@ -802,7 +803,7 @@ struct KnnSearchResults {
     inline void add (idx_t j, float dis) {
         if (C::cmp (heap_sim[0], dis)) {
             heap_pop<C> (k, heap_sim, heap_ids);
-            idx_t id = ids ? ids[j] : (key << 32 | j);
+            idx_t id = ids ? ids[j] : lo_build (key, j);
             heap_push<C> (k, heap_sim, heap_ids, dis, id);
             nup++;
         }
@@ -821,7 +822,7 @@ struct RangeSearchResults {
 
     inline void add (idx_t j, float dis) {
         if (C::cmp (radius, dis)) {
-            idx_t id = ids ? ids[j] : (key << 32 | j);
+            idx_t id = ids ? ids[j] : lo_build (key, j);
             rres.add (dis, id);
         }
     }
@@ -834,7 +835,7 @@ struct RangeSearchResults {
  * The scanning functions call their favorite precompute_*
  * function to precompute the tables they need.
  *****************************************************/
-template <typename IDType, MetricType METRIC_TYPE>
+template <typename IDType, MetricType METRIC_TYPE, class PQDecoder>
 struct IVFPQScannerT: QueryTables {
 
     const uint8_t * list_codes;
@@ -844,7 +845,6 @@ struct IVFPQScannerT: QueryTables {
     IVFPQScannerT (const IndexIVFPQ & ivfpq, const IVFSearchParameters *params):
         QueryTables (ivfpq, params)
     {
-        FAISS_THROW_IF_NOT (pq.nbits == 8);
         assert(METRIC_TYPE == metric_type);
     }
 
@@ -872,12 +872,13 @@ struct IVFPQScannerT: QueryTables {
                                SearchResultType & res) const
     {
         for (size_t j = 0; j < ncode; j++) {
-
+            PQDecoder decoder(codes, pq.nbits);
+            codes += pq.code_size;
             float dis = dis0;
             const float *tab = sim_table;
 
             for (size_t m = 0; m < pq.M; m++) {
-                dis += tab[*codes++];
+                dis += tab[decoder.decode()];
                 tab += pq.ksub;
             }
 
@@ -893,12 +894,14 @@ struct IVFPQScannerT: QueryTables {
                                  SearchResultType & res) const
     {
         for (size_t j = 0; j < ncode; j++) {
+            PQDecoder decoder(codes, pq.nbits);
+            codes += pq.code_size;
 
             float dis = dis0;
             const float *tab = sim_table_2;
 
             for (size_t m = 0; m < pq.M; m++) {
-                int ci = *codes++;
+                int ci = decoder.decode();
                 dis += sim_table_ptrs [m][ci] - 2 * tab [ci];
                 tab += pq.ksub;
             }
@@ -963,12 +966,13 @@ struct IVFPQScannerT: QueryTables {
             int hd = hc.hamming (b_code);
             if (hd < ht) {
                 n_hamming_pass ++;
+                PQDecoder decoder(codes, pq.nbits);
 
                 float dis = dis0;
                 const float *tab = sim_table;
 
                 for (size_t m = 0; m < pq.M; m++) {
-                    dis += tab[*b_code++];
+                    dis += tab[decoder.decode()];
                     tab += pq.ksub;
                 }
 
@@ -1023,16 +1027,18 @@ struct IVFPQScannerT: QueryTables {
  * much we precompute (2 = precompute distance tables, 1 = precompute
  * pointers to distances, 0 = compute distances one by one).
  * Currently only 2 is supported */
-template<MetricType METRIC_TYPE, class C, int precompute_mode>
+template<MetricType METRIC_TYPE, class C, class PQDecoder>
 struct IVFPQScanner:
-    IVFPQScannerT<Index::idx_t, METRIC_TYPE>,
+    IVFPQScannerT<Index::idx_t, METRIC_TYPE, PQDecoder>,
     InvertedListScanner
 {
     bool store_pairs;
+    int precompute_mode;
 
-    IVFPQScanner(const IndexIVFPQ & ivfpq, bool store_pairs):
-        IVFPQScannerT<Index::idx_t, METRIC_TYPE>(ivfpq, nullptr),
-        store_pairs(store_pairs)
+    IVFPQScanner(const IndexIVFPQ & ivfpq, bool store_pairs,
+                 int precompute_mode):
+        IVFPQScannerT<Index::idx_t, METRIC_TYPE, PQDecoder>(ivfpq, nullptr),
+        store_pairs(store_pairs), precompute_mode(precompute_mode)
     {
     }
 
@@ -1048,9 +1054,10 @@ struct IVFPQScanner:
         assert(precompute_mode == 2);
         float dis = this->dis0;
         const float *tab = this->sim_table;
+        PQDecoder decoder(code, this->pq.nbits);
 
         for (size_t m = 0; m < this->pq.M; m++) {
-            dis += tab[*code++];
+            dis += tab[decoder.decode()];
             tab += this->pq.ksub;
         }
         return dis;
@@ -1115,7 +1122,22 @@ struct IVFPQScanner:
     }
 };
 
+template<class PQDecoder>
+InvertedListScanner *get_InvertedListScanner1 (const IndexIVFPQ &index,
+                                               bool store_pairs)
+{
 
+   if (index.metric_type == METRIC_INNER_PRODUCT) {
+        return new IVFPQScanner
+            <METRIC_INNER_PRODUCT, CMin<float, idx_t>, PQDecoder>
+            (index, store_pairs, 2);
+    } else if (index.metric_type == METRIC_L2) {
+        return new IVFPQScanner
+            <METRIC_L2, CMax<float, idx_t>, PQDecoder>
+            (index, store_pairs, 2);
+    }
+    return nullptr;
+}
 
 
 } // anonymous namespace
@@ -1123,12 +1145,13 @@ struct IVFPQScanner:
 InvertedListScanner *
 IndexIVFPQ::get_InvertedListScanner (bool store_pairs) const
 {
-    if (metric_type == METRIC_INNER_PRODUCT) {
-        return new IVFPQScanner<METRIC_INNER_PRODUCT, CMin<float, idx_t>, 2>
-            (*this, store_pairs);
-    } else if (metric_type == METRIC_L2) {
-        return new IVFPQScanner<METRIC_L2, CMax<float, idx_t>, 2>
-            (*this, store_pairs);
+
+    if (pq.nbits == 8) {
+        return get_InvertedListScanner1<PQDecoder8> (*this, store_pairs);
+    } else if (pq.nbits == 16) {
+        return get_InvertedListScanner1<PQDecoder16> (*this, store_pairs);
+    } else {
+        return get_InvertedListScanner1<PQDecoderGeneric> (*this, store_pairs);
     }
     return nullptr;
 

data/vendor/faiss/IndexIVFPQ.h

@@ -42,14 +42,14 @@ struct IndexIVFPQ: IndexIVF {
     int polysemous_ht;             ///< Hamming thresh for polysemous filtering
 
     /** Precompute table that speed up query preprocessing at some
-     * memory cost
+     * memory cost (used only for by_residual with L2 metric)
      * =-1: force disable
      * =0: decide heuristically (default: use tables only if they are
      *     < precomputed_tables_max_bytes)
      * =1: tables that work for all quantizers (size 256 * nlist * M)
      * =2: specific version for MultiIndexQuantizer (much more compact)
      */
-    int use_precomputed_table;     ///< if by_residual, build precompute tables
+    int use_precomputed_table;
     static size_t precomputed_table_max_bytes;
 
     /// if use_precompute_table
@@ -58,7 +58,7 @@ struct IndexIVFPQ: IndexIVF {
 
     IndexIVFPQ (
             Index * quantizer, size_t d, size_t nlist,
-            size_t M, size_t nbits_per_idx);
+            size_t M, size_t nbits_per_idx, MetricType metric = METRIC_L2);
 
     void add_with_ids(idx_t n, const float* x, const idx_t* xids = nullptr)
         override;
@@ -93,9 +93,9 @@ struct IndexIVFPQ: IndexIVF {
      * the duplicates are returned in pre-allocated arrays (see the
      * max sizes).
      *
-     * @params lims   limits between groups of duplicates
+     * @param lims   limits between groups of duplicates
      *                (max size ntotal / 2 + 1)
-     * @params ids    ids[lims[i]] : ids[lims[i+1]-1] is a group of
+     * @param ids    ids[lims[i]] : ids[lims[i+1]-1] is a group of
      *                duplicates (max size ntotal)
      * @return n      number of groups found
      */
@@ -135,15 +135,14 @@ struct IndexIVFPQ: IndexIVF {
 /// statistics are robust to internal threading, but not if
 /// IndexIVFPQ::search_preassigned is called by multiple threads
 struct IndexIVFPQStats {
-    size_t nrefine;  // nb of refines (IVFPQR)
+    size_t nrefine;  ///< nb of refines (IVFPQR)
 
     size_t n_hamming_pass;
-    // nb of passed Hamming distance tests (for polysemous)
+    ///< nb of passed Hamming distance tests (for polysemous)
 
-    // timings measured with the CPU RTC
-    // on all threads
+    // timings measured with the CPU RTC on all threads
     size_t search_cycles;
-    size_t refine_cycles; // only for IVFPQR
+    size_t refine_cycles; ///< only for IVFPQR
 
     IndexIVFPQStats () {reset (); }
     void reset ();