/root/doris/contrib/faiss/faiss/IndexIDMap.cpp
Line | Count | Source |
1 | | /* |
2 | | * Copyright (c) Meta Platforms, Inc. and 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 | | // -*- c++ -*- |
9 | | |
10 | | #include <faiss/IndexIDMap.h> |
11 | | |
12 | | #include <cinttypes> |
13 | | #include <cstdint> |
14 | | #include <cstdio> |
15 | | |
16 | | #include <faiss/impl/AuxIndexStructures.h> |
17 | | #include <faiss/impl/FaissAssert.h> |
18 | | #include <faiss/utils/Heap.h> |
19 | | #include <faiss/utils/WorkerThread.h> |
20 | | |
21 | | namespace faiss { |
22 | | |
23 | | namespace { |
24 | | |
25 | | // IndexBinary needs to update the code_size when d is set... |
26 | | |
27 | 0 | void sync_d(Index* index) {} |
28 | | |
29 | 0 | void sync_d(IndexBinary* index) { |
30 | 0 | FAISS_THROW_IF_NOT(index->d % 8 == 0); |
31 | 0 | index->code_size = index->d / 8; |
32 | 0 | } |
33 | | |
34 | | } // anonymous namespace |
35 | | |
36 | | /***************************************************** |
37 | | * IndexIDMap implementation |
38 | | *******************************************************/ |
39 | | |
40 | | template <typename IndexT> |
41 | 0 | IndexIDMapTemplate<IndexT>::IndexIDMapTemplate(IndexT* index) : index(index) { |
42 | 0 | FAISS_THROW_IF_NOT_MSG(index->ntotal == 0, "index must be empty on input"); |
43 | 0 | this->is_trained = index->is_trained; |
44 | 0 | this->metric_type = index->metric_type; |
45 | 0 | this->verbose = index->verbose; |
46 | 0 | this->d = index->d; |
47 | 0 | sync_d(this); |
48 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEEC2EPS1_ Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEEC2EPS1_ |
49 | | |
50 | | template <typename IndexT> |
51 | | void IndexIDMapTemplate<IndexT>::add( |
52 | | idx_t, |
53 | 0 | const typename IndexT::component_t*) { |
54 | 0 | FAISS_THROW_MSG( |
55 | 0 | "add does not make sense with IndexIDMap, " |
56 | 0 | "use add_with_ids"); |
57 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE3addElPKf Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE3addElPKh |
58 | | |
59 | | template <typename IndexT> |
60 | | void IndexIDMapTemplate<IndexT>::train( |
61 | | idx_t n, |
62 | 0 | const typename IndexT::component_t* x) { |
63 | 0 | index->train(n, x); |
64 | 0 | this->is_trained = index->is_trained; |
65 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE5trainElPKf Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE5trainElPKh |
66 | | |
67 | | template <typename IndexT> |
68 | 0 | void IndexIDMapTemplate<IndexT>::reset() { |
69 | 0 | index->reset(); |
70 | 0 | id_map.clear(); |
71 | 0 | this->ntotal = 0; |
72 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE5resetEv Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE5resetEv |
73 | | |
74 | | template <typename IndexT> |
75 | | void IndexIDMapTemplate<IndexT>::add_with_ids( |
76 | | idx_t n, |
77 | | const typename IndexT::component_t* x, |
78 | 0 | const idx_t* xids) { |
79 | 0 | index->add(n, x); |
80 | 0 | for (idx_t i = 0; i < n; i++) |
81 | 0 | id_map.push_back(xids[i]); |
82 | 0 | this->ntotal = index->ntotal; |
83 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE12add_with_idsElPKfPKl Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE12add_with_idsElPKhPKl |
84 | | |
85 | | template <typename IndexT> |
86 | 0 | size_t IndexIDMapTemplate<IndexT>::sa_code_size() const { |
87 | 0 | return index->sa_code_size(); |
88 | 0 | } Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_5IndexEE12sa_code_sizeEv Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_11IndexBinaryEE12sa_code_sizeEv |
89 | | |
90 | | template <typename IndexT> |
91 | | void IndexIDMapTemplate<IndexT>::add_sa_codes( |
92 | | idx_t n, |
93 | | const uint8_t* codes, |
94 | 0 | const idx_t* xids) { |
95 | 0 | index->add_sa_codes(n, codes, xids); |
96 | 0 | for (idx_t i = 0; i < n; i++) { |
97 | 0 | id_map.push_back(xids[i]); |
98 | 0 | } |
99 | 0 | this->ntotal = index->ntotal; |
100 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE12add_sa_codesElPKhPKl Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE12add_sa_codesElPKhPKl |
101 | | |
102 | | namespace { |
103 | | |
104 | | /// RAII object to reset the IDSelector in the params object |
105 | | struct ScopedSelChange { |
106 | | SearchParameters* params = nullptr; |
107 | | IDSelector* old_sel = nullptr; |
108 | | |
109 | 0 | void set(SearchParameters* params_2, IDSelector* new_sel) { |
110 | 0 | this->params = params_2; |
111 | 0 | old_sel = params_2->sel; |
112 | 0 | params_2->sel = new_sel; |
113 | 0 | } |
114 | 0 | ~ScopedSelChange() { |
115 | 0 | if (params) { |
116 | 0 | params->sel = old_sel; |
117 | 0 | } |
118 | 0 | } |
119 | | }; |
120 | | |
121 | | } // namespace |
122 | | |
123 | | template <typename IndexT> |
124 | | void IndexIDMapTemplate<IndexT>::search( |
125 | | idx_t n, |
126 | | const typename IndexT::component_t* x, |
127 | | idx_t k, |
128 | | typename IndexT::distance_t* distances, |
129 | | idx_t* labels, |
130 | 0 | const SearchParameters* params) const { |
131 | 0 | IDSelectorTranslated this_idtrans(this->id_map, nullptr); |
132 | 0 | ScopedSelChange sel_change; |
133 | |
|
134 | 0 | if (params && params->sel) { |
135 | 0 | auto idtrans = dynamic_cast<const IDSelectorTranslated*>(params->sel); |
136 | |
|
137 | 0 | if (!idtrans) { |
138 | | /* |
139 | | FAISS_THROW_IF_NOT_MSG( |
140 | | idtrans, |
141 | | "IndexIDMap requires an IDSelectorTranslated on input"); |
142 | | */ |
143 | | // then make an idtrans and force it into the SearchParameters |
144 | | // (hence the const_cast) |
145 | 0 | auto params_non_const = const_cast<SearchParameters*>(params); |
146 | 0 | this_idtrans.sel = params->sel; |
147 | 0 | sel_change.set(params_non_const, &this_idtrans); |
148 | 0 | } |
149 | 0 | } |
150 | 0 | index->search(n, x, k, distances, labels, params); |
151 | 0 | idx_t* li = labels; |
152 | 0 | #pragma omp parallel for |
153 | 0 | for (idx_t i = 0; i < n * k; i++) { |
154 | 0 | li[i] = li[i] < 0 ? li[i] : id_map[li[i]]; |
155 | 0 | } Unexecuted instantiation: IndexIDMap.cpp:_ZNK5faiss18IndexIDMapTemplateINS_5IndexEE6searchElPKflPfPlPKNS_16SearchParametersE.omp_outlined_debug__ Unexecuted instantiation: IndexIDMap.cpp:_ZNK5faiss18IndexIDMapTemplateINS_11IndexBinaryEE6searchElPKhlPiPlPKNS_16SearchParametersE.omp_outlined_debug__ |
156 | 0 | } Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_5IndexEE6searchElPKflPfPlPKNS_16SearchParametersE Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_11IndexBinaryEE6searchElPKhlPiPlPKNS_16SearchParametersE |
157 | | |
158 | | template <typename IndexT> |
159 | | void IndexIDMapTemplate<IndexT>::range_search( |
160 | | idx_t n, |
161 | | const typename IndexT::component_t* x, |
162 | | typename IndexT::distance_t radius, |
163 | | RangeSearchResult* result, |
164 | 0 | const SearchParameters* params) const { |
165 | 0 | if (params) { |
166 | 0 | SearchParameters internal_search_parameters; |
167 | 0 | IDSelectorTranslated id_selector_translated(id_map, params->sel); |
168 | 0 | internal_search_parameters.sel = &id_selector_translated; |
169 | |
|
170 | 0 | index->range_search(n, x, radius, result, &internal_search_parameters); |
171 | 0 | } else { |
172 | 0 | index->range_search(n, x, radius, result); |
173 | 0 | } |
174 | |
|
175 | 0 | #pragma omp parallel for |
176 | 0 | for (idx_t i = 0; i < result->lims[result->nq]; i++) { |
177 | 0 | result->labels[i] = result->labels[i] < 0 ? result->labels[i] |
178 | 0 | : id_map[result->labels[i]]; |
179 | 0 | } Unexecuted instantiation: IndexIDMap.cpp:_ZNK5faiss18IndexIDMapTemplateINS_5IndexEE12range_searchElPKffPNS_17RangeSearchResultEPKNS_16SearchParametersE.omp_outlined_debug__ Unexecuted instantiation: IndexIDMap.cpp:_ZNK5faiss18IndexIDMapTemplateINS_11IndexBinaryEE12range_searchElPKhiPNS_17RangeSearchResultEPKNS_16SearchParametersE.omp_outlined_debug__ |
180 | 0 | } Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_5IndexEE12range_searchElPKffPNS_17RangeSearchResultEPKNS_16SearchParametersE Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_11IndexBinaryEE12range_searchElPKhiPNS_17RangeSearchResultEPKNS_16SearchParametersE |
181 | | |
182 | | template <typename IndexT> |
183 | 0 | size_t IndexIDMapTemplate<IndexT>::remove_ids(const IDSelector& sel) { |
184 | | // remove in sub-index first |
185 | 0 | IDSelectorTranslated sel2(id_map, &sel); |
186 | 0 | size_t nremove = index->remove_ids(sel2); |
187 | |
|
188 | 0 | int64_t j = 0; |
189 | 0 | for (idx_t i = 0; i < this->ntotal; i++) { |
190 | 0 | if (sel.is_member(id_map[i])) { |
191 | | // remove |
192 | 0 | } else { |
193 | 0 | id_map[j] = id_map[i]; |
194 | 0 | j++; |
195 | 0 | } |
196 | 0 | } |
197 | 0 | FAISS_ASSERT(j == index->ntotal); |
198 | 0 | this->ntotal = j; |
199 | 0 | id_map.resize(this->ntotal); |
200 | 0 | return nremove; |
201 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE10remove_idsERKNS_10IDSelectorE Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE10remove_idsERKNS_10IDSelectorE |
202 | | |
203 | | template <typename IndexT> |
204 | | void IndexIDMapTemplate<IndexT>::check_compatible_for_merge( |
205 | 0 | const IndexT& otherIndex) const { |
206 | 0 | auto other = dynamic_cast<const IndexIDMapTemplate<IndexT>*>(&otherIndex); |
207 | 0 | FAISS_THROW_IF_NOT(other); |
208 | 0 | index->check_compatible_for_merge(*other->index); |
209 | 0 | } Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_5IndexEE26check_compatible_for_mergeERKS1_ Unexecuted instantiation: _ZNK5faiss18IndexIDMapTemplateINS_11IndexBinaryEE26check_compatible_for_mergeERKS1_ |
210 | | |
211 | | template <typename IndexT> |
212 | 0 | void IndexIDMapTemplate<IndexT>::merge_from(IndexT& otherIndex, idx_t add_id) { |
213 | 0 | check_compatible_for_merge(otherIndex); |
214 | 0 | auto other = static_cast<IndexIDMapTemplate<IndexT>*>(&otherIndex); |
215 | 0 | index->merge_from(*other->index); |
216 | 0 | for (size_t i = 0; i < other->id_map.size(); i++) { |
217 | 0 | id_map.push_back(other->id_map[i] + add_id); |
218 | 0 | } |
219 | 0 | other->id_map.resize(0); |
220 | 0 | this->ntotal = index->ntotal; |
221 | 0 | other->ntotal = 0; |
222 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEE10merge_fromERS1_l Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEE10merge_fromERS1_l |
223 | | |
224 | | template <typename IndexT> |
225 | 0 | IndexIDMapTemplate<IndexT>::~IndexIDMapTemplate() { |
226 | 0 | if (own_fields) |
227 | 0 | delete index; |
228 | 0 | } Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_5IndexEED2Ev Unexecuted instantiation: _ZN5faiss18IndexIDMapTemplateINS_11IndexBinaryEED2Ev |
229 | | |
230 | | /***************************************************** |
231 | | * IndexIDMap2 implementation |
232 | | *******************************************************/ |
233 | | |
234 | | template <typename IndexT> |
235 | | IndexIDMap2Template<IndexT>::IndexIDMap2Template(IndexT* index) |
236 | 0 | : IndexIDMapTemplate<IndexT>(index) {}Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_5IndexEEC2EPS1_ Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_11IndexBinaryEEC2EPS1_ |
237 | | |
238 | | template <typename IndexT> |
239 | | void IndexIDMap2Template<IndexT>::add_with_ids( |
240 | | idx_t n, |
241 | | const typename IndexT::component_t* x, |
242 | 0 | const idx_t* xids) { |
243 | 0 | size_t prev_ntotal = this->ntotal; |
244 | 0 | IndexIDMapTemplate<IndexT>::add_with_ids(n, x, xids); |
245 | 0 | for (size_t i = prev_ntotal; i < this->ntotal; i++) { |
246 | 0 | rev_map[this->id_map[i]] = i; |
247 | 0 | } |
248 | 0 | } Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_5IndexEE12add_with_idsElPKfPKl Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_11IndexBinaryEE12add_with_idsElPKhPKl |
249 | | |
250 | | template <typename IndexT> |
251 | 0 | void IndexIDMap2Template<IndexT>::check_consistency() const { |
252 | 0 | FAISS_THROW_IF_NOT(rev_map.size() == this->id_map.size()); |
253 | 0 | FAISS_THROW_IF_NOT(this->id_map.size() == this->ntotal); |
254 | 0 | for (size_t i = 0; i < this->ntotal; i++) { |
255 | 0 | idx_t ii = rev_map.at(this->id_map[i]); |
256 | 0 | FAISS_THROW_IF_NOT(ii == i); |
257 | 0 | } |
258 | 0 | } Unexecuted instantiation: _ZNK5faiss19IndexIDMap2TemplateINS_5IndexEE17check_consistencyEv Unexecuted instantiation: _ZNK5faiss19IndexIDMap2TemplateINS_11IndexBinaryEE17check_consistencyEv |
259 | | |
260 | | template <typename IndexT> |
261 | 0 | void IndexIDMap2Template<IndexT>::merge_from(IndexT& otherIndex, idx_t add_id) { |
262 | 0 | size_t prev_ntotal = this->ntotal; |
263 | 0 | IndexIDMapTemplate<IndexT>::merge_from(otherIndex, add_id); |
264 | 0 | for (size_t i = prev_ntotal; i < this->ntotal; i++) { |
265 | 0 | rev_map[this->id_map[i]] = i; |
266 | 0 | } |
267 | 0 | static_cast<IndexIDMap2Template<IndexT>&>(otherIndex).rev_map.clear(); |
268 | 0 | } Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_5IndexEE10merge_fromERS1_l Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_11IndexBinaryEE10merge_fromERS1_l |
269 | | |
270 | | template <typename IndexT> |
271 | 0 | void IndexIDMap2Template<IndexT>::construct_rev_map() { |
272 | 0 | rev_map.clear(); |
273 | 0 | for (size_t i = 0; i < this->ntotal; i++) { |
274 | 0 | rev_map[this->id_map[i]] = i; |
275 | 0 | } |
276 | 0 | } Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_5IndexEE17construct_rev_mapEv Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_11IndexBinaryEE17construct_rev_mapEv |
277 | | |
278 | | template <typename IndexT> |
279 | 0 | size_t IndexIDMap2Template<IndexT>::remove_ids(const IDSelector& sel) { |
280 | | // This is quite inefficient |
281 | 0 | size_t nremove = IndexIDMapTemplate<IndexT>::remove_ids(sel); |
282 | 0 | construct_rev_map(); |
283 | 0 | return nremove; |
284 | 0 | } Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_5IndexEE10remove_idsERKNS_10IDSelectorE Unexecuted instantiation: _ZN5faiss19IndexIDMap2TemplateINS_11IndexBinaryEE10remove_idsERKNS_10IDSelectorE |
285 | | |
286 | | template <typename IndexT> |
287 | | void IndexIDMap2Template<IndexT>::reconstruct( |
288 | | idx_t key, |
289 | 0 | typename IndexT::component_t* recons) const { |
290 | 0 | try { |
291 | 0 | this->index->reconstruct(rev_map.at(key), recons); |
292 | 0 | } catch (const std::out_of_range&) { |
293 | 0 | FAISS_THROW_FMT("key %" PRId64 " not found", key); |
294 | 0 | } |
295 | 0 | } Unexecuted instantiation: _ZNK5faiss19IndexIDMap2TemplateINS_5IndexEE11reconstructElPf Unexecuted instantiation: _ZNK5faiss19IndexIDMap2TemplateINS_11IndexBinaryEE11reconstructElPh |
296 | | |
297 | | // explicit template instantiations |
298 | | |
299 | | template struct IndexIDMapTemplate<Index>; |
300 | | template struct IndexIDMapTemplate<IndexBinary>; |
301 | | template struct IndexIDMap2Template<Index>; |
302 | | template struct IndexIDMap2Template<IndexBinary>; |
303 | | |
304 | | } // namespace faiss |