12 #include <gtest/gtest.h>
14 #include <faiss/IndexIVFFlat.h>
15 #include <faiss/index_io.h>
16 #include <faiss/AuxIndexStructures.h>
17 #include <faiss/AutoTune.h>
18 #include <faiss/VectorTransform.h>
19 #include <faiss/utils.h>
20 #include <faiss/IVFlib.h>
31 using namespace faiss;
36 std::vector<float> get_data (
size_t nb,
int seed) {
37 std::vector<float> x (nb * d);
38 float_randn (x.data(), nb * d, seed);
46 void test_index_type(
const char *factory_string) {
55 std::unique_ptr<Index> trained (
index_factory (d, factory_string));
58 auto xt = get_data (nt, 123);
59 trained->train (nt, xt.data());
63 auto xq = get_data (nq, 818);
69 std::unique_ptr<Index> src_index (clone_index (trained.get()));
72 auto xb = get_data (nb, 245);
73 src_index->add (nb, xb.data());
79 std::vector<idx_t> Iref (nq * k);
80 std::vector<float> Dref (nq * k);
81 src_index->search (nq, xq.data(), k, Dref.data(), Iref.data());
87 std::unique_ptr<Index> dst_index (clone_index (trained.get()));
91 auto xb = get_data (nb2, 366);
92 dst_index->add (nb2, xb.data());
95 std::vector<idx_t> Inew (nq * k);
96 std::vector<float> Dnew (nq * k);
101 for (
int sl = 0; sl < nslice; sl++) {
104 dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
105 EXPECT_TRUE (Iref != Inew);
106 EXPECT_TRUE (Dref != Dnew);
109 long i0 = sl * nlist / nslice;
110 long i1 = (sl + 1) * nlist / nslice;
112 std::vector<uint8_t> data_to_transfer;
114 std::unique_ptr<ArrayInvertedLists> il
115 (ivflib::get_invlist_range (src_index.get(), i0, i1));
118 write_InvertedLists (il.get(), &wr);
119 data_to_transfer.swap (wr.data);
126 reader.data.swap (data_to_transfer);
129 std::unique_ptr<ArrayInvertedLists> il
130 (dynamic_cast<ArrayInvertedLists *>
131 (read_InvertedLists (&reader)));
135 ivflib::set_invlist_range (dst_index.get(), i0, i1, il.get());
140 EXPECT_EQ (dst_index->ntotal, src_index->ntotal);
143 dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
144 EXPECT_TRUE (Iref == Inew);
145 EXPECT_TRUE (Dref == Dnew);
150 TEST(TRANS, IVFFlat) {
151 test_index_type (
"IVF40,Flat");
154 TEST(TRANS, IVFFlatPreproc) {
155 test_index_type (
"PCAR32,IVF40,Flat");
long idx_t
all indices are this type
virtual void set_index_parameter(Index *index, const std::string &name, double val) const
set one of the parameters
Index * index_factory(int d, const char *description_in, MetricType metric)