faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/IndexHNSW.h

@@ -21,63 +21,30 @@ 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 */
 
 struct IndexHNSW : Index {
     typedef HNSW::storage_idx_t storage_idx_t;
 
-    // the link strcuture
+    // the link structure
     HNSW hnsw;
 
     // the sequential storage
-    bool own_fields;
-    Index* storage;
+    bool own_fields = false;
+    Index* storage = nullptr;
 
-    ReconstructFromNeighbors* reconstruct_from_neighbors;
+    // When set to false, level 0 in the knn graph is not initialized.
+    // This option is used by GpuIndexCagra::copyTo(IndexHNSWCagra*)
+    // as level 0 knn graph is copied over from the index built by
+    // GpuIndexCagra.
+    bool init_level0 = true;
+
+    // When set to true, all neighbors in level 0 are filled up
+    // to the maximum size allowed (2 * M). This option is used by
+    // IndexHHNSWCagra to create a full base layer graph that is
+    // used when GpuIndexCagra::copyFrom(IndexHNSWCagra*) is invoked.
+    bool keep_max_size_level0 = false;
 
     explicit IndexHNSW(int d = 0, int M = 32, MetricType metric = METRIC_L2);
     explicit IndexHNSW(Index* storage, int M = 32);
@@ -98,6 +65,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;
@@ -119,7 +93,8 @@ struct IndexHNSW : Index {
             float* distances,
             idx_t* labels,
             int nprobe = 1,
-            int search_type = 1) const;
+            int search_type = 1,
+            const SearchParameters* params = nullptr) const;
 
     /// alternative graph building
     void init_level_0_from_knngraph(int k, const float* D, const idx_t* I);
@@ -134,6 +109,10 @@ struct IndexHNSW : Index {
     void reorder_links();
 
     void link_singletons();
+
+    void permute_entries(const idx_t* perm);
+
+    DistanceComputer* get_distance_computer() const override;
 };
 
 /** Flat index topped with with a HNSW structure to access elements
@@ -150,7 +129,12 @@ 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,
+            MetricType metric = METRIC_L2);
     void train(idx_t n, const float* x) override;
 };
 
@@ -184,4 +168,33 @@ struct IndexHNSW2Level : IndexHNSW {
             const SearchParameters* params = nullptr) const override;
 };
 
+struct IndexHNSWCagra : IndexHNSW {
+    IndexHNSWCagra();
+    IndexHNSWCagra(int d, int M, MetricType metric = METRIC_L2);
+
+    /// When set to true, the index is immutable.
+    /// This option is used to copy the knn graph from GpuIndexCagra
+    /// to the base level of IndexHNSWCagra without adding upper levels.
+    /// Doing so enables to search the HNSW index, but removes the
+    /// ability to add vectors.
+    bool base_level_only = false;
+
+    /// When `base_level_only` is set to `True`, the search function
+    /// searches only the base level knn graph of the HNSW index.
+    /// This parameter selects the entry point by randomly selecting
+    /// some points and using the best one.
+    int num_base_level_search_entrypoints = 32;
+
+    void add(idx_t n, const float* x) override;
+
+    /// entry point for search
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+};
+
 } // namespace faiss

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