faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp

@@ -199,60 +199,6 @@ void RangeSearchPartialResult::merge(
     result->lims[0] = 0;
 }
 
-/***********************************************************************
- * IDSelectorRange
- ***********************************************************************/
-
-IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax)
-        : imin(imin), imax(imax) {}
-
-bool IDSelectorRange::is_member(idx_t id) const {
-    return id >= imin && id < imax;
-}
-
-/***********************************************************************
- * 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 > (1L << nbits))
-        nbits++;
-    nbits += 5;
-    // for n = 1M, nbits = 25 is optimal, see P56659518
-
-    mask = (1L << nbits) - 1;
-    bloom.resize(1UL << (nbits - 3), 0);
-    for (long 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);
-}
-
 /***********************************************************
  * Interrupt callback
  ***********************************************************/

data/vendor/faiss/faiss/impl/AuxIndexStructures.h

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 // Auxiliary index structures, that are used in indexes but that can
 // be forward-declared
 
@@ -18,7 +16,6 @@
 #include <cstring>
 #include <memory>
 #include <mutex>
-#include <unordered_set>
 #include <vector>
 
 #include <faiss/Index.h>
@@ -52,55 +49,6 @@ struct RangeSearchResult {
     virtual ~RangeSearchResult();
 };
 
-/** Encapsulates a set of ids to remove. */
-struct IDSelector {
-    typedef Index::idx_t idx_t;
-    virtual bool is_member(idx_t id) const = 0;
-    virtual ~IDSelector() {}
-};
-
-/** remove ids between [imni, imax) */
-struct IDSelectorRange : IDSelector {
-    idx_t imin, imax;
-
-    IDSelectorRange(idx_t imin, idx_t imax);
-    bool is_member(idx_t id) const override;
-    ~IDSelectorRange() override {}
-};
-
-/** simple list of elements to remove
- *
- * this is inefficient in most cases, except for IndexIVF with
- * maintain_direct_map
- */
-struct IDSelectorArray : IDSelector {
-    size_t n;
-    const idx_t* ids;
-
-    IDSelectorArray(size_t n, const idx_t* ids);
-    bool is_member(idx_t id) const override;
-    ~IDSelectorArray() override {}
-};
-
-/** Remove 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;
-
-    typedef unsigned char uint8_t;
-    std::vector<uint8_t> bloom; // assumes low bits of id are a good hash value
-    int nbits;
-    idx_t mask;
-
-    IDSelectorBatch(size_t n, const idx_t* indices);
-    bool is_member(idx_t id) const override;
-    ~IDSelectorBatch() override {}
-};
-
 /****************************************************************
  * Result structures for range search.
  *
@@ -186,30 +134,6 @@ struct RangeSearchPartialResult : BufferList {
             bool do_delete = true);
 };
 
-/***********************************************************
- * The distance computer maintains a current query and computes
- * distances to elements in an index that supports random access.
- *
- * The DistanceComputer is not intended to be thread-safe (eg. because
- * it maintains counters) so the distance functions are not const,
- * instantiate one from each thread if needed.
- ***********************************************************/
-struct DistanceComputer {
-    using idx_t = Index::idx_t;
-
-    /// called before computing distances. Pointer x should remain valid
-    /// while operator () is called
-    virtual void set_query(const float* x) = 0;
-
-    /// compute distance of vector i to current query
-    virtual float operator()(idx_t i) = 0;
-
-    /// compute distance between two stored vectors
-    virtual float symmetric_dis(idx_t i, idx_t j) = 0;
-
-    virtual ~DistanceComputer() {}
-};
-
 /***********************************************************
  * Interrupt callback
  ***********************************************************/

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -0,0 +1,64 @@
+/**
+ * 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>
+
+namespace faiss {
+
+/***********************************************************
+ * The distance computer maintains a current query and computes
+ * distances to elements in an index that supports random access.
+ *
+ * The DistanceComputer is not intended to be thread-safe (eg. because
+ * it maintains counters) so the distance functions are not const,
+ * instantiate one from each thread if needed.
+ *
+ * Note that the equivalent for IVF indexes is the InvertedListScanner,
+ * that has additional methods to handle the inverted list context.
+ ***********************************************************/
+struct DistanceComputer {
+    using idx_t = Index::idx_t;
+
+    /// called before computing distances. Pointer x should remain valid
+    /// while operator () is called
+    virtual void set_query(const float* x) = 0;
+
+    /// compute distance of vector i to current query
+    virtual float operator()(idx_t i) = 0;
+
+    /// compute distance between two stored vectors
+    virtual float symmetric_dis(idx_t i, idx_t j) = 0;
+
+    virtual ~DistanceComputer() {}
+};
+
+/*************************************************************
+ * Specialized version of the DistanceComputer when we know that codes are
+ * laid out in a flat index.
+ */
+struct FlatCodesDistanceComputer : DistanceComputer {
+    const uint8_t* codes;
+    size_t code_size;
+
+    FlatCodesDistanceComputer(const uint8_t* codes, size_t code_size)
+            : codes(codes), code_size(code_size) {}
+
+    FlatCodesDistanceComputer() : codes(nullptr), code_size(0) {}
+
+    float operator()(idx_t i) final {
+        return distance_to_code(codes + i * code_size);
+    }
+
+    /// compute distance of current query to an encoded vector
+    virtual float distance_to_code(const uint8_t* code) = 0;
+
+    virtual ~FlatCodesDistanceComputer() {}
+};
+
+} // namespace faiss

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 {
 
@@ -434,17 +436,22 @@ void HNSW::add_links_starting_from(
 
     ::faiss::shrink_neighbor_list(ptdis, link_targets, M);
 
+    std::vector<storage_idx_t> neighbors;
+    neighbors.reserve(link_targets.size());
     while (!link_targets.empty()) {
-        int other_id = link_targets.top().id;
+        storage_idx_t other_id = link_targets.top().id;
+        add_link(*this, ptdis, pt_id, other_id, level);
+        neighbors.push_back(other_id);
+        link_targets.pop();
+    }
 
+    omp_unset_lock(&locks[pt_id]);
+    for (storage_idx_t other_id : neighbors) {
         omp_set_lock(&locks[other_id]);
         add_link(*this, ptdis, other_id, pt_id, level);
         omp_unset_lock(&locks[other_id]);
-
-        add_link(*this, ptdis, pt_id, other_id, level);
-
-        link_targets.pop();
     }
+    omp_set_lock(&locks[pt_id]);
 }
 
 /**************************************************************
@@ -496,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,
@@ -507,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) {
@@ -541,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)) {
@@ -553,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);
         }
@@ -578,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;
@@ -609,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;
@@ -646,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;
@@ -664,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();
@@ -713,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(),
@@ -732,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