faiss 0.2.4 → 0.2.6

data/vendor/faiss/faiss/IndexIDMap.cpp

@@ -0,0 +1,247 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IndexIDMap.h>
+
+#include <stdint.h>
+#include <cinttypes>
+#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 {
+
+/*****************************************************
+ * IndexIDMap implementation
+ *******************************************************/
+
+template <typename IndexT>
+IndexIDMapTemplate<IndexT>::IndexIDMapTemplate(IndexT* index)
+        : index(index), own_fields(false) {
+    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;
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::add(
+        idx_t,
+        const typename IndexT::component_t*) {
+    FAISS_THROW_MSG(
+            "add does not make sense with IndexIDMap, "
+            "use add_with_ids");
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::train(
+        idx_t n,
+        const typename IndexT::component_t* x) {
+    index->train(n, x);
+    this->is_trained = index->is_trained;
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::reset() {
+    index->reset();
+    id_map.clear();
+    this->ntotal = 0;
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::add_with_ids(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        const typename IndexT::idx_t* xids) {
+    index->add(n, x);
+    for (idx_t i = 0; i < n; i++)
+        id_map.push_back(xids[i]);
+    this->ntotal = index->ntotal;
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::search(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        idx_t k,
+        typename IndexT::distance_t* distances,
+        typename IndexT::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);
+    idx_t* li = labels;
+#pragma omp parallel for
+    for (idx_t i = 0; i < n * k; i++) {
+        li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
+    }
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::range_search(
+        typename IndexT::idx_t n,
+        const typename IndexT::component_t* x,
+        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);
+#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]
+                                                  : id_map[result->labels[i]];
+    }
+}
+
+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);
+    size_t nremove = index->remove_ids(sel2);
+
+    int64_t j = 0;
+    for (idx_t i = 0; i < this->ntotal; i++) {
+        if (sel.is_member(id_map[i])) {
+            // remove
+        } else {
+            id_map[j] = id_map[i];
+            j++;
+        }
+    }
+    FAISS_ASSERT(j == index->ntotal);
+    this->ntotal = j;
+    id_map.resize(this->ntotal);
+    return nremove;
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::check_compatible_for_merge(
+        const IndexT& otherIndex) const {
+    auto other = dynamic_cast<const IndexIDMapTemplate<IndexT>*>(&otherIndex);
+    FAISS_THROW_IF_NOT(other);
+    index->check_compatible_for_merge(*other->index);
+}
+
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::merge_from(IndexT& otherIndex, idx_t add_id) {
+    check_compatible_for_merge(otherIndex);
+    auto other = static_cast<IndexIDMapTemplate<IndexT>*>(&otherIndex);
+    index->merge_from(*other->index);
+    for (size_t i = 0; i < other->id_map.size(); i++) {
+        id_map.push_back(other->id_map[i] + add_id);
+    }
+    other->id_map.resize(0);
+    this->ntotal = index->ntotal;
+    other->ntotal = 0;
+}
+
+template <typename IndexT>
+IndexIDMapTemplate<IndexT>::~IndexIDMapTemplate() {
+    if (own_fields)
+        delete index;
+}
+
+/*****************************************************
+ * IndexIDMap2 implementation
+ *******************************************************/
+
+template <typename IndexT>
+IndexIDMap2Template<IndexT>::IndexIDMap2Template(IndexT* index)
+        : IndexIDMapTemplate<IndexT>(index) {}
+
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::add_with_ids(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        const typename IndexT::idx_t* xids) {
+    size_t prev_ntotal = this->ntotal;
+    IndexIDMapTemplate<IndexT>::add_with_ids(n, x, xids);
+    for (size_t i = prev_ntotal; i < this->ntotal; i++) {
+        rev_map[this->id_map[i]] = i;
+    }
+}
+
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::check_consistency() const {
+    FAISS_THROW_IF_NOT(rev_map.size() == this->id_map.size());
+    FAISS_THROW_IF_NOT(this->id_map.size() == this->ntotal);
+    for (size_t i = 0; i < this->ntotal; i++) {
+        idx_t ii = rev_map.at(this->id_map[i]);
+        FAISS_THROW_IF_NOT(ii == i);
+    }
+}
+
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::merge_from(IndexT& otherIndex, idx_t add_id) {
+    size_t prev_ntotal = this->ntotal;
+    IndexIDMapTemplate<IndexT>::merge_from(otherIndex, add_id);
+    for (size_t i = prev_ntotal; i < this->ntotal; i++) {
+        rev_map[this->id_map[i]] = i;
+    }
+    static_cast<IndexIDMap2Template<IndexT>&>(otherIndex).rev_map.clear();
+}
+
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::construct_rev_map() {
+    rev_map.clear();
+    for (size_t i = 0; i < this->ntotal; i++) {
+        rev_map[this->id_map[i]] = i;
+    }
+}
+
+template <typename IndexT>
+size_t IndexIDMap2Template<IndexT>::remove_ids(const IDSelector& sel) {
+    // This is quite inefficient
+    size_t nremove = IndexIDMapTemplate<IndexT>::remove_ids(sel);
+    construct_rev_map();
+    return nremove;
+}
+
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::reconstruct(
+        idx_t key,
+        typename IndexT::component_t* recons) const {
+    try {
+        this->index->reconstruct(rev_map.at(key), recons);
+    } catch (const std::out_of_range& e) {
+        FAISS_THROW_FMT("key %" PRId64 " not found", key);
+    }
+}
+
+// explicit template instantiations
+
+template struct IndexIDMapTemplate<Index>;
+template struct IndexIDMapTemplate<IndexBinary>;
+template struct IndexIDMap2Template<Index>;
+template struct IndexIDMap2Template<IndexBinary>;
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIDMap.h

@@ -0,0 +1,107 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/Index.h>
+#include <faiss/IndexBinary.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace faiss {
+
+/** Index that translates search results to ids */
+template <typename IndexT>
+struct IndexIDMapTemplate : IndexT {
+    using idx_t = typename IndexT::idx_t;
+    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
+    std::vector<idx_t> id_map;
+
+    explicit IndexIDMapTemplate(IndexT* index);
+
+    /// @param xids if non-null, ids to store for the vectors (size n)
+    void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
+            override;
+
+    /// this will fail. Use add_with_ids
+    void add(idx_t n, const component_t* x) override;
+
+    void search(
+            idx_t n,
+            const component_t* x,
+            idx_t k,
+            distance_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void train(idx_t n, const component_t* x) override;
+
+    void reset() override;
+
+    /// remove ids adapted to IndexFlat
+    size_t remove_ids(const IDSelector& sel) override;
+
+    void range_search(
+            idx_t n,
+            const component_t* x,
+            distance_t radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
+    void merge_from(IndexT& otherIndex, idx_t add_id = 0) override;
+    void check_compatible_for_merge(const IndexT& otherIndex) const override;
+
+    ~IndexIDMapTemplate() override;
+    IndexIDMapTemplate() {
+        own_fields = false;
+        index = nullptr;
+    }
+};
+
+using IndexIDMap = IndexIDMapTemplate<Index>;
+using IndexBinaryIDMap = IndexIDMapTemplate<IndexBinary>;
+
+/** same as IndexIDMap but also provides an efficient reconstruction
+ *  implementation via a 2-way index */
+template <typename IndexT>
+struct IndexIDMap2Template : IndexIDMapTemplate<IndexT> {
+    using idx_t = typename IndexT::idx_t;
+    using component_t = typename IndexT::component_t;
+    using distance_t = typename IndexT::distance_t;
+
+    std::unordered_map<idx_t, idx_t> rev_map;
+
+    explicit IndexIDMap2Template(IndexT* index);
+
+    /// make the rev_map from scratch
+    void construct_rev_map();
+
+    void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
+            override;
+
+    size_t remove_ids(const IDSelector& sel) override;
+
+    void reconstruct(idx_t key, component_t* recons) const override;
+
+    /// check that the rev_map and the id_map are in sync
+    void check_consistency() const;
+
+    void merge_from(IndexT& otherIndex, idx_t add_id = 0) override;
+
+    ~IndexIDMap2Template() override {}
+    IndexIDMap2Template() {}
+};
+
+using IndexIDMap2 = IndexIDMap2Template<Index>;
+using IndexBinaryIDMap2 = IndexIDMap2Template<IndexBinary>;
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -23,6 +23,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 
 namespace faiss {
 
@@ -303,14 +304,20 @@ void IndexIVF::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
-
-    const size_t nprobe = std::min(nlist, this->nprobe);
+    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);
 
     // search function for a subset of queries
-    auto sub_search_func = [this, k, nprobe](
+    auto sub_search_func = [this, k, nprobe, params](
                                    idx_t n,
                                    const float* x,
                                    float* distances,
@@ -320,7 +327,13 @@ void IndexIVF::search(
         std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
 
         double t0 = getmillisecs();
-        quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
+        quantizer->search(
+                n,
+                x,
+                nprobe,
+                coarse_dis.get(),
+                idx.get(),
+                params ? params->quantizer_params : nullptr);
 
         double t1 = getmillisecs();
         invlists->prefetch_lists(idx.get(), n * nprobe);
@@ -334,7 +347,7 @@ void IndexIVF::search(
                 distances,
                 labels,
                 false,
-                nullptr,
+                params,
                 ivf_stats);
         double t2 = getmillisecs();
         ivf_stats->quantization_time += t1 - t0;
@@ -400,6 +413,19 @@ void IndexIVF::search_preassigned(
     FAISS_THROW_IF_NOT(nprobe > 0);
 
     idx_t max_codes = params ? params->max_codes : this->max_codes;
+    IDSelector* sel = params ? params->sel : nullptr;
+    const IDSelectorRange* selr = dynamic_cast<const IDSelectorRange*>(sel);
+    if (selr) {
+        if (selr->assume_sorted) {
+            sel = nullptr; // use special IDSelectorRange processing
+        } else {
+            selr = nullptr; // use generic processing
+        }
+    }
+
+    FAISS_THROW_IF_NOT_MSG(
+            !(sel && store_pairs),
+            "selector and store_pairs cannot be combined");
 
     size_t nlistv = 0, ndis = 0, nheap = 0;
 
@@ -421,7 +447,8 @@ void IndexIVF::search_preassigned(
 
 #pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis, nheap)
     {
-        InvertedListScanner* scanner = get_InvertedListScanner(store_pairs);
+        InvertedListScanner* scanner =
+                get_InvertedListScanner(store_pairs, sel);
         ScopeDeleter1<InvertedListScanner> del(scanner);
 
         /*****************************************************
@@ -492,6 +519,7 @@ void IndexIVF::search_preassigned(
 
             try {
                 InvertedLists::ScopedCodes scodes(invlists, key);
+                const uint8_t* codes = scodes.get();
 
                 std::unique_ptr<InvertedLists::ScopedIds> sids;
                 const Index::idx_t* ids = nullptr;
@@ -501,8 +529,20 @@ void IndexIVF::search_preassigned(
                     ids = sids->get();
                 }
 
+                if (selr) { // IDSelectorRange
+                    // restrict search to a section of the inverted list
+                    size_t jmin, jmax;
+                    selr->find_sorted_ids_bounds(list_size, ids, &jmin, &jmax);
+                    list_size = jmax - jmin;
+                    if (list_size == 0) {
+                        return (size_t)0;
+                    }
+                    codes += jmin * code_size;
+                    ids += jmin;
+                }
+
                 nheap += scanner->scan_codes(
-                        list_size, scodes.get(), ids, simi, idxi, k);
+                        list_size, codes, ids, simi, idxi, k);
 
             } catch (const std::exception& e) {
                 std::lock_guard<std::mutex> lock(exception_mutex);
@@ -651,13 +691,23 @@ void IndexIVF::range_search(
         idx_t nx,
         const float* x,
         float radius,
-        RangeSearchResult* result) const {
-    const size_t nprobe = std::min(nlist, this->nprobe);
+        RangeSearchResult* result,
+        const SearchParameters* params_in) const {
+    const IVFSearchParameters* params = nullptr;
+    const SearchParameters* quantizer_params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+        quantizer_params = params->quantizer_params;
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
     std::unique_ptr<idx_t[]> keys(new idx_t[nx * nprobe]);
     std::unique_ptr<float[]> coarse_dis(new float[nx * nprobe]);
 
     double t0 = getmillisecs();
-    quantizer->search(nx, x, nprobe, coarse_dis.get(), keys.get());
+    quantizer->search(
+            nx, x, nprobe, coarse_dis.get(), keys.get(), quantizer_params);
     indexIVF_stats.quantization_time += getmillisecs() - t0;
 
     t0 = getmillisecs();
@@ -671,7 +721,7 @@ void IndexIVF::range_search(
             coarse_dis.get(),
             result,
             false,
-            nullptr,
+            params,
             &indexIVF_stats);
 
     indexIVF_stats.search_time += getmillisecs() - t0;
@@ -689,7 +739,10 @@ void IndexIVF::range_search_preassigned(
         IndexIVFStats* stats) const {
     idx_t nprobe = params ? params->nprobe : this->nprobe;
     nprobe = std::min((idx_t)nlist, nprobe);
+    FAISS_THROW_IF_NOT(nprobe > 0);
+
     idx_t max_codes = params ? params->max_codes : this->max_codes;
+    IDSelector* sel = params ? params->sel : nullptr;
 
     size_t nlistv = 0, ndis = 0;
 
@@ -711,7 +764,7 @@ void IndexIVF::range_search_preassigned(
     {
         RangeSearchPartialResult pres(result);
         std::unique_ptr<InvertedListScanner> scanner(
-                get_InvertedListScanner(store_pairs));
+                get_InvertedListScanner(store_pairs, sel));
         FAISS_THROW_IF_NOT(scanner.get());
         all_pres[omp_get_thread_num()] = &pres;
 
@@ -774,7 +827,6 @@ void IndexIVF::range_search_preassigned(
                 }
             }
         } else if (parallel_mode == 2) {
-            std::vector<RangeQueryResult*> all_qres(nx);
             RangeQueryResult* qres = nullptr;
 
 #pragma omp for schedule(dynamic)
@@ -782,7 +834,6 @@ void IndexIVF::range_search_preassigned(
                 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);
                     scanner->set_query(x + i * d);
                 }
@@ -818,7 +869,8 @@ void IndexIVF::range_search_preassigned(
 }
 
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
-        bool /*store_pairs*/) const {
+        bool /*store_pairs*/,
+        const IDSelector* /* sel */) const {
     return nullptr;
 }
 
@@ -846,6 +898,21 @@ void IndexIVF::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
     }
 }
 
+bool IndexIVF::check_ids_sorted() const {
+    size_t nflip = 0;
+
+    for (size_t i = 0; i < nlist; i++) {
+        size_t list_size = invlists->list_size(i);
+        InvertedLists::ScopedIds ids(invlists, i);
+        for (size_t j = 0; j + 1 < list_size; j++) {
+            if (ids[j + 1] < ids[j]) {
+                nflip++;
+            }
+        }
+    }
+    return nflip == 0;
+}
+
 /* standalone codec interface */
 size_t IndexIVF::sa_code_size() const {
     size_t coarse_size = coarse_code_size();
@@ -865,10 +932,15 @@ void IndexIVF::search_and_reconstruct(
         idx_t k,
         float* distances,
         idx_t* labels,
-        float* recons) const {
-    FAISS_THROW_IF_NOT(k > 0);
-
-    const size_t nprobe = std::min(nlist, this->nprobe);
+        float* recons,
+        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);
 
     idx_t* idx = new idx_t[n * nprobe];
@@ -890,7 +962,8 @@ void IndexIVF::search_and_reconstruct(
             coarse_dis,
             distances,
             labels,
-            true /* store_pairs */);
+            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;
@@ -976,26 +1049,41 @@ void IndexIVF::train_residual(idx_t /*n*/, const float* /*x*/) {
     // does nothing by default
 }
 
-void IndexIVF::check_compatible_for_merge(const IndexIVF& other) const {
+bool check_compatible_for_merge_expensive_check = true;
+
+void IndexIVF::check_compatible_for_merge(const Index& otherIndex) const {
     // 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);
+    const IndexIVF* other = dynamic_cast<const IndexIVF*>(&otherIndex);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->d == d);
+    FAISS_THROW_IF_NOT(other->nlist == nlist);
+    FAISS_THROW_IF_NOT(quantizer->ntotal == other->quantizer->ntotal);
+    FAISS_THROW_IF_NOT(other->code_size == code_size);
     FAISS_THROW_IF_NOT_MSG(
-            typeid(*this) == typeid(other),
+            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(),
+            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);
+    if (check_compatible_for_merge_expensive_check) {
+        std::vector<float> v(d), v2(d);
+        for (size_t i = 0; i < nlist; i++) {
+            quantizer->reconstruct(i, v.data());
+            other->quantizer->reconstruct(i, v2.data());
+            FAISS_THROW_IF_NOT_MSG(
+                    v == v2, "coarse quantizers should be the same");
+        }
+    }
+}
 
-    invlists->merge_from(other.invlists, add_id);
+void IndexIVF::merge_from(Index& otherIndex, idx_t add_id) {
+    check_compatible_for_merge(otherIndex);
+    IndexIVF* other = static_cast<IndexIVF*>(&otherIndex);
+    invlists->merge_from(other->invlists, add_id);
 
-    ntotal += other.ntotal;
-    other.ntotal = 0;
+    ntotal += other->ntotal;
+    other->ntotal = 0;
 }
 
 void IndexIVF::replace_invlists(InvertedLists* il, bool own) {