faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexIVF.h

@@ -11,9 +11,6 @@
 #define FAISS_INDEX_IVF_H
 
 #include <stdint.h>
-#include <memory>
-#include <unordered_map>
-#include <vector>
 
 #include <faiss/Clustering.h>
 #include <faiss/Index.h>
@@ -61,7 +58,7 @@ struct Level1Quantizer {
     void encode_listno(idx_t list_no, uint8_t* code) const;
     idx_t decode_listno(const uint8_t* code) const;
 
-    Level1Quantizer(Index* quantizer, size_t nlist);
+    Level1Quantizer(Index* quantizer_in, size_t nlist_in);
 
     Level1Quantizer();
 
@@ -71,6 +68,25 @@ struct Level1Quantizer {
 struct SearchParametersIVF : SearchParameters {
     size_t nprobe = 1;    ///< number of probes at query time
     size_t max_codes = 0; ///< max nb of codes to visit to do a query
+
+    /// FastScan k-NN only: maximum number of inverted lists to visit.
+    /// 0 means unlimited, i.e. bounded only by nprobe. When set together
+    /// with max_codes, either budget may stop the scan. With
+    /// ensure_topk_full, this limit is treated as at least k lists.
+    size_t max_lists_num = 0;
+
+    /// For k-NN search, make small early-stop budgets less aggressive:
+    /// max_codes is treated as at least k post-IDSelector scans. Supported
+    /// by generic IVF in parallel_mode 0 and 3, and by FastScan k-NN
+    /// implementations 10 and 11.
+    bool ensure_topk_full = false;
+
+    /// Range-search only: stop after this many consecutive probed lists add
+    /// no in-radius results. 0 disables the heuristic. This trades recall
+    /// for less work. Supported in parallel_mode 0; FastScan range search
+    /// uses implementation 10 for this option.
+    size_t max_empty_result_buckets = 0;
+
     SearchParameters* quantizer_params = nullptr;
     /// context object to pass to InvertedLists
     void* inverted_list_context = nullptr;
@@ -89,8 +105,10 @@ struct IndexIVFInterface : Level1Quantizer {
     size_t nprobe = 1;    ///< number of probes at query time
     size_t max_codes = 0; ///< max nb of codes to visit to do a query
 
-    explicit IndexIVFInterface(Index* quantizer = nullptr, size_t nlist = 0)
-            : Level1Quantizer(quantizer, nlist) {}
+    explicit IndexIVFInterface(
+            Index* quantizer_in = nullptr,
+            size_t nlist_in = 0)
+            : Level1Quantizer(quantizer_in, nlist_in) {}
 
     /** search a set of vectors, that are pre-quantized by the IVF
      *  quantizer. Fill in the corresponding heaps with the query
@@ -206,12 +224,12 @@ struct IndexIVF : Index, IndexIVFInterface {
      * identifier.
      */
     IndexIVF(
-            Index* quantizer,
-            size_t d,
-            size_t nlist,
-            size_t code_size,
+            Index* quantizer_in,
+            size_t d_in,
+            size_t nlist_in,
+            size_t code_size_in,
             MetricType metric = METRIC_L2,
-            bool own_invlists = true);
+            bool own_invlists_in = true);
 
     void reset() override;
 
@@ -325,6 +343,12 @@ struct IndexIVF : Index, IndexIVFInterface {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) const override;
 
+    /** search one vector with a custom result handler */
+    void search1(
+            const float* x,
+            ResultHandler& handler,
+            SearchParameters* params = nullptr) const override;
+
     /** Get a scanner for this index (store_pairs means ignore labels)
      *
      * The default search implementation uses this to compute the distances.
@@ -481,9 +505,9 @@ struct InvertedListScanner {
     const IDSelector* sel;
 
     InvertedListScanner(
-            bool store_pairs = false,
-            const IDSelector* sel = nullptr)
-            : store_pairs(store_pairs), sel(sel) {}
+            bool store_pairs_in = false,
+            const IDSelector* sel_in = nullptr)
+            : store_pairs(store_pairs_in), sel(sel_in) {}
 
     /// used in default implementation of scan_codes
     size_t code_size = 0;
@@ -492,7 +516,7 @@ struct InvertedListScanner {
     virtual void set_query(const float* query_vector) = 0;
 
     /// following codes come from this inverted list
-    virtual void set_list(idx_t list_no, float coarse_dis) = 0;
+    virtual void set_list(idx_t list_no, float coarse_dis);
 
     /// compute a single query-to-code distance
     virtual float distance_to_code(const uint8_t* code) const = 0;
@@ -509,7 +533,7 @@ struct InvertedListScanner {
      * @param k          heap size
      * @return number of heap updates performed
      */
-    virtual size_t scan_codes(
+    size_t scan_codes(
             size_t n,
             const uint8_t* codes,
             const idx_t* ids,
@@ -543,6 +567,13 @@ struct InvertedListScanner {
             RangeQueryResult& result,
             size_t& list_size) const;
 
+    // accumulate results with a ResultHandler
+    virtual size_t scan_codes(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const;
+
     virtual ~InvertedListScanner() {}
 };
 

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp

@@ -24,23 +24,23 @@ namespace faiss {
  **************************************************************************************/
 
 IndexIVFAdditiveQuantizer::IndexIVFAdditiveQuantizer(
-        AdditiveQuantizer* aq,
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        AdditiveQuantizer* aq_,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         MetricType metric,
-        bool own_invlists)
-        : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists), aq(aq) {
+        bool own_invlists_)
+        : IndexIVF(quantizer_, d_, nlist_, 0, metric, own_invlists_), aq(aq_) {
     by_residual = true;
 }
 
-IndexIVFAdditiveQuantizer::IndexIVFAdditiveQuantizer(AdditiveQuantizer* aq)
-        : IndexIVF(), aq(aq) {}
+IndexIVFAdditiveQuantizer::IndexIVFAdditiveQuantizer(AdditiveQuantizer* aq_)
+        : IndexIVF(), aq(aq_) {}
 
 void IndexIVFAdditiveQuantizer::train_encoder(
         idx_t n,
         const float* x,
-        const idx_t* assign) {
+        const idx_t* /*assign*/) {
     aq->train(n, x);
 }
 
@@ -107,7 +107,7 @@ void IndexIVFAdditiveQuantizer::decode_vectors(
             if (by_residual) {
                 int64_t list_no = listnos[i];
                 quantizer->reconstruct(list_no, residual.data());
-                for (size_t j = 0; j < d; j++) {
+                for (int j = 0; j < d; j++) {
                     xi[j] += residual[j];
                 }
             }
@@ -133,7 +133,7 @@ void IndexIVFAdditiveQuantizer::sa_decode(
             aq->decode(code + coarse_size, xi, 1);
             if (by_residual) {
                 quantizer->reconstruct(list_no, residual.data());
-                for (size_t j = 0; j < d; j++) {
+                for (int j = 0; j < d; j++) {
                     xi[j] += residual[j];
                 }
             }
@@ -171,27 +171,29 @@ struct AQInvertedListScanner : InvertedListScanner {
     const AdditiveQuantizer& aq;
     std::vector<float> tmp;
 
-    AQInvertedListScanner(const IndexIVFAdditiveQuantizer& ia, bool store_pairs)
-            : ia(ia), aq(*ia.aq) {
-        this->store_pairs = store_pairs;
-        this->code_size = ia.code_size;
+    AQInvertedListScanner(
+            const IndexIVFAdditiveQuantizer& ia_,
+            bool store_pairs_)
+            : ia(ia_), aq(*ia_.aq) {
+        this->store_pairs = store_pairs_;
+        this->code_size = ia_.code_size;
         keep_max = is_similarity_metric(ia.metric_type);
         tmp.resize(ia.d);
     }
 
-    const float* q0;
+    const float* q0 = nullptr;
 
     /// from now on we handle this query.
     void set_query(const float* query_vector) override {
         q0 = query_vector;
     }
 
-    const float* q;
+    const float* q = nullptr;
     /// following codes come from this inverted list
-    void set_list(idx_t list_no, float coarse_dis) override {
-        this->list_no = list_no;
+    void set_list(idx_t list_no_, float /*coarse_dis*/) override {
+        this->list_no = list_no_;
         if (ia.metric_type == METRIC_L2 && ia.by_residual) {
-            ia.quantizer->compute_residual(q0, tmp.data(), list_no);
+            ia.quantizer->compute_residual(q0, tmp.data(), list_no_);
             q = tmp.data();
         } else {
             q = q0;
@@ -204,15 +206,15 @@ struct AQInvertedListScanner : InvertedListScanner {
 template <bool is_IP>
 struct AQInvertedListScannerDecompress : AQInvertedListScanner {
     AQInvertedListScannerDecompress(
-            const IndexIVFAdditiveQuantizer& ia,
-            bool store_pairs)
-            : AQInvertedListScanner(ia, store_pairs) {}
+            const IndexIVFAdditiveQuantizer& ia_,
+            bool store_pairs_)
+            : AQInvertedListScanner(ia_, store_pairs_) {}
 
     float coarse_dis = 0;
 
     /// following codes come from this inverted list
-    void set_list(idx_t list_no, float coarse_dis_2) override {
-        AQInvertedListScanner::set_list(list_no, coarse_dis_2);
+    void set_list(idx_t list_no_, float coarse_dis_2) override {
+        AQInvertedListScanner::set_list(list_no_, coarse_dis_2);
         if (ia.by_residual) {
             this->coarse_dis = coarse_dis_2;
         }
@@ -238,9 +240,9 @@ struct AQInvertedListScannerLUT : AQInvertedListScanner {
     float distance_bias;
 
     AQInvertedListScannerLUT(
-            const IndexIVFAdditiveQuantizer& ia,
-            bool store_pairs)
-            : AQInvertedListScanner(ia, store_pairs) {
+            const IndexIVFAdditiveQuantizer& ia_,
+            bool store_pairs_)
+            : AQInvertedListScanner(ia_, store_pairs_) {
         LUT.resize(aq.total_codebook_size);
         tmp.resize(ia.d);
         distance_bias = 0;
@@ -255,8 +257,8 @@ struct AQInvertedListScannerLUT : AQInvertedListScanner {
     }
 
     /// following codes come from this inverted list
-    void set_list(idx_t list_no, float coarse_dis) override {
-        AQInvertedListScanner::set_list(list_no, coarse_dis);
+    void set_list(idx_t list_no_, float coarse_dis) override {
+        AQInvertedListScanner::set_list(list_no_, coarse_dis);
         // TODO find a way to provide the nprobes together to do a matmul
         // +  precompute tables
         aq.compute_LUT(1, q, LUT.data());
@@ -310,6 +312,7 @@ InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
             case AdditiveQuantizer::ST_norm_rq2x4:
                 A(ST_norm_cqint8)
 #undef A
+            case AdditiveQuantizer::ST_count:
             default:
                 FAISS_THROW_FMT(
                         "search type %d not supported", aq->search_type);
@@ -322,21 +325,21 @@ InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
  **************************************************************************************/
 
 IndexIVFResidualQuantizer::IndexIVFResidualQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         const std::vector<size_t>& nbits,
         MetricType metric,
         Search_type_t search_type,
-        bool own_invlists)
+        bool own_invlists_)
         : IndexIVFAdditiveQuantizer(
                   &rq,
-                  quantizer,
-                  d,
-                  nlist,
+                  quantizer_,
+                  d_,
+                  nlist_,
                   metric,
-                  own_invlists),
-          rq(d, nbits, search_type) {
+                  own_invlists_),
+          rq(d_, nbits, search_type) {
     code_size = rq.code_size;
     if (invlists) {
         invlists->code_size = code_size;
@@ -347,22 +350,22 @@ IndexIVFResidualQuantizer::IndexIVFResidualQuantizer()
         : IndexIVFAdditiveQuantizer(&rq) {}
 
 IndexIVFResidualQuantizer::IndexIVFResidualQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         size_t M,     /* number of subquantizers */
         size_t nbits, /* number of bit per subvector index */
         MetricType metric,
         Search_type_t search_type,
-        bool own_invlists)
+        bool own_invlists_)
         : IndexIVFResidualQuantizer(
-                  quantizer,
-                  d,
-                  nlist,
+                  quantizer_,
+                  d_,
+                  nlist_,
                   std::vector<size_t>(M, nbits),
                   metric,
                   search_type,
-                  own_invlists) {}
+                  own_invlists_) {}
 
 IndexIVFResidualQuantizer::~IndexIVFResidualQuantizer() = default;
 
@@ -371,22 +374,22 @@ IndexIVFResidualQuantizer::~IndexIVFResidualQuantizer() = default;
  **************************************************************************************/
 
 IndexIVFLocalSearchQuantizer::IndexIVFLocalSearchQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         size_t M,     /* number of subquantizers */
         size_t nbits, /* number of bit per subvector index */
         MetricType metric,
         Search_type_t search_type,
-        bool own_invlists)
+        bool own_invlists_)
         : IndexIVFAdditiveQuantizer(
                   &lsq,
-                  quantizer,
-                  d,
-                  nlist,
+                  quantizer_,
+                  d_,
+                  nlist_,
                   metric,
-                  own_invlists),
-          lsq(d, M, nbits, search_type) {
+                  own_invlists_),
+          lsq(d_, M, nbits, search_type) {
     code_size = lsq.code_size;
     if (invlists) {
         invlists->code_size = code_size;
@@ -403,23 +406,23 @@ IndexIVFLocalSearchQuantizer::~IndexIVFLocalSearchQuantizer() = default;
  **************************************************************************************/
 
 IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         size_t nsplits,
         size_t Msub,
         size_t nbits,
         MetricType metric,
         Search_type_t search_type,
-        bool own_invlists)
+        bool own_invlists_)
         : IndexIVFAdditiveQuantizer(
                   &prq,
-                  quantizer,
-                  d,
-                  nlist,
+                  quantizer_,
+                  d_,
+                  nlist_,
                   metric,
-                  own_invlists),
-          prq(d, nsplits, Msub, nbits, search_type) {
+                  own_invlists_),
+          prq(d_, nsplits, Msub, nbits, search_type) {
     code_size = prq.code_size;
     if (invlists) {
         invlists->code_size = code_size;
@@ -436,23 +439,23 @@ IndexIVFProductResidualQuantizer::~IndexIVFProductResidualQuantizer() = default;
  **************************************************************************************/
 
 IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         size_t nsplits,
         size_t Msub,
         size_t nbits,
         MetricType metric,
         Search_type_t search_type,
-        bool own_invlists)
+        bool own_invlists_)
         : IndexIVFAdditiveQuantizer(
                   &plsq,
-                  quantizer,
-                  d,
-                  nlist,
+                  quantizer_,
+                  d_,
+                  nlist_,
                   metric,
-                  own_invlists),
-          plsq(d, nsplits, Msub, nbits, search_type) {
+                  own_invlists_),
+          plsq(d_, nsplits, Msub, nbits, search_type) {
     code_size = plsq.code_size;
     if (invlists) {
         invlists->code_size = code_size;