faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/IndexIVF.h

@@ -11,11 +11,13 @@
 #define FAISS_INDEX_IVF_H
 
 #include <stdint.h>
+#include <memory>
 #include <unordered_map>
 #include <vector>
 
 #include <faiss/Clustering.h>
 #include <faiss/Index.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/invlists/DirectMap.h>
 #include <faiss/invlists/InvertedLists.h>
@@ -62,13 +64,18 @@ struct Level1Quantizer {
     ~Level1Quantizer();
 };
 
-struct IVFSearchParameters {
+struct SearchParametersIVF : SearchParameters {
     size_t nprobe;    ///< number of probes at query time
     size_t max_codes; ///< max nb of codes to visit to do a query
-    IVFSearchParameters() : nprobe(1), max_codes(0) {}
-    virtual ~IVFSearchParameters() {}
+    SearchParameters* quantizer_params = nullptr;
+
+    SearchParametersIVF() : nprobe(1), max_codes(0) {}
+    virtual ~SearchParametersIVF() {}
 };
 
+// the new convention puts the index type after SearchParameters
+using IVFSearchParameters = SearchParametersIVF;
+
 struct InvertedListScanner;
 struct IndexIVFStats;
 
@@ -218,13 +225,15 @@ struct IndexIVF : Index, Level1Quantizer {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void range_search(
             idx_t n,
             const float* x,
             float radius,
-            RangeSearchResult* result) const override;
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
 
     void range_search_preassigned(
             idx_t nx,
@@ -242,7 +251,8 @@ struct IndexIVF : Index, Level1Quantizer {
      * The default search implementation uses this to compute the distances
      */
     virtual InvertedListScanner* get_InvertedListScanner(
-            bool store_pairs = false) const;
+            bool store_pairs = false,
+            const IDSelector* sel = nullptr) const;
 
     /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2
      */
@@ -284,7 +294,8 @@ struct IndexIVF : Index, Level1Quantizer {
             idx_t k,
             float* distances,
             idx_t* labels,
-            float* recons) const override;
+            float* recons,
+            const SearchParameters* params = nullptr) const override;
 
     /** Reconstruct a vector given the location in terms of (inv list index +
      * inv list offset) instead of the id.
@@ -302,15 +313,9 @@ struct IndexIVF : Index, Level1Quantizer {
 
     size_t remove_ids(const IDSelector& sel) override;
 
-    /** check that the two indexes are compatible (ie, they are
-     * trained in the same way and have the same
-     * parameters). Otherwise throw. */
-    void check_compatible_for_merge(const IndexIVF& other) const;
+    void check_compatible_for_merge(const Index& otherIndex) const override;
 
-    /** moves the entries from another dataset to self. On output,
-     * other is empty. add_id is added to all moved ids (for
-     * sequential ids, this would be this->ntotal */
-    virtual void merge_from(IndexIVF& other, idx_t add_id);
+    virtual void merge_from(Index& otherIndex, idx_t add_id) override;
 
     /** copy a subset of the entries index to the other index
      *
@@ -331,6 +336,9 @@ struct IndexIVF : Index, Level1Quantizer {
         return invlists->list_size(list_no);
     }
 
+    /// are the ids sorted?
+    bool check_ids_sorted() const;
+
     /** intialize a direct map
      *
      * @param new_maintain_direct_map    if true, create a direct map,
@@ -363,7 +371,15 @@ struct InvertedListScanner {
     idx_t list_no = -1;    ///< remember current list
     bool keep_max = false; ///< keep maximum instead of minimum
     /// store positions in invlists rather than labels
-    bool store_pairs = false;
+    bool store_pairs;
+
+    /// search in this subset of ids
+    const IDSelector* sel;
+
+    InvertedListScanner(
+            bool store_pairs = false,
+            const IDSelector* sel = nullptr)
+            : store_pairs(store_pairs), sel(sel) {}
 
     /// used in default implementation of scan_codes
     size_t code_size = 0;
@@ -411,6 +427,9 @@ struct InvertedListScanner {
     virtual ~InvertedListScanner() {}
 };
 
+// whether to check that coarse quantizers are the same
+FAISS_API extern bool check_compatible_for_merge_expensive_check;
+
 struct IndexIVFStats {
     size_t nq;                // nb of queries run
     size_t nlist;             // nb of inverted lists scanned

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp

@@ -5,9 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// quiet the noise
-// XXclang-format off
-
 #include <faiss/IndexIVFAdditiveQuantizer.h>
 
 #include <algorithm>
@@ -44,10 +41,14 @@ void IndexIVFAdditiveQuantizer::train_residual(idx_t n, const float* x) {
     const float* x_in = x;
 
     size_t max_train_points = 1024 * ((size_t)1 << aq->nbits[0]);
+    // we need more data to train LSQ
+    if (dynamic_cast<LocalSearchQuantizer*>(aq)) {
+        max_train_points = 1024 * aq->M * ((size_t)1 << aq->nbits[0]);
+    }
 
     x = fvecs_maybe_subsample(
             d, (size_t*)&n, max_train_points, x, verbose, 1234);
-    ScopeDeleter1<float> del_x(x_in == x ? nullptr : x);
+    ScopeDeleter<float> del_x(x_in == x ? nullptr : x);
 
     if (by_residual) {
         std::vector<Index::idx_t> idx(n);
@@ -76,7 +77,7 @@ void IndexIVFAdditiveQuantizer::encode_vectors(
         // subtract centroids
         std::vector<float> residuals(n * d);
 
-#pragma omp parallel if (n > 10000)
+#pragma omp parallel for if (n > 10000)
         for (idx_t i = 0; i < n; i++) {
             quantizer->compute_residual(
                     x + i * d,
@@ -99,6 +100,32 @@ void IndexIVFAdditiveQuantizer::encode_vectors(
     }
 }
 
+void IndexIVFAdditiveQuantizer::sa_decode(
+        idx_t n,
+        const uint8_t* codes,
+        float* x) const {
+    const size_t coarse_size = coarse_code_size();
+
+#pragma omp parallel if (n > 1000)
+    {
+        std::vector<float> residual(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+            aq->decode(code + coarse_size, xi, 1);
+            if (by_residual) {
+                quantizer->reconstruct(list_no, residual.data());
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] += residual[j];
+                }
+            }
+        }
+    }
+}
+
 IndexIVFAdditiveQuantizer::~IndexIVFAdditiveQuantizer() {}
 
 /*********************************************
@@ -220,7 +247,9 @@ struct AQInvertedListScannerLUT : AQInvertedListScanner {
 } // anonymous namespace
 
 InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
-        bool store_pairs) const {
+        bool store_pairs,
+        const IDSelector* sel) const {
+    FAISS_THROW_IF_NOT(!sel);
     if (metric_type == METRIC_INNER_PRODUCT) {
         if (aq->search_type == AdditiveQuantizer::ST_decompress) {
             return new AQInvertedListScannerDecompress<true>(
@@ -244,8 +273,10 @@ InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
                 A(ST_norm_float)
                 A(ST_norm_qint8)
                 A(ST_norm_qint4)
-                A(ST_norm_cqint8)
                 A(ST_norm_cqint4)
+            case AdditiveQuantizer::ST_norm_lsq2x4:
+            case AdditiveQuantizer::ST_norm_rq2x4:
+                A(ST_norm_cqint8)
 #undef A
             default:
                 FAISS_THROW_FMT(
@@ -313,4 +344,50 @@ IndexIVFLocalSearchQuantizer::IndexIVFLocalSearchQuantizer()
 
 IndexIVFLocalSearchQuantizer::~IndexIVFLocalSearchQuantizer() {}
 
+/**************************************************************************************
+ * IndexIVFProductResidualQuantizer
+ **************************************************************************************/
+
+IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        size_t nsplits,
+        size_t Msub,
+        size_t nbits,
+        MetricType metric,
+        Search_type_t search_type)
+        : IndexIVFAdditiveQuantizer(&prq, quantizer, d, nlist, metric),
+          prq(d, nsplits, Msub, nbits, search_type) {
+    code_size = invlists->code_size = prq.code_size;
+}
+
+IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer()
+        : IndexIVFAdditiveQuantizer(&prq) {}
+
+IndexIVFProductResidualQuantizer::~IndexIVFProductResidualQuantizer() {}
+
+/**************************************************************************************
+ * IndexIVFProductLocalSearchQuantizer
+ **************************************************************************************/
+
+IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        size_t nsplits,
+        size_t Msub,
+        size_t nbits,
+        MetricType metric,
+        Search_type_t search_type)
+        : IndexIVFAdditiveQuantizer(&plsq, quantizer, d, nlist, metric),
+          plsq(d, nsplits, Msub, nbits, search_type) {
+    code_size = invlists->code_size = plsq.code_size;
+}
+
+IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer()
+        : IndexIVFAdditiveQuantizer(&plsq) {}
+
+IndexIVFProductLocalSearchQuantizer::~IndexIVFProductLocalSearchQuantizer() {}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h

@@ -15,6 +15,7 @@
 
 #include <faiss/IndexIVF.h>
 #include <faiss/impl/LocalSearchQuantizer.h>
+#include <faiss/impl/ProductAdditiveQuantizer.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/platform_macros.h>
 
@@ -49,7 +50,10 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
             bool include_listnos = false) const override;
 
     InvertedListScanner* get_InvertedListScanner(
-            bool store_pairs) const override;
+            bool store_pairs,
+            const IDSelector* sel) const override;
+
+    void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
 
     ~IndexIVFAdditiveQuantizer() override;
 };
@@ -116,6 +120,64 @@ struct IndexIVFLocalSearchQuantizer : IndexIVFAdditiveQuantizer {
     virtual ~IndexIVFLocalSearchQuantizer();
 };
 
+/** IndexIVF based on a product residual quantizer. Stored vectors are
+ * approximated by product residual quantization codes.
+ */
+struct IndexIVFProductResidualQuantizer : IndexIVFAdditiveQuantizer {
+    /// The product residual quantizer used to encode the vectors
+    ProductResidualQuantizer prq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param nsplits  number of residual quantizers
+     * @param Msub   number of subquantizers per RQ
+     * @param nbits  number of bit per subvector index
+     */
+    IndexIVFProductResidualQuantizer(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t nsplits,
+            size_t Msub,
+            size_t nbits,
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+
+    IndexIVFProductResidualQuantizer();
+
+    virtual ~IndexIVFProductResidualQuantizer();
+};
+
+/** IndexIVF based on a product local search quantizer. Stored vectors are
+ * approximated by product local search quantization codes.
+ */
+struct IndexIVFProductLocalSearchQuantizer : IndexIVFAdditiveQuantizer {
+    /// The product local search quantizer used to encode the vectors
+    ProductLocalSearchQuantizer plsq;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param nsplits  number of local search quantizers
+     * @param Msub   number of subquantizers per LSQ
+     * @param nbits  number of bit per subvector index
+     */
+    IndexIVFProductLocalSearchQuantizer(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t nsplits,
+            size_t Msub,
+            size_t nbits,
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+
+    IndexIVFProductLocalSearchQuantizer();
+
+    virtual ~IndexIVFProductLocalSearchQuantizer();
+};
+
 } // namespace faiss
 
 #endif