/root/doris/contrib/faiss/faiss/IndexBinaryFlat.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/IndexBinaryFlat.h> |
11 | | |
12 | | #include <faiss/impl/AuxIndexStructures.h> |
13 | | #include <faiss/impl/FaissAssert.h> |
14 | | #include <faiss/impl/IDSelector.h> |
15 | | #include <faiss/utils/Heap.h> |
16 | | #include <faiss/utils/hamming.h> |
17 | | #include <cstring> |
18 | | |
19 | | namespace faiss { |
20 | | |
21 | 0 | IndexBinaryFlat::IndexBinaryFlat(idx_t d) : IndexBinary(d) {} |
22 | | |
23 | 0 | void IndexBinaryFlat::add(idx_t n, const uint8_t* x) { |
24 | 0 | xb.insert(xb.end(), x, x + n * code_size); |
25 | 0 | ntotal += n; |
26 | 0 | } |
27 | | |
28 | 0 | void IndexBinaryFlat::reset() { |
29 | 0 | xb.clear(); |
30 | 0 | ntotal = 0; |
31 | 0 | } |
32 | | |
33 | | void IndexBinaryFlat::search( |
34 | | idx_t n, |
35 | | const uint8_t* x, |
36 | | idx_t k, |
37 | | int32_t* distances, |
38 | | idx_t* labels, |
39 | 0 | const SearchParameters* params) const { |
40 | | // Extract IDSelector from params if present |
41 | 0 | const IDSelector* sel = params ? params->sel : nullptr; |
42 | 0 | FAISS_THROW_IF_NOT(k > 0); |
43 | | |
44 | 0 | const idx_t block_size = query_batch_size; |
45 | 0 | for (idx_t s = 0; s < n; s += block_size) { |
46 | 0 | idx_t nn = block_size; |
47 | 0 | if (s + block_size > n) { |
48 | 0 | nn = n - s; |
49 | 0 | } |
50 | |
|
51 | 0 | if (use_heap) { |
52 | | // We see the distances and labels as heaps. |
53 | 0 | int_maxheap_array_t res = { |
54 | 0 | size_t(nn), size_t(k), labels + s * k, distances + s * k}; |
55 | |
|
56 | 0 | hammings_knn_hc( |
57 | 0 | &res, |
58 | 0 | x + s * code_size, |
59 | 0 | xb.data(), |
60 | 0 | ntotal, |
61 | 0 | code_size, |
62 | 0 | /* ordered = */ true, |
63 | 0 | approx_topk_mode, |
64 | 0 | sel); |
65 | 0 | } else { |
66 | 0 | hammings_knn_mc( |
67 | 0 | x + s * code_size, |
68 | 0 | xb.data(), |
69 | 0 | nn, |
70 | 0 | ntotal, |
71 | 0 | k, |
72 | 0 | code_size, |
73 | 0 | distances + s * k, |
74 | 0 | labels + s * k, |
75 | 0 | sel); |
76 | 0 | } |
77 | 0 | } |
78 | 0 | } |
79 | | |
80 | 0 | size_t IndexBinaryFlat::remove_ids(const IDSelector& sel) { |
81 | 0 | idx_t j = 0; |
82 | 0 | for (idx_t i = 0; i < ntotal; i++) { |
83 | 0 | if (sel.is_member(i)) { |
84 | | // should be removed |
85 | 0 | } else { |
86 | 0 | if (i > j) { |
87 | 0 | memmove(&xb[code_size * j], |
88 | 0 | &xb[code_size * i], |
89 | 0 | sizeof(xb[0]) * code_size); |
90 | 0 | } |
91 | 0 | j++; |
92 | 0 | } |
93 | 0 | } |
94 | 0 | long nremove = ntotal - j; |
95 | 0 | if (nremove > 0) { |
96 | 0 | ntotal = j; |
97 | 0 | xb.resize(ntotal * code_size); |
98 | 0 | } |
99 | 0 | return nremove; |
100 | 0 | } |
101 | | |
102 | 0 | void IndexBinaryFlat::reconstruct(idx_t key, uint8_t* recons) const { |
103 | 0 | memcpy(recons, &(xb[code_size * key]), sizeof(*recons) * code_size); |
104 | 0 | } |
105 | | |
106 | | void IndexBinaryFlat::range_search( |
107 | | idx_t n, |
108 | | const uint8_t* x, |
109 | | int radius, |
110 | | RangeSearchResult* result, |
111 | 0 | const SearchParameters* params) const { |
112 | 0 | const IDSelector* sel = params ? params->sel : nullptr; |
113 | 0 | hamming_range_search( |
114 | 0 | x, xb.data(), n, ntotal, radius, code_size, result, sel); |
115 | 0 | } |
116 | | |
117 | | } // namespace faiss |