contrib/faiss/faiss/impl/IDSelector.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 | | #include <faiss/impl/FaissAssert.h> |
9 | | #include <faiss/impl/IDSelector.h> |
10 | | |
11 | | namespace faiss { |
12 | | |
13 | | /*********************************************************************** |
14 | | * IDSelectorRange |
15 | | ***********************************************************************/ |
16 | | |
17 | | IDSelectorRange::IDSelectorRange(idx_t imin, idx_t imax, bool assume_sorted) |
18 | 0 | : imin(imin), imax(imax), assume_sorted(assume_sorted) {} |
19 | | |
20 | 0 | bool IDSelectorRange::is_member(idx_t id) const { |
21 | 0 | return id >= imin && id < imax; |
22 | 0 | } |
23 | | |
24 | | void IDSelectorRange::find_sorted_ids_bounds( |
25 | | size_t list_size, |
26 | | const idx_t* ids, |
27 | | size_t* jmin_out, |
28 | 0 | size_t* jmax_out) const { |
29 | 0 | FAISS_ASSERT(assume_sorted); |
30 | 0 | if (list_size == 0 || imax <= ids[0] || imin > ids[list_size - 1]) { |
31 | 0 | *jmin_out = *jmax_out = 0; |
32 | 0 | return; |
33 | 0 | } |
34 | | // bissection to find imin |
35 | 0 | if (ids[0] >= imin) { |
36 | 0 | *jmin_out = 0; |
37 | 0 | } else { |
38 | 0 | size_t j0 = 0, j1 = list_size; |
39 | 0 | while (j1 > j0 + 1) { |
40 | 0 | size_t jmed = (j0 + j1) / 2; |
41 | 0 | if (ids[jmed] >= imin) { |
42 | 0 | j1 = jmed; |
43 | 0 | } else { |
44 | 0 | j0 = jmed; |
45 | 0 | } |
46 | 0 | } |
47 | 0 | *jmin_out = j1; |
48 | 0 | } |
49 | | // bissection to find imax |
50 | 0 | if (*jmin_out == list_size || ids[*jmin_out] >= imax) { |
51 | 0 | *jmax_out = *jmin_out; |
52 | 0 | } else { |
53 | 0 | size_t j0 = *jmin_out, j1 = list_size; |
54 | 0 | while (j1 > j0 + 1) { |
55 | 0 | size_t jmed = (j0 + j1) / 2; |
56 | 0 | if (ids[jmed] >= imax) { |
57 | 0 | j1 = jmed; |
58 | 0 | } else { |
59 | 0 | j0 = jmed; |
60 | 0 | } |
61 | 0 | } |
62 | 0 | *jmax_out = j1; |
63 | 0 | } |
64 | 0 | } |
65 | | |
66 | | /*********************************************************************** |
67 | | * IDSelectorArray |
68 | | ***********************************************************************/ |
69 | | |
70 | 0 | IDSelectorArray::IDSelectorArray(size_t n, const idx_t* ids) : n(n), ids(ids) {} |
71 | | |
72 | 0 | bool IDSelectorArray::is_member(idx_t id) const { |
73 | 0 | for (idx_t i = 0; i < n; i++) { |
74 | 0 | if (ids[i] == id) |
75 | 0 | return true; |
76 | 0 | } |
77 | 0 | return false; |
78 | 0 | } |
79 | | |
80 | | /*********************************************************************** |
81 | | * IDSelectorBatch |
82 | | ***********************************************************************/ |
83 | | |
84 | 0 | IDSelectorBatch::IDSelectorBatch(size_t n, const idx_t* indices) { |
85 | 0 | nbits = 0; |
86 | 0 | while (n > ((idx_t)1 << nbits)) { |
87 | 0 | nbits++; |
88 | 0 | } |
89 | 0 | nbits += 5; |
90 | | // for n = 1M, nbits = 25 is optimal, see P56659518 |
91 | |
|
92 | 0 | mask = ((idx_t)1 << nbits) - 1; |
93 | 0 | bloom.resize((idx_t)1 << (nbits - 3), 0); |
94 | 0 | for (idx_t i = 0; i < n; i++) { |
95 | 0 | idx_t id = indices[i]; |
96 | 0 | set.insert(id); |
97 | 0 | id &= mask; |
98 | 0 | bloom[id >> 3] |= 1 << (id & 7); |
99 | 0 | } |
100 | 0 | } |
101 | | |
102 | 0 | bool IDSelectorBatch::is_member(idx_t i) const { |
103 | 0 | long im = i & mask; |
104 | 0 | if (!(bloom[im >> 3] & (1 << (im & 7)))) { |
105 | 0 | return 0; |
106 | 0 | } |
107 | 0 | return set.count(i); |
108 | 0 | } |
109 | | |
110 | | /*********************************************************************** |
111 | | * IDSelectorBitmap |
112 | | ***********************************************************************/ |
113 | | |
114 | | IDSelectorBitmap::IDSelectorBitmap(size_t n, const uint8_t* bitmap) |
115 | 0 | : n(n), bitmap(bitmap) {} |
116 | | |
117 | 0 | bool IDSelectorBitmap::is_member(idx_t ii) const { |
118 | 0 | uint64_t i = ii; |
119 | 0 | if ((i >> 3) >= n) { |
120 | 0 | return false; |
121 | 0 | } |
122 | 0 | return (bitmap[i >> 3] >> (i & 7)) & 1; |
123 | 0 | } |
124 | | |
125 | | } // namespace faiss |