faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -9,132 +9,88 @@
 
 #include <faiss/IndexFlat.h>
 
-#include <cstring>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
 #include <faiss/utils/utils.h>
-#include <faiss/utils/Heap.h>
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/AuxIndexStructures.h>
-
+#include <cstring>
 
 namespace faiss {
 
-IndexFlat::IndexFlat (idx_t d, MetricType metric):
-            Index(d, metric)
-{
-}
-
+IndexFlat::IndexFlat(idx_t d, MetricType metric)
+        : IndexFlatCodes(sizeof(float) * d, d, metric) {}
 
+void IndexFlat::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const {
+    FAISS_THROW_IF_NOT(k > 0);
 
-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;
-}
-
-
-void IndexFlat::search (idx_t n, const float *x, idx_t k,
-                               float *distances, idx_t *labels) const
-{
     // we see the distances and labels as heaps
 
     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);
+        float_minheap_array_t res = {size_t(n), size_t(k), labels, distances};
+        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);
+        float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
+        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);
+        float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
+        knn_extra_metrics(
+                x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
     }
 }
 
-void IndexFlat::range_search (idx_t n, const float *x, float radius,
-                              RangeSearchResult *result) const
-{
+void IndexFlat::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result) const {
     switch (metric_type) {
-    case METRIC_INNER_PRODUCT:
-        range_search_inner_product (x, xb.data(), d, n, ntotal,
-                                    radius, result);
-        break;
-    case METRIC_L2:
-        range_search_L2sqr (x, xb.data(), d, n, ntotal, radius, result);
-        break;
-    default:
-        FAISS_THROW_MSG("metric type not supported");
+        case METRIC_INNER_PRODUCT:
+            range_search_inner_product(
+                    x, get_xb(), d, n, ntotal, radius, result);
+            break;
+        case METRIC_L2:
+            range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result);
+            break;
+        default:
+            FAISS_THROW_MSG("metric type not supported");
     }
 }
 
-
-void IndexFlat::compute_distance_subset (
-            idx_t n,
-            const float *x,
-            idx_t k,
-            float *distances,
-            const idx_t *labels) const
-{
+void IndexFlat::compute_distance_subset(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        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 {
     size_t d;
     Index::idx_t nb;
-    const float *q;
-    const float *b;
+    const float* q;
+    const float* b;
     size_t ndis;
 
-    float operator () (idx_t i) override {
+    float operator()(idx_t i) override {
         ndis++;
         return fvec_L2sqr(q, b + i * d, d);
     }
@@ -143,14 +99,14 @@ struct FlatL2Dis : DistanceComputer {
         return fvec_L2sqr(b + j * d, b + i * d, d);
     }
 
-    explicit FlatL2Dis(const IndexFlat& storage, const float *q = nullptr)
-        : d(storage.d),
-          nb(storage.ntotal),
-          q(q),
-          b(storage.xb.data()),
-          ndis(0) {}
+    explicit FlatL2Dis(const IndexFlat& storage, const float* q = nullptr)
+            : d(storage.d),
+              nb(storage.ntotal),
+              q(q),
+              b(storage.get_xb()),
+              ndis(0) {}
 
-    void set_query(const float *x) override {
+    void set_query(const float* x) override {
         q = x;
     }
 };
@@ -158,128 +114,106 @@ struct FlatL2Dis : DistanceComputer {
 struct FlatIPDis : DistanceComputer {
     size_t d;
     Index::idx_t nb;
-    const float *q;
-    const float *b;
+    const float* q;
+    const float* b;
     size_t ndis;
 
-    float operator () (idx_t i) override {
+    float operator()(idx_t i) override {
         ndis++;
-        return fvec_inner_product (q, b + i * d, d);
+        return fvec_inner_product(q, b + i * d, d);
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
-        return fvec_inner_product (b + j * d, b + i * d, d);
+        return fvec_inner_product(b + j * d, b + i * d, d);
     }
 
-    explicit FlatIPDis(const IndexFlat& storage, const float *q = nullptr)
-        : d(storage.d),
-          nb(storage.ntotal),
-          q(q),
-          b(storage.xb.data()),
-          ndis(0) {}
+    explicit FlatIPDis(const IndexFlat& storage, const float* q = nullptr)
+            : d(storage.d),
+              nb(storage.ntotal),
+              q(q),
+              b(storage.get_xb()),
+              ndis(0) {}
 
-    void set_query(const float *x) override {
+    void set_query(const float* x) override {
         q = x;
     }
 };
 
+} // namespace
 
-
-
-}  // namespace
-
-
-DistanceComputer * IndexFlat::get_distance_computer() const {
+DistanceComputer* IndexFlat::get_distance_computer() const {
     if (metric_type == METRIC_L2) {
         return new FlatL2Dis(*this);
     } else if (metric_type == METRIC_INNER_PRODUCT) {
         return new FlatIPDis(*this);
     } else {
-        return get_extra_distance_computer (d, metric_type, metric_arg,
-                                            ntotal, xb.data());
+        return get_extra_distance_computer(
+                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);
+void IndexFlat::reconstruct(idx_t key, float* recons) const {
+    memcpy(recons, &(codes[key * code_size]), code_size);
 }
 
-
-/* The standalone codec interface */
-size_t IndexFlat::sa_code_size () const
-{
-    return sizeof(float) * d;
-}
-
-void IndexFlat::sa_encode (idx_t n, const float *x, uint8_t *bytes) const
-{
-    memcpy (bytes, x, sizeof(float) * d * n);
+void IndexFlat::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
+    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);
+void IndexFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
+    if (n > 0) {
+        memcpy(x, bytes, sizeof(float) * d * n);
+    }
 }
 
-
-
-
-
-
 /***************************************************
  * IndexFlat1D
  ***************************************************/
 
-
-IndexFlat1D::IndexFlat1D (bool continuous_update):
-    IndexFlatL2 (1),
-    continuous_update (continuous_update)
-{
-}
+IndexFlat1D::IndexFlat1D(bool continuous_update)
+        : IndexFlatL2(1), continuous_update(continuous_update) {}
 
 /// if not continuous_update, call this between the last add and
 /// the first search
-void IndexFlat1D::update_permutation ()
-{
-    perm.resize (ntotal);
+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());
     }
 }
 
-void IndexFlat1D::add (idx_t n, const float *x)
-{
-    IndexFlatL2::add (n, x);
+void IndexFlat1D::add(idx_t n, const float* x) {
+    IndexFlatL2::add(n, x);
     if (continuous_update)
         update_permutation();
 }
 
-void IndexFlat1D::reset()
-{
+void IndexFlat1D::reset() {
     IndexFlatL2::reset();
     perm.clear();
 }
 
-void IndexFlat1D::search (
-            idx_t n,
-            const float *x,
-            idx_t k,
-            float *distances,
-            idx_t *labels) const
-{
-    FAISS_THROW_IF_NOT_MSG (perm.size() == ntotal,
-                    "Call update_permutation before search");
+void IndexFlat1D::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const {
+    FAISS_THROW_IF_NOT(k > 0);
+
+    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++) {
-
         float q = x[i]; // query
-        float *D = distances + i * k;
-        idx_t *I = labels + i * k;
+        float* D = distances + i * k;
+        idx_t* I = labels + i * k;
 
         // binary search
         idx_t i0 = 0, i1 = ntotal;
@@ -297,8 +231,10 @@ void IndexFlat1D::search (
 
         while (i0 + 1 < i1) {
             idx_t imed = (i0 + i1) / 2;
-            if (xb[perm[imed]] <= q) i0 = imed;
-            else                    i1 = imed;
+            if (xb[perm[imed]] <= q)
+                i0 = imed;
+            else
+                i1 = imed;
         }
 
         // query is between xb[perm[i0]] and xb[perm[i1]]
@@ -311,13 +247,19 @@ void IndexFlat1D::search (
             if (q - xleft < xright - q) {
                 D[wp] = q - xleft;
                 I[wp] = perm[i0];
-                i0--; wp++;
-                if (i0 < 0) { goto finish_right; }
+                i0--;
+                wp++;
+                if (i0 < 0) {
+                    goto finish_right;
+                }
             } else {
                 D[wp] = xright - q;
                 I[wp] = perm[i1];
-                i1++; wp++;
-                if (i1 >= ntotal) { goto finish_left; }
+                i1++;
+                wp++;
+                if (i1 >= ntotal) {
+                    goto finish_left;
+                }
             }
         }
         goto done;
@@ -350,11 +292,8 @@ void IndexFlat1D::search (
             }
             wp++;
         }
-    done:  ;
+    done:;
     }
-
 }
 
-
-
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlat.h

@@ -12,35 +12,26 @@
 
 #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;
-
-    explicit IndexFlat (idx_t d, MetricType metric = METRIC_L2);
-
-    void add(idx_t n, const float* x) override;
-
-    void reset() override;
+struct IndexFlat : IndexFlatCodes {
+    explicit IndexFlat(idx_t d, MetricType metric = METRIC_L2);
 
     void search(
-        idx_t n,
-        const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels) const override;
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
 
     void range_search(
-        idx_t n,
-        const float* x,
-        float radius,
-        RangeSearchResult* result) const override;
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
 
@@ -52,59 +43,52 @@ struct IndexFlat: Index {
      * @param distances
      *                corresponding output distances, size n * k
      */
-    void compute_distance_subset (
+    void compute_distance_subset(
             idx_t n,
-            const float *x,
+            const float* x,
             idx_t k,
-            float *distances,
-            const idx_t *labels) const;
+            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 () {}
+    IndexFlat() {}
 
-    DistanceComputer * get_distance_computer() const override;
+    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;
+    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;
 };
 
-
-
-struct IndexFlatIP:IndexFlat {
-    explicit IndexFlatIP (idx_t d): IndexFlat (d, METRIC_INNER_PRODUCT) {}
-    IndexFlatIP () {}
+struct IndexFlatIP : IndexFlat {
+    explicit IndexFlatIP(idx_t d) : IndexFlat(d, METRIC_INNER_PRODUCT) {}
+    IndexFlatIP() {}
 };
 
-
-struct IndexFlatL2:IndexFlat {
-    explicit IndexFlatL2 (idx_t d): IndexFlat (d, METRIC_L2) {}
-    IndexFlatL2 () {}
+struct IndexFlatL2 : IndexFlat {
+    explicit IndexFlatL2(idx_t d) : IndexFlat(d, METRIC_L2) {}
+    IndexFlatL2() {}
 };
 
-
-
 /// optimized version for 1D "vectors".
-struct IndexFlat1D:IndexFlatL2 {
+struct IndexFlat1D : IndexFlatL2 {
     bool continuous_update; ///< is the permutation updated continuously?
 
     std::vector<idx_t> perm; ///< sorted database indices
 
-    explicit IndexFlat1D (bool continuous_update=true);
+    explicit IndexFlat1D(bool continuous_update = true);
 
     /// if not continuous_update, call this between the last add and
     /// the first search
-    void update_permutation ();
+    void update_permutation();
 
     void add(idx_t n, const float* x) override;
 
@@ -112,14 +96,13 @@ struct IndexFlat1D:IndexFlatL2 {
 
     /// Warn: the distances returned are L1 not L2
     void search(
-        idx_t n,
-        const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels) const override;
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
 };
 
-
-}
+} // namespace faiss
 
 #endif

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