faiss 0.2.3 → 0.2.4

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -19,17 +19,8 @@
 
 namespace faiss {
 
-IndexFlat::IndexFlat(idx_t d, MetricType metric) : Index(d, metric) {}
-
-void IndexFlat::add(idx_t n, const float* x) {
-    xb.insert(xb.end(), x, x + n * d);
-    ntotal += n;
-}
-
-void IndexFlat::reset() {
-    xb.clear();
-    ntotal = 0;
-}
+IndexFlat::IndexFlat(idx_t d, MetricType metric)
+        : IndexFlatCodes(sizeof(float) * d, d, metric) {}
 
 void IndexFlat::search(
         idx_t n,
@@ -43,14 +34,14 @@ void IndexFlat::search(
 
     if (metric_type == METRIC_INNER_PRODUCT) {
         float_minheap_array_t res = {size_t(n), size_t(k), labels, distances};
-        knn_inner_product(x, xb.data(), d, n, ntotal, &res);
+        knn_inner_product(x, get_xb(), d, n, ntotal, &res);
     } else if (metric_type == METRIC_L2) {
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
-        knn_L2sqr(x, xb.data(), d, n, ntotal, &res);
+        knn_L2sqr(x, get_xb(), d, n, ntotal, &res);
     } else {
         float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
         knn_extra_metrics(
-                x, xb.data(), d, n, ntotal, metric_type, metric_arg, &res);
+                x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
     }
 }
 
@@ -62,10 +53,10 @@ void IndexFlat::range_search(
     switch (metric_type) {
         case METRIC_INNER_PRODUCT:
             range_search_inner_product(
-                    x, xb.data(), d, n, ntotal, radius, result);
+                    x, get_xb(), d, n, ntotal, radius, result);
             break;
         case METRIC_L2:
-            range_search_L2sqr(x, xb.data(), d, n, ntotal, radius, result);
+            range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result);
             break;
         default:
             FAISS_THROW_MSG("metric type not supported");
@@ -80,37 +71,16 @@ void IndexFlat::compute_distance_subset(
         const idx_t* labels) const {
     switch (metric_type) {
         case METRIC_INNER_PRODUCT:
-            fvec_inner_products_by_idx(
-                    distances, x, xb.data(), labels, d, n, k);
+            fvec_inner_products_by_idx(distances, x, get_xb(), labels, d, n, k);
             break;
         case METRIC_L2:
-            fvec_L2sqr_by_idx(distances, x, xb.data(), labels, d, n, k);
+            fvec_L2sqr_by_idx(distances, x, get_xb(), labels, d, n, k);
             break;
         default:
             FAISS_THROW_MSG("metric type not supported");
     }
 }
 
-size_t IndexFlat::remove_ids(const IDSelector& sel) {
-    idx_t j = 0;
-    for (idx_t i = 0; i < ntotal; i++) {
-        if (sel.is_member(i)) {
-            // should be removed
-        } else {
-            if (i > j) {
-                memmove(&xb[d * j], &xb[d * i], sizeof(xb[0]) * d);
-            }
-            j++;
-        }
-    }
-    size_t nremove = ntotal - j;
-    if (nremove > 0) {
-        ntotal = j;
-        xb.resize(ntotal * d);
-    }
-    return nremove;
-}
-
 namespace {
 
 struct FlatL2Dis : DistanceComputer {
@@ -133,7 +103,7 @@ struct FlatL2Dis : DistanceComputer {
             : d(storage.d),
               nb(storage.ntotal),
               q(q),
-              b(storage.xb.data()),
+              b(storage.get_xb()),
               ndis(0) {}
 
     void set_query(const float* x) override {
@@ -161,7 +131,7 @@ struct FlatIPDis : DistanceComputer {
             : d(storage.d),
               nb(storage.ntotal),
               q(q),
-              b(storage.xb.data()),
+              b(storage.get_xb()),
               ndis(0) {}
 
     void set_query(const float* x) override {
@@ -178,25 +148,24 @@ DistanceComputer* IndexFlat::get_distance_computer() const {
         return new FlatIPDis(*this);
     } else {
         return get_extra_distance_computer(
-                d, metric_type, metric_arg, ntotal, xb.data());
+                d, metric_type, metric_arg, ntotal, get_xb());
     }
 }
 
 void IndexFlat::reconstruct(idx_t key, float* recons) const {
-    memcpy(recons, &(xb[key * d]), sizeof(*recons) * d);
-}
-
-/* The standalone codec interface */
-size_t IndexFlat::sa_code_size() const {
-    return sizeof(float) * d;
+    memcpy(recons, &(codes[key * code_size]), code_size);
 }
 
 void IndexFlat::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
-    memcpy(bytes, x, sizeof(float) * d * n);
+    if (n > 0) {
+        memcpy(bytes, x, sizeof(float) * d * n);
+    }
 }
 
 void IndexFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
-    memcpy(x, bytes, sizeof(float) * d * n);
+    if (n > 0) {
+        memcpy(x, bytes, sizeof(float) * d * n);
+    }
 }
 
 /***************************************************
@@ -211,9 +180,9 @@ IndexFlat1D::IndexFlat1D(bool continuous_update)
 void IndexFlat1D::update_permutation() {
     perm.resize(ntotal);
     if (ntotal < 1000000) {
-        fvec_argsort(ntotal, xb.data(), (size_t*)perm.data());
+        fvec_argsort(ntotal, get_xb(), (size_t*)perm.data());
     } else {
-        fvec_argsort_parallel(ntotal, xb.data(), (size_t*)perm.data());
+        fvec_argsort_parallel(ntotal, get_xb(), (size_t*)perm.data());
     }
 }
 
@@ -238,6 +207,7 @@ void IndexFlat1D::search(
 
     FAISS_THROW_IF_NOT_MSG(
             perm.size() == ntotal, "Call update_permutation before search");
+    const float* xb = get_xb();
 
 #pragma omp parallel for
     for (idx_t i = 0; i < n; i++) {

data/vendor/faiss/faiss/IndexFlat.h

@@ -12,21 +12,14 @@
 
 #include <vector>
 
-#include <faiss/Index.h>
+#include <faiss/IndexFlatCodes.h>
 
 namespace faiss {
 
 /** Index that stores the full vectors and performs exhaustive search */
-struct IndexFlat : Index {
-    /// database vectors, size ntotal * d
-    std::vector<float> xb;
-
+struct IndexFlat : IndexFlatCodes {
     explicit IndexFlat(idx_t d, MetricType metric = METRIC_L2);
 
-    void add(idx_t n, const float* x) override;
-
-    void reset() override;
-
     void search(
             idx_t n,
             const float* x,
@@ -57,18 +50,19 @@ struct IndexFlat : Index {
             float* distances,
             const idx_t* labels) const;
 
-    /** remove some ids. NB that Because of the structure of the
-     * indexing structure, the semantics of this operation are
-     * different from the usual ones: the new ids are shifted */
-    size_t remove_ids(const IDSelector& sel) override;
+    // get pointer to the floating point data
+    float* get_xb() {
+        return (float*)codes.data();
+    }
+    const float* get_xb() const {
+        return (const float*)codes.data();
+    }
 
     IndexFlat() {}
 
     DistanceComputer* get_distance_computer() const override;
 
     /* The stanadlone codec interface (just memcopies in this case) */
-    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;

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -0,0 +1,67 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/IndexFlatCodes.h>
+
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss {
+
+IndexFlatCodes::IndexFlatCodes(size_t code_size, idx_t d, MetricType metric)
+        : Index(d, metric), code_size(code_size) {}
+
+IndexFlatCodes::IndexFlatCodes() : code_size(0) {}
+
+void IndexFlatCodes::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT(is_trained);
+    codes.resize((ntotal + n) * code_size);
+    sa_encode(n, x, &codes[ntotal * code_size]);
+    ntotal += n;
+}
+
+void IndexFlatCodes::reset() {
+    codes.clear();
+    ntotal = 0;
+}
+
+size_t IndexFlatCodes::sa_code_size() const {
+    return code_size;
+}
+
+size_t IndexFlatCodes::remove_ids(const IDSelector& sel) {
+    idx_t j = 0;
+    for (idx_t i = 0; i < ntotal; i++) {
+        if (sel.is_member(i)) {
+            // should be removed
+        } else {
+            if (i > j) {
+                memmove(&codes[code_size * j],
+                        &codes[code_size * i],
+                        code_size);
+            }
+            j++;
+        }
+    }
+    size_t nremove = ntotal - j;
+    if (nremove > 0) {
+        ntotal = j;
+        codes.resize(ntotal * code_size);
+    }
+    return nremove;
+}
+
+void IndexFlatCodes::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
+    FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
+    sa_decode(ni, codes.data() + i0 * code_size, recons);
+}
+
+void IndexFlatCodes::reconstruct(idx_t key, float* recons) const {
+    reconstruct_n(key, 1, recons);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.h

@@ -0,0 +1,47 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#pragma once
+
+#include <faiss/Index.h>
+#include <vector>
+
+namespace faiss {
+
+/** Index that encodes all vectors as fixed-size codes (size code_size). Storage
+ * is in the codes vector */
+struct IndexFlatCodes : Index {
+    size_t code_size;
+
+    /// encoded dataset, size ntotal * code_size
+    std::vector<uint8_t> codes;
+
+    IndexFlatCodes();
+
+    IndexFlatCodes(size_t code_size, idx_t d, MetricType metric = METRIC_L2);
+
+    /// default add uses sa_encode
+    void add(idx_t n, const float* x) override;
+
+    void reset() override;
+
+    /// reconstruction using the codec interface
+    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    size_t sa_code_size() const override;
+
+    /** remove some ids. NB that Because of the structure of the
+     * indexing structure, the semantics of this operation are
+     * different from the usual ones: the new ids are shifted */
+    size_t remove_ids(const IDSelector& sel) override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -107,8 +107,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 +197,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,
@@ -1068,6 +1089,10 @@ IndexIVF::~IndexIVF() {
     }
 }
 
+/*************************************************************************
+ * IndexIVFStats
+ *************************************************************************/
+
 void IndexIVFStats::reset() {
     memset((void*)this, 0, sizeof(*this));
 }
@@ -1083,13 +1108,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

@@ -38,7 +38,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
@@ -121,8 +121,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 +170,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);
@@ -231,7 +237,10 @@ 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;
 
@@ -351,6 +360,14 @@ 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 = false;
+
+    /// 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 +378,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 +395,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
@@ -396,7 +414,7 @@ struct InvertedListScanner {
 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)