Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
test_transfer_invlists.cpp
1 /**
2  * Copyright (c) 2015-present, Facebook, Inc.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD+Patents license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #include <cstdio>
10 #include <cstdlib>
11 
12 #include <gtest/gtest.h>
13 
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>
21 
22 
23 // parameters to use for the test
24 int d = 64;
25 size_t nb = 1000;
26 size_t nq = 100;
27 size_t nt = 500;
28 int k = 10;
29 int nlist = 40;
30 
31 using namespace faiss;
32 
33 typedef faiss::Index::idx_t idx_t;
34 
35 
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);
39  return x;
40 }
41 
42 
43 
44 
45 
46 void test_index_type(const char *factory_string) {
47 
48  // transfer inverted lists in nslice slices
49  int nslice = 3;
50 
51  /****************************************************************
52  * trained reference index
53  ****************************************************************/
54 
55  std::unique_ptr<Index> trained (index_factory (d, factory_string));
56 
57  {
58  auto xt = get_data (nt, 123);
59  trained->train (nt, xt.data());
60  }
61 
62  // sample nq query vectors to check if results are the same
63  auto xq = get_data (nq, 818);
64 
65 
66  /****************************************************************
67  * source index
68  ***************************************************************/
69  std::unique_ptr<Index> src_index (clone_index (trained.get()));
70 
71  { // add some data to source index
72  auto xb = get_data (nb, 245);
73  src_index->add (nb, xb.data());
74  }
75 
76  ParameterSpace().set_index_parameter (src_index.get(), "nprobe", 4);
77 
78  // remember reference search result on source index
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());
82 
83  /****************************************************************
84  * destination index -- should be replaced by source index
85  ***************************************************************/
86 
87  std::unique_ptr<Index> dst_index (clone_index (trained.get()));
88 
89  { // initial state: filled in with some garbage
90  int nb2 = nb + 10;
91  auto xb = get_data (nb2, 366);
92  dst_index->add (nb2, xb.data());
93  }
94 
95  std::vector<idx_t> Inew (nq * k);
96  std::vector<float> Dnew (nq * k);
97 
98  ParameterSpace().set_index_parameter (dst_index.get(), "nprobe", 4);
99 
100  // transfer from source to destination in nslice slices
101  for (int sl = 0; sl < nslice; sl++) {
102 
103  // so far, the indexes are different
104  dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
105  EXPECT_TRUE (Iref != Inew);
106  EXPECT_TRUE (Dref != Dnew);
107 
108  // range of inverted list indices to transfer
109  long i0 = sl * nlist / nslice;
110  long i1 = (sl + 1) * nlist / nslice;
111 
112  std::vector<uint8_t> data_to_transfer;
113  {
114  std::unique_ptr<ArrayInvertedLists> il
115  (ivflib::get_invlist_range (src_index.get(), i0, i1));
116  // serialize inverted lists
117  VectorIOWriter wr;
118  write_InvertedLists (il.get(), &wr);
119  data_to_transfer.swap (wr.data);
120  }
121 
122  // transfer data here from source machine to dest machine
123 
124  {
125  VectorIOReader reader;
126  reader.data.swap (data_to_transfer);
127 
128  // deserialize inverted lists
129  std::unique_ptr<ArrayInvertedLists> il
130  (dynamic_cast<ArrayInvertedLists *>
131  (read_InvertedLists (&reader)));
132 
133  // swap inverted lists. Block searches here!
134  {
135  ivflib::set_invlist_range (dst_index.get(), i0, i1, il.get());
136  }
137  }
138 
139  }
140  EXPECT_EQ (dst_index->ntotal, src_index->ntotal);
141 
142  // now, the indexes are the same
143  dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
144  EXPECT_TRUE (Iref == Inew);
145  EXPECT_TRUE (Dref == Dnew);
146 
147 }
148 
149 
150 TEST(TRANS, IVFFlat) {
151  test_index_type ("IVF40,Flat");
152 }
153 
154 TEST(TRANS, IVFFlatPreproc) {
155  test_index_type ("PCAR32,IVF40,Flat");
156 }
long idx_t
all indices are this type
Definition: Index.h:64
virtual void set_index_parameter(Index *index, const std::string &name, double val) const
set one of the parameters
Definition: AutoTune.cpp:452
Index * index_factory(int d, const char *description_in, MetricType metric)
Definition: AutoTune.cpp:722