faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/impl/HNSW.cpp

@@ -12,6 +12,8 @@
 #include <string>
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/IDSelector.h>
 
 namespace faiss {
 
@@ -501,9 +503,19 @@ void HNSW::add_with_locks(
     }
 }
 
+/**************************************************************
+ * Searching
+ **************************************************************/
+
+namespace {
+
+using idx_t = HNSW::idx_t;
+using MinimaxHeap = HNSW::MinimaxHeap;
+using Node = HNSW::Node;
 /** Do a BFS on the candidates list */
 
-int HNSW::search_from_candidates(
+int search_from_candidates(
+        const HNSW& hnsw,
         DistanceComputer& qdis,
         int k,
         idx_t* I,
@@ -512,22 +524,31 @@ int HNSW::search_from_candidates(
         VisitedTable& vt,
         HNSWStats& stats,
         int level,
-        int nres_in) const {
+        int nres_in = 0,
+        const SearchParametersHNSW* params = nullptr) {
     int nres = nres_in;
     int ndis = 0;
+
+    // can be overridden by search params
+    bool do_dis_check = params ? params->check_relative_distance
+                               : hnsw.check_relative_distance;
+    int efSearch = params ? params->efSearch : hnsw.efSearch;
+    const IDSelector* sel = params ? params->sel : nullptr;
+
     for (int i = 0; i < candidates.size(); i++) {
         idx_t v1 = candidates.ids[i];
         float d = candidates.dis[i];
         FAISS_ASSERT(v1 >= 0);
-        if (nres < k) {
-            faiss::maxheap_push(++nres, D, I, d, v1);
-        } else if (d < D[0]) {
-            faiss::maxheap_replace_top(nres, D, I, d, v1);
+        if (!sel || sel->is_member(v1)) {
+            if (nres < k) {
+                faiss::maxheap_push(++nres, D, I, d, v1);
+            } else if (d < D[0]) {
+                faiss::maxheap_replace_top(nres, D, I, d, v1);
+            }
         }
         vt.set(v1);
     }
 
-    bool do_dis_check = check_relative_distance;
     int nstep = 0;
 
     while (candidates.size() > 0) {
@@ -546,10 +567,10 @@ int HNSW::search_from_candidates(
         }
 
         size_t begin, end;
-        neighbor_range(v0, level, &begin, &end);
+        hnsw.neighbor_range(v0, level, &begin, &end);
 
         for (size_t j = begin; j < end; j++) {
-            int v1 = neighbors[j];
+            int v1 = hnsw.neighbors[j];
             if (v1 < 0)
                 break;
             if (vt.get(v1)) {
@@ -558,10 +579,12 @@ int HNSW::search_from_candidates(
             vt.set(v1);
             ndis++;
             float d = qdis(v1);
-            if (nres < k) {
-                faiss::maxheap_push(++nres, D, I, d, v1);
-            } else if (d < D[0]) {
-                faiss::maxheap_replace_top(nres, D, I, d, v1);
+            if (!sel || sel->is_member(v1)) {
+                if (nres < k) {
+                    faiss::maxheap_push(++nres, D, I, d, v1);
+                } else if (d < D[0]) {
+                    faiss::maxheap_replace_top(nres, D, I, d, v1);
+                }
             }
             candidates.push(v1, d);
         }
@@ -583,16 +606,13 @@ int HNSW::search_from_candidates(
     return nres;
 }
 
-/**************************************************************
- * Searching
- **************************************************************/
-
-std::priority_queue<HNSW::Node> HNSW::search_from_candidate_unbounded(
+std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
+        const HNSW& hnsw,
         const Node& node,
         DistanceComputer& qdis,
         int ef,
         VisitedTable* vt,
-        HNSWStats& stats) const {
+        HNSWStats& stats) {
     int ndis = 0;
     std::priority_queue<Node> top_candidates;
     std::priority_queue<Node, std::vector<Node>, std::greater<Node>> candidates;
@@ -614,10 +634,10 @@ std::priority_queue<HNSW::Node> HNSW::search_from_candidate_unbounded(
         candidates.pop();
 
         size_t begin, end;
-        neighbor_range(v0, 0, &begin, &end);
+        hnsw.neighbor_range(v0, 0, &begin, &end);
 
         for (size_t j = begin; j < end; ++j) {
-            int v1 = neighbors[j];
+            int v1 = hnsw.neighbors[j];
 
             if (v1 < 0) {
                 break;
@@ -651,14 +671,19 @@ std::priority_queue<HNSW::Node> HNSW::search_from_candidate_unbounded(
     return top_candidates;
 }
 
+} // anonymous namespace
+
 HNSWStats HNSW::search(
         DistanceComputer& qdis,
         int k,
         idx_t* I,
         float* D,
-        VisitedTable& vt) const {
+        VisitedTable& vt,
+        const SearchParametersHNSW* params) const {
     HNSWStats stats;
-
+    if (entry_point == -1) {
+        return stats;
+    }
     if (upper_beam == 1) {
         //  greedy search on upper levels
         storage_idx_t nearest = entry_point;
@@ -669,16 +694,22 @@ HNSWStats HNSW::search(
         }
 
         int ef = std::max(efSearch, k);
-        if (search_bounded_queue) {
+        if (search_bounded_queue) { // this is the most common branch
             MinimaxHeap candidates(ef);
 
             candidates.push(nearest, d_nearest);
 
-            search_from_candidates(qdis, k, I, D, candidates, vt, stats, 0);
+            search_from_candidates(
+                    *this, qdis, k, I, D, candidates, vt, stats, 0, 0, params);
         } else {
             std::priority_queue<Node> top_candidates =
                     search_from_candidate_unbounded(
-                            Node(d_nearest, nearest), qdis, ef, &vt, stats);
+                            *this,
+                            Node(d_nearest, nearest),
+                            qdis,
+                            ef,
+                            &vt,
+                            stats);
 
             while (top_candidates.size() > k) {
                 top_candidates.pop();
@@ -718,9 +749,10 @@ HNSWStats HNSW::search(
 
             if (level == 0) {
                 nres = search_from_candidates(
-                        qdis, k, I, D, candidates, vt, stats, 0);
+                        *this, qdis, k, I, D, candidates, vt, stats, 0);
             } else {
                 nres = search_from_candidates(
+                        *this,
                         qdis,
                         candidates_size,
                         I_to_next.data(),
@@ -737,6 +769,70 @@ HNSWStats HNSW::search(
     return stats;
 }
 
+void HNSW::search_level_0(
+        DistanceComputer& qdis,
+        int k,
+        idx_t* idxi,
+        float* simi,
+        idx_t nprobe,
+        const storage_idx_t* nearest_i,
+        const float* nearest_d,
+        int search_type,
+        HNSWStats& search_stats,
+        VisitedTable& vt) const {
+    const HNSW& hnsw = *this;
+
+    if (search_type == 1) {
+        int nres = 0;
+
+        for (int j = 0; j < nprobe; j++) {
+            storage_idx_t cj = nearest_i[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[j]);
+
+            nres = search_from_candidates(
+                    hnsw,
+                    qdis,
+                    k,
+                    idxi,
+                    simi,
+                    candidates,
+                    vt,
+                    search_stats,
+                    0,
+                    nres);
+        }
+    } else if (search_type == 2) {
+        int candidates_size = std::max(hnsw.efSearch, int(k));
+        candidates_size = std::max(candidates_size, int(nprobe));
+
+        MinimaxHeap candidates(candidates_size);
+        for (int j = 0; j < nprobe; j++) {
+            storage_idx_t cj = nearest_i[j];
+
+            if (cj < 0)
+                break;
+            candidates.push(cj, nearest_d[j]);
+        }
+
+        search_from_candidates(
+                hnsw, qdis, k, idxi, simi, candidates, vt, search_stats, 0);
+    }
+}
+
+/**************************************************************
+ * MinimaxHeap
+ **************************************************************/
+
 void HNSW::MinimaxHeap::push(storage_idx_t i, float v) {
     if (k == n) {
         if (v >= dis[0])

data/vendor/faiss/faiss/impl/HNSW.h

@@ -43,6 +43,13 @@ struct VisitedTable;
 struct DistanceComputer; // from AuxIndexStructures
 struct HNSWStats;
 
+struct SearchParametersHNSW : SearchParameters {
+    int efSearch = 16;
+    bool check_relative_distance = true;
+
+    ~SearchParametersHNSW() {}
+};
+
 struct HNSW {
     /// internal storage of vectors (32 bits: this is expensive)
     typedef int storage_idx_t;
@@ -188,30 +195,26 @@ struct HNSW {
             std::vector<omp_lock_t>& locks,
             VisitedTable& vt);
 
-    int search_from_candidates(
+    /// search interface for 1 point, single thread
+    HNSWStats search(
             DistanceComputer& qdis,
             int k,
             idx_t* I,
             float* D,
-            MinimaxHeap& candidates,
             VisitedTable& vt,
-            HNSWStats& stats,
-            int level,
-            int nres_in = 0) const;
+            const SearchParametersHNSW* params = nullptr) const;
 
-    std::priority_queue<Node> search_from_candidate_unbounded(
-            const Node& node,
-            DistanceComputer& qdis,
-            int ef,
-            VisitedTable* vt,
-            HNSWStats& stats) const;
-
-    /// search interface
-    HNSWStats search(
+    /// search only in level 0 from a given vertex
+    void search_level_0(
             DistanceComputer& qdis,
             int k,
-            idx_t* I,
-            float* D,
+            idx_t* idxi,
+            float* simi,
+            idx_t nprobe,
+            const storage_idx_t* nearest_i,
+            const float* nearest_d,
+            int search_type,
+            HNSWStats& search_stats,
             VisitedTable& vt) const;
 
     void reset();

data/vendor/faiss/faiss/impl/IDSelector.cpp

@@ -0,0 +1,125 @@
+/**
+ * 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.
+ */
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
+
+namespace faiss {
+
+/***********************************************************************
+ * IDSelectorRange
+ ***********************************************************************/
+
+IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax, bool assume_sorted)
+        : imin(imin), imax(imax), assume_sorted(assume_sorted) {}
+
+bool IDSelectorRange::is_member(idx_t id) const {
+    return id >= imin && id < imax;
+}
+
+void IDSelectorRange::find_sorted_ids_bounds(
+        size_t list_size,
+        const idx_t* ids,
+        size_t* jmin_out,
+        size_t* jmax_out) const {
+    FAISS_ASSERT(assume_sorted);
+    if (list_size == 0 || imax <= ids[0] || imin > ids[list_size - 1]) {
+        *jmin_out = *jmax_out = 0;
+        return;
+    }
+    // bissection to find imin
+    if (ids[0] >= imin) {
+        *jmin_out = 0;
+    } else {
+        size_t j0 = 0, j1 = list_size;
+        while (j1 > j0 + 1) {
+            size_t jmed = (j0 + j1) / 2;
+            if (ids[jmed] >= imin) {
+                j1 = jmed;
+            } else {
+                j0 = jmed;
+            }
+        }
+        *jmin_out = j1;
+    }
+    // bissection to find imax
+    if (*jmin_out == list_size || ids[*jmin_out] >= imax) {
+        *jmax_out = *jmin_out;
+    } else {
+        size_t j0 = *jmin_out, j1 = list_size;
+        while (j1 > j0 + 1) {
+            size_t jmed = (j0 + j1) / 2;
+            if (ids[jmed] >= imax) {
+                j1 = jmed;
+            } else {
+                j0 = jmed;
+            }
+        }
+        *jmax_out = j1;
+    }
+}
+
+/***********************************************************************
+ * IDSelectorArray
+ ***********************************************************************/
+
+IDSelectorArray::IDSelectorArray(size_t n, const idx_t* ids) : n(n), ids(ids) {}
+
+bool IDSelectorArray::is_member(idx_t id) const {
+    for (idx_t i = 0; i < n; i++) {
+        if (ids[i] == id)
+            return true;
+    }
+    return false;
+}
+
+/***********************************************************************
+ * IDSelectorBatch
+ ***********************************************************************/
+
+IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) {
+    nbits = 0;
+    while (n > ((idx_t)1 << nbits)) {
+        nbits++;
+    }
+    nbits += 5;
+    // for n = 1M, nbits = 25 is optimal, see P56659518
+
+    mask = ((idx_t)1 << nbits) - 1;
+    bloom.resize((idx_t)1 << (nbits - 3), 0);
+    for (idx_t i = 0; i < n; i++) {
+        Index::idx_t id = indices[i];
+        set.insert(id);
+        id &= mask;
+        bloom[id >> 3] |= 1 << (id & 7);
+    }
+}
+
+bool IDSelectorBatch::is_member(idx_t i) const {
+    long im = i & mask;
+    if (!(bloom[im >> 3] & (1 << (im & 7)))) {
+        return 0;
+    }
+    return set.count(i);
+}
+
+/***********************************************************************
+ * IDSelectorBitmap
+ ***********************************************************************/
+
+IDSelectorBitmap::IDSelectorBitmap(size_t n, const uint8_t* bitmap)
+        : n(n), bitmap(bitmap) {}
+
+bool IDSelectorBitmap::is_member(idx_t ii) const {
+    uint64_t i = ii;
+    if ((i >> 3) >= n) {
+        return false;
+    }
+    return (bitmap[i >> 3] >> (i & 7)) & 1;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/IDSelector.h

@@ -0,0 +1,135 @@
+/**
+ * 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 <unordered_set>
+#include <vector>
+
+#include <faiss/Index.h>
+
+/** IDSelector is intended to define a subset of vectors to handle (for removal
+ * or as subset to search) */
+
+namespace faiss {
+
+/** Encapsulates a set of ids to handle. */
+struct IDSelector {
+    using idx_t = Index::idx_t;
+    virtual bool is_member(idx_t id) const = 0;
+    virtual ~IDSelector() {}
+};
+
+/** ids between [imin, imax) */
+struct IDSelectorRange : IDSelector {
+    idx_t imin, imax;
+
+    /// Assume that the ids to handle are sorted. In some cases this can speed
+    /// up processing
+    bool assume_sorted;
+
+    IDSelectorRange(idx_t imin, idx_t imax, bool assume_sorted = false);
+
+    bool is_member(idx_t id) const final;
+
+    /// for sorted ids, find the range of list indices where the valid ids are
+    /// stored
+    void find_sorted_ids_bounds(
+            size_t list_size,
+            const idx_t* ids,
+            size_t* jmin,
+            size_t* jmax) const;
+
+    ~IDSelectorRange() override {}
+};
+
+/** Simple array of elements
+ *
+ * is_member calls are very inefficient, but some operations can use the ids
+ * directly.
+ */
+struct IDSelectorArray : IDSelector {
+    size_t n;
+    const idx_t* ids;
+
+    /** Construct with an array of ids to process
+     *
+     * @param n number of ids to store
+     * @param ids elements to store. The pointer should remain valid during
+     *            IDSelectorArray's lifetime
+     */
+    IDSelectorArray(size_t n, const idx_t* ids);
+    bool is_member(idx_t id) const final;
+    ~IDSelectorArray() override {}
+};
+
+/** Ids from a set.
+ *
+ * Repetitions of ids in the indices set passed to the constructor does not hurt
+ * performance.
+ *
+ * The hash function used for the bloom filter and GCC's implementation of
+ * unordered_set are just the least significant bits of the id. This works fine
+ * for random ids or ids in sequences but will produce many hash collisions if
+ * lsb's are always the same
+ */
+struct IDSelectorBatch : IDSelector {
+    std::unordered_set<idx_t> set;
+
+    // Bloom filter to avoid accessing the unordered set if it is unlikely
+    // to be true
+    std::vector<uint8_t> bloom;
+    int nbits;
+    idx_t mask;
+
+    /** Construct with an array of ids to process
+     *
+     * @param n number of ids to store
+     * @param ids elements to store. The pointer can be released after
+     *            construction
+     */
+    IDSelectorBatch(size_t n, const idx_t* indices);
+    bool is_member(idx_t id) const final;
+    ~IDSelectorBatch() override {}
+};
+
+/** One bit per element. Constructed with a bitmap, size ceil(n / 8).
+ */
+struct IDSelectorBitmap : IDSelector {
+    size_t n;
+    const uint8_t* bitmap;
+
+    /** Construct with a binary mask
+     *
+     * @param n size of the bitmap array
+     * @param bitmap id will be selected iff id / 8 < n and bit number
+     *               (i%8) of bitmap[floor(i / 8)] is 1.
+     */
+    IDSelectorBitmap(size_t n, const uint8_t* bitmap);
+    bool is_member(idx_t id) const final;
+    ~IDSelectorBitmap() override {}
+};
+
+/** reverts the membership test of another selector */
+struct IDSelectorNot : IDSelector {
+    const IDSelector* sel;
+    IDSelectorNot(const IDSelector* sel) : sel(sel) {}
+    bool is_member(idx_t id) const final {
+        return !sel->is_member(id);
+    }
+    virtual ~IDSelectorNot() {}
+};
+
+/// selects all entries (useful for benchmarking)
+struct IDSelectorAll : IDSelector {
+    bool is_member(idx_t id) const final {
+        return true;
+    }
+    virtual ~IDSelectorAll() {}
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp

@@ -15,7 +15,6 @@
 
 #include <algorithm>
 
-#include <faiss/Clustering.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
@@ -151,9 +150,6 @@ LocalSearchQuantizer::LocalSearchQuantizer(
         size_t nbits,
         Search_type_t search_type)
         : AdditiveQuantizer(d, std::vector<size_t>(M, nbits), search_type) {
-    is_trained = false;
-    verbose = false;
-
     K = (1 << nbits);
 
     train_iters = 25;
@@ -182,7 +178,7 @@ LocalSearchQuantizer::LocalSearchQuantizer() : LocalSearchQuantizer(0, 0, 0) {}
 
 void LocalSearchQuantizer::train(size_t n, const float* x) {
     FAISS_THROW_IF_NOT(K == (1 << nbits[0]));
-    FAISS_THROW_IF_NOT(nperts <= M);
+    nperts = std::min(nperts, M);
 
     lsq_timer.reset();
     LSQTimerScope scope(&lsq_timer, "train");
@@ -264,26 +260,7 @@ void LocalSearchQuantizer::train(size_t n, const float* x) {
         decode_unpacked(codes.data(), x_recons.data(), n);
         fvec_norms_L2sqr(norms.data(), x_recons.data(), d, n);
 
-        norm_min = HUGE_VALF;
-        norm_max = -HUGE_VALF;
-        for (idx_t i = 0; i < n; i++) {
-            if (norms[i] < norm_min) {
-                norm_min = norms[i];
-            }
-            if (norms[i] > norm_max) {
-                norm_max = norms[i];
-            }
-        }
-
-        if (search_type == ST_norm_cqint8 || search_type == ST_norm_cqint4) {
-            size_t k = (1 << 8);
-            if (search_type == ST_norm_cqint4) {
-                k = (1 << 4);
-            }
-            Clustering1D clus(k);
-            clus.train_exact(n, norms.data());
-            qnorm.add(clus.k, clus.centroids.data());
-        }
+        train_norm(n, norms.data());
     }
 
     if (verbose) {
@@ -318,10 +295,11 @@ void LocalSearchQuantizer::perturb_codebooks(
     }
 }
 
-void LocalSearchQuantizer::compute_codes(
+void LocalSearchQuantizer::compute_codes_add_centroids(
         const float* x,
         uint8_t* codes_out,
-        size_t n) const {
+        size_t n,
+        const float* centroids) const {
     FAISS_THROW_IF_NOT_MSG(is_trained, "LSQ is not trained yet.");
 
     lsq_timer.reset();
@@ -335,7 +313,7 @@ void LocalSearchQuantizer::compute_codes(
     random_int32(codes, 0, K - 1, gen);
 
     icm_encode(codes.data(), x, n, encode_ils_iters, gen);
-    pack_codes(n, codes.data(), codes_out);
+    pack_codes(n, codes.data(), codes_out, -1, nullptr, centroids);
 
     if (verbose) {
         scope.finish();

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h

@@ -83,8 +83,13 @@ struct LocalSearchQuantizer : AdditiveQuantizer {
      * @param x      vectors to encode, size n * d
      * @param codes  output codes, size n * code_size
      * @param n      number of vectors
+     * @param centroids  centroids to be added to x, size n * d
      */
-    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
+    void compute_codes_add_centroids(
+            const float* x,
+            uint8_t* codes,
+            size_t n,
+            const float* centroids = nullptr) const override;
 
     /** Update codebooks given encodings
      *