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