faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -202,7 +202,10 @@ void hnsw_add_vertices(
                         verbose && omp_get_thread_num() == 0 ? 0 : -1;
                 size_t counter = 0;
 
-#pragma omp for schedule(dynamic)
+                // here we should do schedule(dynamic) but this segfaults for
+                // some versions of LLVM. The performance impact should not be
+                // too large when (i1 - i0) / num_threads >> 1
+#pragma omp for schedule(static)
                 for (int i = i0; i < i1; i++) {
                     storage_idx_t pt_id = order[i];
                     dis->set_query(x + (pt_id - n0) * d);
@@ -219,7 +222,6 @@ void hnsw_add_vertices(
                         printf("  %d / %d\r", i - i0, i1 - i0);
                         fflush(stdout);
                     }
-
                     if (counter % check_period == 0) {
                         if (InterruptCallback::is_interrupted()) {
                             interrupt = true;
@@ -284,18 +286,24 @@ void IndexHNSW::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);
-
     FAISS_THROW_IF_NOT_MSG(
             storage,
             "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
+    const SearchParametersHNSW* params = nullptr;
+
+    int efSearch = hnsw.efSearch;
+    if (params_in) {
+        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "params type invalid");
+        efSearch = params->efSearch;
+    }
     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);
+    idx_t check_period =
+            InterruptCallback::get_period_hint(hnsw.max_level * d * efSearch);
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
@@ -314,7 +322,7 @@ void IndexHNSW::search(
                 dis->set_query(x + i * d);
 
                 maxheap_heapify(k, simi, idxi);
-                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt);
+                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt, params);
                 n1 += stats.n1;
                 n2 += stats.n2;
                 n3 += stats.n3;
@@ -423,16 +431,15 @@ void IndexHNSW::search_level_0(
     FAISS_THROW_IF_NOT(nprobe > 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);
-        ScopeDeleter1<DistanceComputer> del(qdis);
-
+        std::unique_ptr<DistanceComputer> qdis(
+                storage_distance_computer(storage));
+        HNSWStats search_stats;
         VisitedTable vt(ntotal);
 
-#pragma omp for reduction(+ : n1, n2, n3, ndis, nreorder)
+#pragma omp for
         for (idx_t i = 0; i < n; i++) {
             idx_t* idxi = labels + i * k;
             float* simi = distances + i * k;
@@ -440,69 +447,24 @@ void IndexHNSW::search_level_0(
             qdis->set_query(x + i * d);
             maxheap_heapify(k, simi, idxi);
 
-            if (search_type == 1) {
-                int nres = 0;
-
-                for (int j = 0; j < nprobe; j++) {
-                    storage_idx_t cj = nearest[i * nprobe + j];
-
-                    if (cj < 0)
-                        break;
-
-                    if (vt.get(cj))
-                        continue;
-
-                    int candidates_size = std::max(hnsw.efSearch, int(k));
-                    MinimaxHeap candidates(candidates_size);
-
-                    candidates.push(cj, nearest_d[i * nprobe + j]);
-
-                    HNSWStats search_stats;
-                    nres = hnsw.search_from_candidates(
-                            *qdis,
-                            k,
-                            idxi,
-                            simi,
-                            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) {
-                int candidates_size = std::max(hnsw.efSearch, int(k));
-                candidates_size = std::max(candidates_size, nprobe);
-
-                MinimaxHeap candidates(candidates_size);
-                for (int j = 0; j < nprobe; j++) {
-                    storage_idx_t cj = nearest[i * nprobe + j];
-
-                    if (cj < 0)
-                        break;
-                    candidates.push(cj, nearest_d[i * nprobe + j]);
-                }
+            hnsw.search_level_0(
+                    *qdis.get(),
+                    k,
+                    idxi,
+                    simi,
+                    nprobe,
+                    nearest + i * nprobe,
+                    nearest_d + i * nprobe,
+                    search_type,
+                    search_stats,
+                    vt);
 
-                HNSWStats search_stats;
-                hnsw.search_from_candidates(
-                        *qdis, k, idxi, simi, 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();
-
             maxheap_reorder(k, simi, idxi);
         }
+#pragma omp critical
+        { hnsw_stats.combine(search_stats); }
     }
-
-    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
 }
 
 void IndexHNSW::init_level_0_from_knngraph(
@@ -1035,8 +997,11 @@ void IndexHNSW2Level::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
 
     if (dynamic_cast<const Index2Layer*>(storage)) {
         IndexHNSW::search(n, x, k, distances, labels);
@@ -1095,74 +1060,37 @@ void IndexHNSW2Level::search(
                 }
 
                 candidates.clear();
-                // copy the upper_beam elements to candidates list
-
-                int search_policy = 2;
-
-                if (search_policy == 1) {
-                    for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
-                        if (idxi[j] < 0)
-                            break;
-                        candidates.push(idxi[j], simi[j]);
-                        // search_from_candidates adds them back
-                        idxi[j] = -1;
-                        simi[j] = HUGE_VAL;
-                    }
 
-                    // 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,
-                            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();
-
-                } else if (search_policy == 2) {
-                    for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
-                        if (idxi[j] < 0)
-                            break;
-                        candidates.push(idxi[j], simi[j]);
-                    }
-
-                    // 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,
-                            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();
+                for (int j = 0; j < hnsw.upper_beam && j < k; j++) {
+                    if (idxi[j] < 0)
+                        break;
+                    candidates.push(idxi[j], simi[j]);
                 }
 
+                // 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,
+                        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();
+
                 maxheap_reorder(k, simi, idxi);
             }
         }

data/vendor/faiss/faiss/IndexHNSW.h

@@ -96,7 +96,8 @@ struct IndexHNSW : Index {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
 
@@ -180,7 +181,8 @@ struct IndexHNSW2Level : IndexHNSW {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 };
 
 } // namespace faiss

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