faiss 0.2.7 → 0.3.1

data/vendor/faiss/faiss/IndexHNSW.h

@@ -21,49 +21,6 @@ namespace faiss {
 
 struct IndexHNSW;
 
-struct ReconstructFromNeighbors {
-    typedef HNSW::storage_idx_t storage_idx_t;
-
-    const IndexHNSW& index;
-    size_t M;   // number of neighbors
-    size_t k;   // number of codebook entries
-    size_t nsq; // number of subvectors
-    size_t code_size;
-    int k_reorder; // nb to reorder. -1 = all
-
-    std::vector<float> codebook; // size nsq * k * (M + 1)
-
-    std::vector<uint8_t> codes; // size ntotal * code_size
-    size_t ntotal;
-    size_t d, dsub; // derived values
-
-    explicit ReconstructFromNeighbors(
-            const IndexHNSW& index,
-            size_t k = 256,
-            size_t nsq = 1);
-
-    /// codes must be added in the correct order and the IndexHNSW
-    /// must be populated and sorted
-    void add_codes(size_t n, const float* x);
-
-    size_t compute_distances(
-            size_t n,
-            const idx_t* shortlist,
-            const float* query,
-            float* distances) const;
-
-    /// called by add_codes
-    void estimate_code(const float* x, storage_idx_t i, uint8_t* code) const;
-
-    /// called by compute_distances
-    void reconstruct(storage_idx_t i, float* x, float* tmp) const;
-
-    void reconstruct_n(storage_idx_t n0, storage_idx_t ni, float* x) const;
-
-    /// get the M+1 -by-d table for neighbor coordinates for vector i
-    void get_neighbor_table(storage_idx_t i, float* out) const;
-};
-
 /** The HNSW index is a normal random-access index with a HNSW
  * link structure built on top */
 
@@ -74,10 +31,8 @@ struct IndexHNSW : Index {
     HNSW hnsw;
 
     // the sequential storage
-    bool own_fields;
-    Index* storage;
-
-    ReconstructFromNeighbors* reconstruct_from_neighbors;
+    bool own_fields = false;
+    Index* storage = nullptr;
 
     explicit IndexHNSW(int d = 0, int M = 32, MetricType metric = METRIC_L2);
     explicit IndexHNSW(Index* storage, int M = 32);
@@ -98,6 +53,13 @@ struct IndexHNSW : Index {
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
     void reconstruct(idx_t key, float* recons) const override;
 
     void reset() override;
@@ -134,6 +96,8 @@ struct IndexHNSW : Index {
     void reorder_links();
 
     void link_singletons();
+
+    void permute_entries(const idx_t* perm);
 };
 
 /** Flat index topped with with a HNSW structure to access elements
@@ -150,7 +114,7 @@ struct IndexHNSWFlat : IndexHNSW {
  */
 struct IndexHNSWPQ : IndexHNSW {
     IndexHNSWPQ();
-    IndexHNSWPQ(int d, int pq_m, int M);
+    IndexHNSWPQ(int d, int pq_m, int M, int pq_nbits = 8);
     void train(idx_t n, const float* x) override;
 };
 

data/vendor/faiss/faiss/IndexIDMap.cpp

@@ -9,31 +9,43 @@
 
 #include <faiss/IndexIDMap.h>
 
-#include <stdint.h>
 #include <cinttypes>
+#include <cstdint>
 #include <cstdio>
 #include <limits>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/IDSelector.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/WorkerThread.h>
 
 namespace faiss {
 
+namespace {
+
+// IndexBinary needs to update the code_size when d is set...
+
+void sync_d(Index* index) {}
+
+void sync_d(IndexBinary* index) {
+    FAISS_THROW_IF_NOT(index->d % 8 == 0);
+    index->code_size = index->d / 8;
+}
+
+} // anonymous namespace
+
 /*****************************************************
  * IndexIDMap implementation
  *******************************************************/
 
 template <typename IndexT>
-IndexIDMapTemplate<IndexT>::IndexIDMapTemplate(IndexT* index)
-        : index(index), own_fields(false) {
+IndexIDMapTemplate<IndexT>::IndexIDMapTemplate(IndexT* index) : index(index) {
     FAISS_THROW_IF_NOT_MSG(index->ntotal == 0, "index must be empty on input");
     this->is_trained = index->is_trained;
     this->metric_type = index->metric_type;
     this->verbose = index->verbose;
     this->d = index->d;
+    sync_d(this);
 }
 
 template <typename IndexT>
@@ -71,6 +83,27 @@ void IndexIDMapTemplate<IndexT>::add_with_ids(
     this->ntotal = index->ntotal;
 }
 
+namespace {
+
+/// RAII object to reset the IDSelector in the params object
+struct ScopedSelChange {
+    SearchParameters* params = nullptr;
+    IDSelector* old_sel = nullptr;
+
+    void set(SearchParameters* params_2, IDSelector* new_sel) {
+        this->params = params_2;
+        old_sel = params_2->sel;
+        params_2->sel = new_sel;
+    }
+    ~ScopedSelChange() {
+        if (params) {
+            params->sel = old_sel;
+        }
+    }
+};
+
+} // namespace
+
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::search(
         idx_t n,
@@ -79,9 +112,26 @@ void IndexIDMapTemplate<IndexT>::search(
         typename IndexT::distance_t* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
-    index->search(n, x, k, distances, labels);
+    IDSelectorTranslated this_idtrans(this->id_map, nullptr);
+    ScopedSelChange sel_change;
+
+    if (params && params->sel) {
+        auto idtrans = dynamic_cast<const IDSelectorTranslated*>(params->sel);
+
+        if (!idtrans) {
+            /*
+            FAISS_THROW_IF_NOT_MSG(
+                    idtrans,
+                    "IndexIDMap requires an IDSelectorTranslated on input");
+            */
+            // then make an idtrans and force it into the SearchParameters
+            // (hence the const_cast)
+            auto params_non_const = const_cast<SearchParameters*>(params);
+            this_idtrans.sel = params->sel;
+            sel_change.set(params_non_const, &this_idtrans);
+        }
+    }
+    index->search(n, x, k, distances, labels, params);
     idx_t* li = labels;
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; i++) {
@@ -96,9 +146,16 @@ void IndexIDMapTemplate<IndexT>::range_search(
         typename IndexT::distance_t radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
-    index->range_search(n, x, radius, result);
+    if (params) {
+        SearchParameters internal_search_parameters;
+        IDSelectorTranslated id_selector_translated(id_map, params->sel);
+        internal_search_parameters.sel = &id_selector_translated;
+
+        index->range_search(n, x, radius, result, &internal_search_parameters);
+    } else {
+        index->range_search(n, x, radius, result);
+    }
+
 #pragma omp parallel for
     for (idx_t i = 0; i < result->lims[result->nq]; i++) {
         result->labels[i] = result->labels[i] < 0 ? result->labels[i]
@@ -106,26 +163,10 @@ void IndexIDMapTemplate<IndexT>::range_search(
     }
 }
 
-namespace {
-
-struct IDTranslatedSelector : IDSelector {
-    const std::vector<int64_t>& id_map;
-    const IDSelector& sel;
-    IDTranslatedSelector(
-            const std::vector<int64_t>& id_map,
-            const IDSelector& sel)
-            : id_map(id_map), sel(sel) {}
-    bool is_member(idx_t id) const override {
-        return sel.is_member(id_map[id]);
-    }
-};
-
-} // namespace
-
 template <typename IndexT>
 size_t IndexIDMapTemplate<IndexT>::remove_ids(const IDSelector& sel) {
     // remove in sub-index first
-    IDTranslatedSelector sel2(id_map, sel);
+    IDSelectorTranslated sel2(id_map, &sel);
     size_t nremove = index->remove_ids(sel2);
 
     int64_t j = 0;
@@ -232,7 +273,7 @@ void IndexIDMap2Template<IndexT>::reconstruct(
         typename IndexT::component_t* recons) const {
     try {
         this->index->reconstruct(rev_map.at(key), recons);
-    } catch (const std::out_of_range& e) {
+    } catch (const std::out_of_range&) {
         FAISS_THROW_FMT("key %" PRId64 " not found", key);
     }
 }

data/vendor/faiss/faiss/IndexIDMap.h

@@ -9,6 +9,7 @@
 
 #include <faiss/Index.h>
 #include <faiss/IndexBinary.h>
+#include <faiss/impl/IDSelector.h>
 
 #include <unordered_map>
 #include <vector>
@@ -21,8 +22,8 @@ struct IndexIDMapTemplate : IndexT {
     using component_t = typename IndexT::component_t;
     using distance_t = typename IndexT::distance_t;
 
-    IndexT* index;   ///! the sub-index
-    bool own_fields; ///! whether pointers are deleted in destructo
+    IndexT* index = nullptr; ///! the sub-index
+    bool own_fields = false; ///! whether pointers are deleted in destructo
     std::vector<idx_t> id_map;
 
     explicit IndexIDMapTemplate(IndexT* index);
@@ -102,4 +103,25 @@ struct IndexIDMap2Template : IndexIDMapTemplate<IndexT> {
 using IndexIDMap2 = IndexIDMap2Template<Index>;
 using IndexBinaryIDMap2 = IndexIDMap2Template<IndexBinary>;
 
+// IDSelector that translates the ids using an IDMap
+struct IDSelectorTranslated : IDSelector {
+    const std::vector<int64_t>& id_map;
+    const IDSelector* sel;
+
+    IDSelectorTranslated(
+            const std::vector<int64_t>& id_map,
+            const IDSelector* sel)
+            : id_map(id_map), sel(sel) {}
+
+    IDSelectorTranslated(IndexBinaryIDMap& index_idmap, const IDSelector* sel)
+            : id_map(index_idmap.id_map), sel(sel) {}
+
+    IDSelectorTranslated(IndexIDMap& index_idmap, const IDSelector* sel)
+            : id_map(index_idmap.id_map), sel(sel) {}
+
+    bool is_member(idx_t id) const override {
+        return sel->is_member(id_map[id]);
+    }
+};
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -11,6 +11,7 @@
 
 #include <omp.h>
 #include <cstdint>
+#include <memory>
 #include <mutex>
 
 #include <algorithm>
@@ -45,7 +46,7 @@ Level1Quantizer::Level1Quantizer(Index* quantizer, size_t nlist)
     cp.niter = 10;
 }
 
-Level1Quantizer::Level1Quantizer() {}
+Level1Quantizer::Level1Quantizer() = default;
 
 Level1Quantizer::~Level1Quantizer() {
     if (own_fields) {
@@ -172,7 +173,7 @@ IndexIVF::IndexIVF(
     }
 }
 
-IndexIVF::IndexIVF() {}
+IndexIVF::IndexIVF() = default;
 
 void IndexIVF::add(idx_t n, const float* x) {
     add_with_ids(n, x, nullptr);
@@ -202,7 +203,8 @@ void IndexIVF::add_core(
         idx_t n,
         const float* x,
         const idx_t* xids,
-        const idx_t* coarse_idx) {
+        const idx_t* coarse_idx,
+        void* inverted_list_context) {
     // do some blocking to avoid excessive allocs
     idx_t bs = 65536;
     if (n > bs) {
@@ -217,7 +219,8 @@ void IndexIVF::add_core(
                     i1 - i0,
                     x + i0 * d,
                     xids ? xids + i0 : nullptr,
-                    coarse_idx + i0);
+                    coarse_idx + i0,
+                    inverted_list_context);
         }
         return;
     }
@@ -248,7 +251,10 @@ void IndexIVF::add_core(
             if (list_no >= 0 && list_no % nt == rank) {
                 idx_t id = xids ? xids[i] : ntotal + i;
                 size_t ofs = invlists->add_entry(
-                        list_no, id, flat_codes.get() + i * code_size);
+                        list_no,
+                        id,
+                        flat_codes.get() + i * code_size,
+                        inverted_list_context);
 
                 dm_adder.add(i, list_no, ofs);
 
@@ -375,7 +381,7 @@ void IndexIVF::search(
             indexIVF_stats.add(stats[slice]);
         }
     } else {
-        // handle paralellization at level below (or don't run in parallel at
+        // handle parallelization at level below (or don't run in parallel at
         // all)
         sub_search_func(n, x, distances, labels, &indexIVF_stats);
     }
@@ -444,11 +450,13 @@ void IndexIVF::search_preassigned(
                      : pmode == 1 ? nprobe > 1
                                   : nprobe * n > 1);
 
+    void* inverted_list_context =
+            params ? params->inverted_list_context : nullptr;
+
 #pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis, nheap)
     {
-        InvertedListScanner* scanner =
-                get_InvertedListScanner(store_pairs, sel);
-        ScopeDeleter1<InvertedListScanner> del(scanner);
+        std::unique_ptr<InvertedListScanner> scanner(
+                get_InvertedListScanner(store_pairs, sel));
 
         /*****************************************************
          * Depending on parallel_mode, there are two possible ways
@@ -507,7 +515,7 @@ void IndexIVF::search_preassigned(
                     nlist);
 
             // don't waste time on empty lists
-            if (invlists->is_empty(key)) {
+            if (invlists->is_empty(key, inverted_list_context)) {
                 return (size_t)0;
             }
 
@@ -520,7 +528,7 @@ void IndexIVF::search_preassigned(
                     size_t list_size = 0;
 
                     std::unique_ptr<InvertedListsIterator> it(
-                            invlists->get_iterator(key));
+                            invlists->get_iterator(key, inverted_list_context));
 
                     nheap += scanner->iterate_codes(
                             it.get(), simi, idxi, k, list_size);
@@ -539,7 +547,8 @@ void IndexIVF::search_preassigned(
                     const idx_t* ids = nullptr;
 
                     if (!store_pairs) {
-                        sids.reset(new InvertedLists::ScopedIds(invlists, key));
+                        sids = std::make_unique<InvertedLists::ScopedIds>(
+                                invlists, key);
                         ids = sids->get();
                     }
 
@@ -659,7 +668,6 @@ void IndexIVF::search_preassigned(
 #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());
@@ -696,12 +704,13 @@ void IndexIVF::search_preassigned(
         }
     }
 
-    if (ivf_stats) {
-        ivf_stats->nq += n;
-        ivf_stats->nlist += nlistv;
-        ivf_stats->ndis += ndis;
-        ivf_stats->nheap_updates += nheap;
+    if (ivf_stats == nullptr) {
+        ivf_stats = &indexIVF_stats;
     }
+    ivf_stats->nq += n;
+    ivf_stats->nlist += nlistv;
+    ivf_stats->ndis += ndis;
+    ivf_stats->nheap_updates += nheap;
 }
 
 void IndexIVF::range_search(
@@ -781,6 +790,9 @@ void IndexIVF::range_search_preassigned(
                      : pmode == 1 ? nprobe > 1
                                   : nprobe * nx > 1);
 
+    void* inverted_list_context =
+            params ? params->inverted_list_context : nullptr;
+
 #pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis)
     {
         RangeSearchPartialResult pres(result);
@@ -802,7 +814,7 @@ void IndexIVF::range_search_preassigned(
                     ik,
                     nlist);
 
-            if (invlists->is_empty(key)) {
+            if (invlists->is_empty(key, inverted_list_context)) {
                 return;
             }
 
@@ -811,7 +823,7 @@ void IndexIVF::range_search_preassigned(
                 scanner->set_list(key, coarse_dis[i * nprobe + ik]);
                 if (invlists->use_iterator) {
                     std::unique_ptr<InvertedListsIterator> it(
-                            invlists->get_iterator(key));
+                            invlists->get_iterator(key, inverted_list_context));
 
                     scanner->iterate_codes_range(
                             it.get(), radius, qres, list_size);
@@ -891,17 +903,18 @@ void IndexIVF::range_search_preassigned(
         }
     }
 
-    if (stats) {
-        stats->nq += nx;
-        stats->nlist += nlistv;
-        stats->ndis += ndis;
+    if (stats == nullptr) {
+        stats = &indexIVF_stats;
     }
+    stats->nq += nx;
+    stats->nlist += nlistv;
+    stats->ndis += ndis;
 }
 
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
         bool /*store_pairs*/,
         const IDSelector* /* sel */) const {
-    return nullptr;
+    FAISS_THROW_MSG("get_InvertedListScanner not implemented");
 }
 
 void IndexIVF::reconstruct(idx_t key, float* recons) const {
@@ -973,14 +986,12 @@ void IndexIVF::search_and_reconstruct(
             std::min(nlist, params ? params->nprobe : this->nprobe);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
-    idx_t* idx = new idx_t[n * nprobe];
-    ScopeDeleter<idx_t> del(idx);
-    float* coarse_dis = new float[n * nprobe];
-    ScopeDeleter<float> del2(coarse_dis);
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
 
-    quantizer->search(n, x, nprobe, coarse_dis, idx);
+    quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
 
-    invlists->prefetch_lists(idx, n * nprobe);
+    invlists->prefetch_lists(idx.get(), n * nprobe);
 
     // search_preassigned() with `store_pairs` enabled to obtain the list_no
     // and offset into `codes` for reconstruction
@@ -988,29 +999,94 @@ void IndexIVF::search_and_reconstruct(
             n,
             x,
             k,
-            idx,
-            coarse_dis,
+            idx.get(),
+            coarse_dis.get(),
             distances,
             labels,
             true /* store_pairs */,
             params);
-    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];
-            float* reconstructed = recons + ij * d;
-            if (key < 0) {
-                // Fill with NaNs
-                memset(reconstructed, -1, sizeof(*reconstructed) * d);
-            } else {
-                int list_no = lo_listno(key);
-                int offset = lo_offset(key);
+#pragma omp parallel for if (n * k > 1000)
+    for (idx_t ij = 0; ij < n * k; ij++) {
+        idx_t key = labels[ij];
+        float* reconstructed = recons + ij * d;
+        if (key < 0) {
+            // Fill with NaNs
+            memset(reconstructed, -1, sizeof(*reconstructed) * d);
+        } else {
+            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);
+
+            reconstruct_from_offset(list_no, offset, reconstructed);
+        }
+    }
+}
+
+void IndexIVF::search_and_return_codes(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        uint8_t* codes,
+        bool include_listno,
+        const SearchParameters* params_in) const {
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
+    FAISS_THROW_IF_NOT(nprobe > 0);
+
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[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,
+            true /* store_pairs */,
+            params);
+
+    size_t code_size_1 = code_size;
+    if (include_listno) {
+        code_size_1 += coarse_code_size();
+    }
 
-                // Update label to the actual id
-                labels[ij] = invlists->get_single_id(list_no, offset);
+#pragma omp parallel for if (n * k > 1000)
+    for (idx_t ij = 0; ij < n * k; ij++) {
+        idx_t key = labels[ij];
+        uint8_t* code1 = codes + ij * code_size_1;
 
-                reconstruct_from_offset(list_no, offset, reconstructed);
+        if (key < 0) {
+            // Fill with 0xff
+            memset(code1, -1, code_size_1);
+        } else {
+            int list_no = lo_listno(key);
+            int offset = lo_offset(key);
+            const uint8_t* cc = invlists->get_single_code(list_no, offset);
+
+            labels[ij] = invlists->get_single_id(list_no, offset);
+
+            if (include_listno) {
+                encode_listno(list_no, code1);
+                code1 += code_size_1 - code_size;
             }
+            memcpy(code1, cc, code_size);
         }
     }
 }
@@ -1061,22 +1137,52 @@ void IndexIVF::update_vectors(int n, const idx_t* new_ids, const float* x) {
 }
 
 void IndexIVF::train(idx_t n, const float* x) {
-    if (verbose)
+    if (verbose) {
         printf("Training level-1 quantizer\n");
+    }
 
     train_q1(n, x, verbose, metric_type);
 
-    if (verbose)
+    if (verbose) {
         printf("Training IVF residual\n");
+    }
+
+    // optional subsampling
+    idx_t max_nt = train_encoder_num_vectors();
+    if (max_nt <= 0) {
+        max_nt = (size_t)1 << 35;
+    }
+
+    TransformedVectors tv(
+            x, fvecs_maybe_subsample(d, (size_t*)&n, max_nt, x, verbose));
+
+    if (by_residual) {
+        std::vector<idx_t> assign(n);
+        quantizer->assign(n, tv.x, assign.data());
+
+        std::vector<float> residuals(n * d);
+        quantizer->compute_residual_n(n, tv.x, residuals.data(), assign.data());
+
+        train_encoder(n, residuals.data(), assign.data());
+    } else {
+        train_encoder(n, tv.x, nullptr);
+    }
 
-    train_residual(n, x);
     is_trained = true;
 }
 
-void IndexIVF::train_residual(idx_t /*n*/, const float* /*x*/) {
-    if (verbose)
-        printf("IndexIVF: no residual training\n");
+idx_t IndexIVF::train_encoder_num_vectors() const {
+    return 0;
+}
+
+void IndexIVF::train_encoder(
+        idx_t /*n*/,
+        const float* /*x*/,
+        const idx_t* assign) {
     // does nothing by default
+    if (verbose) {
+        printf("IndexIVF: no residual training\n");
+    }
 }
 
 bool check_compatible_for_merge_expensive_check = true;