faiss 0.1.0 → 0.1.1

data/vendor/faiss/tests/test_sliding_ivf.cpp

@@ -0,0 +1,240 @@
+/**
+ * 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 <cstdio>
+#include <cstdlib>
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/AutoTune.h>
+#include <faiss/index_factory.h>
+#include <faiss/clone_index.h>
+#include <faiss/IVFlib.h>
+
+using namespace faiss;
+
+typedef Index::idx_t idx_t;
+
+
+// dimension of the vectors to index
+int d = 32;
+
+// nb of training vectors
+size_t nt = 5000;
+
+// size of the database points per window step
+size_t nb = 1000;
+
+// nb of queries
+size_t nq = 200;
+
+
+int total_size = 40;
+int window_size = 10;
+
+
+
+
+
+std::vector<float> make_data(size_t n)
+{
+    std::vector <float> database (n * d);
+    for (size_t i = 0; i < n * d; i++) {
+        database[i] = drand48();
+    }
+    return database;
+}
+
+std::unique_ptr<Index> make_trained_index(const char *index_type)
+{
+    auto index = std::unique_ptr<Index>(index_factory(d, index_type));
+    auto xt = make_data(nt * d);
+    index->train(nt, xt.data());
+    ParameterSpace().set_index_parameter (index.get(), "nprobe", 4);
+    return index;
+}
+
+std::vector<idx_t> search_index(Index *index, const float *xq) {
+    int k = 10;
+    std::vector<idx_t> I(k * nq);
+    std::vector<float> D(k * nq);
+    index->search (nq, xq, k, D.data(), I.data());
+    return I;
+}
+
+
+
+
+
+/*************************************************************
+ * Test functions for a given index type
+ *************************************************************/
+
+
+// make a few slices of indexes that can be merged
+void make_index_slices (const Index* trained_index,
+                        std::vector<std::unique_ptr<Index> > & sub_indexes) {
+
+    for (int i = 0; i < total_size; i++) {
+        sub_indexes.emplace_back (clone_index (trained_index));
+
+        printf ("preparing sub-index # %d\n", i);
+
+        Index * index = sub_indexes.back().get();
+
+        auto xb = make_data(nb * d);
+        std::vector<faiss::Index::idx_t> ids (nb);
+        for (int j = 0; j < nb; j++) {
+            ids[j] = lrand48();
+        }
+        index->add_with_ids (nb, xb.data(), ids.data());
+    }
+
+}
+
+// build merged index explicitly at sliding window position i
+Index *make_merged_index(
+         const Index* trained_index,
+         const std::vector<std::unique_ptr<Index> > & sub_indexes,
+         int i) {
+
+    Index * merged_index = clone_index (trained_index);
+    for (int j = i - window_size + 1; j <= i; j++) {
+        if (j < 0 || j >= total_size) continue;
+        std::unique_ptr<Index> sub_index (
+                clone_index (sub_indexes[j].get()));
+        IndexIVF *ivf0 = ivflib::extract_index_ivf (merged_index);
+        IndexIVF *ivf1 = ivflib::extract_index_ivf (sub_index.get());
+        ivf0->merge_from (*ivf1, 0);
+        merged_index->ntotal = ivf0->ntotal;
+    }
+    return merged_index;
+}
+
+int test_sliding_window (const char *index_key) {
+
+    std::unique_ptr<Index> trained_index = make_trained_index(index_key);
+
+    // make the index slices
+    std::vector<std::unique_ptr<Index> > sub_indexes;
+
+    make_index_slices (trained_index.get(), sub_indexes);
+
+    // now slide over the windows
+    std::unique_ptr<Index> index (clone_index (trained_index.get()));
+    ivflib::SlidingIndexWindow window (index.get());
+
+    auto xq = make_data (nq * d);
+
+    for (int i = 0; i < total_size + window_size; i++) {
+
+        printf ("doing step %d / %d\n", i, total_size + window_size);
+
+        // update the index
+        window.step (i < total_size ? sub_indexes[i].get() : nullptr,
+                     i >= window_size);
+        printf ("   current n_slice = %d\n", window.n_slice);
+
+        auto new_res = search_index (index.get(), xq.data());
+
+        std::unique_ptr<Index> merged_index (
+             make_merged_index (trained_index.get(), sub_indexes, i));
+
+        auto ref_res = search_index (merged_index.get(), xq.data ());
+
+        EXPECT_EQ (ref_res.size(), new_res.size());
+
+        EXPECT_EQ (ref_res, new_res);
+    }
+    return 0;
+}
+
+
+int test_sliding_invlists (const char *index_key) {
+
+    std::unique_ptr<Index> trained_index = make_trained_index(index_key);
+
+    // make the index slices
+    std::vector<std::unique_ptr<Index> > sub_indexes;
+
+    make_index_slices (trained_index.get(), sub_indexes);
+
+    // now slide over the windows
+    std::unique_ptr<Index> index (clone_index (trained_index.get()));
+    IndexIVF * index_ivf = ivflib::extract_index_ivf (index.get());
+
+    auto xq = make_data (nq * d);
+
+    for (int i = 0; i < total_size + window_size; i++) {
+
+        printf ("doing step %d / %d\n", i, total_size + window_size);
+
+        // update the index
+        std::vector<const InvertedLists*> ils;
+        for (int j = i - window_size + 1; j <= i; j++) {
+            if (j < 0 || j >= total_size) continue;
+            ils.push_back (ivflib::extract_index_ivf (
+                      sub_indexes[j].get())->invlists);
+        }
+        if (ils.size() == 0) continue;
+
+        ConcatenatedInvertedLists *ci =
+            new ConcatenatedInvertedLists (ils.size(), ils.data());
+
+        // will be deleted by the index
+        index_ivf->replace_invlists (ci, true);
+
+        printf ("   nb invlists = %ld\n", ils.size());
+
+        auto new_res = search_index (index.get(), xq.data());
+
+        std::unique_ptr<Index> merged_index (
+             make_merged_index (trained_index.get(), sub_indexes, i));
+
+        auto ref_res = search_index (merged_index.get(), xq.data ());
+
+        EXPECT_EQ (ref_res.size(), new_res.size());
+
+        size_t ndiff = 0;
+        for (size_t j = 0; j < ref_res.size(); j++) {
+            if (ref_res[j] != new_res[j])
+                ndiff++;
+        }
+        printf("  nb differences: %ld / %ld\n",
+               ndiff, ref_res.size());
+        EXPECT_EQ (ref_res, new_res);
+    }
+    return 0;
+}
+
+
+
+
+
+/*************************************************************
+ * Test entry points
+ *************************************************************/
+
+TEST(SlidingWindow, IVFFlat) {
+    test_sliding_window ("IVF32,Flat");
+}
+
+TEST(SlidingWindow, PCAIVFFlat) {
+    test_sliding_window ("PCA24,IVF32,Flat");
+}
+
+TEST(SlidingInvlists, IVFFlat) {
+    test_sliding_invlists ("IVF32,Flat");
+}
+
+TEST(SlidingInvlists, PCAIVFFlat) {
+    test_sliding_invlists ("PCA24,IVF32,Flat");
+}

data/vendor/faiss/tests/test_threaded_index.cpp

@@ -0,0 +1,253 @@
+/**
+ * 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/impl/ThreadedIndex.h>
+#include <faiss/IndexReplicas.h>
+#include <faiss/IndexShards.h>
+
+#include <chrono>
+#include <gtest/gtest.h>
+#include <memory>
+#include <vector>
+#include <thread>
+
+namespace {
+
+struct TestException : public std::exception { };
+
+struct MockIndex : public faiss::Index {
+  explicit MockIndex(idx_t d) :
+      faiss::Index(d) {
+    resetMock();
+  }
+
+  void resetMock() {
+    flag = false;
+    nCalled = 0;
+    xCalled = nullptr;
+    kCalled = 0;
+    distancesCalled = nullptr;
+    labelsCalled = nullptr;
+  }
+
+  void add(idx_t n, const float* x) override {
+    nCalled = n;
+    xCalled = x;
+  }
+
+  void search(idx_t n,
+              const float* x,
+              idx_t k,
+              float* distances,
+              idx_t* labels) const override {
+    nCalled = n;
+    xCalled = x;
+    kCalled = k;
+    distancesCalled = distances;
+    labelsCalled = labels;
+  }
+
+  void reset() override { }
+
+  bool flag;
+
+  mutable idx_t nCalled;
+  mutable const float* xCalled;
+  mutable idx_t kCalled;
+  mutable float* distancesCalled;
+  mutable idx_t* labelsCalled;
+};
+
+template <typename IndexT>
+struct MockThreadedIndex : public faiss::ThreadedIndex<IndexT> {
+  using idx_t = faiss::Index::idx_t;
+
+  explicit MockThreadedIndex(bool threaded)
+      : faiss::ThreadedIndex<IndexT>(threaded) {
+  }
+
+  void add(idx_t, const float*) override { }
+  void search(idx_t, const float*, idx_t, float*, idx_t*) const override {}
+  void reset() override {}
+};
+
+}
+
+TEST(ThreadedIndex, SingleException) {
+  std::vector<std::unique_ptr<MockIndex>> idxs;
+
+  for (int i = 0; i < 3; ++i) {
+    idxs.emplace_back(new MockIndex(1));
+  }
+
+  auto fn =
+    [](int i, MockIndex* index) {
+      if (i == 1) {
+        throw TestException();
+      } else {
+        std::this_thread::sleep_for(std::chrono::milliseconds(i * 250));
+
+        index->flag = true;
+      }
+    };
+
+  // Try with threading and without
+  for (bool threaded : {true, false}) {
+    // clear flags
+    for (auto& idx : idxs) {
+      idx->resetMock();
+    }
+
+    MockThreadedIndex<MockIndex> ti(threaded);
+    for (auto& idx : idxs) {
+      ti.addIndex(idx.get());
+    }
+
+    // The second index should throw
+    EXPECT_THROW(ti.runOnIndex(fn), TestException);
+
+    // Index 0 and 2 should have processed
+    EXPECT_TRUE(idxs[0]->flag);
+    EXPECT_TRUE(idxs[2]->flag);
+  }
+}
+
+TEST(ThreadedIndex, MultipleException) {
+  std::vector<std::unique_ptr<MockIndex>> idxs;
+
+  for (int i = 0; i < 3; ++i) {
+    idxs.emplace_back(new MockIndex(1));
+  }
+
+  auto fn =
+    [](int i, MockIndex* index) {
+      if (i < 2) {
+        throw TestException();
+      } else {
+        std::this_thread::sleep_for(std::chrono::milliseconds(i * 250));
+
+        index->flag = true;
+      }
+    };
+
+  // Try with threading and without
+  for (bool threaded : {true, false}) {
+    // clear flags
+    for (auto& idx : idxs) {
+      idx->resetMock();
+    }
+
+    MockThreadedIndex<MockIndex> ti(threaded);
+    for (auto& idx : idxs) {
+      ti.addIndex(idx.get());
+    }
+
+    // Multiple indices threw an exception that was aggregated into a
+    // FaissException
+    EXPECT_THROW(ti.runOnIndex(fn), faiss::FaissException);
+
+    // Index 2 should have processed
+    EXPECT_TRUE(idxs[2]->flag);
+  }
+}
+
+TEST(ThreadedIndex, TestReplica) {
+  int numReplicas = 5;
+  int n = 10 * numReplicas;
+  int d = 3;
+  int k = 6;
+
+  // Try with threading and without
+  for (bool threaded : {true, false}) {
+    std::vector<std::unique_ptr<MockIndex>> idxs;
+    faiss::IndexReplicas replica(d);
+
+    for (int i = 0; i < numReplicas; ++i) {
+      idxs.emplace_back(new MockIndex(d));
+      replica.addIndex(idxs.back().get());
+    }
+
+    std::vector<float> x(n * d);
+    std::vector<float> distances(n * k);
+    std::vector<faiss::Index::idx_t> labels(n * k);
+
+    replica.add(n, x.data());
+
+    for (int i = 0; i < idxs.size(); ++i) {
+      EXPECT_EQ(idxs[i]->nCalled, n);
+      EXPECT_EQ(idxs[i]->xCalled, x.data());
+    }
+
+    for (auto& idx : idxs) {
+      idx->resetMock();
+    }
+
+    replica.search(n, x.data(), k, distances.data(), labels.data());
+
+    for (int i = 0; i < idxs.size(); ++i) {
+      auto perReplica = n / idxs.size();
+
+      EXPECT_EQ(idxs[i]->nCalled, perReplica);
+      EXPECT_EQ(idxs[i]->xCalled, x.data() + i * perReplica * d);
+      EXPECT_EQ(idxs[i]->kCalled, k);
+      EXPECT_EQ(idxs[i]->distancesCalled,
+                distances.data() + (i * perReplica) * k);
+      EXPECT_EQ(idxs[i]->labelsCalled,
+                labels.data() + (i * perReplica) * k);
+    }
+  }
+}
+
+TEST(ThreadedIndex, TestShards) {
+  int numShards = 7;
+  int d = 3;
+  int n = 10 * numShards;
+  int k = 6;
+
+  // Try with threading and without
+  for (bool threaded : {true, false}) {
+    std::vector<std::unique_ptr<MockIndex>> idxs;
+    faiss::IndexShards shards(d, threaded);
+
+    for (int i = 0; i < numShards; ++i) {
+      idxs.emplace_back(new MockIndex(d));
+      shards.addIndex(idxs.back().get());
+    }
+
+    std::vector<float> x(n * d);
+    std::vector<float> distances(n * k);
+    std::vector<faiss::Index::idx_t> labels(n * k);
+
+    shards.add(n, x.data());
+
+    for (int i = 0; i < idxs.size(); ++i) {
+      auto perShard = n / idxs.size();
+
+      EXPECT_EQ(idxs[i]->nCalled, perShard);
+      EXPECT_EQ(idxs[i]->xCalled, x.data() + i * perShard * d);
+    }
+
+    for (auto& idx : idxs) {
+      idx->resetMock();
+    }
+
+    shards.search(n, x.data(), k, distances.data(), labels.data());
+
+    for (int i = 0; i < idxs.size(); ++i) {
+      auto perShard = n / idxs.size();
+
+      EXPECT_EQ(idxs[i]->nCalled, n);
+      EXPECT_EQ(idxs[i]->xCalled, x.data());
+      EXPECT_EQ(idxs[i]->kCalled, k);
+      // There is a temporary buffer used for shards
+      EXPECT_EQ(idxs[i]->distancesCalled,
+                idxs[0]->distancesCalled + i * k * n);
+      EXPECT_EQ(idxs[i]->labelsCalled,
+                idxs[0]->labelsCalled + i * k * n);
+    }
+  }
+}