faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/IndexRowwiseMinMax.h

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include <faiss/Index.h>
+#include <faiss/impl/platform_macros.h>
+
+namespace faiss {
+
+/// Index wrapper that performs rowwise normalization to [0,1], preserving
+/// the coefficients. This is a vector codec index only.
+///
+/// Basically, this index performs a rowwise scaling to [0,1] of every row
+/// in an input dataset before calling subindex::train() and
+/// subindex::sa_encode(). sa_encode() call stores the scaling coefficients
+///  (scaler and minv) in the very beginning of every output code. The format:
+///     [scaler][minv][subindex::sa_encode() output]
+/// The de-scaling in sa_decode() is done using:
+///     output_rescaled = scaler * output + minv
+///
+/// An additional ::train_inplace() function is provided in order to do
+/// an inplace scaling before calling subindex::train() and, thus, avoiding
+/// the cloning of the input dataset, but modifying the input dataset because
+/// of the scaling and the scaling back. It is up to user to call
+/// this function instead of ::train()
+///
+/// Derived classes provide different data types for scaling coefficients.
+/// Currently, versions with fp16 and fp32 scaling coefficients are available.
+/// * fp16 version adds 4 extra bytes per encoded vector
+/// * fp32 version adds 8 extra bytes per encoded vector
+
+/// Provides base functions for rowwise normalizing indices.
+struct IndexRowwiseMinMaxBase : Index {
+    /// sub-index
+    Index* index;
+
+    /// whether the subindex needs to be freed in the destructor.
+    bool own_fields;
+
+    explicit IndexRowwiseMinMaxBase(Index* index);
+
+    IndexRowwiseMinMaxBase();
+    ~IndexRowwiseMinMaxBase() override;
+
+    void add(idx_t n, const float* x) override;
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void reset() override;
+
+    virtual void train_inplace(idx_t n, float* x) = 0;
+};
+
+/// Stores scaling coefficients as fp16 values.
+struct IndexRowwiseMinMaxFP16 : IndexRowwiseMinMaxBase {
+    explicit IndexRowwiseMinMaxFP16(Index* index);
+
+    IndexRowwiseMinMaxFP16();
+
+    void train(idx_t n, const float* x) override;
+    void train_inplace(idx_t n, float* x) override;
+
+    size_t sa_code_size() const override;
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+};
+
+/// Stores scaling coefficients as fp32 values.
+struct IndexRowwiseMinMax : IndexRowwiseMinMaxBase {
+    explicit IndexRowwiseMinMax(Index* index);
+
+    IndexRowwiseMinMax();
+
+    void train(idx_t n, const float* x) override;
+    void train_inplace(idx_t n, float* x) override;
+
+    size_t sa_code_size() const override;
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+};
+
+/// block size for performing sa_encode and sa_decode
+FAISS_API extern int rowwise_minmax_sa_encode_bs;
+FAISS_API extern int rowwise_minmax_sa_decode_bs;
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexScalarQuantizer.cpp

@@ -16,6 +16,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ScalarQuantizer.h>
 #include <faiss/utils/utils.h>
 
@@ -48,9 +49,11 @@ void IndexScalarQuantizer::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
-    FAISS_THROW_IF_NOT(k > 0);
+        idx_t* labels,
+        const SearchParameters* params) const {
+    const IDSelector* sel = params ? params->sel : nullptr;
 
+    FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT(is_trained);
     FAISS_THROW_IF_NOT(
             metric_type == METRIC_L2 || metric_type == METRIC_INNER_PRODUCT);
@@ -58,7 +61,8 @@ void IndexScalarQuantizer::search(
 #pragma omp parallel
     {
         InvertedListScanner* scanner =
-                sq.select_InvertedListScanner(metric_type, nullptr, true);
+                sq.select_InvertedListScanner(metric_type, nullptr, true, sel);
+
         ScopeDeleter1<InvertedListScanner> del(scanner);
         scanner->list_no = 0; // directly the list number
 
@@ -85,7 +89,8 @@ void IndexScalarQuantizer::search(
     }
 }
 
-DistanceComputer* IndexScalarQuantizer::get_distance_computer() const {
+FlatCodesDistanceComputer* IndexScalarQuantizer::get_FlatCodesDistanceComputer()
+        const {
     ScalarQuantizer::SQDistanceComputer* dc =
             sq.get_distance_computer(metric_type);
     dc->code_size = sq.code_size;
@@ -140,7 +145,7 @@ void IndexIVFScalarQuantizer::encode_vectors(
         const idx_t* list_nos,
         uint8_t* codes,
         bool include_listnos) const {
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
+    std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
     size_t coarse_size = include_listnos ? coarse_code_size() : 0;
     memset(codes, 0, (code_size + coarse_size) * n);
 
@@ -169,7 +174,7 @@ void IndexIVFScalarQuantizer::encode_vectors(
 
 void IndexIVFScalarQuantizer::sa_decode(idx_t n, const uint8_t* codes, float* x)
         const {
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
+    std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
     size_t coarse_size = coarse_code_size();
 
 #pragma omp parallel if (n > 1000)
@@ -200,7 +205,7 @@ void IndexIVFScalarQuantizer::add_core(
     FAISS_THROW_IF_NOT(is_trained);
 
     size_t nadd = 0;
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
+    std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
 
     DirectMapAdd dm_add(direct_map, n, xids);
 
@@ -241,22 +246,28 @@ void IndexIVFScalarQuantizer::add_core(
 }
 
 InvertedListScanner* IndexIVFScalarQuantizer::get_InvertedListScanner(
-        bool store_pairs) const {
+        bool store_pairs,
+        const IDSelector* sel) const {
     return sq.select_InvertedListScanner(
-            metric_type, quantizer, store_pairs, by_residual);
+            metric_type, quantizer, store_pairs, sel, by_residual);
 }
 
 void IndexIVFScalarQuantizer::reconstruct_from_offset(
         int64_t list_no,
         int64_t offset,
         float* recons) const {
-    std::vector<float> centroid(d);
-    quantizer->reconstruct(list_no, centroid.data());
-
     const uint8_t* code = invlists->get_single_code(list_no, offset);
-    sq.decode(code, recons, 1);
-    for (int i = 0; i < d; ++i) {
-        recons[i] += centroid[i];
+
+    if (by_residual) {
+        std::vector<float> centroid(d);
+        quantizer->reconstruct(list_no, centroid.data());
+
+        sq.decode(code, recons, 1);
+        for (int i = 0; i < d; ++i) {
+            recons[i] += centroid[i];
+        }
+    } else {
+        sq.decode(code, recons, 1);
     }
 }
 

data/vendor/faiss/faiss/IndexScalarQuantizer.h

@@ -20,11 +20,8 @@
 namespace faiss {
 
 /**
- * The uniform quantizer has a range [vmin, vmax]. The range can be
- * the same for all dimensions (uniform) or specific per dimension
- * (default).
+ * Flat index built on a scalar quantizer.
  */
-
 struct IndexScalarQuantizer : IndexFlatCodes {
     /// Used to encode the vectors
     ScalarQuantizer sq;
@@ -49,9 +46,10 @@ struct IndexScalarQuantizer : IndexFlatCodes {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
-    DistanceComputer* get_distance_computer() const override;
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
 
     /* standalone codec interface */
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
@@ -95,7 +93,8 @@ struct IndexIVFScalarQuantizer : IndexIVF {
             const idx_t* precomputed_idx) override;
 
     InvertedListScanner* get_InvertedListScanner(
-            bool store_pairs) const override;
+            bool store_pairs,
+            const IDSelector* sel) const override;
 
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;

data/vendor/faiss/faiss/IndexShards.cpp

@@ -288,7 +288,10 @@ void IndexShardsTemplate<IndexT>::search(
         const component_t* x,
         idx_t k,
         distance_t* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
     long nshard = this->count();

data/vendor/faiss/faiss/IndexShards.h

@@ -87,7 +87,8 @@ struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
             const component_t* x,
             idx_t k,
             distance_t* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void train(idx_t n, const component_t* x) override;
 

data/vendor/faiss/faiss/MetaIndexes.cpp

@@ -16,188 +16,12 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/WorkerThread.h>
 
 namespace faiss {
 
-namespace {} // namespace
-
-/*****************************************************
- * IndexIDMap implementation
- *******************************************************/
-
-template <typename IndexT>
-IndexIDMapTemplate<IndexT>::IndexIDMapTemplate(IndexT* index)
-        : index(index), own_fields(false) {
-    FAISS_THROW_IF_NOT_MSG(index->ntotal == 0, "index must be empty on input");
-    this->is_trained = index->is_trained;
-    this->metric_type = index->metric_type;
-    this->verbose = index->verbose;
-    this->d = index->d;
-}
-
-template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::add(
-        idx_t,
-        const typename IndexT::component_t*) {
-    FAISS_THROW_MSG(
-            "add does not make sense with IndexIDMap, "
-            "use add_with_ids");
-}
-
-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;
-}
-
-template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::reset() {
-    index->reset();
-    id_map.clear();
-    this->ntotal = 0;
-}
-
-template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::add_with_ids(
-        idx_t n,
-        const typename IndexT::component_t* x,
-        const typename IndexT::idx_t* xids) {
-    index->add(n, x);
-    for (idx_t i = 0; i < n; i++)
-        id_map.push_back(xids[i]);
-    this->ntotal = index->ntotal;
-}
-
-template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::search(
-        idx_t n,
-        const typename IndexT::component_t* x,
-        idx_t k,
-        typename IndexT::distance_t* distances,
-        typename IndexT::idx_t* labels) const {
-    index->search(n, x, k, distances, labels);
-    idx_t* li = labels;
-#pragma omp parallel for
-    for (idx_t i = 0; i < n * k; i++) {
-        li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
-    }
-}
-
-template <typename IndexT>
-void IndexIDMapTemplate<IndexT>::range_search(
-        typename IndexT::idx_t n,
-        const typename IndexT::component_t* x,
-        typename IndexT::distance_t radius,
-        RangeSearchResult* result) const {
-    index->range_search(n, x, radius, result);
-#pragma omp parallel for
-    for (idx_t i = 0; i < result->lims[result->nq]; i++) {
-        result->labels[i] = result->labels[i] < 0 ? result->labels[i]
-                                                  : id_map[result->labels[i]];
-    }
-}
-
-namespace {
-
-struct IDTranslatedSelector : IDSelector {
-    const std::vector<int64_t>& id_map;
-    const IDSelector& sel;
-    IDTranslatedSelector(
-            const std::vector<int64_t>& id_map,
-            const IDSelector& sel)
-            : id_map(id_map), sel(sel) {}
-    bool is_member(idx_t id) const override {
-        return sel.is_member(id_map[id]);
-    }
-};
-
-} // namespace
-
-template <typename IndexT>
-size_t IndexIDMapTemplate<IndexT>::remove_ids(const IDSelector& sel) {
-    // remove in sub-index first
-    IDTranslatedSelector sel2(id_map, sel);
-    size_t nremove = index->remove_ids(sel2);
-
-    int64_t j = 0;
-    for (idx_t i = 0; i < this->ntotal; i++) {
-        if (sel.is_member(id_map[i])) {
-            // remove
-        } else {
-            id_map[j] = id_map[i];
-            j++;
-        }
-    }
-    FAISS_ASSERT(j == index->ntotal);
-    this->ntotal = j;
-    id_map.resize(this->ntotal);
-    return nremove;
-}
-
-template <typename IndexT>
-IndexIDMapTemplate<IndexT>::~IndexIDMapTemplate() {
-    if (own_fields)
-        delete index;
-}
-
-/*****************************************************
- * IndexIDMap2 implementation
- *******************************************************/
-
-template <typename IndexT>
-IndexIDMap2Template<IndexT>::IndexIDMap2Template(IndexT* index)
-        : IndexIDMapTemplate<IndexT>(index) {}
-
-template <typename IndexT>
-void IndexIDMap2Template<IndexT>::add_with_ids(
-        idx_t n,
-        const typename IndexT::component_t* x,
-        const typename IndexT::idx_t* xids) {
-    size_t prev_ntotal = this->ntotal;
-    IndexIDMapTemplate<IndexT>::add_with_ids(n, x, xids);
-    for (size_t i = prev_ntotal; i < this->ntotal; i++) {
-        rev_map[this->id_map[i]] = i;
-    }
-}
-
-template <typename IndexT>
-void IndexIDMap2Template<IndexT>::construct_rev_map() {
-    rev_map.clear();
-    for (size_t i = 0; i < this->ntotal; i++) {
-        rev_map[this->id_map[i]] = i;
-    }
-}
-
-template <typename IndexT>
-size_t IndexIDMap2Template<IndexT>::remove_ids(const IDSelector& sel) {
-    // This is quite inefficient
-    size_t nremove = IndexIDMapTemplate<IndexT>::remove_ids(sel);
-    construct_rev_map();
-    return nremove;
-}
-
-template <typename IndexT>
-void IndexIDMap2Template<IndexT>::reconstruct(
-        idx_t key,
-        typename IndexT::component_t* recons) const {
-    try {
-        this->index->reconstruct(rev_map.at(key), recons);
-    } catch (const std::out_of_range& e) {
-        FAISS_THROW_FMT("key %" PRId64 " not found", key);
-    }
-}
-
-// explicit template instantiations
-
-template struct IndexIDMapTemplate<Index>;
-template struct IndexIDMapTemplate<IndexBinary>;
-template struct IndexIDMap2Template<Index>;
-template struct IndexIDMap2Template<IndexBinary>;
-
 /*****************************************************
  * IndexSplitVectors implementation
  *******************************************************/
@@ -235,7 +59,10 @@ void IndexSplitVectors::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT_MSG(k == 1, "search implemented only for k=1");
     FAISS_THROW_IF_NOT_MSG(
             sum_d == d, "not enough indexes compared to # dimensions");

data/vendor/faiss/faiss/MetaIndexes.h

@@ -11,92 +11,13 @@
 #define META_INDEXES_H
 
 #include <faiss/Index.h>
+#include <faiss/IndexIDMap.h>
 #include <faiss/IndexReplicas.h>
 #include <faiss/IndexShards.h>
-#include <unordered_map>
 #include <vector>
 
 namespace faiss {
 
-/** Index that translates search results to ids */
-template <typename IndexT>
-struct IndexIDMapTemplate : IndexT {
-    using idx_t = typename IndexT::idx_t;
-    using component_t = typename IndexT::component_t;
-    using distance_t = typename IndexT::distance_t;
-
-    IndexT* index;   ///! the sub-index
-    bool own_fields; ///! whether pointers are deleted in destructo
-    std::vector<idx_t> id_map;
-
-    explicit IndexIDMapTemplate(IndexT* index);
-
-    /// @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;
-
-    /// this will fail. Use add_with_ids
-    void add(idx_t n, const component_t* x) override;
-
-    void search(
-            idx_t n,
-            const component_t* x,
-            idx_t k,
-            distance_t* distances,
-            idx_t* labels) const override;
-
-    void train(idx_t n, const component_t* x) override;
-
-    void reset() override;
-
-    /// remove ids adapted to IndexFlat
-    size_t remove_ids(const IDSelector& sel) override;
-
-    void range_search(
-            idx_t n,
-            const component_t* x,
-            distance_t radius,
-            RangeSearchResult* result) const override;
-
-    ~IndexIDMapTemplate() override;
-    IndexIDMapTemplate() {
-        own_fields = false;
-        index = nullptr;
-    }
-};
-
-using IndexIDMap = IndexIDMapTemplate<Index>;
-using IndexBinaryIDMap = IndexIDMapTemplate<IndexBinary>;
-
-/** same as IndexIDMap but also provides an efficient reconstruction
- *  implementation via a 2-way index */
-template <typename IndexT>
-struct IndexIDMap2Template : IndexIDMapTemplate<IndexT> {
-    using idx_t = typename IndexT::idx_t;
-    using component_t = typename IndexT::component_t;
-    using distance_t = typename IndexT::distance_t;
-
-    std::unordered_map<idx_t, idx_t> rev_map;
-
-    explicit IndexIDMap2Template(IndexT* index);
-
-    /// make the rev_map from scratch
-    void construct_rev_map();
-
-    void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
-            override;
-
-    size_t remove_ids(const IDSelector& sel) override;
-
-    void reconstruct(idx_t key, component_t* recons) const override;
-
-    ~IndexIDMap2Template() override {}
-    IndexIDMap2Template() {}
-};
-
-using IndexIDMap2 = IndexIDMap2Template<Index>;
-using IndexBinaryIDMap2 = IndexIDMap2Template<IndexBinary>;
-
 /** splits input vectors in segments and assigns each segment to a sub-index
  * used to distribute a MultiIndexQuantizer
  */
@@ -118,7 +39,8 @@ struct IndexSplitVectors : Index {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void train(idx_t n, const float* x) override;
 

data/vendor/faiss/faiss/VectorTransform.cpp

@@ -149,6 +149,10 @@ void VectorTransform::reverse_transform(idx_t, const float*, float*) const {
     FAISS_THROW_MSG("reverse transform not implemented");
 }
 
+void VectorTransform::check_identical(const VectorTransform& other) const {
+    FAISS_THROW_IF_NOT(other.d_in == d_in && other.d_in == d_in);
+}
+
 /*********************************************
  * LinearTransform
  *********************************************/
@@ -308,6 +312,13 @@ void LinearTransform::print_if_verbose(
     printf("]\n");
 }
 
+void LinearTransform::check_identical(const VectorTransform& other_in) const {
+    VectorTransform::check_identical(other_in);
+    auto other = dynamic_cast<const LinearTransform*>(&other_in);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->A == A && other->b == b);
+}
+
 /*********************************************
  * RandomRotationMatrix
  *********************************************/
@@ -966,6 +977,14 @@ void ITQTransform::apply_noalloc(Index::idx_t n, const float* x, float* xt)
     pca_then_itq.apply_noalloc(n, x_norm.get(), xt);
 }
 
+void ITQTransform::check_identical(const VectorTransform& other_in) const {
+    VectorTransform::check_identical(other_in);
+    auto other = dynamic_cast<const ITQTransform*>(&other_in);
+    FAISS_THROW_IF_NOT(other);
+    pca_then_itq.check_identical(other->pca_then_itq);
+    FAISS_THROW_IF_NOT(other->mean == mean);
+}
+
 /*********************************************
  * OPQMatrix
  *********************************************/
@@ -1226,6 +1245,14 @@ void NormalizationTransform::reverse_transform(
     memcpy(x, xt, sizeof(xt[0]) * n * d_in);
 }
 
+void NormalizationTransform::check_identical(
+        const VectorTransform& other_in) const {
+    VectorTransform::check_identical(other_in);
+    auto other = dynamic_cast<const NormalizationTransform*>(&other_in);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->norm == norm);
+}
+
 /*********************************************
  * CenteringTransform
  *********************************************/
@@ -1271,6 +1298,14 @@ void CenteringTransform::reverse_transform(idx_t n, const float* xt, float* x)
     }
 }
 
+void CenteringTransform::check_identical(
+        const VectorTransform& other_in) const {
+    VectorTransform::check_identical(other_in);
+    auto other = dynamic_cast<const CenteringTransform*>(&other_in);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->mean == mean);
+}
+
 /*********************************************
  * RemapDimensionsTransform
  *********************************************/
@@ -1335,3 +1370,11 @@ void RemapDimensionsTransform::reverse_transform(
         xt += d_out;
     }
 }
+
+void RemapDimensionsTransform::check_identical(
+        const VectorTransform& other_in) const {
+    VectorTransform::check_identical(other_in);
+    auto other = dynamic_cast<const RemapDimensionsTransform*>(&other_in);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->map == map);
+}