faiss 0.4.1 → 0.4.2

data/vendor/faiss/faiss/IndexIDMap.cpp

@@ -12,11 +12,11 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstdio>
+#include "faiss/Index.h"
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/Heap.h>
-#include <faiss/utils/WorkerThread.h>
 
 namespace faiss {
 
@@ -33,6 +33,17 @@ void sync_d(IndexBinary* index) {
 
 } // anonymous namespace
 
+template <typename componentT>
+NumericType component_t_to_numeric() {
+    if constexpr (std::is_same<componentT, float>::value) {
+        return NumericType::Float32;
+    } else if constexpr (std::is_same<componentT, uint8_t>::value) {
+        return NumericType::UInt8;
+    } else {
+        FAISS_THROW_MSG("Unsupported component_t");
+    }
+}
+
 /*****************************************************
  * IndexIDMap implementation
  *******************************************************/
@@ -47,6 +58,16 @@ IndexIDMapTemplate<IndexT>::IndexIDMapTemplate(IndexT* index) : index(index) {
     sync_d(this);
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::addEx(
+        idx_t,
+        const void*,
+        NumericType numeric_type) {
+    FAISS_THROW_MSG(
+            "add does not make sense with IndexIDMap, "
+            "use add_with_ids");
+}
+
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::add(
         idx_t,
@@ -56,12 +77,22 @@ void IndexIDMapTemplate<IndexT>::add(
             "use add_with_ids");
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::trainEx(
+        idx_t n,
+        const void* x,
+        NumericType numeric_type) {
+    index->trainEx(n, x, numeric_type);
+    this->is_trained = index->is_trained;
+}
+
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::train(
         idx_t n,
         const typename IndexT::component_t* x) {
-    index->train(n, x);
-    this->is_trained = index->is_trained;
+    trainEx(n,
+            static_cast<const void*>(x),
+            component_t_to_numeric<typename IndexT::component_t>());
 }
 
 template <typename IndexT>
@@ -72,16 +103,30 @@ void IndexIDMapTemplate<IndexT>::reset() {
 }
 
 template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::add_with_ids(
+void IndexIDMapTemplate<IndexT>::add_with_idsEx(
         idx_t n,
-        const typename IndexT::component_t* x,
+        const void* x,
+        NumericType numeric_type,
         const idx_t* xids) {
-    index->add(n, x);
-    for (idx_t i = 0; i < n; i++)
+    index->addEx(n, x, numeric_type);
+    for (idx_t i = 0; i < n; i++) {
         id_map.push_back(xids[i]);
+    }
     this->ntotal = index->ntotal;
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::add_with_ids(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        const idx_t* xids) {
+    add_with_idsEx(
+            n,
+            static_cast<const void*>(x),
+            component_t_to_numeric<typename IndexT::component_t>(),
+            xids);
+}
+
 template <typename IndexT>
 size_t IndexIDMapTemplate<IndexT>::sa_code_size() const {
     return index->sa_code_size();
@@ -106,10 +151,10 @@ 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;
+    void set(SearchParameters* target_params, IDSelector* new_sel) {
+        this->params = target_params;
+        old_sel = target_params->sel;
+        target_params->sel = new_sel;
     }
     ~ScopedSelChange() {
         if (params) {
@@ -121,9 +166,10 @@ struct ScopedSelChange {
 } // namespace
 
 template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::search(
+void IndexIDMapTemplate<IndexT>::searchEx(
         idx_t n,
-        const typename IndexT::component_t* x,
+        const void* x,
+        NumericType numeric_type,
         idx_t k,
         typename IndexT::distance_t* distances,
         idx_t* labels,
@@ -147,7 +193,7 @@ void IndexIDMapTemplate<IndexT>::search(
             sel_change.set(params_non_const, &this_idtrans);
         }
     }
-    index->search(n, x, k, distances, labels, params);
+    index->searchEx(n, x, numeric_type, k, distances, labels, params);
     idx_t* li = labels;
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; i++) {
@@ -155,6 +201,24 @@ void IndexIDMapTemplate<IndexT>::search(
     }
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::search(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        idx_t k,
+        typename IndexT::distance_t* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    searchEx(
+            n,
+            static_cast<const void*>(x),
+            component_t_to_numeric<typename IndexT::component_t>(),
+            k,
+            distances,
+            labels,
+            params);
+}
+
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::range_search(
         idx_t n,
@@ -223,8 +287,9 @@ void IndexIDMapTemplate<IndexT>::merge_from(IndexT& otherIndex, idx_t add_id) {
 
 template <typename IndexT>
 IndexIDMapTemplate<IndexT>::~IndexIDMapTemplate() {
-    if (own_fields)
+    if (own_fields) {
         delete index;
+    }
 }
 
 /*****************************************************
@@ -236,17 +301,30 @@ IndexIDMap2Template<IndexT>::IndexIDMap2Template(IndexT* index)
         : IndexIDMapTemplate<IndexT>(index) {}
 
 template <typename IndexT>
-void IndexIDMap2Template<IndexT>::add_with_ids(
+void IndexIDMap2Template<IndexT>::add_with_idsEx(
         idx_t n,
-        const typename IndexT::component_t* x,
+        const void* x,
+        NumericType numeric_type,
         const idx_t* xids) {
     size_t prev_ntotal = this->ntotal;
-    IndexIDMapTemplate<IndexT>::add_with_ids(n, x, xids);
+    IndexIDMapTemplate<IndexT>::add_with_idsEx(n, x, numeric_type, xids);
     for (size_t i = prev_ntotal; i < this->ntotal; i++) {
         rev_map[this->id_map[i]] = i;
     }
 }
 
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::add_with_ids(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        const idx_t* xids) {
+    add_with_idsEx(
+            n,
+            static_cast<const void*>(x),
+            component_t_to_numeric<typename IndexT::component_t>(),
+            xids);
+}
+
 template <typename IndexT>
 void IndexIDMap2Template<IndexT>::check_consistency() const {
     FAISS_THROW_IF_NOT(rev_map.size() == this->id_map.size());

data/vendor/faiss/faiss/IndexIDMap.h

@@ -31,9 +31,15 @@ struct IndexIDMapTemplate : IndexT {
     /// @param xids if non-null, ids to store for the vectors (size n)
     void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
             override;
+    void add_with_idsEx(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            const idx_t* xids) override;
 
     /// this will fail. Use add_with_ids
     void add(idx_t n, const component_t* x) override;
+    void addEx(idx_t n, const void* x, NumericType numeric_type) override;
 
     void search(
             idx_t n,
@@ -42,8 +48,17 @@ struct IndexIDMapTemplate : IndexT {
             distance_t* distances,
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
+    void searchEx(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            idx_t k,
+            distance_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void train(idx_t n, const component_t* x) override;
+    void trainEx(idx_t n, const void* x, NumericType numeric_type) override;
 
     void reset() override;
 
@@ -89,6 +104,11 @@ struct IndexIDMap2Template : IndexIDMapTemplate<IndexT> {
 
     void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
             override;
+    void add_with_idsEx(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            const idx_t* xids) override;
 
     size_t remove_ids(const IDSelector& sel) override;
 

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -60,19 +60,22 @@ void Level1Quantizer::train_q1(
         MetricType metric_type) {
     size_t d = quantizer->d;
     if (quantizer->is_trained && (quantizer->ntotal == nlist)) {
-        if (verbose)
+        if (verbose) {
             printf("IVF quantizer does not need training.\n");
+        }
     } else if (quantizer_trains_alone == 1) {
-        if (verbose)
+        if (verbose) {
             printf("IVF quantizer trains alone...\n");
+        }
         quantizer->verbose = verbose;
         quantizer->train(n, x);
         FAISS_THROW_IF_NOT_MSG(
                 quantizer->ntotal == nlist,
                 "nlist not consistent with quantizer size");
     } else if (quantizer_trains_alone == 0) {
-        if (verbose)
+        if (verbose) {
             printf("Training level-1 quantizer on %zd vectors in %zdD\n", n, d);
+        }
 
         Clustering clus(d, nlist, cp);
         quantizer->reset();
@@ -158,11 +161,14 @@ IndexIVF::IndexIVF(
         size_t d,
         size_t nlist,
         size_t code_size,
-        MetricType metric)
+        MetricType metric,
+        bool own_invlists)
         : Index(d, metric),
           IndexIVFInterface(quantizer, nlist),
-          invlists(new ArrayInvertedLists(nlist, code_size)),
-          own_invlists(true),
+          invlists(
+                  own_invlists ? new ArrayInvertedLists(nlist, code_size)
+                               : nullptr),
+          own_invlists(own_invlists),
           code_size(code_size) {
     FAISS_THROW_IF_NOT(d == quantizer->d);
     is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);
@@ -230,8 +236,9 @@ void IndexIVF::add_core(
     size_t nadd = 0, nminus1 = 0;
 
     for (size_t i = 0; i < n; i++) {
-        if (coarse_idx[i] < 0)
+        if (coarse_idx[i] < 0) {
             nminus1++;
+        }
     }
 
     std::unique_ptr<uint8_t[]> flat_codes(new uint8_t[n * code_size]);
@@ -466,8 +473,9 @@ void IndexIVF::search_preassigned(
         // initialize + reorder a result heap
 
         auto init_result = [&](float* simi, idx_t* idxi) {
-            if (!do_heap_init)
+            if (!do_heap_init) {
                 return;
+            }
             if (metric_type == METRIC_INNER_PRODUCT) {
                 heap_heapify<HeapForIP>(k, simi, idxi);
             } else {
@@ -487,8 +495,9 @@ void IndexIVF::search_preassigned(
         };
 
         auto reorder_result = [&](float* simi, idx_t* idxi) {
-            if (!do_heap_init)
+            if (!do_heap_init) {
                 return;
+            }
             if (metric_type == METRIC_INNER_PRODUCT) {
                 heap_reorder<HeapForIP>(k, simi, idxi);
             } else {
@@ -804,8 +813,9 @@ void IndexIVF::range_search_preassigned(
 
         auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) {
             idx_t key = keys[i * nprobe + ik]; /* select the list  */
-            if (key < 0)
+            if (key < 0) {
                 return;
+            }
             FAISS_THROW_IF_NOT_FMT(
                     key < (idx_t)nlist,
                     "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
@@ -956,6 +966,14 @@ bool IndexIVF::check_ids_sorted() const {
     return nflip == 0;
 }
 
+void IndexIVF::decode_vectors(
+        idx_t /*n*/,
+        const uint8_t* /*codes*/,
+        const idx_t* /*list_nos*/,
+        float* /*x*/) const {
+    FAISS_THROW_MSG("decode_vectors not implemented");
+}
+
 /* standalone codec interface */
 size_t IndexIVF::sa_code_size() const {
     size_t coarse_size = coarse_code_size();

data/vendor/faiss/faiss/IndexIVF.h

@@ -210,7 +210,8 @@ struct IndexIVF : Index, IndexIVFInterface {
             size_t d,
             size_t nlist,
             size_t code_size,
-            MetricType metric = METRIC_L2);
+            MetricType metric = METRIC_L2,
+            bool own_invlists = true);
 
     void reset() override;
 
@@ -253,6 +254,20 @@ struct IndexIVF : Index, IndexIVFInterface {
             uint8_t* codes,
             bool include_listno = false) const = 0;
 
+    /** Decodes a set of vectors as they would appear in a given set of inverted
+     * lists (inverse of encode_vectors)
+     *
+     * @param codes      input codes, size n * code_size
+     * @param x          output decoded vectors
+     * @param list_nos   input listnos, size n
+     *
+     */
+    virtual void decode_vectors(
+            idx_t n,
+            const uint8_t* codes,
+            const idx_t* list_nos,
+            float* x) const;
+
     /** Add vectors that are computed with the standalone codec
      *
      * @param codes  codes to add size n * sa_code_size()

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp

@@ -28,8 +28,9 @@ IndexIVFAdditiveQuantizer::IndexIVFAdditiveQuantizer(
         Index* quantizer,
         size_t d,
         size_t nlist,
-        MetricType metric)
-        : IndexIVF(quantizer, d, nlist, 0, metric), aq(aq) {
+        MetricType metric,
+        bool own_invlists)
+        : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists), aq(aq) {
     by_residual = true;
 }
 
@@ -89,6 +90,31 @@ void IndexIVFAdditiveQuantizer::encode_vectors(
     }
 }
 
+void IndexIVFAdditiveQuantizer::decode_vectors(
+        idx_t n,
+        const uint8_t* codes,
+        const idx_t* listnos,
+        float* x) const {
+#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);
+            float* xi = x + i * d;
+            aq->decode(code, xi, 1);
+            if (by_residual) {
+                int64_t list_no = listnos[i];
+                quantizer->reconstruct(list_no, residual.data());
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] += residual[j];
+                }
+            }
+        }
+    }
+}
+
 void IndexIVFAdditiveQuantizer::sa_decode(
         idx_t n,
         const uint8_t* codes,
@@ -301,10 +327,20 @@ IndexIVFResidualQuantizer::IndexIVFResidualQuantizer(
         size_t nlist,
         const std::vector<size_t>& nbits,
         MetricType metric,
-        Search_type_t search_type)
-        : IndexIVFAdditiveQuantizer(&rq, quantizer, d, nlist, metric),
+        Search_type_t search_type,
+        bool own_invlists)
+        : IndexIVFAdditiveQuantizer(
+                  &rq,
+                  quantizer,
+                  d,
+                  nlist,
+                  metric,
+                  own_invlists),
           rq(d, nbits, search_type) {
-    code_size = invlists->code_size = rq.code_size;
+    code_size = rq.code_size;
+    if (invlists) {
+        invlists->code_size = code_size;
+    }
 }
 
 IndexIVFResidualQuantizer::IndexIVFResidualQuantizer()
@@ -317,14 +353,16 @@ IndexIVFResidualQuantizer::IndexIVFResidualQuantizer(
         size_t M,     /* number of subquantizers */
         size_t nbits, /* number of bit per subvector index */
         MetricType metric,
-        Search_type_t search_type)
+        Search_type_t search_type,
+        bool own_invlists)
         : IndexIVFResidualQuantizer(
                   quantizer,
                   d,
                   nlist,
                   std::vector<size_t>(M, nbits),
                   metric,
-                  search_type) {}
+                  search_type,
+                  own_invlists) {}
 
 IndexIVFResidualQuantizer::~IndexIVFResidualQuantizer() = default;
 
@@ -339,10 +377,20 @@ IndexIVFLocalSearchQuantizer::IndexIVFLocalSearchQuantizer(
         size_t M,     /* number of subquantizers */
         size_t nbits, /* number of bit per subvector index */
         MetricType metric,
-        Search_type_t search_type)
-        : IndexIVFAdditiveQuantizer(&lsq, quantizer, d, nlist, metric),
+        Search_type_t search_type,
+        bool own_invlists)
+        : IndexIVFAdditiveQuantizer(
+                  &lsq,
+                  quantizer,
+                  d,
+                  nlist,
+                  metric,
+                  own_invlists),
           lsq(d, M, nbits, search_type) {
-    code_size = invlists->code_size = lsq.code_size;
+    code_size = lsq.code_size;
+    if (invlists) {
+        invlists->code_size = code_size;
+    }
 }
 
 IndexIVFLocalSearchQuantizer::IndexIVFLocalSearchQuantizer()
@@ -362,10 +410,20 @@ IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer(
         size_t Msub,
         size_t nbits,
         MetricType metric,
-        Search_type_t search_type)
-        : IndexIVFAdditiveQuantizer(&prq, quantizer, d, nlist, metric),
+        Search_type_t search_type,
+        bool own_invlists)
+        : IndexIVFAdditiveQuantizer(
+                  &prq,
+                  quantizer,
+                  d,
+                  nlist,
+                  metric,
+                  own_invlists),
           prq(d, nsplits, Msub, nbits, search_type) {
-    code_size = invlists->code_size = prq.code_size;
+    code_size = prq.code_size;
+    if (invlists) {
+        invlists->code_size = code_size;
+    }
 }
 
 IndexIVFProductResidualQuantizer::IndexIVFProductResidualQuantizer()
@@ -385,10 +443,20 @@ IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer(
         size_t Msub,
         size_t nbits,
         MetricType metric,
-        Search_type_t search_type)
-        : IndexIVFAdditiveQuantizer(&plsq, quantizer, d, nlist, metric),
+        Search_type_t search_type,
+        bool own_invlists)
+        : IndexIVFAdditiveQuantizer(
+                  &plsq,
+                  quantizer,
+                  d,
+                  nlist,
+                  metric,
+                  own_invlists),
           plsq(d, nsplits, Msub, nbits, search_type) {
-    code_size = invlists->code_size = plsq.code_size;
+    code_size = plsq.code_size;
+    if (invlists) {
+        invlists->code_size = code_size;
+    }
 }
 
 IndexIVFProductLocalSearchQuantizer::IndexIVFProductLocalSearchQuantizer()

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h

@@ -35,7 +35,8 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
             Index* quantizer,
             size_t d,
             size_t nlist,
-            MetricType metric = METRIC_L2);
+            MetricType metric = METRIC_L2,
+            bool own_invlists = true);
 
     explicit IndexIVFAdditiveQuantizer(AdditiveQuantizer* aq);
 
@@ -50,6 +51,12 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
             uint8_t* codes,
             bool include_listnos = false) const override;
 
+    void decode_vectors(
+            idx_t n,
+            const uint8_t* codes,
+            const idx_t* list_nos,
+            float* x) const override;
+
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
             const IDSelector* sel,
@@ -82,7 +89,8 @@ struct IndexIVFResidualQuantizer : IndexIVFAdditiveQuantizer {
             size_t nlist,
             const std::vector<size_t>& nbits,
             MetricType metric = METRIC_L2,
-            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress,
+            bool own_invlists = true);
 
     IndexIVFResidualQuantizer(
             Index* quantizer,
@@ -91,7 +99,8 @@ struct IndexIVFResidualQuantizer : IndexIVFAdditiveQuantizer {
             size_t M,     /* number of subquantizers */
             size_t nbits, /* number of bit per subvector index */
             MetricType metric = METRIC_L2,
-            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress,
+            bool own_invlists = true);
 
     IndexIVFResidualQuantizer();
 
@@ -118,7 +127,8 @@ struct IndexIVFLocalSearchQuantizer : IndexIVFAdditiveQuantizer {
             size_t M,     /* number of subquantizers */
             size_t nbits, /* number of bit per subvector index */
             MetricType metric = METRIC_L2,
-            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress,
+            bool own_invlists = true);
 
     IndexIVFLocalSearchQuantizer();
 
@@ -147,7 +157,8 @@ struct IndexIVFProductResidualQuantizer : IndexIVFAdditiveQuantizer {
             size_t Msub,
             size_t nbits,
             MetricType metric = METRIC_L2,
-            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress,
+            bool own_invlists = true);
 
     IndexIVFProductResidualQuantizer();
 
@@ -176,7 +187,8 @@ struct IndexIVFProductLocalSearchQuantizer : IndexIVFAdditiveQuantizer {
             size_t Msub,
             size_t nbits,
             MetricType metric = METRIC_L2,
-            Search_type_t search_type = AdditiveQuantizer::ST_decompress);
+            Search_type_t search_type = AdditiveQuantizer::ST_decompress,
+            bool own_invlists = true);
 
     IndexIVFProductLocalSearchQuantizer();