Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/hoss/faiss/IndexReplicas.cpp
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8 #include "IndexReplicas.h"
9 #include "FaissAssert.h"
10
11 namespace faiss {
12
13 template <typename IndexT>
14 IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(bool threaded)
15 : ThreadedIndex<IndexT>(threaded) {
16 }
17
18 template <typename IndexT>
19 IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(idx_t d, bool threaded)
20 : ThreadedIndex<IndexT>(d, threaded) {
21 }
22
23 template <typename IndexT>
24 IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(int d, bool threaded)
25 : ThreadedIndex<IndexT>(d, threaded) {
26 }
27
28 template <typename IndexT>
29 void
30 IndexReplicasTemplate<IndexT>::onAfterAddIndex(IndexT* index) {
31 // Make sure that the parameters are the same for all prior indices, unless
32 // we're the first index to be added
33 if (this->count() > 0 && this->at(0) != index) {
34 auto existing = this->at(0);
35
36 FAISS_THROW_IF_NOT_FMT(index->ntotal == existing->ntotal,
37 "IndexReplicas: newly added index does "
38 "not have same number of vectors as prior index; "
39 "prior index has %ld vectors, new index has %ld",
40 existing->ntotal, index->ntotal);
41
42 FAISS_THROW_IF_NOT_MSG(index->is_trained == existing->is_trained,
43 "IndexReplicas: newly added index does "
44 "not have same train status as prior index");
45 } else {
46 // Set our parameters based on the first index we're adding
47 // (dimension is handled in ThreadedIndex)
48 this->ntotal = index->ntotal;
49 this->verbose = index->verbose;
50 this->is_trained = index->is_trained;
51 this->metric_type = index->metric_type;
52 }
53 }
54
55 template <typename IndexT>
56 void
57 IndexReplicasTemplate<IndexT>::train(idx_t n, const component_t* x) {
58 this->runOnIndex([n, x](int, IndexT* index){ index->train(n, x); });
59 }
60
61 template <typename IndexT>
62 void
63 IndexReplicasTemplate<IndexT>::add(idx_t n, const component_t* x) {
64 this->runOnIndex([n, x](int, IndexT* index){ index->add(n, x); });
65 this->ntotal += n;
66 }
67
68 template <typename IndexT>
69 void
70 IndexReplicasTemplate<IndexT>::reconstruct(idx_t n, component_t* x) const {
71 FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
72
73 // Just pass to the first replica
74 this->at(0)->reconstruct(n, x);
75 }
76
77 template <typename IndexT>
78 void
79 IndexReplicasTemplate<IndexT>::search(idx_t n,
80 const component_t* x,
81 idx_t k,
82 distance_t* distances,
83 idx_t* labels) const {
84 FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
85
86 if (n == 0) {
87 return;
88 }
89
90 auto dim = this->d;
91 size_t componentsPerVec =
92 sizeof(component_t) == 1 ? (dim + 7) / 8 : dim;
93
94 // Partition the query by the number of indices we have
95 faiss::Index::idx_t queriesPerIndex =
96 (faiss::Index::idx_t) (n + this->count() - 1) /
97 (faiss::Index::idx_t) this->count();
98 FAISS_ASSERT(n / queriesPerIndex <= this->count());
99
100 auto fn =
101 [queriesPerIndex, componentsPerVec,
102 n, x, k, distances, labels](int i, const IndexT* index) {
103 faiss::Index::idx_t base = (faiss::Index::idx_t) i * queriesPerIndex;
104
105 if (base < n) {
106 auto numForIndex = std::min(queriesPerIndex, n - base);
107
108 index->search(numForIndex,
109 x + base * componentsPerVec,
110 k,
111 distances + base * k,
112 labels + base * k);
113 }
114 };
115
116 this->runOnIndex(fn);
117 }
118
119 // explicit instantiations
120 template struct IndexReplicasTemplate<Index>;
121 template struct IndexReplicasTemplate<IndexBinary>;
122
123 } // namespace
faiss::IndexReplicasTemplate::train
void train(idx_t n, const component_t *x) override
Definition: IndexReplicas.cpp:57
faiss::IndexReplicasTemplate::IndexReplicasTemplate
IndexReplicasTemplate(bool threaded=true)
Definition: IndexReplicas.cpp:14
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::IndexReplicasTemplate::onAfterAddIndex
void onAfterAddIndex(IndexT *index) override
Called just after an index is added.
Definition: IndexReplicas.cpp:30
faiss::ThreadedIndex
Definition: ThreadedIndex.h:21
faiss::IndexReplicasTemplate::search
void search(idx_t n, const component_t *x, idx_t k, distance_t *distances, idx_t *labels) const override
Definition: IndexReplicas.cpp:79
faiss::IndexReplicasTemplate::reconstruct
void reconstruct(idx_t, component_t *v) const override
reconstructs from the first index
Definition: IndexReplicas.cpp:70
faiss::IndexReplicasTemplate::add
void add(idx_t n, const component_t *x) override
Definition: IndexReplicas.cpp:63
faiss::IndexReplicasTemplate
Definition: IndexReplicas.h:21
faiss::IndexIVF::search
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexIVF.cpp:228