usearch 0.1.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.clang-format +10 -0
- package/.github/workflows/package.json +82 -0
- package/.github/workflows/release.yml +111 -0
- package/.gitmodules +6 -0
- package/.vscode/launch.json +116 -0
- package/.vscode/settings.json +109 -0
- package/.vscode/tasks.json +35 -0
- package/CMakeLists.txt +84 -0
- package/CMakeListsInstructuons.md +109 -0
- package/README.md +69 -0
- package/VERSION +1 -0
- package/assets/unum.png +0 -0
- package/binding.gyp +28 -0
- package/build.rs +14 -0
- package/cargo.toml +15 -0
- package/cmake/clipp.cmake +8 -0
- package/cmake/jemalloc.cmake +80 -0
- package/docs/Makefile +20 -0
- package/docs/_static/custom.css +26 -0
- package/docs/_static/custom.js +7 -0
- package/docs/conf.py +48 -0
- package/docs/index.rst +10 -0
- package/fp16/CMakeLists.txt +176 -0
- package/fp16/LICENSE +11 -0
- package/fp16/README.md +20 -0
- package/fp16/bench/alt-element.cc +89 -0
- package/fp16/bench/from-alt-array.cc +175 -0
- package/fp16/bench/from-ieee-array.cc +399 -0
- package/fp16/bench/ieee-element.cc +272 -0
- package/fp16/bench/to-alt-array.cc +91 -0
- package/fp16/bench/to-ieee-array.cc +303 -0
- package/fp16/cmake/DownloadGoogleBenchmark.cmake +15 -0
- package/fp16/cmake/DownloadGoogleTest.cmake +15 -0
- package/fp16/cmake/DownloadPSimd.cmake +15 -0
- package/fp16/configure.py +65 -0
- package/fp16/confu.yaml +8 -0
- package/fp16/include/fp16/__init__.py +0 -0
- package/fp16/include/fp16/avx.py +50 -0
- package/fp16/include/fp16/avx2.py +53 -0
- package/fp16/include/fp16/bitcasts.h +92 -0
- package/fp16/include/fp16/fp16.h +451 -0
- package/fp16/include/fp16/psimd.h +131 -0
- package/fp16/include/fp16.h +11 -0
- package/fp16/jni/Android.mk +6 -0
- package/fp16/jni/Application.mk +4 -0
- package/fp16/test/alt-from-fp32-value.cc +498 -0
- package/fp16/test/alt-to-fp32-bits.cc +263 -0
- package/fp16/test/alt-to-fp32-psimd.cc +145 -0
- package/fp16/test/alt-to-fp32-value.cc +348 -0
- package/fp16/test/alt-to-fp32x2-psimd.cc +245 -0
- package/fp16/test/bitcasts.cc +56 -0
- package/fp16/test/ieee-from-fp32-value.cc +523 -0
- package/fp16/test/ieee-to-fp32-bits.cc +333 -0
- package/fp16/test/ieee-to-fp32-psimd.cc +312 -0
- package/fp16/test/ieee-to-fp32-value.cc +420 -0
- package/fp16/test/ieee-to-fp32x2-psimd.cc +541 -0
- package/fp16/test/peachpy/alt-xmm-to-fp32-xmm-avx.cc +96 -0
- package/fp16/test/peachpy/alt-xmm-to-fp32-ymm-avx2.cc +96 -0
- package/fp16/test/peachpy/stubs.py +38 -0
- package/fp16/test/tables.cc +530 -0
- package/fp16/test/tables.h +13 -0
- package/fp16/third-party/THHalf.h +92 -0
- package/fp16/third-party/eigen-half.h +170 -0
- package/fp16/third-party/float16-compressor.h +80 -0
- package/fp16/third-party/half.hpp +2916 -0
- package/fp16/third-party/npy-halffloat.h +328 -0
- package/include/usearch/usearch.hpp +1345 -0
- package/package.json +19 -0
- package/pyproject.toml +33 -0
- package/setup.py +63 -0
- package/simsimd/.clang-format +10 -0
- package/simsimd/.vscode/launch.json +24 -0
- package/simsimd/.vscode/settings.json +73 -0
- package/simsimd/.vscode/tasks.json +17 -0
- package/simsimd/CMakeLists.txt +82 -0
- package/simsimd/LICENSE +201 -0
- package/simsimd/README.md +73 -0
- package/simsimd/bench.cxx +82 -0
- package/simsimd/include/simsimd/simsimd.h +322 -0
- package/simsimd/include/simsimd/simsimd.hpp +443 -0
- package/src/advanced.hpp +489 -0
- package/src/bench.cpp +415 -0
- package/src/bench.py +89 -0
- package/src/javascript.cpp +198 -0
- package/src/main.rs +42 -0
- package/src/python.cpp +226 -0
- package/src/rust.cpp +23 -0
- package/src/rust.h +22 -0
- package/src/test.js +17 -0
- package/src/test.py +20 -0
|
@@ -0,0 +1,1345 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* @file usearch.hpp
|
|
3
|
+
* @author Ashot Vardanian
|
|
4
|
+
* @brief Single-header Vector Search.
|
|
5
|
+
* @date 2023-04-26
|
|
6
|
+
*
|
|
7
|
+
* @copyright Copyright (c) 2023
|
|
8
|
+
*/
|
|
9
|
+
#pragma once
|
|
10
|
+
#include <atomic>
|
|
11
|
+
#include <bitset> // `std::bitset`
|
|
12
|
+
#include <climits> // `CHAR_BIT`
|
|
13
|
+
#include <cmath> // `std::sqrt`
|
|
14
|
+
#include <cstring> // `std::memset`
|
|
15
|
+
#include <mutex>
|
|
16
|
+
#include <random>
|
|
17
|
+
#include <utility> // `std::exchange`
|
|
18
|
+
#include <vector>
|
|
19
|
+
|
|
20
|
+
#include <fcntl.h> // `fallocate`
|
|
21
|
+
#include <stdlib.h> // `posix_memalign`
|
|
22
|
+
#include <sys/mman.h> // `mmap`
|
|
23
|
+
#include <sys/stat.h> // `fstat` for file size
|
|
24
|
+
#include <unistd.h> // `open`, `close`
|
|
25
|
+
|
|
26
|
+
#if defined(__GNUC__)
|
|
27
|
+
// https://gcc.gnu.org/onlinedocs/gcc/Other-Builtins.html
|
|
28
|
+
// Zero means we are only going to read from that memory.
|
|
29
|
+
// Three means high temporal locality and suggests to keep
|
|
30
|
+
// the data in all layers of cache.
|
|
31
|
+
#define prefetch_m(ptr) __builtin_prefetch((void*)(ptr), 0, 3)
|
|
32
|
+
#elif defined(__x86_64__)
|
|
33
|
+
#define prefetch_m(ptr) _mm_prefetch((void*)(ptr), _MM_HINT_T0)
|
|
34
|
+
#else
|
|
35
|
+
#define prefetch_m(ptr)
|
|
36
|
+
#endif
|
|
37
|
+
|
|
38
|
+
#ifdef NDEBUG
|
|
39
|
+
#define assert_m(must_be_true, message)
|
|
40
|
+
#else
|
|
41
|
+
#define assert_m(must_be_true, message) \
|
|
42
|
+
if (!(must_be_true)) { \
|
|
43
|
+
throw std::runtime_error(message); \
|
|
44
|
+
}
|
|
45
|
+
#endif
|
|
46
|
+
|
|
47
|
+
#define usearch_align_m __attribute__((aligned(64)))
|
|
48
|
+
#define usearch_pack_m __attribute__((packed))
|
|
49
|
+
|
|
50
|
+
namespace unum {
|
|
51
|
+
namespace usearch {
|
|
52
|
+
|
|
53
|
+
using f64_t = double;
|
|
54
|
+
using f32_t = float;
|
|
55
|
+
|
|
56
|
+
/**
|
|
57
|
+
* GCC supports half-precision types.
|
|
58
|
+
* https://gcc.gnu.org/onlinedocs/gcc/Half-Precision.html
|
|
59
|
+
*/
|
|
60
|
+
enum class f16_t : std::uint16_t {};
|
|
61
|
+
|
|
62
|
+
template <typename at> constexpr at angle_to_radians(at angle) noexcept { return angle * M_PI / 180; }
|
|
63
|
+
|
|
64
|
+
template <typename at> constexpr at square(at value) noexcept { return value * value; }
|
|
65
|
+
|
|
66
|
+
inline std::size_t ceil2(std::size_t v) noexcept {
|
|
67
|
+
v--;
|
|
68
|
+
v |= v >> 1;
|
|
69
|
+
v |= v >> 2;
|
|
70
|
+
v |= v >> 4;
|
|
71
|
+
v |= v >> 8;
|
|
72
|
+
v |= v >> 16;
|
|
73
|
+
v |= v >> 32;
|
|
74
|
+
v++;
|
|
75
|
+
return v;
|
|
76
|
+
}
|
|
77
|
+
|
|
78
|
+
template <typename scalar_at> struct ip_gt {
|
|
79
|
+
using scalar_t = scalar_at;
|
|
80
|
+
using result_type = scalar_t;
|
|
81
|
+
|
|
82
|
+
inline scalar_t operator()(scalar_t const* a, scalar_t const* b, std::size_t dim, std::size_t = 0) const noexcept {
|
|
83
|
+
result_type ab{};
|
|
84
|
+
#if defined(__GNUC__)
|
|
85
|
+
#pragma GCC ivdep
|
|
86
|
+
#elif defined(__clang__)
|
|
87
|
+
#pragma clang loop vectorize(enable)
|
|
88
|
+
#elif defined(_OPENMP)
|
|
89
|
+
#pragma omp simd reduction(+ : ab)
|
|
90
|
+
#endif
|
|
91
|
+
for (std::size_t i = 0; i != dim; ++i)
|
|
92
|
+
ab += a[i] * b[i];
|
|
93
|
+
return 1 - ab;
|
|
94
|
+
}
|
|
95
|
+
};
|
|
96
|
+
|
|
97
|
+
template <typename scalar_at> struct cos_gt {
|
|
98
|
+
using scalar_t = scalar_at;
|
|
99
|
+
using result_type = scalar_t;
|
|
100
|
+
|
|
101
|
+
inline scalar_t operator()(scalar_t const* a, scalar_t const* b, std::size_t dim, std::size_t = 0) const noexcept {
|
|
102
|
+
result_type ab{}, a2{}, b2{};
|
|
103
|
+
#if defined(__GNUC__)
|
|
104
|
+
#pragma GCC ivdep
|
|
105
|
+
#elif defined(__clang__)
|
|
106
|
+
#pragma clang loop vectorize(enable)
|
|
107
|
+
#elif defined(_OPENMP)
|
|
108
|
+
#pragma omp simd reduction(+ : ab, a2, b2)
|
|
109
|
+
#endif
|
|
110
|
+
for (std::size_t i = 0; i != dim; ++i)
|
|
111
|
+
ab += a[i] * b[i], a2 += a[i] * a[i], b2 += b[i] * b[i];
|
|
112
|
+
return ab / (std::sqrt(a2) * std::sqrt(b2));
|
|
113
|
+
}
|
|
114
|
+
};
|
|
115
|
+
|
|
116
|
+
template <typename scalar_at> struct l2_squared_gt {
|
|
117
|
+
using scalar_t = scalar_at;
|
|
118
|
+
using result_type = scalar_t;
|
|
119
|
+
|
|
120
|
+
inline scalar_t operator()(scalar_t const* a, scalar_t const* b, std::size_t dim, std::size_t = 0) const noexcept {
|
|
121
|
+
result_type ab_deltas_sq{};
|
|
122
|
+
#if defined(__GNUC__)
|
|
123
|
+
#pragma GCC ivdep
|
|
124
|
+
#elif defined(__clang__)
|
|
125
|
+
#pragma clang loop vectorize(enable)
|
|
126
|
+
#elif defined(_OPENMP)
|
|
127
|
+
#pragma omp simd reduction(+ : ab_deltas_sq)
|
|
128
|
+
#endif
|
|
129
|
+
for (std::size_t i = 0; i != dim; ++i)
|
|
130
|
+
ab_deltas_sq += (a[i] - b[i]) * (a[i] - b[i]);
|
|
131
|
+
return ab_deltas_sq;
|
|
132
|
+
}
|
|
133
|
+
};
|
|
134
|
+
|
|
135
|
+
/**
|
|
136
|
+
* @brief Hamming distance computes the number of differing bits in
|
|
137
|
+
* two arrays of integers. An example would be a textual document,
|
|
138
|
+
* tokenized and hashed into a fixed-capacity bitset.
|
|
139
|
+
*/
|
|
140
|
+
template <typename scalar_at> struct bit_hamming_gt {
|
|
141
|
+
using scalar_t = scalar_at;
|
|
142
|
+
using result_type = std::size_t;
|
|
143
|
+
static_assert(std::is_unsigned<scalar_t>::value, "Hamming distance requires unsigned integral words");
|
|
144
|
+
static constexpr std::size_t bits_per_word_k = sizeof(scalar_t) * CHAR_BIT;
|
|
145
|
+
|
|
146
|
+
inline std::size_t operator()(scalar_t const* a, scalar_t const* b, std::size_t words,
|
|
147
|
+
std::size_t = 0) const noexcept {
|
|
148
|
+
std::size_t matches{};
|
|
149
|
+
#if defined(__GNUC__)
|
|
150
|
+
#pragma GCC ivdep
|
|
151
|
+
#elif defined(__clang__)
|
|
152
|
+
#pragma clang loop vectorize(enable)
|
|
153
|
+
#elif defined(_OPENMP)
|
|
154
|
+
#pragma omp simd reduction(+ : matches)
|
|
155
|
+
#endif
|
|
156
|
+
for (std::size_t i = 0; i != words; ++i)
|
|
157
|
+
matches += std::bitset<bits_per_word_k>(a[i] ^ b[i]).count();
|
|
158
|
+
return matches;
|
|
159
|
+
}
|
|
160
|
+
};
|
|
161
|
+
|
|
162
|
+
/**
|
|
163
|
+
* @brief Counts the number of matching elements in two unique sorted sets.
|
|
164
|
+
* Can be used to compute the similarity between two textual documents
|
|
165
|
+
* using the IDs of tokens present in them.
|
|
166
|
+
*/
|
|
167
|
+
template <typename scalar_at> struct jaccard_gt {
|
|
168
|
+
using scalar_t = scalar_at;
|
|
169
|
+
using result_type = f32_t;
|
|
170
|
+
static_assert(!std::is_floating_point<scalar_t>::value, "Jaccard distance requires integral scalars");
|
|
171
|
+
|
|
172
|
+
inline result_type operator()( //
|
|
173
|
+
scalar_t const* a, scalar_t const* b, std::size_t a_length, std::size_t b_length) const noexcept {
|
|
174
|
+
std::size_t intersection{};
|
|
175
|
+
std::size_t i{};
|
|
176
|
+
std::size_t j{};
|
|
177
|
+
while (i != a_length && j != b_length) {
|
|
178
|
+
intersection += a[i] == b[j];
|
|
179
|
+
i += a[i] < b[j];
|
|
180
|
+
j += a[i] >= b[j];
|
|
181
|
+
}
|
|
182
|
+
return 1.f - intersection * 1.f / (a_length + b_length - intersection);
|
|
183
|
+
}
|
|
184
|
+
};
|
|
185
|
+
|
|
186
|
+
/**
|
|
187
|
+
* @brief Counts the number of matching elements in two unique sorted sets.
|
|
188
|
+
* Can be used to compute the similarity between two textual documents
|
|
189
|
+
* using the IDs of tokens present in them.
|
|
190
|
+
*/
|
|
191
|
+
template <typename scalar_at> struct pearson_correlation_gt {
|
|
192
|
+
using scalar_t = scalar_at;
|
|
193
|
+
using result_type = scalar_t;
|
|
194
|
+
|
|
195
|
+
inline result_type operator()( //
|
|
196
|
+
scalar_t const* a, scalar_t const* b, std::size_t dim, std::size_t = 0) const noexcept {
|
|
197
|
+
scalar_t a_sum{}, b_sum{}, ab_sum{};
|
|
198
|
+
scalar_t a_sq_sum{}, b_sq_sum{};
|
|
199
|
+
#if defined(__GNUC__)
|
|
200
|
+
#pragma GCC ivdep
|
|
201
|
+
#elif defined(__clang__)
|
|
202
|
+
#pragma clang loop vectorize(enable)
|
|
203
|
+
#elif defined(_OPENMP)
|
|
204
|
+
#pragma omp simd reduction(+ : a_sum, b_sum, ab_sum, a_sq_sum, b_sq_sum)
|
|
205
|
+
#endif
|
|
206
|
+
for (std::size_t i = 0; i != dim; ++i) {
|
|
207
|
+
a_sum += a[i];
|
|
208
|
+
b_sum += b[i];
|
|
209
|
+
ab_sum += a[i] * b[i];
|
|
210
|
+
a_sq_sum += a[i] * a[i];
|
|
211
|
+
b_sq_sum += b[i] * b[i];
|
|
212
|
+
}
|
|
213
|
+
result_type denom = std::sqrt((dim * a_sq_sum - a_sum * a_sum) * (dim * b_sq_sum - b_sum * b_sum));
|
|
214
|
+
result_type corr = (dim * ab_sum - a_sum * b_sum) / denom;
|
|
215
|
+
return -corr;
|
|
216
|
+
}
|
|
217
|
+
};
|
|
218
|
+
|
|
219
|
+
/**
|
|
220
|
+
* @brief Haversine distance for the shortest distance between two nodes on
|
|
221
|
+
* the surface of a 3D sphere, defined with latitude and longitude.
|
|
222
|
+
*/
|
|
223
|
+
template <typename scalar_at> struct haversine_gt {
|
|
224
|
+
using scalar_t = scalar_at;
|
|
225
|
+
using result_type = scalar_t;
|
|
226
|
+
static_assert(std::is_floating_point<scalar_t>::value, "Latitude and longitude must be floating-node");
|
|
227
|
+
|
|
228
|
+
inline scalar_t operator()(scalar_t const* a, scalar_t const* b, std::size_t = 2, std::size_t = 2) const noexcept {
|
|
229
|
+
scalar_t lat_a = a[0], lon_a = a[1];
|
|
230
|
+
scalar_t lat_b = b[0], lon_b = b[1];
|
|
231
|
+
|
|
232
|
+
scalar_t lat_delta = angle_to_radians(lat_b - lat_a);
|
|
233
|
+
scalar_t lon_delta = angle_to_radians(lon_b - lon_a);
|
|
234
|
+
|
|
235
|
+
scalar_t converted_lat_a = angle_to_radians(lat_a);
|
|
236
|
+
scalar_t converted_lat_b = angle_to_radians(lat_b);
|
|
237
|
+
|
|
238
|
+
scalar_t x = square(std::sin(lat_delta / 2)) +
|
|
239
|
+
std::cos(converted_lat_a) * std::cos(converted_lat_b) * square(std::sin(lon_delta / 2));
|
|
240
|
+
|
|
241
|
+
return std::atan2(std::sqrt(x), std::sqrt(1 - x));
|
|
242
|
+
}
|
|
243
|
+
};
|
|
244
|
+
|
|
245
|
+
template <std::size_t multiple_ak> inline std::size_t divide_round_up(std::size_t num) noexcept {
|
|
246
|
+
return (num + multiple_ak - 1) / multiple_ak;
|
|
247
|
+
}
|
|
248
|
+
|
|
249
|
+
template <typename allocator_at = std::allocator<char>> class visits_bitset_gt {
|
|
250
|
+
using allocator_t = allocator_at;
|
|
251
|
+
using byte_t = typename allocator_t::value_type;
|
|
252
|
+
static_assert(sizeof(byte_t) == 1, "Allocator must allocate separate addressable bytes");
|
|
253
|
+
|
|
254
|
+
using slot_t = std::uint64_t;
|
|
255
|
+
|
|
256
|
+
std::uint64_t* u64s_{};
|
|
257
|
+
std::size_t slots_{};
|
|
258
|
+
|
|
259
|
+
public:
|
|
260
|
+
visits_bitset_gt() noexcept {}
|
|
261
|
+
visits_bitset_gt(std::size_t capacity) noexcept {
|
|
262
|
+
slots_ = divide_round_up<64>(capacity);
|
|
263
|
+
u64s_ = (slot_t*)allocator_t{}.allocate(slots_ * sizeof(slot_t));
|
|
264
|
+
}
|
|
265
|
+
~visits_bitset_gt() noexcept { allocator_t{}.deallocate((byte_t*)u64s_, slots_ * sizeof(slot_t)), u64s_ = nullptr; }
|
|
266
|
+
void clear() noexcept { std::memset(u64s_, 0, slots_ * sizeof(slot_t)); }
|
|
267
|
+
inline bool test(std::size_t i) const noexcept { return u64s_[i / 64] & (1ul << (i & 63ul)); }
|
|
268
|
+
inline void set(std::size_t i) noexcept { u64s_[i / 64] |= (1ul << (i & 63ul)); }
|
|
269
|
+
|
|
270
|
+
visits_bitset_gt(visits_bitset_gt&& other) noexcept {
|
|
271
|
+
u64s_ = other.u64s_;
|
|
272
|
+
slots_ = other.slots_;
|
|
273
|
+
other.u64s_ = nullptr;
|
|
274
|
+
other.slots_ = 0;
|
|
275
|
+
}
|
|
276
|
+
|
|
277
|
+
visits_bitset_gt& operator=(visits_bitset_gt&& other) noexcept {
|
|
278
|
+
std::swap(u64s_, other.u64s_);
|
|
279
|
+
std::swap(slots_, other.slots_);
|
|
280
|
+
return *this;
|
|
281
|
+
}
|
|
282
|
+
};
|
|
283
|
+
|
|
284
|
+
using visits_bitset_t = visits_bitset_gt<>;
|
|
285
|
+
|
|
286
|
+
/**
|
|
287
|
+
* @brief Similar to `std::priority_queue`, but allows raw access to underlying
|
|
288
|
+
* memory, in case you want to shuffle it or sort.
|
|
289
|
+
*/
|
|
290
|
+
template <typename element_at, //
|
|
291
|
+
typename comparator_at = std::less<void>, // <void> is needed before C++14.
|
|
292
|
+
typename allocator_at = std::allocator<element_at>> //
|
|
293
|
+
class max_heap_gt {
|
|
294
|
+
public:
|
|
295
|
+
using element_t = element_at;
|
|
296
|
+
using comparator_t = comparator_at;
|
|
297
|
+
using allocator_t = allocator_at;
|
|
298
|
+
|
|
299
|
+
using value_type = element_t;
|
|
300
|
+
|
|
301
|
+
private:
|
|
302
|
+
element_t* elements_;
|
|
303
|
+
std::size_t size_;
|
|
304
|
+
std::size_t capacity_;
|
|
305
|
+
std::size_t max_capacity_;
|
|
306
|
+
|
|
307
|
+
public:
|
|
308
|
+
max_heap_gt(std::size_t max_capacity = 0) noexcept
|
|
309
|
+
: elements_(nullptr), size_(0), capacity_(0), max_capacity_(max_capacity) {}
|
|
310
|
+
|
|
311
|
+
bool empty() const noexcept { return !size_; }
|
|
312
|
+
std::size_t size() const noexcept { return size_; }
|
|
313
|
+
std::size_t capacity() const noexcept { return capacity_; }
|
|
314
|
+
element_t const& top() const noexcept { return elements_[0]; }
|
|
315
|
+
element_t const& max() const noexcept { return elements_[0]; }
|
|
316
|
+
|
|
317
|
+
void clear() noexcept {
|
|
318
|
+
while (size_) {
|
|
319
|
+
size_--;
|
|
320
|
+
elements_[size_].~element_t();
|
|
321
|
+
}
|
|
322
|
+
}
|
|
323
|
+
|
|
324
|
+
bool reserve(std::size_t n) noexcept {
|
|
325
|
+
if (n < capacity_)
|
|
326
|
+
return true;
|
|
327
|
+
if (max_capacity_ && capacity_ == max_capacity_)
|
|
328
|
+
return false;
|
|
329
|
+
|
|
330
|
+
auto new_capacity = std::max<std::size_t>(capacity_ * 2u, 16u);
|
|
331
|
+
if (max_capacity_)
|
|
332
|
+
new_capacity = std::min(new_capacity, max_capacity_);
|
|
333
|
+
|
|
334
|
+
auto allocator = allocator_t{};
|
|
335
|
+
auto new_elements = allocator.allocate(new_capacity);
|
|
336
|
+
if (!new_elements)
|
|
337
|
+
return false;
|
|
338
|
+
|
|
339
|
+
if (elements_) {
|
|
340
|
+
std::uninitialized_copy_n(elements_, size_, new_elements);
|
|
341
|
+
allocator.deallocate(elements_, capacity_);
|
|
342
|
+
}
|
|
343
|
+
elements_ = new_elements;
|
|
344
|
+
capacity_ = new_capacity;
|
|
345
|
+
return new_elements;
|
|
346
|
+
}
|
|
347
|
+
|
|
348
|
+
template <typename... args_at> //
|
|
349
|
+
bool emplace(args_at&&... args) noexcept {
|
|
350
|
+
if (!reserve(size_ + 1))
|
|
351
|
+
return false;
|
|
352
|
+
|
|
353
|
+
new (&elements_[size_]) element_t({std::forward<args_at>(args)...});
|
|
354
|
+
size_++;
|
|
355
|
+
shift_up(size_ - 1);
|
|
356
|
+
return true;
|
|
357
|
+
}
|
|
358
|
+
|
|
359
|
+
element_t pop() noexcept {
|
|
360
|
+
element_t result = top();
|
|
361
|
+
std::swap(elements_[0], elements_[size_ - 1]);
|
|
362
|
+
size_--;
|
|
363
|
+
elements_[size_].~element_t();
|
|
364
|
+
shift_down(0);
|
|
365
|
+
return result;
|
|
366
|
+
}
|
|
367
|
+
|
|
368
|
+
/** @brief Invalidates the "max-heap" property, transforming into ascending range. */
|
|
369
|
+
void sort_ascending() noexcept { std::sort_heap(elements_, elements_ + size_, &less); }
|
|
370
|
+
/** @brief Invalidates the "max-heap" property, transforming into descending range. */
|
|
371
|
+
void sort_descending() noexcept { sort_ascending(), std::reverse(elements_, elements_ + size_); }
|
|
372
|
+
void sort_heap() noexcept { std::make_heap(elements_, elements_ + size_, &less); }
|
|
373
|
+
void shrink(std::size_t n) noexcept { size_ = n; }
|
|
374
|
+
|
|
375
|
+
element_t* data() noexcept { return elements_; }
|
|
376
|
+
element_t const* data() const noexcept { return elements_; }
|
|
377
|
+
|
|
378
|
+
private:
|
|
379
|
+
inline std::size_t parent_idx(std::size_t i) const noexcept { return (i - 1u) / 2u; }
|
|
380
|
+
inline std::size_t left_child_idx(std::size_t i) const noexcept { return (i * 2u) + 1u; }
|
|
381
|
+
inline std::size_t right_child_idx(std::size_t i) const noexcept { return (i * 2u) + 2u; }
|
|
382
|
+
static bool less(element_t const& a, element_t const& b) noexcept { return comparator_t{}(a, b); }
|
|
383
|
+
|
|
384
|
+
void shift_up(std::size_t i) noexcept {
|
|
385
|
+
for (; i && less(elements_[parent_idx(i)], elements_[i]); i = parent_idx(i))
|
|
386
|
+
std::swap(elements_[parent_idx(i)], elements_[i]);
|
|
387
|
+
}
|
|
388
|
+
|
|
389
|
+
void shift_down(std::size_t i) noexcept {
|
|
390
|
+
std::size_t max_idx = i;
|
|
391
|
+
|
|
392
|
+
std::size_t left = left_child_idx(i);
|
|
393
|
+
if (left < size_ && less(elements_[max_idx], elements_[left]))
|
|
394
|
+
max_idx = left;
|
|
395
|
+
|
|
396
|
+
std::size_t right = right_child_idx(i);
|
|
397
|
+
if (right < size_ && less(elements_[max_idx], elements_[right]))
|
|
398
|
+
max_idx = right;
|
|
399
|
+
|
|
400
|
+
if (i != max_idx) {
|
|
401
|
+
std::swap(elements_[i], elements_[max_idx]);
|
|
402
|
+
shift_down(max_idx);
|
|
403
|
+
}
|
|
404
|
+
}
|
|
405
|
+
};
|
|
406
|
+
|
|
407
|
+
/**
|
|
408
|
+
*
|
|
409
|
+
*/
|
|
410
|
+
class mutex_t {
|
|
411
|
+
using slot_t = std::int32_t;
|
|
412
|
+
slot_t flag_;
|
|
413
|
+
|
|
414
|
+
public:
|
|
415
|
+
inline mutex_t(slot_t flag = 0) noexcept : flag_(flag) {}
|
|
416
|
+
inline ~mutex_t() noexcept {}
|
|
417
|
+
|
|
418
|
+
inline bool try_lock() noexcept {
|
|
419
|
+
slot_t raw = 0;
|
|
420
|
+
return __atomic_compare_exchange_n(&flag_, &raw, 1, true, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED);
|
|
421
|
+
}
|
|
422
|
+
|
|
423
|
+
inline void lock() noexcept {
|
|
424
|
+
slot_t raw;
|
|
425
|
+
lock_again:
|
|
426
|
+
raw = 0;
|
|
427
|
+
if (!__atomic_compare_exchange_n(&flag_, &raw, 1, true, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED))
|
|
428
|
+
goto lock_again;
|
|
429
|
+
}
|
|
430
|
+
|
|
431
|
+
inline void unlock() noexcept { __atomic_store_n(&flag_, 0, __ATOMIC_RELEASE); }
|
|
432
|
+
};
|
|
433
|
+
|
|
434
|
+
static_assert(sizeof(mutex_t) == sizeof(std::int32_t), "Mutex is larger than expected");
|
|
435
|
+
|
|
436
|
+
using lock_t = std::unique_lock<mutex_t>;
|
|
437
|
+
|
|
438
|
+
/**
|
|
439
|
+
* @brief Five-byte integer type to address node clouds with over 4B entries.
|
|
440
|
+
*/
|
|
441
|
+
class usearch_pack_m uint40_t {
|
|
442
|
+
unsigned char octets[5];
|
|
443
|
+
|
|
444
|
+
public:
|
|
445
|
+
inline uint40_t() noexcept { std::memset(octets, 0, 5); }
|
|
446
|
+
inline uint40_t(std::uint32_t n) noexcept { std::memcpy(octets + 1, (char*)&n, 4), octets[0] = 0; }
|
|
447
|
+
inline uint40_t(std::uint64_t n) noexcept { std::memcpy(octets, (char*)&n + 3, 5); }
|
|
448
|
+
#if defined(__clang__)
|
|
449
|
+
inline uint40_t(std::size_t n) noexcept { std::memcpy(octets, (char*)&n + 3, 5); }
|
|
450
|
+
#endif
|
|
451
|
+
|
|
452
|
+
uint40_t(uint40_t&&) = default;
|
|
453
|
+
uint40_t(uint40_t const&) = default;
|
|
454
|
+
uint40_t& operator=(uint40_t&&) = default;
|
|
455
|
+
uint40_t& operator=(uint40_t const&) = default;
|
|
456
|
+
|
|
457
|
+
inline uint40_t& operator+=(std::uint32_t i) noexcept {
|
|
458
|
+
std::uint32_t& tail = *reinterpret_cast<std::uint32_t*>(octets + 1);
|
|
459
|
+
octets[0] += static_cast<unsigned char>((tail + i) < tail);
|
|
460
|
+
tail += i;
|
|
461
|
+
return *this;
|
|
462
|
+
}
|
|
463
|
+
|
|
464
|
+
inline operator std::size_t() const noexcept {
|
|
465
|
+
std::size_t result = 0;
|
|
466
|
+
std::memcpy((char*)&result + 3, octets, 5);
|
|
467
|
+
return result;
|
|
468
|
+
}
|
|
469
|
+
|
|
470
|
+
inline uint40_t& operator++() noexcept { return *this += 1; }
|
|
471
|
+
|
|
472
|
+
inline uint40_t operator++(int) noexcept {
|
|
473
|
+
uint40_t old = *this;
|
|
474
|
+
*this += 1;
|
|
475
|
+
return old;
|
|
476
|
+
}
|
|
477
|
+
};
|
|
478
|
+
|
|
479
|
+
static_assert(sizeof(uint40_t) == 5, "uint40_t must be exactly 5 bytes");
|
|
480
|
+
|
|
481
|
+
template <typename scalar_at> class span_gt {
|
|
482
|
+
scalar_at* data_;
|
|
483
|
+
std::size_t size_;
|
|
484
|
+
|
|
485
|
+
public:
|
|
486
|
+
span_gt(scalar_at* begin, scalar_at* end) noexcept : data_(begin), size_(end - begin) {}
|
|
487
|
+
span_gt(scalar_at* begin, std::size_t size) noexcept : data_(begin), size_(size) {}
|
|
488
|
+
scalar_at* data() const noexcept { return data_; }
|
|
489
|
+
std::size_t size() const noexcept { return size_; }
|
|
490
|
+
operator scalar_at*() const noexcept { return data(); }
|
|
491
|
+
};
|
|
492
|
+
|
|
493
|
+
/**
|
|
494
|
+
* @brief
|
|
495
|
+
* Configuration settings for the index construction.
|
|
496
|
+
* Includes the main `::connectivity` parameter (`M` in the paper)
|
|
497
|
+
* and two expansion factors - for construction and search.
|
|
498
|
+
*/
|
|
499
|
+
struct config_t {
|
|
500
|
+
|
|
501
|
+
/// @brief Number of neighbors per graph node.
|
|
502
|
+
/// Defaults to 32 in FAISS and 16 in hnswlib.
|
|
503
|
+
/// > It is called `M` in the paper.
|
|
504
|
+
static constexpr std::size_t connectivity_default_k = 16;
|
|
505
|
+
|
|
506
|
+
/// @brief Hyper-parameter controlling the quality of indexing.
|
|
507
|
+
/// Defaults to 40 in FAISS and 200 in hnswlib.
|
|
508
|
+
/// > It is called `efConstruction` in the paper.
|
|
509
|
+
static constexpr std::size_t expansion_construction_default_k = 128;
|
|
510
|
+
|
|
511
|
+
/// @brief Hyper-parameter controlling the quality of search.
|
|
512
|
+
/// Defaults to 16 in FAISS and 10 in hnswlib.
|
|
513
|
+
/// > It is called `ef` in the paper.
|
|
514
|
+
static constexpr std::size_t expansion_search_default_k = 64;
|
|
515
|
+
|
|
516
|
+
std::size_t connectivity = connectivity_default_k;
|
|
517
|
+
std::size_t expansion_add = expansion_construction_default_k;
|
|
518
|
+
std::size_t expansion_search = expansion_search_default_k;
|
|
519
|
+
|
|
520
|
+
///
|
|
521
|
+
std::size_t max_elements = 0;
|
|
522
|
+
std::size_t max_threads_add = 64;
|
|
523
|
+
std::size_t max_threads_search = 64;
|
|
524
|
+
};
|
|
525
|
+
|
|
526
|
+
/**
|
|
527
|
+
* @brief
|
|
528
|
+
* Approximate Nearest Neighbors Search index using the
|
|
529
|
+
* Hierarchical Navigable Small World graph algorithm.
|
|
530
|
+
*
|
|
531
|
+
* @section Features
|
|
532
|
+
* - Search for vectors of different dimensionality.
|
|
533
|
+
* - Thread-safe.
|
|
534
|
+
* - Bring your threads!
|
|
535
|
+
*
|
|
536
|
+
* @tparam metric_at
|
|
537
|
+
* A function vector responsible for computing the distance between two vectors.
|
|
538
|
+
* Must overload the call operator with the following signature:
|
|
539
|
+
* - `result_type (*) (scalar_t const *, scalar_t const *, std::size_t, std::size_t)`
|
|
540
|
+
*
|
|
541
|
+
* @tparam label_at
|
|
542
|
+
* The type of unique labels to assign to vectors.
|
|
543
|
+
*
|
|
544
|
+
* @tparam id_at
|
|
545
|
+
* The smallest unsigned integer type to address indexed elements.
|
|
546
|
+
* Can be a built-in `uint32_t`, `uint64_t`, or our custom `uint40_t`.
|
|
547
|
+
*
|
|
548
|
+
* @tparam allocator_at
|
|
549
|
+
* Dynamic memory allocator.
|
|
550
|
+
*/
|
|
551
|
+
template <typename metric_at = ip_gt<float>, //
|
|
552
|
+
typename label_at = std::size_t, //
|
|
553
|
+
typename id_at = std::uint32_t, //
|
|
554
|
+
typename scalar_at = float, //
|
|
555
|
+
typename allocator_at = std::allocator<char>> //
|
|
556
|
+
class index_gt {
|
|
557
|
+
public:
|
|
558
|
+
using metric_t = metric_at;
|
|
559
|
+
using scalar_t = scalar_at;
|
|
560
|
+
using label_t = label_at;
|
|
561
|
+
using id_t = id_at;
|
|
562
|
+
using allocator_t = allocator_at;
|
|
563
|
+
|
|
564
|
+
using distance_t =
|
|
565
|
+
typename std::result_of<metric_t(scalar_t const*, scalar_t const*, std::size_t, std::size_t)>::type;
|
|
566
|
+
|
|
567
|
+
private:
|
|
568
|
+
using neighbors_count_t = id_t;
|
|
569
|
+
using dim_t = std::uint32_t;
|
|
570
|
+
using level_t = std::int32_t;
|
|
571
|
+
|
|
572
|
+
using allocator_traits_t = std::allocator_traits<allocator_t>;
|
|
573
|
+
using byte_t = typename allocator_t::value_type;
|
|
574
|
+
static_assert(sizeof(byte_t) == 1, "Allocator must allocate separate addressable bytes");
|
|
575
|
+
static constexpr std::size_t base_level_multiple_k = 2;
|
|
576
|
+
|
|
577
|
+
using visits_bitset_t = visits_bitset_gt<allocator_t>;
|
|
578
|
+
|
|
579
|
+
struct precomputed_constants_t {
|
|
580
|
+
double inverse_log_connectivity{};
|
|
581
|
+
std::size_t connectivity_max_base{};
|
|
582
|
+
std::size_t neighbors_bytes{};
|
|
583
|
+
std::size_t neighbors_base_bytes{};
|
|
584
|
+
std::size_t mutex_bytes{};
|
|
585
|
+
};
|
|
586
|
+
struct distance_and_id_t {
|
|
587
|
+
distance_t first;
|
|
588
|
+
id_t second;
|
|
589
|
+
};
|
|
590
|
+
struct compare_by_distance_t {
|
|
591
|
+
inline bool operator()(distance_and_id_t a, distance_and_id_t b) const noexcept { return a.first < b.first; }
|
|
592
|
+
};
|
|
593
|
+
|
|
594
|
+
using distances_and_ids_allocator_t = typename allocator_traits_t::template rebind_alloc<distance_and_id_t>;
|
|
595
|
+
using distances_and_ids_t = max_heap_gt<distance_and_id_t, compare_by_distance_t, distances_and_ids_allocator_t>;
|
|
596
|
+
|
|
597
|
+
struct neighbors_ref_t {
|
|
598
|
+
neighbors_count_t& count;
|
|
599
|
+
id_t* neighbors{};
|
|
600
|
+
|
|
601
|
+
inline neighbors_ref_t(byte_t* tape) noexcept
|
|
602
|
+
: count(*(neighbors_count_t*)tape), neighbors((neighbors_count_t*)tape + 1) {}
|
|
603
|
+
};
|
|
604
|
+
|
|
605
|
+
struct usearch_pack_m node_head_t {
|
|
606
|
+
label_t label;
|
|
607
|
+
dim_t dim;
|
|
608
|
+
level_t level;
|
|
609
|
+
// Variable length array, that has multiple similarly structured segments.
|
|
610
|
+
// Each starts with a `neighbors_count_t` and is followed by such number of `id_t`s.
|
|
611
|
+
byte_t neighbors[1];
|
|
612
|
+
};
|
|
613
|
+
static constexpr std::size_t head_bytes_k = sizeof(label_t) + sizeof(dim_t) + sizeof(level_t);
|
|
614
|
+
|
|
615
|
+
struct node_t {
|
|
616
|
+
byte_t* tape_;
|
|
617
|
+
scalar_t* vector_;
|
|
618
|
+
};
|
|
619
|
+
|
|
620
|
+
class node_ref_t {
|
|
621
|
+
mutex_t* mutex_{};
|
|
622
|
+
|
|
623
|
+
public:
|
|
624
|
+
node_head_t& head{};
|
|
625
|
+
scalar_t* vector{};
|
|
626
|
+
|
|
627
|
+
inline node_ref_t(mutex_t& m, node_head_t& h, scalar_t* s) noexcept : mutex_(&m), head(h), vector(s) {}
|
|
628
|
+
inline lock_t lock() const noexcept { return mutex_ ? lock_t{*mutex_} : lock_t{}; }
|
|
629
|
+
inline operator node_t() const noexcept { return {mutex_ ? (byte_t*)mutex_ : (byte_t*)&head, vector}; }
|
|
630
|
+
};
|
|
631
|
+
|
|
632
|
+
struct usearch_align_m thread_context_t {
|
|
633
|
+
distances_and_ids_t top_candidates;
|
|
634
|
+
distances_and_ids_t candidates_set;
|
|
635
|
+
distances_and_ids_t candidates_filter;
|
|
636
|
+
visits_bitset_t visits;
|
|
637
|
+
std::default_random_engine level_generator;
|
|
638
|
+
metric_t metric;
|
|
639
|
+
};
|
|
640
|
+
|
|
641
|
+
config_t config_{};
|
|
642
|
+
precomputed_constants_t pre_;
|
|
643
|
+
int viewed_file_descriptor_{};
|
|
644
|
+
|
|
645
|
+
usearch_align_m mutable std::atomic<std::size_t> capacity_{};
|
|
646
|
+
usearch_align_m mutable std::atomic<std::size_t> size_{};
|
|
647
|
+
|
|
648
|
+
mutex_t global_mutex_{};
|
|
649
|
+
level_t max_level_{};
|
|
650
|
+
id_t entry_id_{};
|
|
651
|
+
|
|
652
|
+
using node_allocator_t = typename allocator_traits_t::template rebind_alloc<node_t>;
|
|
653
|
+
std::vector<node_t, node_allocator_t> nodes_{};
|
|
654
|
+
|
|
655
|
+
using thread_context_allocator_t = typename allocator_traits_t::template rebind_alloc<thread_context_t>;
|
|
656
|
+
mutable std::vector<thread_context_t, thread_context_allocator_t> thread_contexts_;
|
|
657
|
+
|
|
658
|
+
public:
|
|
659
|
+
std::size_t connectivity() const noexcept { return config_.connectivity; }
|
|
660
|
+
std::size_t capacity() const noexcept { return capacity_; }
|
|
661
|
+
std::size_t size() const noexcept { return size_; }
|
|
662
|
+
std::size_t max_threads_add() const noexcept { return config_.max_threads_add; }
|
|
663
|
+
bool is_immutable() const noexcept { return viewed_file_descriptor_ != 0; }
|
|
664
|
+
bool synchronize() const noexcept { return config_.max_threads_add > 1; }
|
|
665
|
+
|
|
666
|
+
index_gt(config_t config = {}, metric_t metric = {}, allocator_t = {}) : config_(config) {
|
|
667
|
+
|
|
668
|
+
// Externally defined hyper-parameters:
|
|
669
|
+
config_.expansion_add = std::max(config_.expansion_add, config_.connectivity);
|
|
670
|
+
pre_ = precompute(config);
|
|
671
|
+
|
|
672
|
+
// Configure initial empty state:
|
|
673
|
+
size_ = 0u;
|
|
674
|
+
max_level_ = -1;
|
|
675
|
+
entry_id_ = 0u;
|
|
676
|
+
viewed_file_descriptor_ = 0;
|
|
677
|
+
|
|
678
|
+
// Dynamic memory:
|
|
679
|
+
thread_contexts_.resize(std::max(config.max_threads_search, config.max_threads_add));
|
|
680
|
+
for (thread_context_t& context : thread_contexts_)
|
|
681
|
+
context.metric = metric;
|
|
682
|
+
reserve(config.max_elements);
|
|
683
|
+
}
|
|
684
|
+
|
|
685
|
+
~index_gt() noexcept { clear(); }
|
|
686
|
+
|
|
687
|
+
#pragma region Adjusting Configuration
|
|
688
|
+
|
|
689
|
+
void clear() noexcept {
|
|
690
|
+
std::size_t n = size_;
|
|
691
|
+
for (std::size_t i = 0; i != n; ++i)
|
|
692
|
+
node_free(i);
|
|
693
|
+
size_ = 0;
|
|
694
|
+
max_level_ = -1;
|
|
695
|
+
entry_id_ = 0u;
|
|
696
|
+
}
|
|
697
|
+
|
|
698
|
+
void reserve(std::size_t new_capacity) noexcept(false) {
|
|
699
|
+
|
|
700
|
+
assert_m(new_capacity >= size_, "Can't drop existing values");
|
|
701
|
+
nodes_.resize(new_capacity);
|
|
702
|
+
for (thread_context_t& context : thread_contexts_)
|
|
703
|
+
context.visits = visits_bitset_t(new_capacity);
|
|
704
|
+
|
|
705
|
+
capacity_ = new_capacity;
|
|
706
|
+
}
|
|
707
|
+
|
|
708
|
+
static config_t optimize(config_t const& config) noexcept {
|
|
709
|
+
precomputed_constants_t pre = precompute(config);
|
|
710
|
+
std::size_t bytes_per_node_base = head_bytes_k + pre.neighbors_base_bytes + pre.mutex_bytes;
|
|
711
|
+
std::size_t rounded_size = divide_round_up<64>(bytes_per_node_base) * 64;
|
|
712
|
+
std::size_t added_connections = (rounded_size - rounded_size) / sizeof(id_t);
|
|
713
|
+
config_t result = config;
|
|
714
|
+
result.connectivity = config.connectivity + added_connections / base_level_multiple_k;
|
|
715
|
+
return result;
|
|
716
|
+
}
|
|
717
|
+
|
|
718
|
+
#pragma endregion
|
|
719
|
+
|
|
720
|
+
#pragma region Construction and Search
|
|
721
|
+
|
|
722
|
+
id_t add( //
|
|
723
|
+
label_t new_label, span_gt<scalar_t const> new_span, //
|
|
724
|
+
std::size_t thread_idx = 0, bool store_vector = true) {
|
|
725
|
+
|
|
726
|
+
assert_m(!is_immutable(), "Can't add to an immutable index");
|
|
727
|
+
id_t new_id = static_cast<id_t>(size_.fetch_add(1));
|
|
728
|
+
scalar_t const* new_vector = new_span.data();
|
|
729
|
+
std::size_t new_dim = new_span.size();
|
|
730
|
+
|
|
731
|
+
// Determining how much memory to allocate depends on the target level.
|
|
732
|
+
lock_t new_level_lock(global_mutex_);
|
|
733
|
+
level_t max_level = max_level_;
|
|
734
|
+
thread_context_t& context = thread_contexts_[thread_idx];
|
|
735
|
+
level_t new_target_level = choose_random_level(context.level_generator);
|
|
736
|
+
if (new_target_level <= max_level)
|
|
737
|
+
new_level_lock.unlock();
|
|
738
|
+
|
|
739
|
+
// Allocate the neighbors
|
|
740
|
+
node_ref_t new_node = node_malloc(new_label, new_vector, new_dim, new_target_level, store_vector);
|
|
741
|
+
lock_t new_lock = new_node.lock();
|
|
742
|
+
nodes_[new_id] = new_node;
|
|
743
|
+
|
|
744
|
+
// Do nothing for the first element
|
|
745
|
+
if (!new_id) {
|
|
746
|
+
max_level_ = new_target_level;
|
|
747
|
+
return new_id;
|
|
748
|
+
}
|
|
749
|
+
|
|
750
|
+
// Go down the level, tracking only the closest match
|
|
751
|
+
id_t closest_id = entry_id_;
|
|
752
|
+
distance_t closest_dist =
|
|
753
|
+
context.metric(new_vector, node(closest_id).vector, new_dim, node(closest_id).head.dim);
|
|
754
|
+
for (level_t level = max_level; level > new_target_level; level--) {
|
|
755
|
+
bool changed;
|
|
756
|
+
do {
|
|
757
|
+
changed = false;
|
|
758
|
+
node_ref_t closest_node = node(closest_id);
|
|
759
|
+
lock_t closest_lock = closest_node.lock();
|
|
760
|
+
neighbors_ref_t closest_header = neighbors_non_base(closest_node, level);
|
|
761
|
+
iterate_through_neighbors(closest_header, [&](id_t candidate_id) noexcept {
|
|
762
|
+
node_ref_t candidate_node = node(candidate_id);
|
|
763
|
+
distance_t candidate_dist =
|
|
764
|
+
context.metric(new_vector, candidate_node.vector, new_dim, candidate_node.head.dim);
|
|
765
|
+
if (candidate_dist < closest_dist) {
|
|
766
|
+
closest_dist = candidate_dist;
|
|
767
|
+
closest_id = candidate_id;
|
|
768
|
+
changed = true;
|
|
769
|
+
}
|
|
770
|
+
});
|
|
771
|
+
} while (changed);
|
|
772
|
+
}
|
|
773
|
+
|
|
774
|
+
// From `new_target_level` down perform proper extensive search.
|
|
775
|
+
for (level_t level = std::min(new_target_level, max_level); level >= 0; level--) {
|
|
776
|
+
search_to_insert(closest_id, new_vector, new_dim, level, context);
|
|
777
|
+
closest_id = connect_new_element(new_id, level, context);
|
|
778
|
+
}
|
|
779
|
+
|
|
780
|
+
// Releasing lock for the maximum level
|
|
781
|
+
if (new_target_level > max_level) {
|
|
782
|
+
entry_id_ = new_id;
|
|
783
|
+
max_level_ = new_target_level;
|
|
784
|
+
}
|
|
785
|
+
return new_id;
|
|
786
|
+
}
|
|
787
|
+
|
|
788
|
+
template <typename label_and_distance_callback_at>
|
|
789
|
+
void search( //
|
|
790
|
+
span_gt<scalar_t const> query_span, std::size_t wanted, //
|
|
791
|
+
label_and_distance_callback_at&& callback, std::size_t thread_idx = 0) const {
|
|
792
|
+
|
|
793
|
+
if (!size_)
|
|
794
|
+
return;
|
|
795
|
+
|
|
796
|
+
scalar_t const* query_vec = query_span.data();
|
|
797
|
+
std::size_t query_dim = query_span.size();
|
|
798
|
+
|
|
799
|
+
// Go down the level, tracking only the closest match
|
|
800
|
+
thread_context_t& context = thread_contexts_[thread_idx];
|
|
801
|
+
id_t closest_id = entry_id_;
|
|
802
|
+
distance_t closest_dist =
|
|
803
|
+
context.metric(query_vec, node(closest_id).vector, query_dim, node(closest_id).head.dim);
|
|
804
|
+
for (level_t level = max_level_; level > 0; level--) {
|
|
805
|
+
bool changed;
|
|
806
|
+
do {
|
|
807
|
+
changed = false;
|
|
808
|
+
node_ref_t closest_node = node(closest_id);
|
|
809
|
+
neighbors_ref_t closest_header = neighbors_non_base(closest_node, level);
|
|
810
|
+
iterate_through_neighbors(closest_header, [&](id_t candidate_id) noexcept {
|
|
811
|
+
node_ref_t candidate_node = node(candidate_id);
|
|
812
|
+
distance_t candidate_dist = context.metric( //
|
|
813
|
+
query_vec, candidate_node.vector, query_dim, candidate_node.head.dim);
|
|
814
|
+
if (candidate_dist < closest_dist) {
|
|
815
|
+
closest_dist = candidate_dist;
|
|
816
|
+
closest_id = candidate_id;
|
|
817
|
+
changed = true;
|
|
818
|
+
}
|
|
819
|
+
});
|
|
820
|
+
|
|
821
|
+
} while (changed);
|
|
822
|
+
}
|
|
823
|
+
|
|
824
|
+
// For bottom layer we need a more optimized procedure
|
|
825
|
+
search_to_find_in_base( //
|
|
826
|
+
closest_id, query_vec, query_dim, std::max(config_.expansion_search, wanted), context);
|
|
827
|
+
while (context.top_candidates.size() > wanted)
|
|
828
|
+
context.top_candidates.pop();
|
|
829
|
+
|
|
830
|
+
while (context.top_candidates.size()) {
|
|
831
|
+
distance_and_id_t top = context.top_candidates.top();
|
|
832
|
+
callback(node(top.second).head.label, top.first);
|
|
833
|
+
context.top_candidates.pop();
|
|
834
|
+
}
|
|
835
|
+
}
|
|
836
|
+
|
|
837
|
+
std::size_t search( //
|
|
838
|
+
span_gt<scalar_t const> query_span, std::size_t wanted, //
|
|
839
|
+
label_t* matches, distance_t* distances, //
|
|
840
|
+
std::size_t thread_idx = 0) const {
|
|
841
|
+
|
|
842
|
+
std::size_t found = 0;
|
|
843
|
+
auto callback = [&](label_t label, distance_t distance) noexcept {
|
|
844
|
+
if (matches)
|
|
845
|
+
matches[found] = label;
|
|
846
|
+
if (distances)
|
|
847
|
+
distances[found] = distance;
|
|
848
|
+
++found;
|
|
849
|
+
};
|
|
850
|
+
search(query_span, wanted, callback, thread_idx);
|
|
851
|
+
std::reverse(matches, matches + found * (matches != nullptr));
|
|
852
|
+
std::reverse(distances, distances + found * (distances != nullptr));
|
|
853
|
+
return found;
|
|
854
|
+
}
|
|
855
|
+
|
|
856
|
+
#pragma endregion
|
|
857
|
+
|
|
858
|
+
#pragma region Serialization
|
|
859
|
+
|
|
860
|
+
struct state_t {
|
|
861
|
+
// Check compatibility
|
|
862
|
+
std::uint64_t bytes_per_label{};
|
|
863
|
+
std::uint64_t bytes_per_id{};
|
|
864
|
+
|
|
865
|
+
// Describe state
|
|
866
|
+
std::uint64_t connectivity{};
|
|
867
|
+
std::uint64_t size{};
|
|
868
|
+
std::uint64_t entry_id{};
|
|
869
|
+
std::uint64_t max_level{};
|
|
870
|
+
};
|
|
871
|
+
|
|
872
|
+
/**
|
|
873
|
+
* Compatibility: Linux, MacOS, Windows.
|
|
874
|
+
*/
|
|
875
|
+
void save(char const* file_path) const noexcept(false) {
|
|
876
|
+
|
|
877
|
+
state_t state;
|
|
878
|
+
state.connectivity = config_.connectivity;
|
|
879
|
+
state.size = size_;
|
|
880
|
+
state.entry_id = entry_id_;
|
|
881
|
+
state.max_level = max_level_;
|
|
882
|
+
|
|
883
|
+
std::FILE* file = std::fopen(file_path, "w");
|
|
884
|
+
if (!file)
|
|
885
|
+
throw std::runtime_error(std::strerror(errno));
|
|
886
|
+
|
|
887
|
+
// Write the header
|
|
888
|
+
{
|
|
889
|
+
std::size_t written = std::fwrite(&state, sizeof(state), 1, file);
|
|
890
|
+
if (!written) {
|
|
891
|
+
std::fclose(file);
|
|
892
|
+
throw std::runtime_error(std::strerror(errno));
|
|
893
|
+
}
|
|
894
|
+
}
|
|
895
|
+
|
|
896
|
+
// Serialize nodes one by one
|
|
897
|
+
for (std::size_t i = 0; i != state.size; ++i) {
|
|
898
|
+
node_ref_t node_ref = node(static_cast<id_t>(i));
|
|
899
|
+
std::size_t bytes_to_dump = node_dump_size(node_ref.head.dim, node_ref.head.level);
|
|
900
|
+
std::size_t written = std::fwrite(&node_ref.head, bytes_to_dump - node_ref.head.dim, 1, file);
|
|
901
|
+
if (!written) {
|
|
902
|
+
std::fclose(file);
|
|
903
|
+
throw std::runtime_error(std::strerror(errno));
|
|
904
|
+
}
|
|
905
|
+
written = std::fwrite(node_ref.vector, node_ref.head.dim, 1, file);
|
|
906
|
+
if (!written) {
|
|
907
|
+
std::fclose(file);
|
|
908
|
+
throw std::runtime_error(std::strerror(errno));
|
|
909
|
+
}
|
|
910
|
+
}
|
|
911
|
+
|
|
912
|
+
std::fclose(file);
|
|
913
|
+
}
|
|
914
|
+
|
|
915
|
+
/**
|
|
916
|
+
* Compatibility: Linux, MacOS, Windows.
|
|
917
|
+
*/
|
|
918
|
+
void load(char const* file_path) noexcept(false) {
|
|
919
|
+
state_t state;
|
|
920
|
+
std::FILE* file = std::fopen(file_path, "r");
|
|
921
|
+
if (!file)
|
|
922
|
+
throw std::runtime_error(std::strerror(errno));
|
|
923
|
+
|
|
924
|
+
// Read the header
|
|
925
|
+
{
|
|
926
|
+
std::size_t read = std::fread(&state, sizeof(state), 1, file);
|
|
927
|
+
if (!read) {
|
|
928
|
+
std::fclose(file);
|
|
929
|
+
throw std::runtime_error(std::strerror(errno));
|
|
930
|
+
}
|
|
931
|
+
config_.connectivity = state.connectivity;
|
|
932
|
+
config_.max_elements = state.size;
|
|
933
|
+
pre_ = precompute(config_);
|
|
934
|
+
reserve(state.size);
|
|
935
|
+
size_ = state.size;
|
|
936
|
+
max_level_ = state.max_level;
|
|
937
|
+
entry_id_ = state.entry_id;
|
|
938
|
+
}
|
|
939
|
+
|
|
940
|
+
// Load nodes one by one
|
|
941
|
+
for (std::size_t i = 0; i != state.size; ++i) {
|
|
942
|
+
node_head_t head;
|
|
943
|
+
std::size_t read = std::fread(&head, head_bytes_k, 1, file);
|
|
944
|
+
if (!read) {
|
|
945
|
+
std::fclose(file);
|
|
946
|
+
throw std::runtime_error(std::strerror(errno));
|
|
947
|
+
}
|
|
948
|
+
|
|
949
|
+
std::size_t bytes_to_dump = node_dump_size(head.dim, head.level);
|
|
950
|
+
node_ref_t node_ref = node_malloc(head.label, nullptr, head.dim, head.level, true);
|
|
951
|
+
read = std::fread((byte_t*)&node_ref.head + head_bytes_k, bytes_to_dump - head_bytes_k, 1, file);
|
|
952
|
+
if (!read) {
|
|
953
|
+
std::fclose(file);
|
|
954
|
+
throw std::runtime_error(std::strerror(errno));
|
|
955
|
+
}
|
|
956
|
+
|
|
957
|
+
nodes_[i] = node_ref;
|
|
958
|
+
}
|
|
959
|
+
|
|
960
|
+
std::fclose(file);
|
|
961
|
+
viewed_file_descriptor_ = 0;
|
|
962
|
+
}
|
|
963
|
+
|
|
964
|
+
/**
|
|
965
|
+
* Compatibility: Linux, MacOS.
|
|
966
|
+
*/
|
|
967
|
+
void view(char const* file_path) noexcept(false) {
|
|
968
|
+
state_t state;
|
|
969
|
+
int open_flags = O_RDONLY;
|
|
970
|
+
#if __linux__
|
|
971
|
+
open_flags |= O_NOATIME;
|
|
972
|
+
#endif
|
|
973
|
+
int descriptor = open(file_path, open_flags);
|
|
974
|
+
|
|
975
|
+
// Estimate the file size
|
|
976
|
+
struct stat file_stat;
|
|
977
|
+
int fstat_status = fstat(descriptor, &file_stat);
|
|
978
|
+
if (fstat_status < 0) {
|
|
979
|
+
close(descriptor);
|
|
980
|
+
throw std::runtime_error(std::strerror(errno));
|
|
981
|
+
}
|
|
982
|
+
|
|
983
|
+
// Map the entire file
|
|
984
|
+
byte_t* file = (byte_t*)mmap(NULL, file_stat.st_size, PROT_READ, MAP_PRIVATE, descriptor, 0);
|
|
985
|
+
if (file == MAP_FAILED) {
|
|
986
|
+
close(descriptor);
|
|
987
|
+
throw std::runtime_error(std::strerror(errno));
|
|
988
|
+
}
|
|
989
|
+
|
|
990
|
+
// Read the header
|
|
991
|
+
{
|
|
992
|
+
std::memcpy(&state, file, sizeof(state));
|
|
993
|
+
config_.connectivity = state.connectivity;
|
|
994
|
+
config_.max_elements = state.size;
|
|
995
|
+
config_.max_threads_add = 0;
|
|
996
|
+
pre_ = precompute(config_);
|
|
997
|
+
reserve(state.size);
|
|
998
|
+
size_ = state.size;
|
|
999
|
+
max_level_ = state.max_level;
|
|
1000
|
+
entry_id_ = state.entry_id;
|
|
1001
|
+
}
|
|
1002
|
+
|
|
1003
|
+
// Locate every node packed into file
|
|
1004
|
+
std::size_t progress = sizeof(state);
|
|
1005
|
+
for (std::size_t i = 0; i != state.size; ++i) {
|
|
1006
|
+
node_head_t const& head = *(node_head_t const*)(file + progress);
|
|
1007
|
+
std::size_t bytes_to_dump = node_dump_size(head.dim, head.level);
|
|
1008
|
+
nodes_[i].tape_ = (byte_t*)(file + progress);
|
|
1009
|
+
nodes_[i].vector_ = (scalar_t*)(file + progress + bytes_to_dump - head.dim);
|
|
1010
|
+
progress += bytes_to_dump;
|
|
1011
|
+
max_level_ = std::max(max_level_, head.level);
|
|
1012
|
+
}
|
|
1013
|
+
|
|
1014
|
+
viewed_file_descriptor_ = descriptor;
|
|
1015
|
+
}
|
|
1016
|
+
|
|
1017
|
+
#pragma endregion
|
|
1018
|
+
|
|
1019
|
+
private:
|
|
1020
|
+
inline static precomputed_constants_t precompute(config_t const& config) noexcept {
|
|
1021
|
+
precomputed_constants_t pre;
|
|
1022
|
+
pre.connectivity_max_base = config.connectivity * base_level_multiple_k;
|
|
1023
|
+
pre.inverse_log_connectivity = 1.0 / std::log(static_cast<double>(config.connectivity));
|
|
1024
|
+
pre.neighbors_bytes = config.connectivity * sizeof(id_t) + sizeof(neighbors_count_t);
|
|
1025
|
+
pre.neighbors_base_bytes = pre.connectivity_max_base * sizeof(id_t) + sizeof(neighbors_count_t);
|
|
1026
|
+
pre.mutex_bytes = sizeof(mutex_t) * (config.max_threads_add > 1);
|
|
1027
|
+
return pre;
|
|
1028
|
+
}
|
|
1029
|
+
|
|
1030
|
+
inline std::size_t node_dump_size(dim_t dim, level_t level) const noexcept {
|
|
1031
|
+
return head_bytes_k + pre_.neighbors_base_bytes + pre_.neighbors_bytes * level + sizeof(scalar_t) * dim;
|
|
1032
|
+
}
|
|
1033
|
+
|
|
1034
|
+
void node_free(std::size_t id) noexcept {
|
|
1035
|
+
|
|
1036
|
+
if (viewed_file_descriptor_)
|
|
1037
|
+
return;
|
|
1038
|
+
|
|
1039
|
+
// This function is rarely called and can be as expensive as needed for higher space-efficiency.
|
|
1040
|
+
node_t& node = nodes_[id];
|
|
1041
|
+
if (!node.tape_)
|
|
1042
|
+
return;
|
|
1043
|
+
|
|
1044
|
+
node_head_t const& head = *(node_head_t const*)(node.tape_ + pre_.mutex_bytes);
|
|
1045
|
+
std::size_t levels_bytes = pre_.neighbors_base_bytes + pre_.neighbors_bytes * head.level;
|
|
1046
|
+
bool store_vector = (byte_t*)(node.tape_ + pre_.mutex_bytes + head_bytes_k + levels_bytes) == //
|
|
1047
|
+
(byte_t*)(node.vector_);
|
|
1048
|
+
std::size_t node_bytes = //
|
|
1049
|
+
pre_.mutex_bytes + // Optional concurrency-control
|
|
1050
|
+
head_bytes_k + levels_bytes + // Obligatory neighborhood index
|
|
1051
|
+
head.dim * store_vector // Optional vector copy
|
|
1052
|
+
;
|
|
1053
|
+
|
|
1054
|
+
allocator_t{}.deallocate(node.tape_, node_bytes);
|
|
1055
|
+
node = {};
|
|
1056
|
+
}
|
|
1057
|
+
|
|
1058
|
+
node_ref_t node_malloc( //
|
|
1059
|
+
label_t label, scalar_t const* vector, std::size_t dim, //
|
|
1060
|
+
level_t level, bool store_vector = true) noexcept(false) {
|
|
1061
|
+
|
|
1062
|
+
// This function is rarely called and can be as expensive as needed for higher space-efficiency.
|
|
1063
|
+
std::size_t levels_bytes = pre_.neighbors_base_bytes + pre_.neighbors_bytes * level;
|
|
1064
|
+
std::size_t node_bytes = //
|
|
1065
|
+
pre_.mutex_bytes + // Optional concurrency-control
|
|
1066
|
+
head_bytes_k + levels_bytes + // Obligatory neighborhood index
|
|
1067
|
+
sizeof(scalar_t) * dim * store_vector // Optional vector copy
|
|
1068
|
+
;
|
|
1069
|
+
|
|
1070
|
+
byte_t* data = (byte_t*)allocator_t{}.allocate(node_bytes);
|
|
1071
|
+
assert_m(data, "Not enough memory for links");
|
|
1072
|
+
|
|
1073
|
+
mutex_t* mutex = synchronize() ? (mutex_t*)data : nullptr;
|
|
1074
|
+
scalar_t* scalars = store_vector //
|
|
1075
|
+
? (scalar_t*)(data + pre_.mutex_bytes + head_bytes_k + levels_bytes)
|
|
1076
|
+
: (scalar_t*)(vector);
|
|
1077
|
+
|
|
1078
|
+
std::memset(data, 0, node_bytes);
|
|
1079
|
+
std::memcpy(scalars, vector, sizeof(scalar_t) * dim * (store_vector && vector));
|
|
1080
|
+
|
|
1081
|
+
node_head_t& head = *(node_head_t*)(data + pre_.mutex_bytes);
|
|
1082
|
+
head.label = label;
|
|
1083
|
+
head.dim = static_cast<dim_t>(dim);
|
|
1084
|
+
head.level = level;
|
|
1085
|
+
|
|
1086
|
+
return {*mutex, head, scalars};
|
|
1087
|
+
}
|
|
1088
|
+
|
|
1089
|
+
inline node_ref_t node(id_t id) const noexcept {
|
|
1090
|
+
|
|
1091
|
+
node_t node = nodes_[id];
|
|
1092
|
+
byte_t* data = node.tape_;
|
|
1093
|
+
mutex_t* mutex = synchronize() ? (mutex_t*)data : nullptr;
|
|
1094
|
+
node_head_t& head = *(node_head_t*)(data + pre_.mutex_bytes);
|
|
1095
|
+
scalar_t* scalars = node.vector_;
|
|
1096
|
+
|
|
1097
|
+
return {*mutex, head, scalars};
|
|
1098
|
+
}
|
|
1099
|
+
|
|
1100
|
+
inline neighbors_ref_t neighbors_base(node_ref_t node) const noexcept { return {node.head.neighbors}; }
|
|
1101
|
+
|
|
1102
|
+
inline neighbors_ref_t neighbors_non_base(node_ref_t node, level_t level) const noexcept {
|
|
1103
|
+
return {node.head.neighbors + pre_.neighbors_base_bytes + (level - 1) * pre_.neighbors_bytes};
|
|
1104
|
+
}
|
|
1105
|
+
|
|
1106
|
+
inline neighbors_ref_t neighbors(node_ref_t node, level_t level) const noexcept {
|
|
1107
|
+
return level ? neighbors_non_base(node, level) : neighbors_base(node);
|
|
1108
|
+
}
|
|
1109
|
+
|
|
1110
|
+
id_t connect_new_element(id_t new_id, level_t level, thread_context_t& context) noexcept(false) {
|
|
1111
|
+
|
|
1112
|
+
node_ref_t new_node = node(new_id);
|
|
1113
|
+
distances_and_ids_t& top_candidates = context.top_candidates;
|
|
1114
|
+
std::size_t connectivity_max = level ? config_.connectivity : pre_.connectivity_max_base;
|
|
1115
|
+
filter_top_candidates_with_heuristic(top_candidates, context.candidates_set, config_.connectivity, context);
|
|
1116
|
+
|
|
1117
|
+
distance_and_id_t const* const top_unordered = top_candidates.data();
|
|
1118
|
+
std::size_t const top_count = top_candidates.size();
|
|
1119
|
+
id_t next_closest_entry_id = top_unordered[0].second;
|
|
1120
|
+
distance_t next_closest_distance = top_unordered[0].first;
|
|
1121
|
+
|
|
1122
|
+
// Outgoing links from `new_id`:
|
|
1123
|
+
{
|
|
1124
|
+
neighbors_ref_t new_neighbors = neighbors(new_node, level);
|
|
1125
|
+
assert_m(!new_neighbors.count, "The newly inserted element should have blank link list");
|
|
1126
|
+
|
|
1127
|
+
new_neighbors.count = static_cast<neighbors_count_t>(top_count);
|
|
1128
|
+
for (std::size_t idx = 0; idx < top_count; idx++) {
|
|
1129
|
+
assert_m(!new_neighbors.neighbors[idx], "Possible memory corruption");
|
|
1130
|
+
assert_m(level <= node(top_unordered[idx].second).head.level, "Linking to missing level");
|
|
1131
|
+
|
|
1132
|
+
new_neighbors.neighbors[idx] = top_unordered[idx].second;
|
|
1133
|
+
if (top_unordered[idx].first < next_closest_distance) {
|
|
1134
|
+
next_closest_entry_id = top_unordered[idx].second;
|
|
1135
|
+
next_closest_distance = top_unordered[idx].first;
|
|
1136
|
+
}
|
|
1137
|
+
}
|
|
1138
|
+
}
|
|
1139
|
+
|
|
1140
|
+
// Reverse links from the neighbors:
|
|
1141
|
+
for (std::size_t idx = 0; idx < top_count; idx++) {
|
|
1142
|
+
id_t close_id = top_unordered[idx].second;
|
|
1143
|
+
node_ref_t close_node = node(close_id);
|
|
1144
|
+
lock_t close_lock = close_node.lock();
|
|
1145
|
+
|
|
1146
|
+
neighbors_ref_t close_header = neighbors(close_node, level);
|
|
1147
|
+
assert_m(close_header.count <= connectivity_max, "Possible corruption");
|
|
1148
|
+
assert_m(close_id != new_id, "Self-loops are impossible");
|
|
1149
|
+
assert_m(level <= close_node.head.level, "Linking to missing level");
|
|
1150
|
+
|
|
1151
|
+
// If `new_id` is already present in the neighboring connections of `close_id`
|
|
1152
|
+
// then no need to modify any connections or run the heuristics.
|
|
1153
|
+
if (close_header.count < connectivity_max) {
|
|
1154
|
+
close_header.neighbors[close_header.count] = new_id;
|
|
1155
|
+
close_header.count++;
|
|
1156
|
+
continue;
|
|
1157
|
+
}
|
|
1158
|
+
|
|
1159
|
+
// To fit a new connection we need to drop an existing one.
|
|
1160
|
+
distances_and_ids_t& candidates = context.candidates_filter;
|
|
1161
|
+
candidates.clear();
|
|
1162
|
+
candidates.emplace( //
|
|
1163
|
+
context.metric( //
|
|
1164
|
+
new_node.vector, close_node.vector, new_node.head.dim, close_node.head.dim),
|
|
1165
|
+
new_id);
|
|
1166
|
+
iterate_through_neighbors(close_header, [&](id_t successor_id) noexcept {
|
|
1167
|
+
node_ref_t successor_node = node(successor_id);
|
|
1168
|
+
candidates.emplace( //
|
|
1169
|
+
context.metric( //
|
|
1170
|
+
successor_node.vector, close_node.vector, successor_node.head.dim, close_node.head.dim),
|
|
1171
|
+
successor_id);
|
|
1172
|
+
});
|
|
1173
|
+
filter_top_candidates_with_heuristic(candidates, context.candidates_set, connectivity_max, context);
|
|
1174
|
+
|
|
1175
|
+
// Export the results:
|
|
1176
|
+
close_header.count = 0u;
|
|
1177
|
+
while (candidates.size()) {
|
|
1178
|
+
close_header.neighbors[close_header.count] = candidates.top().second;
|
|
1179
|
+
close_header.count++;
|
|
1180
|
+
candidates.pop();
|
|
1181
|
+
}
|
|
1182
|
+
}
|
|
1183
|
+
|
|
1184
|
+
return next_closest_entry_id;
|
|
1185
|
+
}
|
|
1186
|
+
|
|
1187
|
+
level_t choose_random_level(std::default_random_engine& level_generator) const noexcept {
|
|
1188
|
+
std::uniform_real_distribution<double> distribution(0.0, 1.0);
|
|
1189
|
+
double r = -std::log(distribution(level_generator)) * pre_.inverse_log_connectivity;
|
|
1190
|
+
return (level_t)r;
|
|
1191
|
+
}
|
|
1192
|
+
|
|
1193
|
+
void search_to_insert( //
|
|
1194
|
+
id_t start_id, scalar_t const* query_vec, dim_t query_dim, //
|
|
1195
|
+
level_t level, thread_context_t& context) noexcept(false) {
|
|
1196
|
+
|
|
1197
|
+
visits_bitset_t& visits = context.visits;
|
|
1198
|
+
distances_and_ids_t& top_candidates = context.top_candidates;
|
|
1199
|
+
distances_and_ids_t& candidates_set = context.candidates_set;
|
|
1200
|
+
|
|
1201
|
+
top_candidates.clear();
|
|
1202
|
+
candidates_set.clear();
|
|
1203
|
+
visits.clear();
|
|
1204
|
+
|
|
1205
|
+
distance_t closest_dist = context.metric(query_vec, node(start_id).vector, query_dim, node(start_id).head.dim);
|
|
1206
|
+
top_candidates.emplace(closest_dist, start_id);
|
|
1207
|
+
candidates_set.emplace(-closest_dist, start_id);
|
|
1208
|
+
visits.set(start_id);
|
|
1209
|
+
|
|
1210
|
+
while (!candidates_set.empty()) {
|
|
1211
|
+
|
|
1212
|
+
distance_and_id_t candidacy = candidates_set.top();
|
|
1213
|
+
if ((-candidacy.first) > closest_dist && top_candidates.size() == config_.expansion_add)
|
|
1214
|
+
break;
|
|
1215
|
+
|
|
1216
|
+
candidates_set.pop();
|
|
1217
|
+
id_t candidate_id = candidacy.second;
|
|
1218
|
+
node_ref_t candidate_node = node(candidate_id);
|
|
1219
|
+
lock_t candidate_lock = candidate_node.lock();
|
|
1220
|
+
neighbors_ref_t candidate_header = neighbors(candidate_node, level);
|
|
1221
|
+
|
|
1222
|
+
iterate_through_neighbors(candidate_header, [&](id_t successor_id) noexcept {
|
|
1223
|
+
if (visits.test(successor_id))
|
|
1224
|
+
return;
|
|
1225
|
+
|
|
1226
|
+
visits.set(successor_id);
|
|
1227
|
+
node_ref_t successor_node = node(successor_id);
|
|
1228
|
+
distance_t successor_dist =
|
|
1229
|
+
context.metric(query_vec, successor_node.vector, query_dim, successor_node.head.dim);
|
|
1230
|
+
if (top_candidates.size() < config_.expansion_add || closest_dist > successor_dist) {
|
|
1231
|
+
candidates_set.emplace(-successor_dist, successor_id);
|
|
1232
|
+
|
|
1233
|
+
top_candidates.emplace(successor_dist, successor_id);
|
|
1234
|
+
if (top_candidates.size() > config_.expansion_add)
|
|
1235
|
+
top_candidates.pop();
|
|
1236
|
+
if (!top_candidates.empty())
|
|
1237
|
+
closest_dist = top_candidates.top().first;
|
|
1238
|
+
}
|
|
1239
|
+
});
|
|
1240
|
+
}
|
|
1241
|
+
}
|
|
1242
|
+
|
|
1243
|
+
void search_to_find_in_base( //
|
|
1244
|
+
id_t start_id, scalar_t const* query_vec, dim_t query_dim, //
|
|
1245
|
+
std::size_t expansion, thread_context_t& context) const noexcept(false) {
|
|
1246
|
+
|
|
1247
|
+
visits_bitset_t& visits = context.visits;
|
|
1248
|
+
distances_and_ids_t& top_candidates = context.top_candidates;
|
|
1249
|
+
distances_and_ids_t& candidates_set = context.candidates_set;
|
|
1250
|
+
|
|
1251
|
+
visits.clear();
|
|
1252
|
+
top_candidates.clear();
|
|
1253
|
+
candidates_set.clear();
|
|
1254
|
+
|
|
1255
|
+
distance_t closest_dist = context.metric(query_vec, node(start_id).vector, query_dim, node(start_id).head.dim);
|
|
1256
|
+
top_candidates.emplace(closest_dist, start_id);
|
|
1257
|
+
candidates_set.emplace(-closest_dist, start_id);
|
|
1258
|
+
visits.set(start_id);
|
|
1259
|
+
|
|
1260
|
+
while (!candidates_set.empty()) {
|
|
1261
|
+
|
|
1262
|
+
distance_and_id_t current_node_node = candidates_set.top();
|
|
1263
|
+
if ((-current_node_node.first) > closest_dist)
|
|
1264
|
+
break;
|
|
1265
|
+
|
|
1266
|
+
candidates_set.pop();
|
|
1267
|
+
|
|
1268
|
+
id_t candidate_id = current_node_node.second;
|
|
1269
|
+
neighbors_ref_t candidate_header = neighbors_base(node(candidate_id));
|
|
1270
|
+
|
|
1271
|
+
iterate_through_neighbors(candidate_header, [&](id_t successor_id) noexcept {
|
|
1272
|
+
if (visits.test(successor_id))
|
|
1273
|
+
return;
|
|
1274
|
+
|
|
1275
|
+
visits.set(successor_id);
|
|
1276
|
+
node_ref_t successor_node = node(successor_id);
|
|
1277
|
+
distance_t successor_dist =
|
|
1278
|
+
context.metric(query_vec, successor_node.vector, query_dim, successor_node.head.dim);
|
|
1279
|
+
|
|
1280
|
+
if (top_candidates.size() < expansion || closest_dist > successor_dist) {
|
|
1281
|
+
candidates_set.emplace(-successor_dist, successor_id);
|
|
1282
|
+
|
|
1283
|
+
top_candidates.emplace(successor_dist, successor_id);
|
|
1284
|
+
if (top_candidates.size() > expansion)
|
|
1285
|
+
top_candidates.pop();
|
|
1286
|
+
if (!top_candidates.empty())
|
|
1287
|
+
closest_dist = top_candidates.top().first;
|
|
1288
|
+
}
|
|
1289
|
+
});
|
|
1290
|
+
}
|
|
1291
|
+
}
|
|
1292
|
+
|
|
1293
|
+
/**
|
|
1294
|
+
* @brief A simple `for`-loop that prefetches vectors of neighbors.
|
|
1295
|
+
*/
|
|
1296
|
+
template <typename neighbor_id_callback_at>
|
|
1297
|
+
inline void iterate_through_neighbors(neighbors_ref_t head, neighbor_id_callback_at&& callback) const noexcept {
|
|
1298
|
+
std::size_t n = head.count;
|
|
1299
|
+
for (std::size_t j = 0; j != n; j++)
|
|
1300
|
+
callback(head.neighbors[j]);
|
|
1301
|
+
}
|
|
1302
|
+
|
|
1303
|
+
void filter_top_candidates_with_heuristic( //
|
|
1304
|
+
distances_and_ids_t& top_candidates, distances_and_ids_t& temporary, std::size_t needed,
|
|
1305
|
+
thread_context_t& context) const noexcept(false) {
|
|
1306
|
+
|
|
1307
|
+
if (top_candidates.size() < needed)
|
|
1308
|
+
return;
|
|
1309
|
+
|
|
1310
|
+
// TODO: Sort ascending, then run an inplace triangular reduction.
|
|
1311
|
+
temporary.clear();
|
|
1312
|
+
while (top_candidates.size()) {
|
|
1313
|
+
temporary.emplace(-top_candidates.top().first, top_candidates.top().second);
|
|
1314
|
+
top_candidates.pop();
|
|
1315
|
+
}
|
|
1316
|
+
|
|
1317
|
+
while (temporary.size() && top_candidates.size() < needed) {
|
|
1318
|
+
|
|
1319
|
+
distance_and_id_t best = temporary.top();
|
|
1320
|
+
distance_t dist_to_query = -best.first;
|
|
1321
|
+
temporary.pop();
|
|
1322
|
+
bool good = true;
|
|
1323
|
+
|
|
1324
|
+
distance_and_id_t const* const top_unordered = top_candidates.data();
|
|
1325
|
+
std::size_t const top_count = top_candidates.size();
|
|
1326
|
+
for (std::size_t idx = 0; idx < top_count; idx++) {
|
|
1327
|
+
distance_and_id_t other = top_unordered[idx];
|
|
1328
|
+
node_ref_t other_node = node(other.second);
|
|
1329
|
+
node_ref_t best_node = node(best.second);
|
|
1330
|
+
distance_t inter_result_dist =
|
|
1331
|
+
context.metric(other_node.vector, best_node.vector, other_node.head.dim, best_node.head.dim);
|
|
1332
|
+
if (inter_result_dist < dist_to_query) {
|
|
1333
|
+
good = false;
|
|
1334
|
+
break;
|
|
1335
|
+
}
|
|
1336
|
+
}
|
|
1337
|
+
|
|
1338
|
+
if (good)
|
|
1339
|
+
top_candidates.emplace(-best.first, best.second);
|
|
1340
|
+
}
|
|
1341
|
+
}
|
|
1342
|
+
};
|
|
1343
|
+
|
|
1344
|
+
} // namespace usearch
|
|
1345
|
+
} // namespace unum
|