Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
test_ondisk_ivf.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 <omp.h>
13
14 #include <unordered_map>
15 #include <pthread.h>
16
17 #include <gtest/gtest.h>
18
19 #include <faiss/OnDiskInvertedLists.h>
20 #include <faiss/IndexIVFFlat.h>
21 #include <faiss/IndexFlat.h>
22 #include <faiss/utils.h>
23 #include <faiss/index_io.h>
24
25
26 namespace {
27
28 struct Tempfilename {
29
30 static pthread_mutex_t mutex;
31
32 std::string filename;
33
34 Tempfilename (const char *prefix = nullptr) {
35 pthread_mutex_lock (&mutex);
36 char *cfname = tempnam (nullptr, prefix);
37 filename = cfname;
38 free(cfname);
39 pthread_mutex_unlock (&mutex);
40 }
41
42 ~Tempfilename () {
43 if (access (filename.c_str(), F_OK)) {
44 unlink (filename.c_str());
45 }
46 }
47
48 const char *c_str() {
49 return filename.c_str();
50 }
51
52 };
53
54 pthread_mutex_t Tempfilename::mutex = PTHREAD_MUTEX_INITIALIZER;
55
56 } // namespace
57
58
59 TEST(ONDISK, make_invlists) {
60 int nlist = 100;
61 int code_size = 32;
62 int nadd = 1000000;
63 std::unordered_map<int, int> listnos;
64
65 Tempfilename filename;
66
67 faiss::OnDiskInvertedLists ivf (
68 nlist, code_size,
69 filename.c_str());
70
71 {
72 std::vector<uint8_t> code(32);
73 for (int i = 0; i < nadd; i++) {
74 double d = drand48();
75 int list_no = int(nlist * d * d); // skewed distribution
76 int * ar = (int*)code.data();
77 ar[0] = i;
78 ar[1] = list_no;
79 ivf.add_entry (list_no, i, code.data());
80 listnos[i] = list_no;
81 }
82 }
83
84 int ntot = 0;
85 for (int i = 0; i < nlist; i++) {
86 int size = ivf.list_size(i);
87 const long *ids = ivf.get_ids (i);
88 const uint8_t *codes = ivf.get_codes (i);
89 for (int j = 0; j < size; j++) {
90 long id = ids[j];
91 const int * ar = (const int*)&codes[code_size * j];
92 EXPECT_EQ (ar[0], id);
93 EXPECT_EQ (ar[1], i);
94 EXPECT_EQ (listnos[id], i);
95 ntot ++;
96 }
97 }
98 EXPECT_EQ (ntot, nadd);
99 };
100
101
102 TEST(ONDISK, test_add) {
103 int d = 8;
104 int nlist = 30, nq = 200, nb = 1500, k = 10;
105 faiss::IndexFlatL2 quantizer(d);
106 {
107 std::vector<float> x(d * nlist);
108 faiss::float_rand(x.data(), d * nlist, 12345);
109 quantizer.add(nlist, x.data());
110 }
111 std::vector<float> xb(d * nb);
112 faiss::float_rand(xb.data(), d * nb, 23456);
113
114 faiss::IndexIVFFlat index(&quantizer, d, nlist);
115 index.add(nb, xb.data());
116
117 std::vector<float> xq(d * nb);
118 faiss::float_rand(xq.data(), d * nq, 34567);
119
120 std::vector<float> ref_D (nq * k);
121 std::vector<faiss::Index::idx_t> ref_I (nq * k);
122
123 index.search (nq, xq.data(), k,
124 ref_D.data(), ref_I.data());
125
126 Tempfilename filename, filename2;
127
128 // test add + search
129 {
130 faiss::IndexIVFFlat index2(&quantizer, d, nlist);
131
132 faiss::OnDiskInvertedLists ivf (
133 index.nlist, index.code_size,
134 filename.c_str());
135
136 index2.replace_invlists(&ivf);
137
138 index2.add(nb, xb.data());
139
140 std::vector<float> new_D (nq * k);
141 std::vector<faiss::Index::idx_t> new_I (nq * k);
142
143 index2.search (nq, xq.data(), k,
144 new_D.data(), new_I.data());
145
146 EXPECT_EQ (ref_D, new_D);
147 EXPECT_EQ (ref_I, new_I);
148
149 write_index(&index2, filename2.c_str());
150
151 }
152
153 // test io
154 {
155 faiss::Index *index3 = faiss::read_index(filename2.c_str());
156
157 std::vector<float> new_D (nq * k);
158 std::vector<faiss::Index::idx_t> new_I (nq * k);
159
160 index3->search (nq, xq.data(), k,
161 new_D.data(), new_I.data());
162
163 EXPECT_EQ (ref_D, new_D);
164 EXPECT_EQ (ref_I, new_I);
165
166 delete index3;
167 }
168
169 };
170
171
172
173 // WARN this thest will run multithreaded only in opt mode
174 TEST(ONDISK, make_invlists_threaded) {
175 int nlist = 100;
176 int code_size = 32;
177 int nadd = 1000000;
178 int nt = 20;
179
180 Tempfilename filename;
181
182 faiss::OnDiskInvertedLists ivf (
183 nlist, code_size,
184 filename.c_str());
185
186 std::vector<int> list_nos (nadd);
187
188 for (int i = 0; i < nadd; i++) {
189 double d = drand48();
190 list_nos[i] = int(nlist * d * d); // skewed distribution
191 }
192
193 #pragma omp parallel
194 {
195 std::vector<uint8_t> code(32);
196 #pragma omp for
197 for (int i = 0; i < nadd; i++) {
198 int list_no = list_nos[i];
199 int * ar = (int*)code.data();
200 ar[0] = i;
201 ar[1] = list_no;
202 ivf.add_entry (list_no, i, code.data());
203 }
204 }
205
206 int ntot = 0;
207 for (int i = 0; i < nlist; i++) {
208 int size = ivf.list_size(i);
209 const long *ids = ivf.get_ids (i);
210 const uint8_t *codes = ivf.get_codes (i);
211 for (int j = 0; j < size; j++) {
212 long id = ids[j];
213 const int * ar = (const int*)&codes[code_size * j];
214 EXPECT_EQ (ar[0], id);
215 EXPECT_EQ (ar[1], i);
216 EXPECT_EQ (list_nos[id], i);
217 ntot ++;
218 }
219 }
220 EXPECT_EQ (ntot, nadd);
221
222 };
faiss::IndexFlatL2
Definition: IndexFlat.h:78
faiss::IndexIVFFlat
Definition: IndexIVFFlat.h:25
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::OnDiskInvertedLists
Definition: OnDiskInvertedLists.h:52
faiss::Index
Definition: Index.h:62