Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
MetaIndexes.cpp
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved
11 // -*- c++ -*-
12
13 #include "MetaIndexes.h"
14
15 #include <pthread.h>
16
17 #include <cstdio>
18
19 #include "FaissAssert.h"
20 #include "Heap.h"
21
22 namespace faiss {
23
24 /*****************************************************
25 * IndexIDMap implementation
26 *******************************************************/
27
28 IndexIDMap::IndexIDMap (Index *index):
29 index (index),
30 own_fields (false)
31 {
32 FAISS_ASSERT (index->ntotal == 0 || !"index must be empty on input");
33 is_trained = index->is_trained;
34 metric_type = index->metric_type;
35 verbose = index->verbose;
36 d = index->d;
37 set_typename ();
38 }
39
40 void IndexIDMap::add (idx_t, const float *)
41 {
42 FAISS_ASSERT (!"add does not make sense with IndexIDMap, "
43 "use add_with_ids");
44 }
45
46
47 void IndexIDMap::train (idx_t n, const float *x)
48 {
49 index->train (n, x);
50 is_trained = index->is_trained;
51 }
52
53 void IndexIDMap::reset ()
54 {
55 index->reset ();
56 ntotal = 0;
57 }
58
59
60 void IndexIDMap::add_with_ids (idx_t n, const float * x, const long *xids)
61 {
62 index->add (n, x);
63 for (idx_t i = 0; i < n; i++)
64 id_map.push_back (xids[i]);
65 ntotal = index->ntotal;
66 }
67
68
69 void IndexIDMap::search (idx_t n, const float *x, idx_t k,
70 float *distances, idx_t *labels) const
71 {
72 index->search (n, x, k, distances, labels);
73 idx_t *li = labels;
74 for (idx_t i = 0; i < n * k; i++) {
75 li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
76 }
77 }
78
79
80
81 IndexIDMap::~IndexIDMap ()
82 {
83 if (own_fields) delete index;
84 }
85
86 void IndexIDMap::set_typename ()
87 {
88 index_typename = "IDMap[" + index->index_typename + "]";
89 }
90
91
92 /*****************************************************
93 * IndexShards implementation
94 *******************************************************/
95
96 // subroutines
97 namespace {
98
99
100 typedef Index::idx_t idx_t;
101
102
103 template<class Job>
104 struct Thread {
105 Job job;
106 pthread_t thread;
107
108 Thread () {}
109
110 explicit Thread (const Job & job): job(job) {}
111
112 void start () {
113 pthread_create (&thread, nullptr, run, this);
114 }
115
116 void wait () {
117 pthread_join (thread, nullptr);
118 }
119
120 static void * run (void *arg) {
121 static_cast<Thread*> (arg)->job.run();
122 return nullptr;
123 }
124
125 };
126
127
128 /// callback + thread management to train 1 shard
129 struct TrainJob {
130 IndexShards *index; // the relevant index
131 int no; // shard number
132 idx_t n; // train points
133 const float *x;
134
135 void run ()
136 {
137 if (index->verbose)
138 printf ("begin train shard %d on %ld points\n", no, n);
139 index->shard_indexes [no]->train(n, x);
140 if (index->verbose)
141 printf ("end train shard %d\n", no);
142 }
143
144 };
145
146 struct AddJob {
147 IndexShards *index; // the relevant index
148 int no; // shard number
149
150 idx_t n;
151 const float *x;
152 const idx_t *ids;
153
154 void run ()
155 {
156 if (index->verbose)
157 printf ("begin add shard %d on %ld points\n", no, n);
158 if (ids)
159 index->shard_indexes[no]->add_with_ids (n, x, ids);
160 else
161 index->shard_indexes[no]->add (n, x);
162 if (index->verbose)
163 printf ("end add shard %d on %ld points\n", no, n);
164 }
165 };
166
167
168
169 /// callback + thread management to query in 1 shard
170 struct QueryJob {
171 const IndexShards *index; // the relevant index
172 int no; // shard number
173
174 // query params
175 idx_t n;
176 const float *x;
177 idx_t k;
178 float *distances;
179 idx_t *labels;
180
181
182 void run ()
183 {
184 if (index->verbose)
185 printf ("begin query shard %d on %ld points\n", no, n);
186 index->shard_indexes [no]->search (n, x, k,
187 distances, labels);
188 if (index->verbose)
189 printf ("end query shard %d\n", no);
190 }
191
192
193 };
194
195
196
197
198 // add translation to all valid labels
199 void translate_labels (long n, idx_t *labels, long translation)
200 {
201 if (translation == 0) return;
202 for (long i = 0; i < n; i++) {
203 if(labels[i] < 0) return;
204 labels[i] += translation;
205 }
206 }
207
208
209 /** merge result tables from several shards.
210 * @param all_distances size nshard * n * k
211 * @param all_labels idem
212 * @param translartions label translations to apply, size nshard
213 */
214
215 template <class C>
216 void merge_tables (long n, long k, long nshard,
217 float *distances, idx_t *labels,
218 const float *all_distances,
219 idx_t *all_labels,
220 const long *translations)
221 {
222 if(k == 0) {
223 return;
224 }
225
226 long stride = n * k;
227 #pragma omp parallel
228 {
229 std::vector<int> buf (2 * nshard);
230 int * pointer = buf.data();
231 int * shard_ids = pointer + nshard;
232 std::vector<float> buf2 (nshard);
233 float * heap_vals = buf2.data();
234 #pragma omp for
235 for (long i = 0; i < n; i++) {
236 // the heap maps values to the shard where they are
237 // produced.
238 const float *D_in = all_distances + i * k;
239 const idx_t *I_in = all_labels + i * k;
240 int heap_size = 0;
241
242 for (long s = 0; s < nshard; s++) {
243 pointer[s] = 0;
244 if (I_in[stride * s] >= 0)
245 heap_push<C> (++heap_size, heap_vals, shard_ids,
246 D_in[stride * s], s);
247 }
248
249 float *D = distances + i * k;
250 idx_t *I = labels + i * k;
251
252 for (int j = 0; j < k; j++) {
253 if (heap_size == 0) {
254 I[j] = -1;
255 D[j] = C::neutral();
256 } else {
257 // pop best element
258 int s = shard_ids[0];
259 int & p = pointer[s];
260 D[j] = heap_vals[0];
261 I[j] = I_in[stride * s + p] + translations[s];
262
263 heap_pop<C> (heap_size--, heap_vals, shard_ids);
264 p++;
265 if (p < k && I_in[stride * s + p] >= 0)
266 heap_push<C> (++heap_size, heap_vals, shard_ids,
267 D_in[stride * s + p], s);
268 }
269 }
270 }
271 }
272 }
273
274
275 };
276
277
278
279
280 IndexShards::IndexShards (idx_t d, bool threaded, bool successive_ids):
281 Index (d), own_fields (false),
282 threaded (threaded), successive_ids (successive_ids)
283 {
284
285 }
286
287 void IndexShards::add_shard (Index *idx)
288 {
289 shard_indexes.push_back (idx);
290 sync_with_shard_indexes ();
291 }
292
293 void IndexShards::sync_with_shard_indexes ()
294 {
295 if (shard_indexes.empty()) return;
296 Index * index0 = shard_indexes[0];
297 d = index0->d;
298 metric_type = index0->metric_type;
299 is_trained = index0->is_trained;
300 ntotal = index0->ntotal;
301 for (int i = 1; i < shard_indexes.size(); i++) {
302 Index * index = shard_indexes[i];
303 FAISS_ASSERT (metric_type == index->metric_type);
304 FAISS_ASSERT (d == index->d);
305 ntotal += index->ntotal;
306 }
307 }
308
309
310 void IndexShards::train (idx_t n, const float *x)
311 {
312
313 // pre-alloc because we don't want reallocs
314 std::vector<Thread<TrainJob > > tss (shard_indexes.size());
315 int nt = 0;
316 for (int i = 0; i < shard_indexes.size(); i++) {
317 if(!shard_indexes[i]->is_trained) {
318 TrainJob ts = {this, i, n, x};
319 if (threaded) {
320 tss[nt] = Thread<TrainJob> (ts);
321 tss[nt++].start();
322 } else {
323 ts.run();
324 }
325 }
326 }
327 for (int i = 0; i < nt; i++) {
328 tss[i].wait();
329 }
330 sync_with_shard_indexes ();
331 }
332
333 void IndexShards::add (idx_t n, const float *x)
334 {
335 add_with_ids (n, x, nullptr);
336 }
337
338 /**
339 * Cases (successive_ids, xids):
340 * - true, non-NULL ERROR: it makes no sense to pass in ids and
341 * request them to be shifted
342 * - true, NULL OK, but should be called only once (calls add()
343 * on sub-indexes).
344 * - false, non-NULL OK: will call add_with_ids with passed in xids
345 * distributed evenly over shards
346 * - false, NULL OK: will call add_with_ids on each sub-index,
347 * starting at ntotal
348 */
349
350 void IndexShards::add_with_ids (idx_t n, const float * x, const long *xids)
351 {
352
353 FAISS_ASSERT(!(successive_ids && xids) ||
354 !"It makes no sense to pass in ids and request them to be shifted");
355
356 if (successive_ids) {
357 FAISS_ASSERT(!xids ||
358 !"It makes no sense to pass in ids and request them to be shifted");
359 FAISS_ASSERT(ntotal == 0 ||
360 !"when adding to IndexShards with sucessive_ids, only add() "
361 "in a single pass is supported");
362 }
363
364 long nshard = shard_indexes.size();
365 const long *ids = xids;
366 if (!ids && !successive_ids) {
367 long *aids = new long[n];
368 for (long i = 0; i < n; i++)
369 aids[i] = ntotal + i;
370 ids = aids;
371 }
372
373 std::vector<Thread<AddJob > > asa (shard_indexes.size());
374 int nt = 0;
375 for (int i = 0; i < nshard; i++) {
376 long i0 = i * n / nshard;
377 long i1 = (i + 1) * n / nshard;
378
379 AddJob as = {this, i,
380 i1 - i0, x + i0 * d,
381 ids ? ids + i0 : nullptr};
382 if (threaded) {
383 asa[nt] = Thread<AddJob>(as);
384 asa[nt++].start();
385 } else {
386 as.run();
387 }
388 }
389 for (int i = 0; i < nt; i++) {
390 asa[i].wait();
391 }
392 if (ids != xids) delete [] ids;
393 ntotal += n;
394 }
395
396
397
398
399
400 void IndexShards::reset ()
401 {
402 for (int i = 0; i < shard_indexes.size(); i++) {
403 shard_indexes[i]->reset ();
404 }
405 sync_with_shard_indexes ();
406 }
407
408 void IndexShards::search (
409 idx_t n, const float *x, idx_t k,
410 float *distances, idx_t *labels) const
411 {
412 long nshard = shard_indexes.size();
413 float *all_distances = new float [nshard * k * n];
414 idx_t *all_labels = new idx_t [nshard * k * n];
415
416 #if 1
417
418 // pre-alloc because we don't want reallocs
419 std::vector<Thread<QueryJob> > qss (nshard);
420 for (int i = 0; i < nshard; i++) {
421 QueryJob qs = {
422 this, i, n, x, k,
423 all_distances + i * k * n,
424 all_labels + i * k * n
425 };
426 if (threaded) {
427 qss[i] = Thread<QueryJob> (qs);
428 qss[i].start();
429 } else {
430 qs.run();
431 }
432 }
433
434 if (threaded) {
435 for (int i = 0; i < qss.size(); i++) {
436 qss[i].wait();
437 }
438 }
439 #else
440
441 // pre-alloc because we don't want reallocs
442 std::vector<QueryJob> qss (nshard);
443 for (int i = 0; i < nshard; i++) {
444 QueryJob qs = {
445 this, i, n, x, k,
446 all_distances + i * k * n,
447 all_labels + i * k * n
448 };
449 if (threaded) {
450 qss[i] = qs;
451 } else {
452 qs.run();
453 }
454 }
455
456 if (threaded) {
457 #pragma omp parallel for
458 for (int i = 0; i < qss.size(); i++) {
459 qss[i].run();
460 }
461 }
462
463 #endif
464 std::vector<long> translations (nshard, 0);
465 if (successive_ids) {
466 translations[0] = 0;
467 for (int s = 0; s + 1 < nshard; s++)
468 translations [s + 1] = translations [s] +
469 shard_indexes [s]->ntotal;
470 }
471
472 if (metric_type == METRIC_L2) {
473 merge_tables< CMin<float, int> > (
474 n, k, nshard, distances, labels,
475 all_distances, all_labels, translations.data ());
476 } else {
477 merge_tables< CMax<float, int> > (
478 n, k, nshard, distances, labels,
479 all_distances, all_labels, translations.data ());
480 }
481 delete [] all_distances;
482 delete [] all_labels;
483 }
484
485
486 void IndexShards::set_typename ()
487 {
488
489 }
490
491 IndexShards::~IndexShards ()
492 {
493 if (own_fields) {
494 for (int s = 0; s < shard_indexes.size(); s++)
495 delete shard_indexes [s];
496 }
497 }
498
499
500 /*****************************************************
501 * IndexSplitVectors implementation
502 *******************************************************/
503
504
505 IndexSplitVectors::IndexSplitVectors (idx_t d, bool threaded):
506 Index (d), own_fields (false),
507 threaded (threaded), sum_d (0)
508 {
509
510 }
511
512 void IndexSplitVectors::add_sub_index (Index *index)
513 {
514 sub_indexes.push_back (index);
515 sync_with_sub_indexes ();
516 }
517
518 void IndexSplitVectors::sync_with_sub_indexes ()
519 {
520 if (sub_indexes.empty()) return;
521 Index * index0 = sub_indexes[0];
522 sum_d = index0->d;
523 metric_type = index0->metric_type;
524 is_trained = index0->is_trained;
525 ntotal = index0->ntotal;
526 for (int i = 1; i < sub_indexes.size(); i++) {
527 Index * index = sub_indexes[i];
528 FAISS_ASSERT (metric_type == index->metric_type);
529 FAISS_ASSERT (ntotal == index->ntotal);
530 sum_d += index->d;
531 }
532
533 }
534
535 void IndexSplitVectors::add (idx_t n, const float *x)
536 {
537 FAISS_ASSERT (!"not implemented");
538 }
539
540 namespace {
541
542 /// callback + thread management to query in 1 shard
543 struct SplitQueryJob {
544 const IndexSplitVectors *index; // the relevant index
545 int no; // shard number
546
547 // query params
548 idx_t n;
549 const float *x;
550 idx_t k;
551 float *distances;
552 idx_t *labels;
553
554
555 void run ()
556 {
557 if (index->verbose)
558 printf ("begin query shard %d on %ld points\n", no, n);
559 const Index * sub_index = index->sub_indexes[no];
560 long sub_d = sub_index->d, d = index->d;
561 idx_t ofs = 0;
562 for (int i = 0; i < no; i++) ofs += index->sub_indexes[i]->d;
563 float *sub_x = new float [sub_d * n];
564 for (idx_t i = 0; i < n; i++)
565 memcpy (sub_x + i * sub_d, x + ofs + i * d, sub_d * sizeof (sub_x));
566 sub_index->search (n, sub_x, k, distances, labels);
567 delete [] sub_x;
568 if (index->verbose)
569 printf ("end query shard %d\n", no);
570 }
571
572 };
573
574
575
576 }
577
578
579
580 void IndexSplitVectors::search (
581 idx_t n, const float *x, idx_t k,
582 float *distances, idx_t *labels) const
583 {
584 FAISS_ASSERT (k == 1 || !"search implemented only for k=1");
585 FAISS_ASSERT (sum_d == d || !"not enough indexes compared to # dimensions");
586
587 long nshard = sub_indexes.size();
588 float *all_distances = new float [nshard * k * n];
589 idx_t *all_labels = new idx_t [nshard * k * n];
590
591 // pre-alloc because we don't want reallocs
592 std::vector<Thread<SplitQueryJob> > qss (nshard);
593 for (int i = 0; i < nshard; i++) {
594 SplitQueryJob qs = {
595 this, i, n, x, k,
596 i == 0 ? distances : all_distances + i * k * n,
597 i == 0 ? labels : all_labels + i * k * n
598 };
599 if (threaded) {
600 qss[i] = Thread<SplitQueryJob> (qs);
601 qss[i].start();
602 } else {
603 qs.run();
604 }
605 }
606
607 if (threaded) {
608 for (int i = 0; i < qss.size(); i++) {
609 qss[i].wait();
610 }
611 }
612
613 long factor = 1;
614 for (int i = 0; i < nshard; i++) {
615 if (i > 0) { // results of 0 are already in the table
616 const float *distances_i = all_distances + i * k * n;
617 const idx_t *labels_i = all_labels + i * k * n;
618 for (long j = 0; j < n; j++) {
619 if (labels[j] >= 0 && labels_i[j] >= 0) {
620 labels[j] += labels_i[j] * factor;
621 distances[j] += distances_i[j];
622 } else {
623 labels[j] = -1;
624 distances[j] = 0.0 / 0.0;
625 }
626 }
627 }
628 factor *= sub_indexes[i]->ntotal;
629 }
630 delete [] all_labels;
631 delete [] all_distances;
632 }
633
634
635 void IndexSplitVectors::train (idx_t n, const float *x)
636 {
637 FAISS_ASSERT (!"not implemented");
638 }
639
640 void IndexSplitVectors::reset ()
641 {
642 FAISS_ASSERT (!"not implemented");
643 }
644
645 void IndexSplitVectors::set_typename ()
646 {}
647
648 IndexSplitVectors::~IndexSplitVectors ()
649 {
650 if (own_fields) {
651 for (int s = 0; s < sub_indexes.size(); s++)
652 delete sub_indexes [s];
653 }
654 }
655
656
657
658
659
660
661 }; // namespace faiss
faiss::IndexShards::train
virtual void train(idx_t n, const float *x) override
Definition: MetaIndexes.cpp:310
faiss::IndexShards::IndexShards
IndexShards(idx_t d, bool threaded=false, bool successive_ids=true)
Definition: MetaIndexes.cpp:280
faiss::IndexIDMap::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: MetaIndexes.cpp:69
faiss::IndexShards::add_with_ids
virtual void add_with_ids(idx_t n, const float *x, const long *xids) override
Definition: MetaIndexes.cpp:350
faiss::IndexShards::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: MetaIndexes.cpp:408
faiss::IndexSplitVectors::reset
virtual void reset() override
removes all elements from the database.
Definition: MetaIndexes.cpp:640
faiss::IndexIDMap::add
virtual void add(idx_t n, const float *x) override
this will fail. Use add_with_ids
Definition: MetaIndexes.cpp:40
faiss::IndexSplitVectors
Definition: MetaIndexes.h:109
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::IndexSplitVectors::train
virtual void train(idx_t n, const float *x) override
Definition: MetaIndexes.cpp:635
faiss::IndexShards::add
virtual void add(idx_t n, const float *x) override
supported only for sub-indices that implement add_with_ids
Definition: MetaIndexes.cpp:333
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::Index::verbose
bool verbose
verbosity level
Definition: Index.h:68
faiss::IndexSplitVectors::add
virtual void add(idx_t n, const float *x) override
Definition: MetaIndexes.cpp:535
faiss::IndexShards::threaded
bool threaded
should the sub-indexes be deleted along with this?
Definition: MetaIndexes.h:62
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::IndexShards::reset
virtual void reset() override
removes all elements from the database.
Definition: MetaIndexes.cpp:400
faiss::IndexIDMap::reset
virtual void reset() override
removes all elements from the database.
Definition: MetaIndexes.cpp:53
faiss::IndexSplitVectors::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: MetaIndexes.cpp:580
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::IndexIDMap::train
virtual void train(idx_t n, const float *x) override
Definition: MetaIndexes.cpp:47
faiss::Index
Definition: Index.h:61
faiss::Index::is_trained
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:71
faiss::IndexIDMap::add_with_ids
virtual void add_with_ids(idx_t n, const float *x, const long *xids) override
Definition: MetaIndexes.cpp:60
faiss::IndexSplitVectors::IndexSplitVectors
IndexSplitVectors(idx_t d, bool threaded=false)
sum of dimensions seen so far
Definition: MetaIndexes.cpp:505