faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -23,6 +23,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 
 namespace faiss {
 
@@ -107,8 +108,15 @@ void Level1Quantizer::train_q1(
         } else {
             clus.train(n, x, *clustering_index);
         }
-        if (verbose)
+        if (verbose) {
             printf("Adding centroids to quantizer\n");
+        }
+        if (!quantizer->is_trained) {
+            if (verbose) {
+                printf("But training it first on centroids table...\n");
+            }
+            quantizer->train(nlist, clus.centroids.data());
+        }
         quantizer->add(nlist, clus.centroids.data());
     }
 }
@@ -190,6 +198,20 @@ void IndexIVF::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
     add_core(n, x, xids, coarse_idx.get());
 }
 
+void IndexIVF::add_sa_codes(idx_t n, const uint8_t* codes, const idx_t* xids) {
+    size_t coarse_size = coarse_code_size();
+    DirectMapAdd dm_adder(direct_map, n, xids);
+
+    for (idx_t i = 0; i < n; i++) {
+        const uint8_t* code = codes + (code_size + coarse_size) * i;
+        idx_t list_no = decode_listno(code);
+        idx_t id = xids ? xids[i] : ntotal + i;
+        size_t ofs = invlists->add_entry(list_no, id, code + coarse_size);
+        dm_adder.add(i, list_no, ofs);
+    }
+    ntotal += n;
+}
+
 void IndexIVF::add_core(
         idx_t n,
         const float* x,
@@ -282,14 +304,20 @@ void IndexIVF::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
-
-    const size_t nprobe = std::min(nlist, this->nprobe);
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
     // search function for a subset of queries
-    auto sub_search_func = [this, k, nprobe](
+    auto sub_search_func = [this, k, nprobe, params](
                                    idx_t n,
                                    const float* x,
                                    float* distances,
@@ -299,7 +327,13 @@ void IndexIVF::search(
         std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
 
         double t0 = getmillisecs();
-        quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
+        quantizer->search(
+                n,
+                x,
+                nprobe,
+                coarse_dis.get(),
+                idx.get(),
+                params ? params->quantizer_params : nullptr);
 
         double t1 = getmillisecs();
         invlists->prefetch_lists(idx.get(), n * nprobe);
@@ -313,7 +347,7 @@ void IndexIVF::search(
                 distances,
                 labels,
                 false,
-                nullptr,
+                params,
                 ivf_stats);
         double t2 = getmillisecs();
         ivf_stats->quantization_time += t1 - t0;
@@ -379,6 +413,19 @@ void IndexIVF::search_preassigned(
     FAISS_THROW_IF_NOT(nprobe > 0);
 
     idx_t max_codes = params ? params->max_codes : this->max_codes;
+    IDSelector* sel = params ? params->sel : nullptr;
+    const IDSelectorRange* selr = dynamic_cast<const IDSelectorRange*>(sel);
+    if (selr) {
+        if (selr->assume_sorted) {
+            sel = nullptr; // use special IDSelectorRange processing
+        } else {
+            selr = nullptr; // use generic processing
+        }
+    }
+
+    FAISS_THROW_IF_NOT_MSG(
+            !(sel && store_pairs),
+            "selector and store_pairs cannot be combined");
 
     size_t nlistv = 0, ndis = 0, nheap = 0;
 
@@ -400,7 +447,8 @@ void IndexIVF::search_preassigned(
 
 #pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis, nheap)
     {
-        InvertedListScanner* scanner = get_InvertedListScanner(store_pairs);
+        InvertedListScanner* scanner =
+                get_InvertedListScanner(store_pairs, sel);
         ScopeDeleter1<InvertedListScanner> del(scanner);
 
         /*****************************************************
@@ -471,6 +519,7 @@ void IndexIVF::search_preassigned(
 
             try {
                 InvertedLists::ScopedCodes scodes(invlists, key);
+                const uint8_t* codes = scodes.get();
 
                 std::unique_ptr<InvertedLists::ScopedIds> sids;
                 const Index::idx_t* ids = nullptr;
@@ -480,8 +529,20 @@ void IndexIVF::search_preassigned(
                     ids = sids->get();
                 }
 
+                if (selr) { // IDSelectorRange
+                    // restrict search to a section of the inverted list
+                    size_t jmin, jmax;
+                    selr->find_sorted_ids_bounds(list_size, ids, &jmin, &jmax);
+                    list_size = jmax - jmin;
+                    if (list_size == 0) {
+                        return (size_t)0;
+                    }
+                    codes += jmin * code_size;
+                    ids += jmin;
+                }
+
                 nheap += scanner->scan_codes(
-                        list_size, scodes.get(), ids, simi, idxi, k);
+                        list_size, codes, ids, simi, idxi, k);
 
             } catch (const std::exception& e) {
                 std::lock_guard<std::mutex> lock(exception_mutex);
@@ -630,13 +691,23 @@ void IndexIVF::range_search(
         idx_t nx,
         const float* x,
         float radius,
-        RangeSearchResult* result) const {
-    const size_t nprobe = std::min(nlist, this->nprobe);
+        RangeSearchResult* result,
+        const SearchParameters* params_in) const {
+    const IVFSearchParameters* params = nullptr;
+    const SearchParameters* quantizer_params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+        quantizer_params = params->quantizer_params;
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
     std::unique_ptr<idx_t[]> keys(new idx_t[nx * nprobe]);
     std::unique_ptr<float[]> coarse_dis(new float[nx * nprobe]);
 
     double t0 = getmillisecs();
-    quantizer->search(nx, x, nprobe, coarse_dis.get(), keys.get());
+    quantizer->search(
+            nx, x, nprobe, coarse_dis.get(), keys.get(), quantizer_params);
     indexIVF_stats.quantization_time += getmillisecs() - t0;
 
     t0 = getmillisecs();
@@ -650,7 +721,7 @@ void IndexIVF::range_search(
             coarse_dis.get(),
             result,
             false,
-            nullptr,
+            params,
             &indexIVF_stats);
 
     indexIVF_stats.search_time += getmillisecs() - t0;
@@ -668,7 +739,10 @@ void IndexIVF::range_search_preassigned(
         IndexIVFStats* stats) const {
     idx_t nprobe = params ? params->nprobe : this->nprobe;
     nprobe = std::min((idx_t)nlist, nprobe);
+    FAISS_THROW_IF_NOT(nprobe > 0);
+
     idx_t max_codes = params ? params->max_codes : this->max_codes;
+    IDSelector* sel = params ? params->sel : nullptr;
 
     size_t nlistv = 0, ndis = 0;
 
@@ -690,7 +764,7 @@ void IndexIVF::range_search_preassigned(
     {
         RangeSearchPartialResult pres(result);
         std::unique_ptr<InvertedListScanner> scanner(
-                get_InvertedListScanner(store_pairs));
+                get_InvertedListScanner(store_pairs, sel));
         FAISS_THROW_IF_NOT(scanner.get());
         all_pres[omp_get_thread_num()] = &pres;
 
@@ -753,7 +827,6 @@ void IndexIVF::range_search_preassigned(
                 }
             }
         } else if (parallel_mode == 2) {
-            std::vector<RangeQueryResult*> all_qres(nx);
             RangeQueryResult* qres = nullptr;
 
 #pragma omp for schedule(dynamic)
@@ -761,7 +834,6 @@ void IndexIVF::range_search_preassigned(
                 idx_t i = iik / (idx_t)nprobe;
                 idx_t ik = iik % (idx_t)nprobe;
                 if (qres == nullptr || qres->qno != i) {
-                    FAISS_ASSERT(!qres || i > qres->qno);
                     qres = &pres.new_result(i);
                     scanner->set_query(x + i * d);
                 }
@@ -797,7 +869,8 @@ void IndexIVF::range_search_preassigned(
 }
 
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
-        bool /*store_pairs*/) const {
+        bool /*store_pairs*/,
+        const IDSelector* /* sel */) const {
     return nullptr;
 }
 
@@ -825,6 +898,21 @@ void IndexIVF::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
     }
 }
 
+bool IndexIVF::check_ids_sorted() const {
+    size_t nflip = 0;
+
+    for (size_t i = 0; i < nlist; i++) {
+        size_t list_size = invlists->list_size(i);
+        InvertedLists::ScopedIds ids(invlists, i);
+        for (size_t j = 0; j + 1 < list_size; j++) {
+            if (ids[j + 1] < ids[j]) {
+                nflip++;
+            }
+        }
+    }
+    return nflip == 0;
+}
+
 /* standalone codec interface */
 size_t IndexIVF::sa_code_size() const {
     size_t coarse_size = coarse_code_size();
@@ -844,10 +932,15 @@ void IndexIVF::search_and_reconstruct(
         idx_t k,
         float* distances,
         idx_t* labels,
-        float* recons) const {
-    FAISS_THROW_IF_NOT(k > 0);
-
-    const size_t nprobe = std::min(nlist, this->nprobe);
+        float* recons,
+        const SearchParameters* params_in) const {
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
     idx_t* idx = new idx_t[n * nprobe];
@@ -869,7 +962,8 @@ void IndexIVF::search_and_reconstruct(
             coarse_dis,
             distances,
             labels,
-            true /* store_pairs */);
+            true /* store_pairs */,
+            params);
     for (idx_t i = 0; i < n; ++i) {
         for (idx_t j = 0; j < k; ++j) {
             idx_t ij = i * k + j;
@@ -955,26 +1049,41 @@ void IndexIVF::train_residual(idx_t /*n*/, const float* /*x*/) {
     // does nothing by default
 }
 
-void IndexIVF::check_compatible_for_merge(const IndexIVF& other) const {
+bool check_compatible_for_merge_expensive_check = true;
+
+void IndexIVF::check_compatible_for_merge(const Index& otherIndex) const {
     // minimal sanity checks
-    FAISS_THROW_IF_NOT(other.d == d);
-    FAISS_THROW_IF_NOT(other.nlist == nlist);
-    FAISS_THROW_IF_NOT(other.code_size == code_size);
+    const IndexIVF* other = dynamic_cast<const IndexIVF*>(&otherIndex);
+    FAISS_THROW_IF_NOT(other);
+    FAISS_THROW_IF_NOT(other->d == d);
+    FAISS_THROW_IF_NOT(other->nlist == nlist);
+    FAISS_THROW_IF_NOT(quantizer->ntotal == other->quantizer->ntotal);
+    FAISS_THROW_IF_NOT(other->code_size == code_size);
     FAISS_THROW_IF_NOT_MSG(
-            typeid(*this) == typeid(other),
+            typeid(*this) == typeid(*other),
             "can only merge indexes of the same type");
     FAISS_THROW_IF_NOT_MSG(
-            this->direct_map.no() && other.direct_map.no(),
+            this->direct_map.no() && other->direct_map.no(),
             "merge direct_map not implemented");
-}
 
-void IndexIVF::merge_from(IndexIVF& other, idx_t add_id) {
-    check_compatible_for_merge(other);
+    if (check_compatible_for_merge_expensive_check) {
+        std::vector<float> v(d), v2(d);
+        for (size_t i = 0; i < nlist; i++) {
+            quantizer->reconstruct(i, v.data());
+            other->quantizer->reconstruct(i, v2.data());
+            FAISS_THROW_IF_NOT_MSG(
+                    v == v2, "coarse quantizers should be the same");
+        }
+    }
+}
 
-    invlists->merge_from(other.invlists, add_id);
+void IndexIVF::merge_from(Index& otherIndex, idx_t add_id) {
+    check_compatible_for_merge(otherIndex);
+    IndexIVF* other = static_cast<IndexIVF*>(&otherIndex);
+    invlists->merge_from(other->invlists, add_id);
 
-    ntotal += other.ntotal;
-    other.ntotal = 0;
+    ntotal += other->ntotal;
+    other->ntotal = 0;
 }
 
 void IndexIVF::replace_invlists(InvertedLists* il, bool own) {
@@ -1068,6 +1177,10 @@ IndexIVF::~IndexIVF() {
     }
 }
 
+/*************************************************************************
+ * IndexIVFStats
+ *************************************************************************/
+
 void IndexIVFStats::reset() {
     memset((void*)this, 0, sizeof(*this));
 }
@@ -1083,13 +1196,60 @@ void IndexIVFStats::add(const IndexIVFStats& other) {
 
 IndexIVFStats indexIVF_stats;
 
+/*************************************************************************
+ * InvertedListScanner
+ *************************************************************************/
+
+size_t InvertedListScanner::scan_codes(
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        float* simi,
+        idx_t* idxi,
+        size_t k) const {
+    size_t nup = 0;
+
+    if (!keep_max) {
+        for (size_t j = 0; j < list_size; j++) {
+            float dis = distance_to_code(codes);
+            if (dis < simi[0]) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                maxheap_replace_top(k, simi, idxi, dis, id);
+                nup++;
+            }
+            codes += code_size;
+        }
+    } else {
+        for (size_t j = 0; j < list_size; j++) {
+            float dis = distance_to_code(codes);
+            if (dis > simi[0]) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                minheap_replace_top(k, simi, idxi, dis, id);
+                nup++;
+            }
+            codes += code_size;
+        }
+    }
+    return nup;
+}
+
 void InvertedListScanner::scan_codes_range(
-        size_t,
-        const uint8_t*,
-        const idx_t*,
-        float,
-        RangeQueryResult&) const {
-    FAISS_THROW_MSG("scan_codes_range not implemented");
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        float radius,
+        RangeQueryResult& res) const {
+    for (size_t j = 0; j < list_size; j++) {
+        float dis = distance_to_code(codes);
+        bool keep = !keep_max
+                ? dis < radius
+                : dis > radius; // TODO templatize to remove this test
+        if (keep) {
+            int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+            res.add(dis, id);
+        }
+        codes += code_size;
+    }
 }
 
 } // namespace faiss

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>
@@ -38,7 +40,7 @@ struct Level1Quantizer {
      * = 2: kmeans training on a flat index + add the centroids to the quantizer
      */
     char quantizer_trains_alone;
-    bool own_fields; ///< whether object owns the quantizer
+    bool own_fields; ///< whether object owns the quantizer (false by default)
 
     ClusteringParameters cp; ///< to override default clustering params
     Index* clustering_index; ///< to override index used during clustering
@@ -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;
 
@@ -121,8 +128,7 @@ struct IndexIVF : Index, Level1Quantizer {
 
     /** The Inverted file takes a quantizer (an Index) on input,
      * which implements the function mapping a vector to a list
-     * identifier. The pointer is borrowed: the quantizer should not
-     * be deleted while the IndexIVF is in use.
+     * identifier.
      */
     IndexIVF(
             Index* quantizer,
@@ -171,6 +177,13 @@ struct IndexIVF : Index, Level1Quantizer {
             uint8_t* codes,
             bool include_listno = false) const = 0;
 
+    /** Add vectors that are computed with the standalone codec
+     *
+     * @param codes  codes to add size n * sa_code_size()
+     * @param xids   corresponding ids, size n
+     */
+    void add_sa_codes(idx_t n, const uint8_t* codes, const idx_t* xids);
+
     /// Sub-classes that encode the residuals can train their encoders here
     /// does nothing by default
     virtual void train_residual(idx_t n, const float* x);
@@ -212,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,
@@ -231,9 +246,13 @@ struct IndexIVF : Index, Level1Quantizer {
             const IVFSearchParameters* params = nullptr,
             IndexIVFStats* stats = nullptr) const;
 
-    /// get a scanner for this index (store_pairs means ignore labels)
+    /** Get a scanner for this index (store_pairs means ignore labels)
+     *
+     * 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
      */
@@ -275,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.
@@ -293,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
      *
@@ -322,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,
@@ -351,6 +368,22 @@ struct RangeQueryResult;
 struct InvertedListScanner {
     using idx_t = Index::idx_t;
 
+    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;
+
+    /// 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;
+
     /// from now on we handle this query.
     virtual void set_query(const float* query_vector) = 0;
 
@@ -361,7 +394,8 @@ struct InvertedListScanner {
     virtual float distance_to_code(const uint8_t* code) const = 0;
 
     /** scan a set of codes, compute distances to current query and
-     * update heap of results if necessary.
+     * update heap of results if necessary. Default implemetation
+     * calls distance_to_code.
      *
      * @param n      number of codes to scan
      * @param codes  codes to scan (n * code_size)
@@ -377,7 +411,7 @@ struct InvertedListScanner {
             const idx_t* ids,
             float* distances,
             idx_t* labels,
-            size_t k) const = 0;
+            size_t k) const;
 
     /** scan a set of codes, compute distances to current query and
      * update results if distances are below radius
@@ -393,10 +427,13 @@ 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
-    size_t ndis;              // nb of distancs computed
+    size_t ndis;              // nb of distances computed
     size_t nheap_updates;     // nb of times the heap was updated
     double quantization_time; // time spent quantizing vectors (in ms)
     double search_time;       // time spent searching lists (in ms)