usearch 2.10.2 → 2.10.5
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/README.md +1 -1
- package/binding.gyp +101 -0
- package/c/CMakeLists.txt +27 -0
- package/c/README.md +34 -0
- package/c/lib.cpp +362 -0
- package/c/test.c +382 -0
- package/c/usearch.h +356 -0
- package/cmake/clipp.cmake +7 -0
- package/cmake/config.cmake.in +15 -0
- package/cmake/emscripten.cmake +8 -0
- package/cmake/jemalloc.cmake +62 -0
- package/cmake/pkg-config.pc.in +4 -0
- package/cmake/usearchConfigVersion.cmake.in +20 -0
- package/cpp/CMakeLists.txt +14 -0
- package/cpp/README.md +161 -0
- package/cpp/bench.cpp +616 -0
- package/cpp/test.cpp +324 -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/index.hpp +3838 -0
- package/include/usearch/index_dense.hpp +2004 -0
- package/include/usearch/index_plugins.hpp +2255 -0
- package/javascript/dist-package-cjs.json +3 -0
- package/javascript/dist-package-esm.json +3 -0
- package/javascript/lib.cpp +350 -0
- package/javascript/node-gyp-build.d.ts +1 -0
- package/javascript/tsconfig-base.json +25 -0
- package/javascript/tsconfig-cjs.json +8 -0
- package/javascript/tsconfig-esm.json +8 -0
- package/javascript/usearch.test.js +84 -0
- package/javascript/usearch.ts +640 -0
- package/package.json +10 -3
- package/prebuilds/darwin-arm64+x64/usearch.node +0 -0
- package/prebuilds/linux-arm64/usearch.node +0 -0
- package/prebuilds/linux-x64/usearch.node +0 -0
- package/prebuilds/win32-ia32/usearch.node +0 -0
- package/prebuilds/win32-x64/usearch.node +0 -0
- package/simsimd/.releaserc +87 -0
- package/simsimd/CONTRIBUTING.md +114 -0
- package/simsimd/Cargo.lock +709 -0
- package/simsimd/Cargo.toml +44 -0
- package/simsimd/LICENSE +201 -0
- package/simsimd/README.md +708 -0
- package/simsimd/binding.gyp +36 -0
- package/simsimd/build.rs +51 -0
- package/simsimd/c/lib.c +110 -0
- package/simsimd/include/simsimd/binary.h +258 -0
- package/simsimd/include/simsimd/dot.h +1237 -0
- package/simsimd/include/simsimd/geospatial.h +37 -0
- package/simsimd/include/simsimd/probability.h +563 -0
- package/simsimd/include/simsimd/simsimd.h +1129 -0
- package/simsimd/include/simsimd/spatial.h +1019 -0
- package/simsimd/include/simsimd/types.h +309 -0
- package/simsimd/javascript/bench.js +127 -0
- package/simsimd/javascript/dist-package-cjs.json +3 -0
- package/simsimd/javascript/dist-package-esm.json +3 -0
- package/simsimd/javascript/fallback.ts +203 -0
- package/simsimd/javascript/lib.c +107 -0
- package/simsimd/javascript/node-gyp-build.d.ts +1 -0
- package/simsimd/javascript/simsimd.ts +135 -0
- package/simsimd/javascript/test.mjs +179 -0
- package/simsimd/javascript/tsconfig-base.json +37 -0
- package/simsimd/javascript/tsconfig-cjs.json +8 -0
- package/simsimd/javascript/tsconfig-esm.json +8 -0
- package/simsimd/package-lock.json +1709 -0
- package/simsimd/package.json +61 -0
- package/simsimd/rust/benches/cosine.rs +27 -0
- package/simsimd/rust/benches/native.rs +30 -0
- package/simsimd/rust/benches/sqeuclidean.rs +27 -0
- package/simsimd/rust/lib.rs +778 -0
|
@@ -0,0 +1,2004 @@
|
|
|
1
|
+
#pragma once
|
|
2
|
+
#include <stdlib.h> // `aligned_alloc`
|
|
3
|
+
|
|
4
|
+
#include <functional> // `std::function`
|
|
5
|
+
#include <numeric> // `std::iota`
|
|
6
|
+
#include <thread> // `std::thread`
|
|
7
|
+
#include <vector> // `std::vector`
|
|
8
|
+
|
|
9
|
+
#include <usearch/index.hpp>
|
|
10
|
+
#include <usearch/index_plugins.hpp>
|
|
11
|
+
|
|
12
|
+
#if defined(USEARCH_DEFINED_CPP17)
|
|
13
|
+
#include <shared_mutex> // `std::shared_mutex`
|
|
14
|
+
#endif
|
|
15
|
+
|
|
16
|
+
namespace unum {
|
|
17
|
+
namespace usearch {
|
|
18
|
+
|
|
19
|
+
template <typename, typename> class index_dense_gt;
|
|
20
|
+
|
|
21
|
+
/**
|
|
22
|
+
* @brief The "magic" sequence helps infer the type of the file.
|
|
23
|
+
* USearch indexes start with the "usearch" string.
|
|
24
|
+
*/
|
|
25
|
+
constexpr char const* default_magic() { return "usearch"; }
|
|
26
|
+
|
|
27
|
+
using index_dense_head_buffer_t = byte_t[64];
|
|
28
|
+
|
|
29
|
+
static_assert(sizeof(index_dense_head_buffer_t) == 64, "File header should be exactly 64 bytes");
|
|
30
|
+
|
|
31
|
+
/**
|
|
32
|
+
* @brief Serialized binary representations of the USearch index start with metadata.
|
|
33
|
+
* Metadata is parsed into a `index_dense_head_t`, containing the USearch package version,
|
|
34
|
+
* and the properties of the index.
|
|
35
|
+
*
|
|
36
|
+
* It uses: 13 bytes for file versioning, 22 bytes for structural information = 35 bytes.
|
|
37
|
+
* The following 24 bytes contain binary size of the graph, of the vectors, and the checksum,
|
|
38
|
+
* leaving 5 bytes at the end vacant.
|
|
39
|
+
*/
|
|
40
|
+
struct index_dense_head_t {
|
|
41
|
+
|
|
42
|
+
// Versioning:
|
|
43
|
+
using magic_t = char[7];
|
|
44
|
+
using version_t = std::uint16_t;
|
|
45
|
+
|
|
46
|
+
// Versioning: 7 + 2 * 3 = 13 bytes
|
|
47
|
+
char const* magic;
|
|
48
|
+
misaligned_ref_gt<version_t> version_major;
|
|
49
|
+
misaligned_ref_gt<version_t> version_minor;
|
|
50
|
+
misaligned_ref_gt<version_t> version_patch;
|
|
51
|
+
|
|
52
|
+
// Structural: 4 * 3 = 12 bytes
|
|
53
|
+
misaligned_ref_gt<metric_kind_t> kind_metric;
|
|
54
|
+
misaligned_ref_gt<scalar_kind_t> kind_scalar;
|
|
55
|
+
misaligned_ref_gt<scalar_kind_t> kind_key;
|
|
56
|
+
misaligned_ref_gt<scalar_kind_t> kind_compressed_slot;
|
|
57
|
+
|
|
58
|
+
// Population: 8 * 3 = 24 bytes
|
|
59
|
+
misaligned_ref_gt<std::uint64_t> count_present;
|
|
60
|
+
misaligned_ref_gt<std::uint64_t> count_deleted;
|
|
61
|
+
misaligned_ref_gt<std::uint64_t> dimensions;
|
|
62
|
+
misaligned_ref_gt<bool> multi;
|
|
63
|
+
|
|
64
|
+
index_dense_head_t(byte_t* ptr) noexcept
|
|
65
|
+
: magic((char const*)exchange(ptr, ptr + sizeof(magic_t))), //
|
|
66
|
+
version_major(exchange(ptr, ptr + sizeof(version_t))), //
|
|
67
|
+
version_minor(exchange(ptr, ptr + sizeof(version_t))), //
|
|
68
|
+
version_patch(exchange(ptr, ptr + sizeof(version_t))), //
|
|
69
|
+
kind_metric(exchange(ptr, ptr + sizeof(metric_kind_t))), //
|
|
70
|
+
kind_scalar(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
|
|
71
|
+
kind_key(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
|
|
72
|
+
kind_compressed_slot(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
|
|
73
|
+
count_present(exchange(ptr, ptr + sizeof(std::uint64_t))), //
|
|
74
|
+
count_deleted(exchange(ptr, ptr + sizeof(std::uint64_t))), //
|
|
75
|
+
dimensions(exchange(ptr, ptr + sizeof(std::uint64_t))), //
|
|
76
|
+
multi(exchange(ptr, ptr + sizeof(bool))) {}
|
|
77
|
+
};
|
|
78
|
+
|
|
79
|
+
struct index_dense_head_result_t {
|
|
80
|
+
|
|
81
|
+
index_dense_head_buffer_t buffer;
|
|
82
|
+
index_dense_head_t head;
|
|
83
|
+
error_t error;
|
|
84
|
+
|
|
85
|
+
explicit operator bool() const noexcept { return !error; }
|
|
86
|
+
index_dense_head_result_t failed(error_t message) noexcept {
|
|
87
|
+
error = std::move(message);
|
|
88
|
+
return std::move(*this);
|
|
89
|
+
}
|
|
90
|
+
};
|
|
91
|
+
|
|
92
|
+
struct index_dense_config_t : public index_config_t {
|
|
93
|
+
std::size_t expansion_add = default_expansion_add();
|
|
94
|
+
std::size_t expansion_search = default_expansion_search();
|
|
95
|
+
bool exclude_vectors = false;
|
|
96
|
+
bool multi = false;
|
|
97
|
+
|
|
98
|
+
/**
|
|
99
|
+
* @brief Allows you to reduce RAM consumption by avoiding
|
|
100
|
+
* reverse-indexing keys-to-vectors, and only keeping
|
|
101
|
+
* the vectors-to-keys mappings.
|
|
102
|
+
*
|
|
103
|
+
* ! This configuration parameter doesn't affect the serialized file,
|
|
104
|
+
* ! and is not preserved between runs. Makes sense for small vector
|
|
105
|
+
* ! representations that fit ina single cache line.
|
|
106
|
+
*/
|
|
107
|
+
bool enable_key_lookups = true;
|
|
108
|
+
|
|
109
|
+
index_dense_config_t(index_config_t base) noexcept : index_config_t(base) {}
|
|
110
|
+
|
|
111
|
+
index_dense_config_t(std::size_t c = default_connectivity(), std::size_t ea = default_expansion_add(),
|
|
112
|
+
std::size_t es = default_expansion_search()) noexcept
|
|
113
|
+
: index_config_t(c), expansion_add(ea ? ea : default_expansion_add()),
|
|
114
|
+
expansion_search(es ? es : default_expansion_search()) {}
|
|
115
|
+
};
|
|
116
|
+
|
|
117
|
+
struct index_dense_clustering_config_t {
|
|
118
|
+
std::size_t min_clusters = 0;
|
|
119
|
+
std::size_t max_clusters = 0;
|
|
120
|
+
enum mode_t {
|
|
121
|
+
merge_smallest_k,
|
|
122
|
+
merge_closest_k,
|
|
123
|
+
} mode = merge_smallest_k;
|
|
124
|
+
};
|
|
125
|
+
|
|
126
|
+
struct index_dense_serialization_config_t {
|
|
127
|
+
bool exclude_vectors = false;
|
|
128
|
+
bool use_64_bit_dimensions = false;
|
|
129
|
+
};
|
|
130
|
+
|
|
131
|
+
struct index_dense_copy_config_t : public index_copy_config_t {
|
|
132
|
+
bool force_vector_copy = true;
|
|
133
|
+
|
|
134
|
+
index_dense_copy_config_t() = default;
|
|
135
|
+
index_dense_copy_config_t(index_copy_config_t base) noexcept : index_copy_config_t(base) {}
|
|
136
|
+
};
|
|
137
|
+
|
|
138
|
+
struct index_dense_metadata_result_t {
|
|
139
|
+
index_dense_serialization_config_t config;
|
|
140
|
+
index_dense_head_buffer_t head_buffer;
|
|
141
|
+
index_dense_head_t head;
|
|
142
|
+
error_t error;
|
|
143
|
+
|
|
144
|
+
explicit operator bool() const noexcept { return !error; }
|
|
145
|
+
index_dense_metadata_result_t failed(error_t message) noexcept {
|
|
146
|
+
error = std::move(message);
|
|
147
|
+
return std::move(*this);
|
|
148
|
+
}
|
|
149
|
+
|
|
150
|
+
index_dense_metadata_result_t() noexcept : config(), head_buffer(), head(head_buffer), error() {}
|
|
151
|
+
|
|
152
|
+
index_dense_metadata_result_t(index_dense_metadata_result_t&& other) noexcept
|
|
153
|
+
: config(), head_buffer(), head(head_buffer), error(std::move(other.error)) {
|
|
154
|
+
std::memcpy(&config, &other.config, sizeof(other.config));
|
|
155
|
+
std::memcpy(&head_buffer, &other.head_buffer, sizeof(other.head_buffer));
|
|
156
|
+
}
|
|
157
|
+
|
|
158
|
+
index_dense_metadata_result_t& operator=(index_dense_metadata_result_t&& other) noexcept {
|
|
159
|
+
std::memcpy(&config, &other.config, sizeof(other.config));
|
|
160
|
+
std::memcpy(&head_buffer, &other.head_buffer, sizeof(other.head_buffer));
|
|
161
|
+
error = std::move(other.error);
|
|
162
|
+
return *this;
|
|
163
|
+
}
|
|
164
|
+
};
|
|
165
|
+
|
|
166
|
+
/**
|
|
167
|
+
* @brief Extracts metadata from a pre-constructed index on disk,
|
|
168
|
+
* without loading it or mapping the whole binary file.
|
|
169
|
+
*/
|
|
170
|
+
inline index_dense_metadata_result_t index_dense_metadata_from_path(char const* file_path) noexcept {
|
|
171
|
+
index_dense_metadata_result_t result;
|
|
172
|
+
std::unique_ptr<std::FILE, int (*)(std::FILE*)> file(std::fopen(file_path, "rb"), &std::fclose);
|
|
173
|
+
if (!file)
|
|
174
|
+
return result.failed(std::strerror(errno));
|
|
175
|
+
|
|
176
|
+
// Read the header
|
|
177
|
+
std::size_t read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get());
|
|
178
|
+
if (!read)
|
|
179
|
+
return result.failed(std::feof(file.get()) ? "End of file reached!" : std::strerror(errno));
|
|
180
|
+
|
|
181
|
+
// Check if the file immediately starts with the index, instead of vectors
|
|
182
|
+
result.config.exclude_vectors = true;
|
|
183
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
184
|
+
return result;
|
|
185
|
+
|
|
186
|
+
if (std::fseek(file.get(), 0L, SEEK_END) != 0)
|
|
187
|
+
return result.failed("Can't infer file size");
|
|
188
|
+
|
|
189
|
+
// Check if it starts with 32-bit
|
|
190
|
+
std::size_t const file_size = std::ftell(file.get());
|
|
191
|
+
|
|
192
|
+
std::uint32_t dimensions_u32[2]{0};
|
|
193
|
+
std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32));
|
|
194
|
+
std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32);
|
|
195
|
+
|
|
196
|
+
std::uint64_t dimensions_u64[2]{0};
|
|
197
|
+
std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64));
|
|
198
|
+
std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64);
|
|
199
|
+
|
|
200
|
+
// Check if it starts with 32-bit
|
|
201
|
+
if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) {
|
|
202
|
+
if (std::fseek(file.get(), static_cast<long>(offset_if_u32), SEEK_SET) != 0)
|
|
203
|
+
return result.failed(std::strerror(errno));
|
|
204
|
+
read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get());
|
|
205
|
+
if (!read)
|
|
206
|
+
return result.failed(std::feof(file.get()) ? "End of file reached!" : std::strerror(errno));
|
|
207
|
+
|
|
208
|
+
result.config.exclude_vectors = false;
|
|
209
|
+
result.config.use_64_bit_dimensions = false;
|
|
210
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
211
|
+
return result;
|
|
212
|
+
}
|
|
213
|
+
|
|
214
|
+
// Check if it starts with 64-bit
|
|
215
|
+
if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) {
|
|
216
|
+
if (std::fseek(file.get(), static_cast<long>(offset_if_u64), SEEK_SET) != 0)
|
|
217
|
+
return result.failed(std::strerror(errno));
|
|
218
|
+
read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get());
|
|
219
|
+
if (!read)
|
|
220
|
+
return result.failed(std::feof(file.get()) ? "End of file reached!" : std::strerror(errno));
|
|
221
|
+
|
|
222
|
+
// Check if it starts with 64-bit
|
|
223
|
+
result.config.exclude_vectors = false;
|
|
224
|
+
result.config.use_64_bit_dimensions = true;
|
|
225
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
226
|
+
return result;
|
|
227
|
+
}
|
|
228
|
+
|
|
229
|
+
return result.failed("Not a dense USearch index!");
|
|
230
|
+
}
|
|
231
|
+
|
|
232
|
+
/**
|
|
233
|
+
* @brief Extracts metadata from a pre-constructed index serialized into an in-memory buffer.
|
|
234
|
+
*/
|
|
235
|
+
inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_mapped_file_t file,
|
|
236
|
+
std::size_t offset = 0) noexcept {
|
|
237
|
+
index_dense_metadata_result_t result;
|
|
238
|
+
|
|
239
|
+
// Read the header
|
|
240
|
+
if (offset + sizeof(index_dense_head_buffer_t) >= file.size())
|
|
241
|
+
return result.failed("End of file reached!");
|
|
242
|
+
|
|
243
|
+
byte_t* const file_data = file.data() + offset;
|
|
244
|
+
std::size_t const file_size = file.size() - offset;
|
|
245
|
+
std::memcpy(&result.head_buffer, file_data, sizeof(index_dense_head_buffer_t));
|
|
246
|
+
|
|
247
|
+
// Check if the file immediately starts with the index, instead of vectors
|
|
248
|
+
result.config.exclude_vectors = true;
|
|
249
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
250
|
+
return result;
|
|
251
|
+
|
|
252
|
+
// Check if it starts with 32-bit
|
|
253
|
+
std::uint32_t dimensions_u32[2]{0};
|
|
254
|
+
std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32));
|
|
255
|
+
std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32);
|
|
256
|
+
|
|
257
|
+
std::uint64_t dimensions_u64[2]{0};
|
|
258
|
+
std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64));
|
|
259
|
+
std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64);
|
|
260
|
+
|
|
261
|
+
// Check if it starts with 32-bit
|
|
262
|
+
if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) {
|
|
263
|
+
std::memcpy(&result.head_buffer, file_data + offset_if_u32, sizeof(index_dense_head_buffer_t));
|
|
264
|
+
result.config.exclude_vectors = false;
|
|
265
|
+
result.config.use_64_bit_dimensions = false;
|
|
266
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
267
|
+
return result;
|
|
268
|
+
}
|
|
269
|
+
|
|
270
|
+
// Check if it starts with 64-bit
|
|
271
|
+
if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) {
|
|
272
|
+
std::memcpy(&result.head_buffer, file_data + offset_if_u64, sizeof(index_dense_head_buffer_t));
|
|
273
|
+
result.config.exclude_vectors = false;
|
|
274
|
+
result.config.use_64_bit_dimensions = true;
|
|
275
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
276
|
+
return result;
|
|
277
|
+
}
|
|
278
|
+
|
|
279
|
+
return result.failed("Not a dense USearch index!");
|
|
280
|
+
}
|
|
281
|
+
|
|
282
|
+
/**
|
|
283
|
+
* @brief Oversimplified type-punned index for equidimensional vectors
|
|
284
|
+
* with automatic @b down-casting, hardware-specific @b SIMD metrics,
|
|
285
|
+
* and ability to @b remove existing vectors, common in Semantic Caching
|
|
286
|
+
* applications.
|
|
287
|
+
*
|
|
288
|
+
* @section Serialization
|
|
289
|
+
*
|
|
290
|
+
* The serialized binary form of `index_dense_gt` is made up of three parts:
|
|
291
|
+
* 1. Binary matrix, aka the `.bbin` part,
|
|
292
|
+
* 2. Metadata about used metrics, number of used vs free slots,
|
|
293
|
+
* 3. The HNSW index in a binary form.
|
|
294
|
+
* The first (1.) generally starts with 2 integers - number of rows (vectors) and @b single-byte columns.
|
|
295
|
+
* The second (2.) starts with @b "usearch"-magic-string, used to infer the file type on open.
|
|
296
|
+
* The third (3.) is implemented by the underlying `index_gt` class.
|
|
297
|
+
*/
|
|
298
|
+
template <typename key_at = default_key_t, typename compressed_slot_at = default_slot_t> //
|
|
299
|
+
class index_dense_gt {
|
|
300
|
+
public:
|
|
301
|
+
using vector_key_t = key_at;
|
|
302
|
+
using key_t = vector_key_t;
|
|
303
|
+
using compressed_slot_t = compressed_slot_at;
|
|
304
|
+
using distance_t = distance_punned_t;
|
|
305
|
+
using metric_t = metric_punned_t;
|
|
306
|
+
|
|
307
|
+
using member_ref_t = member_ref_gt<vector_key_t>;
|
|
308
|
+
using member_cref_t = member_cref_gt<vector_key_t>;
|
|
309
|
+
|
|
310
|
+
using head_t = index_dense_head_t;
|
|
311
|
+
using head_buffer_t = index_dense_head_buffer_t;
|
|
312
|
+
using head_result_t = index_dense_head_result_t;
|
|
313
|
+
|
|
314
|
+
using serialization_config_t = index_dense_serialization_config_t;
|
|
315
|
+
|
|
316
|
+
using dynamic_allocator_t = aligned_allocator_gt<byte_t, 64>;
|
|
317
|
+
using tape_allocator_t = memory_mapping_allocator_gt<64>;
|
|
318
|
+
|
|
319
|
+
private:
|
|
320
|
+
/// @brief Schema: input buffer, bytes in input buffer, output buffer.
|
|
321
|
+
using cast_t = std::function<bool(byte_t const*, std::size_t, byte_t*)>;
|
|
322
|
+
/// @brief Punned index.
|
|
323
|
+
using index_t = index_gt< //
|
|
324
|
+
distance_t, vector_key_t, compressed_slot_t, //
|
|
325
|
+
dynamic_allocator_t, tape_allocator_t>;
|
|
326
|
+
using index_allocator_t = aligned_allocator_gt<index_t, 64>;
|
|
327
|
+
|
|
328
|
+
using member_iterator_t = typename index_t::member_iterator_t;
|
|
329
|
+
using member_citerator_t = typename index_t::member_citerator_t;
|
|
330
|
+
|
|
331
|
+
/// @brief Punned metric object.
|
|
332
|
+
class metric_proxy_t {
|
|
333
|
+
index_dense_gt const* index_ = nullptr;
|
|
334
|
+
|
|
335
|
+
public:
|
|
336
|
+
metric_proxy_t(index_dense_gt const& index) noexcept : index_(&index) {}
|
|
337
|
+
|
|
338
|
+
inline distance_t operator()(byte_t const* a, member_cref_t b) const noexcept { return f(a, v(b)); }
|
|
339
|
+
inline distance_t operator()(member_cref_t a, member_cref_t b) const noexcept { return f(v(a), v(b)); }
|
|
340
|
+
|
|
341
|
+
inline distance_t operator()(byte_t const* a, member_citerator_t b) const noexcept { return f(a, v(b)); }
|
|
342
|
+
inline distance_t operator()(member_citerator_t a, member_citerator_t b) const noexcept {
|
|
343
|
+
return f(v(a), v(b));
|
|
344
|
+
}
|
|
345
|
+
|
|
346
|
+
inline distance_t operator()(byte_t const* a, byte_t const* b) const noexcept { return f(a, b); }
|
|
347
|
+
|
|
348
|
+
inline byte_t const* v(member_cref_t m) const noexcept { return index_->vectors_lookup_[get_slot(m)]; }
|
|
349
|
+
inline byte_t const* v(member_citerator_t m) const noexcept { return index_->vectors_lookup_[get_slot(m)]; }
|
|
350
|
+
inline distance_t f(byte_t const* a, byte_t const* b) const noexcept { return index_->metric_(a, b); }
|
|
351
|
+
};
|
|
352
|
+
|
|
353
|
+
index_dense_config_t config_;
|
|
354
|
+
index_t* typed_ = nullptr;
|
|
355
|
+
|
|
356
|
+
mutable std::vector<byte_t> cast_buffer_;
|
|
357
|
+
struct casts_t {
|
|
358
|
+
cast_t from_b1x8;
|
|
359
|
+
cast_t from_i8;
|
|
360
|
+
cast_t from_f16;
|
|
361
|
+
cast_t from_f32;
|
|
362
|
+
cast_t from_f64;
|
|
363
|
+
|
|
364
|
+
cast_t to_b1x8;
|
|
365
|
+
cast_t to_i8;
|
|
366
|
+
cast_t to_f16;
|
|
367
|
+
cast_t to_f32;
|
|
368
|
+
cast_t to_f64;
|
|
369
|
+
} casts_;
|
|
370
|
+
|
|
371
|
+
/// @brief An instance of a potentially stateful `metric_t` used to initialize copies and forks.
|
|
372
|
+
metric_t metric_;
|
|
373
|
+
|
|
374
|
+
using vectors_tape_allocator_t = memory_mapping_allocator_gt<8>;
|
|
375
|
+
/// @brief Allocator for the copied vectors, aligned to widest double-precision scalars.
|
|
376
|
+
vectors_tape_allocator_t vectors_tape_allocator_;
|
|
377
|
+
|
|
378
|
+
/// @brief For every managed `compressed_slot_t` stores a pointer to the allocated vector copy.
|
|
379
|
+
mutable std::vector<byte_t*> vectors_lookup_;
|
|
380
|
+
|
|
381
|
+
/// @brief Originally forms and array of integers [0, threads], marking all
|
|
382
|
+
mutable std::vector<std::size_t> available_threads_;
|
|
383
|
+
|
|
384
|
+
/// @brief Mutex, controlling concurrent access to `available_threads_`.
|
|
385
|
+
mutable std::mutex available_threads_mutex_;
|
|
386
|
+
|
|
387
|
+
#if defined(USEARCH_DEFINED_CPP17)
|
|
388
|
+
using shared_mutex_t = std::shared_mutex;
|
|
389
|
+
#else
|
|
390
|
+
using shared_mutex_t = unfair_shared_mutex_t;
|
|
391
|
+
#endif
|
|
392
|
+
using shared_lock_t = shared_lock_gt<shared_mutex_t>;
|
|
393
|
+
using unique_lock_t = std::unique_lock<shared_mutex_t>;
|
|
394
|
+
|
|
395
|
+
struct key_and_slot_t {
|
|
396
|
+
vector_key_t key;
|
|
397
|
+
compressed_slot_t slot;
|
|
398
|
+
|
|
399
|
+
bool any_slot() const { return slot == default_free_value<compressed_slot_t>(); }
|
|
400
|
+
static key_and_slot_t any_slot(vector_key_t key) { return {key, default_free_value<compressed_slot_t>()}; }
|
|
401
|
+
};
|
|
402
|
+
|
|
403
|
+
struct lookup_key_hash_t {
|
|
404
|
+
using is_transparent = void;
|
|
405
|
+
std::size_t operator()(key_and_slot_t const& k) const noexcept { return std::hash<vector_key_t>{}(k.key); }
|
|
406
|
+
std::size_t operator()(vector_key_t const& k) const noexcept { return std::hash<vector_key_t>{}(k); }
|
|
407
|
+
};
|
|
408
|
+
|
|
409
|
+
struct lookup_key_same_t {
|
|
410
|
+
using is_transparent = void;
|
|
411
|
+
bool operator()(key_and_slot_t const& a, vector_key_t const& b) const noexcept { return a.key == b; }
|
|
412
|
+
bool operator()(vector_key_t const& a, key_and_slot_t const& b) const noexcept { return a == b.key; }
|
|
413
|
+
bool operator()(key_and_slot_t const& a, key_and_slot_t const& b) const noexcept { return a.key == b.key; }
|
|
414
|
+
};
|
|
415
|
+
|
|
416
|
+
/// @brief Multi-Map from keys to IDs, and allocated vectors.
|
|
417
|
+
flat_hash_multi_set_gt<key_and_slot_t, lookup_key_hash_t, lookup_key_same_t> slot_lookup_;
|
|
418
|
+
|
|
419
|
+
/// @brief Mutex, controlling concurrent access to `slot_lookup_`.
|
|
420
|
+
mutable shared_mutex_t slot_lookup_mutex_;
|
|
421
|
+
|
|
422
|
+
/// @brief Ring-shaped queue of deleted entries, to be reused on future insertions.
|
|
423
|
+
ring_gt<compressed_slot_t> free_keys_;
|
|
424
|
+
|
|
425
|
+
/// @brief Mutex, controlling concurrent access to `free_keys_`.
|
|
426
|
+
mutable std::mutex free_keys_mutex_;
|
|
427
|
+
|
|
428
|
+
/// @brief A constant for the reserved key value, used to mark deleted entries.
|
|
429
|
+
vector_key_t free_key_ = default_free_value<vector_key_t>();
|
|
430
|
+
|
|
431
|
+
public:
|
|
432
|
+
using search_result_t = typename index_t::search_result_t;
|
|
433
|
+
using cluster_result_t = typename index_t::cluster_result_t;
|
|
434
|
+
using add_result_t = typename index_t::add_result_t;
|
|
435
|
+
using stats_t = typename index_t::stats_t;
|
|
436
|
+
using match_t = typename index_t::match_t;
|
|
437
|
+
|
|
438
|
+
index_dense_gt() = default;
|
|
439
|
+
index_dense_gt(index_dense_gt&& other)
|
|
440
|
+
: config_(std::move(other.config_)),
|
|
441
|
+
|
|
442
|
+
typed_(exchange(other.typed_, nullptr)), //
|
|
443
|
+
cast_buffer_(std::move(other.cast_buffer_)), //
|
|
444
|
+
casts_(std::move(other.casts_)), //
|
|
445
|
+
metric_(std::move(other.metric_)), //
|
|
446
|
+
|
|
447
|
+
vectors_tape_allocator_(std::move(other.vectors_tape_allocator_)), //
|
|
448
|
+
vectors_lookup_(std::move(other.vectors_lookup_)), //
|
|
449
|
+
|
|
450
|
+
available_threads_(std::move(other.available_threads_)), //
|
|
451
|
+
slot_lookup_(std::move(other.slot_lookup_)), //
|
|
452
|
+
free_keys_(std::move(other.free_keys_)), //
|
|
453
|
+
free_key_(std::move(other.free_key_)) {} //
|
|
454
|
+
|
|
455
|
+
index_dense_gt& operator=(index_dense_gt&& other) {
|
|
456
|
+
swap(other);
|
|
457
|
+
return *this;
|
|
458
|
+
}
|
|
459
|
+
|
|
460
|
+
/**
|
|
461
|
+
* @brief Swaps the contents of this index with another index.
|
|
462
|
+
* @param other The other index to swap with.
|
|
463
|
+
*/
|
|
464
|
+
void swap(index_dense_gt& other) {
|
|
465
|
+
std::swap(config_, other.config_);
|
|
466
|
+
|
|
467
|
+
std::swap(typed_, other.typed_);
|
|
468
|
+
std::swap(cast_buffer_, other.cast_buffer_);
|
|
469
|
+
std::swap(casts_, other.casts_);
|
|
470
|
+
std::swap(metric_, other.metric_);
|
|
471
|
+
|
|
472
|
+
std::swap(vectors_tape_allocator_, other.vectors_tape_allocator_);
|
|
473
|
+
std::swap(vectors_lookup_, other.vectors_lookup_);
|
|
474
|
+
|
|
475
|
+
std::swap(available_threads_, other.available_threads_);
|
|
476
|
+
std::swap(slot_lookup_, other.slot_lookup_);
|
|
477
|
+
std::swap(free_keys_, other.free_keys_);
|
|
478
|
+
std::swap(free_key_, other.free_key_);
|
|
479
|
+
}
|
|
480
|
+
|
|
481
|
+
~index_dense_gt() {
|
|
482
|
+
if (typed_)
|
|
483
|
+
typed_->~index_t();
|
|
484
|
+
index_allocator_t{}.deallocate(typed_, 1);
|
|
485
|
+
typed_ = nullptr;
|
|
486
|
+
}
|
|
487
|
+
|
|
488
|
+
/**
|
|
489
|
+
* @brief Constructs an instance of ::index_dense_gt.
|
|
490
|
+
* @param[in] metric One of the provided or an @b ad-hoc metric, type-punned.
|
|
491
|
+
* @param[in] config The index configuration (optional).
|
|
492
|
+
* @param[in] free_key The key used for freed vectors (optional).
|
|
493
|
+
* @return An instance of ::index_dense_gt.
|
|
494
|
+
*/
|
|
495
|
+
static index_dense_gt make( //
|
|
496
|
+
metric_t metric, //
|
|
497
|
+
index_dense_config_t config = {}, //
|
|
498
|
+
vector_key_t free_key = default_free_value<vector_key_t>()) {
|
|
499
|
+
|
|
500
|
+
scalar_kind_t scalar_kind = metric.scalar_kind();
|
|
501
|
+
std::size_t hardware_threads = std::thread::hardware_concurrency();
|
|
502
|
+
|
|
503
|
+
index_dense_gt result;
|
|
504
|
+
result.config_ = config;
|
|
505
|
+
result.cast_buffer_.resize(hardware_threads * metric.bytes_per_vector());
|
|
506
|
+
result.casts_ = make_casts_(scalar_kind);
|
|
507
|
+
result.metric_ = metric;
|
|
508
|
+
result.free_key_ = free_key;
|
|
509
|
+
|
|
510
|
+
// Fill the thread IDs.
|
|
511
|
+
result.available_threads_.resize(hardware_threads);
|
|
512
|
+
std::iota(result.available_threads_.begin(), result.available_threads_.end(), 0ul);
|
|
513
|
+
|
|
514
|
+
// Available since C11, but only C++17, so we use the C version.
|
|
515
|
+
index_t* raw = index_allocator_t{}.allocate(1);
|
|
516
|
+
new (raw) index_t(config);
|
|
517
|
+
result.typed_ = raw;
|
|
518
|
+
return result;
|
|
519
|
+
}
|
|
520
|
+
|
|
521
|
+
static index_dense_gt make(char const* path, bool view = false) {
|
|
522
|
+
index_dense_metadata_result_t meta = index_dense_metadata_from_path(path);
|
|
523
|
+
if (!meta)
|
|
524
|
+
return {};
|
|
525
|
+
metric_punned_t metric(meta.head.dimensions, meta.head.kind_metric, meta.head.kind_scalar);
|
|
526
|
+
index_dense_gt result = make(metric);
|
|
527
|
+
if (!result)
|
|
528
|
+
return result;
|
|
529
|
+
if (view)
|
|
530
|
+
result.view(path);
|
|
531
|
+
else
|
|
532
|
+
result.load(path);
|
|
533
|
+
return result;
|
|
534
|
+
}
|
|
535
|
+
|
|
536
|
+
explicit operator bool() const { return typed_; }
|
|
537
|
+
std::size_t connectivity() const { return typed_->connectivity(); }
|
|
538
|
+
std::size_t size() const { return typed_->size() - free_keys_.size(); }
|
|
539
|
+
std::size_t capacity() const { return typed_->capacity(); }
|
|
540
|
+
std::size_t max_level() const noexcept { return typed_->max_level(); }
|
|
541
|
+
index_dense_config_t const& config() const { return config_; }
|
|
542
|
+
index_limits_t const& limits() const { return typed_->limits(); }
|
|
543
|
+
bool multi() const { return config_.multi; }
|
|
544
|
+
|
|
545
|
+
// The metric and its properties
|
|
546
|
+
metric_t const& metric() const { return metric_; }
|
|
547
|
+
void change_metric(metric_t metric) { metric_ = std::move(metric); }
|
|
548
|
+
|
|
549
|
+
scalar_kind_t scalar_kind() const noexcept { return metric_.scalar_kind(); }
|
|
550
|
+
std::size_t bytes_per_vector() const noexcept { return metric_.bytes_per_vector(); }
|
|
551
|
+
std::size_t scalar_words() const noexcept { return metric_.scalar_words(); }
|
|
552
|
+
std::size_t dimensions() const noexcept { return metric_.dimensions(); }
|
|
553
|
+
|
|
554
|
+
// Fetching and changing search criteria
|
|
555
|
+
std::size_t expansion_add() const { return config_.expansion_add; }
|
|
556
|
+
std::size_t expansion_search() const { return config_.expansion_search; }
|
|
557
|
+
void change_expansion_add(std::size_t n) { config_.expansion_add = n; }
|
|
558
|
+
void change_expansion_search(std::size_t n) { config_.expansion_search = n; }
|
|
559
|
+
|
|
560
|
+
member_citerator_t cbegin() const { return typed_->cbegin(); }
|
|
561
|
+
member_citerator_t cend() const { return typed_->cend(); }
|
|
562
|
+
member_citerator_t begin() const { return typed_->begin(); }
|
|
563
|
+
member_citerator_t end() const { return typed_->end(); }
|
|
564
|
+
member_iterator_t begin() { return typed_->begin(); }
|
|
565
|
+
member_iterator_t end() { return typed_->end(); }
|
|
566
|
+
|
|
567
|
+
stats_t stats() const { return typed_->stats(); }
|
|
568
|
+
stats_t stats(std::size_t level) const { return typed_->stats(level); }
|
|
569
|
+
stats_t stats(stats_t* stats_per_level, std::size_t max_level) const {
|
|
570
|
+
return typed_->stats(stats_per_level, max_level);
|
|
571
|
+
}
|
|
572
|
+
|
|
573
|
+
dynamic_allocator_t const& allocator() const { return typed_->dynamic_allocator(); }
|
|
574
|
+
vector_key_t const& free_key() const { return free_key_; }
|
|
575
|
+
|
|
576
|
+
/**
|
|
577
|
+
* @brief A relatively accurate lower bound on the amount of memory consumed by the system.
|
|
578
|
+
* In practice it's error will be below 10%.
|
|
579
|
+
*
|
|
580
|
+
* @see `serialized_length` for the length of the binary serialized representation.
|
|
581
|
+
*/
|
|
582
|
+
std::size_t memory_usage() const {
|
|
583
|
+
return //
|
|
584
|
+
typed_->memory_usage(0) + //
|
|
585
|
+
typed_->tape_allocator().total_wasted() + //
|
|
586
|
+
typed_->tape_allocator().total_reserved() + //
|
|
587
|
+
vectors_tape_allocator_.total_allocated();
|
|
588
|
+
}
|
|
589
|
+
|
|
590
|
+
static constexpr std::size_t any_thread() { return std::numeric_limits<std::size_t>::max(); }
|
|
591
|
+
static constexpr distance_t infinite_distance() { return std::numeric_limits<distance_t>::max(); }
|
|
592
|
+
|
|
593
|
+
struct aggregated_distances_t {
|
|
594
|
+
std::size_t count = 0;
|
|
595
|
+
distance_t mean = infinite_distance();
|
|
596
|
+
distance_t min = infinite_distance();
|
|
597
|
+
distance_t max = infinite_distance();
|
|
598
|
+
};
|
|
599
|
+
|
|
600
|
+
// clang-format off
|
|
601
|
+
add_result_t add(vector_key_t key, b1x8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_b1x8); }
|
|
602
|
+
add_result_t add(vector_key_t key, i8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_i8); }
|
|
603
|
+
add_result_t add(vector_key_t key, f16_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f16); }
|
|
604
|
+
add_result_t add(vector_key_t key, f32_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f32); }
|
|
605
|
+
add_result_t add(vector_key_t key, f64_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f64); }
|
|
606
|
+
|
|
607
|
+
search_result_t search(b1x8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_b1x8); }
|
|
608
|
+
search_result_t search(i8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_i8); }
|
|
609
|
+
search_result_t search(f16_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f16); }
|
|
610
|
+
search_result_t search(f32_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f32); }
|
|
611
|
+
search_result_t search(f64_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f64); }
|
|
612
|
+
|
|
613
|
+
std::size_t get(vector_key_t key, b1x8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_b1x8); }
|
|
614
|
+
std::size_t get(vector_key_t key, i8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_i8); }
|
|
615
|
+
std::size_t get(vector_key_t key, f16_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f16); }
|
|
616
|
+
std::size_t get(vector_key_t key, f32_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f32); }
|
|
617
|
+
std::size_t get(vector_key_t key, f64_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f64); }
|
|
618
|
+
|
|
619
|
+
cluster_result_t cluster(b1x8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_b1x8); }
|
|
620
|
+
cluster_result_t cluster(i8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_i8); }
|
|
621
|
+
cluster_result_t cluster(f16_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f16); }
|
|
622
|
+
cluster_result_t cluster(f32_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f32); }
|
|
623
|
+
cluster_result_t cluster(f64_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f64); }
|
|
624
|
+
|
|
625
|
+
aggregated_distances_t distance_between(vector_key_t key, b1x8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_b1x8); }
|
|
626
|
+
aggregated_distances_t distance_between(vector_key_t key, i8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_i8); }
|
|
627
|
+
aggregated_distances_t distance_between(vector_key_t key, f16_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f16); }
|
|
628
|
+
aggregated_distances_t distance_between(vector_key_t key, f32_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f32); }
|
|
629
|
+
aggregated_distances_t distance_between(vector_key_t key, f64_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f64); }
|
|
630
|
+
// clang-format on
|
|
631
|
+
|
|
632
|
+
/**
|
|
633
|
+
* @brief Computes the distance between two managed entities.
|
|
634
|
+
* If either key maps into more than one vector, will aggregate results
|
|
635
|
+
* exporting the mean, maximum, and minimum values.
|
|
636
|
+
*/
|
|
637
|
+
aggregated_distances_t distance_between(vector_key_t a, vector_key_t b, std::size_t = any_thread()) const {
|
|
638
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
639
|
+
aggregated_distances_t result;
|
|
640
|
+
if (!multi()) {
|
|
641
|
+
auto a_it = slot_lookup_.find(key_and_slot_t::any_slot(a));
|
|
642
|
+
auto b_it = slot_lookup_.find(key_and_slot_t::any_slot(b));
|
|
643
|
+
bool a_missing = a_it == slot_lookup_.end();
|
|
644
|
+
bool b_missing = b_it == slot_lookup_.end();
|
|
645
|
+
if (a_missing || b_missing)
|
|
646
|
+
return result;
|
|
647
|
+
|
|
648
|
+
key_and_slot_t a_key_and_slot = *a_it;
|
|
649
|
+
byte_t const* a_vector = vectors_lookup_[a_key_and_slot.slot];
|
|
650
|
+
key_and_slot_t b_key_and_slot = *b_it;
|
|
651
|
+
byte_t const* b_vector = vectors_lookup_[b_key_and_slot.slot];
|
|
652
|
+
distance_t a_b_distance = metric_(a_vector, b_vector);
|
|
653
|
+
|
|
654
|
+
result.mean = result.min = result.max = a_b_distance;
|
|
655
|
+
result.count = 1;
|
|
656
|
+
return result;
|
|
657
|
+
}
|
|
658
|
+
|
|
659
|
+
auto a_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(a));
|
|
660
|
+
auto b_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(b));
|
|
661
|
+
bool a_missing = a_range.first == a_range.second;
|
|
662
|
+
bool b_missing = b_range.first == b_range.second;
|
|
663
|
+
if (a_missing || b_missing)
|
|
664
|
+
return result;
|
|
665
|
+
|
|
666
|
+
result.min = std::numeric_limits<distance_t>::max();
|
|
667
|
+
result.max = std::numeric_limits<distance_t>::min();
|
|
668
|
+
result.mean = 0;
|
|
669
|
+
result.count = 0;
|
|
670
|
+
|
|
671
|
+
while (a_range.first != a_range.second) {
|
|
672
|
+
key_and_slot_t a_key_and_slot = *a_range.first;
|
|
673
|
+
byte_t const* a_vector = vectors_lookup_[a_key_and_slot.slot];
|
|
674
|
+
while (b_range.first != b_range.second) {
|
|
675
|
+
key_and_slot_t b_key_and_slot = *b_range.first;
|
|
676
|
+
byte_t const* b_vector = vectors_lookup_[b_key_and_slot.slot];
|
|
677
|
+
distance_t a_b_distance = metric_(a_vector, b_vector);
|
|
678
|
+
|
|
679
|
+
result.mean += a_b_distance;
|
|
680
|
+
result.min = (std::min)(result.min, a_b_distance);
|
|
681
|
+
result.max = (std::max)(result.max, a_b_distance);
|
|
682
|
+
result.count++;
|
|
683
|
+
|
|
684
|
+
//
|
|
685
|
+
++b_range.first;
|
|
686
|
+
}
|
|
687
|
+
++a_range.first;
|
|
688
|
+
}
|
|
689
|
+
|
|
690
|
+
result.mean /= result.count;
|
|
691
|
+
return result;
|
|
692
|
+
}
|
|
693
|
+
|
|
694
|
+
/**
|
|
695
|
+
* @brief Identifies a node in a given `level`, that is the closest to the `key`.
|
|
696
|
+
*/
|
|
697
|
+
cluster_result_t cluster(vector_key_t key, std::size_t level, std::size_t thread = any_thread()) const {
|
|
698
|
+
|
|
699
|
+
// Check if such `key` is even present.
|
|
700
|
+
shared_lock_t slots_lock(slot_lookup_mutex_);
|
|
701
|
+
auto key_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
|
|
702
|
+
cluster_result_t result;
|
|
703
|
+
if (key_range.first == key_range.second)
|
|
704
|
+
return result.failed("Key missing!");
|
|
705
|
+
|
|
706
|
+
index_cluster_config_t cluster_config;
|
|
707
|
+
thread_lock_t lock = thread_lock_(thread);
|
|
708
|
+
cluster_config.thread = lock.thread_id;
|
|
709
|
+
cluster_config.expansion = config_.expansion_search;
|
|
710
|
+
metric_proxy_t metric{*this};
|
|
711
|
+
auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
|
|
712
|
+
|
|
713
|
+
// Find the closest cluster for any vector under that key.
|
|
714
|
+
while (key_range.first != key_range.second) {
|
|
715
|
+
key_and_slot_t key_and_slot = *key_range.first;
|
|
716
|
+
byte_t const* vector_data = vectors_lookup_[key_and_slot.slot];
|
|
717
|
+
cluster_result_t new_result = typed_->cluster(vector_data, level, metric, cluster_config, allow);
|
|
718
|
+
if (!new_result)
|
|
719
|
+
return new_result;
|
|
720
|
+
if (new_result.cluster.distance < result.cluster.distance)
|
|
721
|
+
result = std::move(new_result);
|
|
722
|
+
|
|
723
|
+
++key_range.first;
|
|
724
|
+
}
|
|
725
|
+
return result;
|
|
726
|
+
}
|
|
727
|
+
|
|
728
|
+
/**
|
|
729
|
+
* @brief Reserves memory for the index and the keyed lookup.
|
|
730
|
+
* @return `true` if the memory reservation was successful, `false` otherwise.
|
|
731
|
+
*/
|
|
732
|
+
bool reserve(index_limits_t limits) {
|
|
733
|
+
{
|
|
734
|
+
unique_lock_t lock(slot_lookup_mutex_);
|
|
735
|
+
slot_lookup_.reserve(limits.members);
|
|
736
|
+
vectors_lookup_.resize(limits.members);
|
|
737
|
+
}
|
|
738
|
+
return typed_->reserve(limits);
|
|
739
|
+
}
|
|
740
|
+
|
|
741
|
+
/**
|
|
742
|
+
* @brief Erases all the vectors from the index.
|
|
743
|
+
*
|
|
744
|
+
* Will change `size()` to zero, but will keep the same `capacity()`.
|
|
745
|
+
* Will keep the number of available threads/contexts the same as it was.
|
|
746
|
+
*/
|
|
747
|
+
void clear() {
|
|
748
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
749
|
+
|
|
750
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
751
|
+
typed_->clear();
|
|
752
|
+
slot_lookup_.clear();
|
|
753
|
+
vectors_lookup_.clear();
|
|
754
|
+
free_keys_.clear();
|
|
755
|
+
vectors_tape_allocator_.reset();
|
|
756
|
+
}
|
|
757
|
+
|
|
758
|
+
/**
|
|
759
|
+
* @brief Erases all members from index, closing files, and returning RAM to OS.
|
|
760
|
+
*
|
|
761
|
+
* Will change both `size()` and `capacity()` to zero.
|
|
762
|
+
* Will deallocate all threads/contexts.
|
|
763
|
+
* If the index is memory-mapped - releases the mapping and the descriptor.
|
|
764
|
+
*/
|
|
765
|
+
void reset() {
|
|
766
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
767
|
+
|
|
768
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
769
|
+
std::unique_lock<std::mutex> available_threads_lock(available_threads_mutex_);
|
|
770
|
+
typed_->reset();
|
|
771
|
+
slot_lookup_.clear();
|
|
772
|
+
vectors_lookup_.clear();
|
|
773
|
+
free_keys_.clear();
|
|
774
|
+
vectors_tape_allocator_.reset();
|
|
775
|
+
|
|
776
|
+
// Reset the thread IDs.
|
|
777
|
+
available_threads_.resize(std::thread::hardware_concurrency());
|
|
778
|
+
std::iota(available_threads_.begin(), available_threads_.end(), 0ul);
|
|
779
|
+
}
|
|
780
|
+
|
|
781
|
+
/**
|
|
782
|
+
* @brief Saves serialized binary index representation to a stream.
|
|
783
|
+
*/
|
|
784
|
+
template <typename output_callback_at, typename progress_at = dummy_progress_t>
|
|
785
|
+
serialization_result_t save_to_stream(output_callback_at&& output, //
|
|
786
|
+
serialization_config_t config = {}, //
|
|
787
|
+
progress_at&& progress = {}) const {
|
|
788
|
+
|
|
789
|
+
serialization_result_t result;
|
|
790
|
+
std::uint64_t matrix_rows = 0;
|
|
791
|
+
std::uint64_t matrix_cols = 0;
|
|
792
|
+
|
|
793
|
+
// We may not want to put the vectors into the same file
|
|
794
|
+
if (!config.exclude_vectors) {
|
|
795
|
+
// Save the matrix size
|
|
796
|
+
if (!config.use_64_bit_dimensions) {
|
|
797
|
+
std::uint32_t dimensions[2];
|
|
798
|
+
dimensions[0] = static_cast<std::uint32_t>(typed_->size());
|
|
799
|
+
dimensions[1] = static_cast<std::uint32_t>(metric_.bytes_per_vector());
|
|
800
|
+
if (!output(&dimensions, sizeof(dimensions)))
|
|
801
|
+
return result.failed("Failed to serialize into stream");
|
|
802
|
+
matrix_rows = dimensions[0];
|
|
803
|
+
matrix_cols = dimensions[1];
|
|
804
|
+
} else {
|
|
805
|
+
std::uint64_t dimensions[2];
|
|
806
|
+
dimensions[0] = static_cast<std::uint64_t>(typed_->size());
|
|
807
|
+
dimensions[1] = static_cast<std::uint64_t>(metric_.bytes_per_vector());
|
|
808
|
+
if (!output(&dimensions, sizeof(dimensions)))
|
|
809
|
+
return result.failed("Failed to serialize into stream");
|
|
810
|
+
matrix_rows = dimensions[0];
|
|
811
|
+
matrix_cols = dimensions[1];
|
|
812
|
+
}
|
|
813
|
+
|
|
814
|
+
// Dump the vectors one after another
|
|
815
|
+
for (std::uint64_t i = 0; i != matrix_rows; ++i) {
|
|
816
|
+
byte_t* vector = vectors_lookup_[i];
|
|
817
|
+
if (!output(vector, matrix_cols))
|
|
818
|
+
return result.failed("Failed to serialize into stream");
|
|
819
|
+
}
|
|
820
|
+
}
|
|
821
|
+
|
|
822
|
+
// Augment metadata
|
|
823
|
+
{
|
|
824
|
+
index_dense_head_buffer_t buffer;
|
|
825
|
+
std::memset(buffer, 0, sizeof(buffer));
|
|
826
|
+
index_dense_head_t head{buffer};
|
|
827
|
+
std::memcpy(buffer, default_magic(), std::strlen(default_magic()));
|
|
828
|
+
|
|
829
|
+
// Describe software version
|
|
830
|
+
using version_t = index_dense_head_t::version_t;
|
|
831
|
+
head.version_major = static_cast<version_t>(USEARCH_VERSION_MAJOR);
|
|
832
|
+
head.version_minor = static_cast<version_t>(USEARCH_VERSION_MINOR);
|
|
833
|
+
head.version_patch = static_cast<version_t>(USEARCH_VERSION_PATCH);
|
|
834
|
+
|
|
835
|
+
// Describes types used
|
|
836
|
+
head.kind_metric = metric_.metric_kind();
|
|
837
|
+
head.kind_scalar = metric_.scalar_kind();
|
|
838
|
+
head.kind_key = unum::usearch::scalar_kind<vector_key_t>();
|
|
839
|
+
head.kind_compressed_slot = unum::usearch::scalar_kind<compressed_slot_t>();
|
|
840
|
+
|
|
841
|
+
head.count_present = size();
|
|
842
|
+
head.count_deleted = typed_->size() - size();
|
|
843
|
+
head.dimensions = dimensions();
|
|
844
|
+
head.multi = multi();
|
|
845
|
+
|
|
846
|
+
if (!output(&buffer, sizeof(buffer)))
|
|
847
|
+
return result.failed("Failed to serialize into stream");
|
|
848
|
+
}
|
|
849
|
+
|
|
850
|
+
// Save the actual proximity graph
|
|
851
|
+
return typed_->save_to_stream(std::forward<output_callback_at>(output), std::forward<progress_at>(progress));
|
|
852
|
+
}
|
|
853
|
+
|
|
854
|
+
/**
|
|
855
|
+
* @brief Estimate the binary length (in bytes) of the serialized index.
|
|
856
|
+
*/
|
|
857
|
+
std::size_t serialized_length(serialization_config_t config = {}) const noexcept {
|
|
858
|
+
std::size_t dimensions_length = 0;
|
|
859
|
+
std::size_t matrix_length = 0;
|
|
860
|
+
if (!config.exclude_vectors) {
|
|
861
|
+
dimensions_length = config.use_64_bit_dimensions ? sizeof(std::uint64_t) * 2 : sizeof(std::uint32_t) * 2;
|
|
862
|
+
matrix_length = typed_->size() * metric_.bytes_per_vector();
|
|
863
|
+
}
|
|
864
|
+
return dimensions_length + matrix_length + sizeof(index_dense_head_buffer_t) + typed_->serialized_length();
|
|
865
|
+
}
|
|
866
|
+
|
|
867
|
+
/**
|
|
868
|
+
* @brief Parses the index from file to RAM.
|
|
869
|
+
* @param[in] path The path to the file.
|
|
870
|
+
* @param[in] config Configuration parameters for imports.
|
|
871
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
872
|
+
*/
|
|
873
|
+
template <typename input_callback_at, typename progress_at = dummy_progress_t>
|
|
874
|
+
serialization_result_t load_from_stream(input_callback_at&& input, //
|
|
875
|
+
serialization_config_t config = {}, //
|
|
876
|
+
progress_at&& progress = {}) {
|
|
877
|
+
|
|
878
|
+
// Discard all previous memory allocations of `vectors_tape_allocator_`
|
|
879
|
+
reset();
|
|
880
|
+
|
|
881
|
+
// Infer the new index size
|
|
882
|
+
serialization_result_t result;
|
|
883
|
+
std::uint64_t matrix_rows = 0;
|
|
884
|
+
std::uint64_t matrix_cols = 0;
|
|
885
|
+
|
|
886
|
+
// We may not want to load the vectors from the same file, or allow attaching them afterwards
|
|
887
|
+
if (!config.exclude_vectors) {
|
|
888
|
+
// Save the matrix size
|
|
889
|
+
if (!config.use_64_bit_dimensions) {
|
|
890
|
+
std::uint32_t dimensions[2];
|
|
891
|
+
if (!input(&dimensions, sizeof(dimensions)))
|
|
892
|
+
return result.failed("Failed to read 32-bit dimensions of the matrix");
|
|
893
|
+
matrix_rows = dimensions[0];
|
|
894
|
+
matrix_cols = dimensions[1];
|
|
895
|
+
} else {
|
|
896
|
+
std::uint64_t dimensions[2];
|
|
897
|
+
if (!input(&dimensions, sizeof(dimensions)))
|
|
898
|
+
return result.failed("Failed to read 64-bit dimensions of the matrix");
|
|
899
|
+
matrix_rows = dimensions[0];
|
|
900
|
+
matrix_cols = dimensions[1];
|
|
901
|
+
}
|
|
902
|
+
// Load the vectors one after another
|
|
903
|
+
vectors_lookup_.resize(matrix_rows);
|
|
904
|
+
for (std::uint64_t slot = 0; slot != matrix_rows; ++slot) {
|
|
905
|
+
byte_t* vector = vectors_tape_allocator_.allocate(matrix_cols);
|
|
906
|
+
if (!input(vector, matrix_cols))
|
|
907
|
+
return result.failed("Failed to read vectors");
|
|
908
|
+
vectors_lookup_[slot] = vector;
|
|
909
|
+
}
|
|
910
|
+
}
|
|
911
|
+
|
|
912
|
+
// Load metadata and choose the right metric
|
|
913
|
+
{
|
|
914
|
+
index_dense_head_buffer_t buffer;
|
|
915
|
+
if (!input(buffer, sizeof(buffer)))
|
|
916
|
+
return result.failed("Failed to read the index ");
|
|
917
|
+
|
|
918
|
+
index_dense_head_t head{buffer};
|
|
919
|
+
if (std::memcmp(buffer, default_magic(), std::strlen(default_magic())) != 0)
|
|
920
|
+
return result.failed("Magic header mismatch - the file isn't an index");
|
|
921
|
+
|
|
922
|
+
// Validate the software version
|
|
923
|
+
if (head.version_major != USEARCH_VERSION_MAJOR)
|
|
924
|
+
return result.failed("File format may be different, please rebuild");
|
|
925
|
+
|
|
926
|
+
// Check the types used
|
|
927
|
+
if (head.kind_key != unum::usearch::scalar_kind<vector_key_t>())
|
|
928
|
+
return result.failed("Key type doesn't match, consider rebuilding");
|
|
929
|
+
if (head.kind_compressed_slot != unum::usearch::scalar_kind<compressed_slot_t>())
|
|
930
|
+
return result.failed("Slot type doesn't match, consider rebuilding");
|
|
931
|
+
|
|
932
|
+
config_.multi = head.multi;
|
|
933
|
+
metric_ = metric_t(head.dimensions, head.kind_metric, head.kind_scalar);
|
|
934
|
+
cast_buffer_.resize(available_threads_.size() * metric_.bytes_per_vector());
|
|
935
|
+
casts_ = make_casts_(head.kind_scalar);
|
|
936
|
+
}
|
|
937
|
+
|
|
938
|
+
// Pull the actual proximity graph
|
|
939
|
+
result = typed_->load_from_stream(std::forward<input_callback_at>(input), std::forward<progress_at>(progress));
|
|
940
|
+
if (!result)
|
|
941
|
+
return result;
|
|
942
|
+
if (typed_->size() != static_cast<std::size_t>(matrix_rows))
|
|
943
|
+
return result.failed("Index size and the number of vectors doesn't match");
|
|
944
|
+
|
|
945
|
+
reindex_keys_();
|
|
946
|
+
return result;
|
|
947
|
+
}
|
|
948
|
+
|
|
949
|
+
/**
|
|
950
|
+
* @brief Parses the index from file, without loading it into RAM.
|
|
951
|
+
* @param[in] path The path to the file.
|
|
952
|
+
* @param[in] config Configuration parameters for imports.
|
|
953
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
954
|
+
*/
|
|
955
|
+
template <typename progress_at = dummy_progress_t>
|
|
956
|
+
serialization_result_t view(memory_mapped_file_t file, //
|
|
957
|
+
std::size_t offset = 0, serialization_config_t config = {}, //
|
|
958
|
+
progress_at&& progress = {}) {
|
|
959
|
+
|
|
960
|
+
// Discard all previous memory allocations of `vectors_tape_allocator_`
|
|
961
|
+
reset();
|
|
962
|
+
|
|
963
|
+
serialization_result_t result = file.open_if_not();
|
|
964
|
+
if (!result)
|
|
965
|
+
return result;
|
|
966
|
+
|
|
967
|
+
// Infer the new index size
|
|
968
|
+
std::uint64_t matrix_rows = 0;
|
|
969
|
+
std::uint64_t matrix_cols = 0;
|
|
970
|
+
span_punned_t vectors_buffer;
|
|
971
|
+
|
|
972
|
+
// We may not want to fetch the vectors from the same file, or allow attaching them afterwards
|
|
973
|
+
if (!config.exclude_vectors) {
|
|
974
|
+
// Save the matrix size
|
|
975
|
+
if (!config.use_64_bit_dimensions) {
|
|
976
|
+
std::uint32_t dimensions[2];
|
|
977
|
+
if (file.size() - offset < sizeof(dimensions))
|
|
978
|
+
return result.failed("File is corrupted and lacks matrix dimensions");
|
|
979
|
+
std::memcpy(&dimensions, file.data() + offset, sizeof(dimensions));
|
|
980
|
+
matrix_rows = dimensions[0];
|
|
981
|
+
matrix_cols = dimensions[1];
|
|
982
|
+
offset += sizeof(dimensions);
|
|
983
|
+
} else {
|
|
984
|
+
std::uint64_t dimensions[2];
|
|
985
|
+
if (file.size() - offset < sizeof(dimensions))
|
|
986
|
+
return result.failed("File is corrupted and lacks matrix dimensions");
|
|
987
|
+
std::memcpy(&dimensions, file.data() + offset, sizeof(dimensions));
|
|
988
|
+
matrix_rows = dimensions[0];
|
|
989
|
+
matrix_cols = dimensions[1];
|
|
990
|
+
offset += sizeof(dimensions);
|
|
991
|
+
}
|
|
992
|
+
vectors_buffer = {file.data() + offset, static_cast<std::size_t>(matrix_rows * matrix_cols)};
|
|
993
|
+
offset += vectors_buffer.size();
|
|
994
|
+
}
|
|
995
|
+
|
|
996
|
+
// Load metadata and choose the right metric
|
|
997
|
+
{
|
|
998
|
+
index_dense_head_buffer_t buffer;
|
|
999
|
+
if (file.size() - offset < sizeof(buffer))
|
|
1000
|
+
return result.failed("File is corrupted and lacks a header");
|
|
1001
|
+
|
|
1002
|
+
std::memcpy(buffer, file.data() + offset, sizeof(buffer));
|
|
1003
|
+
|
|
1004
|
+
index_dense_head_t head{buffer};
|
|
1005
|
+
if (std::memcmp(buffer, default_magic(), std::strlen(default_magic())) != 0)
|
|
1006
|
+
return result.failed("Magic header mismatch - the file isn't an index");
|
|
1007
|
+
|
|
1008
|
+
// Validate the software version
|
|
1009
|
+
if (head.version_major != USEARCH_VERSION_MAJOR)
|
|
1010
|
+
return result.failed("File format may be different, please rebuild");
|
|
1011
|
+
|
|
1012
|
+
// Check the types used
|
|
1013
|
+
if (head.kind_key != unum::usearch::scalar_kind<vector_key_t>())
|
|
1014
|
+
return result.failed("Key type doesn't match, consider rebuilding");
|
|
1015
|
+
if (head.kind_compressed_slot != unum::usearch::scalar_kind<compressed_slot_t>())
|
|
1016
|
+
return result.failed("Slot type doesn't match, consider rebuilding");
|
|
1017
|
+
|
|
1018
|
+
config_.multi = head.multi;
|
|
1019
|
+
metric_ = metric_t(head.dimensions, head.kind_metric, head.kind_scalar);
|
|
1020
|
+
cast_buffer_.resize(available_threads_.size() * metric_.bytes_per_vector());
|
|
1021
|
+
casts_ = make_casts_(head.kind_scalar);
|
|
1022
|
+
offset += sizeof(buffer);
|
|
1023
|
+
}
|
|
1024
|
+
|
|
1025
|
+
// Pull the actual proximity graph
|
|
1026
|
+
result = typed_->view(std::move(file), offset, std::forward<progress_at>(progress));
|
|
1027
|
+
if (!result)
|
|
1028
|
+
return result;
|
|
1029
|
+
if (typed_->size() != static_cast<std::size_t>(matrix_rows))
|
|
1030
|
+
return result.failed("Index size and the number of vectors doesn't match");
|
|
1031
|
+
|
|
1032
|
+
// Address the vectors
|
|
1033
|
+
vectors_lookup_.resize(matrix_rows);
|
|
1034
|
+
if (!config.exclude_vectors)
|
|
1035
|
+
for (std::uint64_t slot = 0; slot != matrix_rows; ++slot)
|
|
1036
|
+
vectors_lookup_[slot] = (byte_t*)vectors_buffer.data() + matrix_cols * slot;
|
|
1037
|
+
|
|
1038
|
+
reindex_keys_();
|
|
1039
|
+
return result;
|
|
1040
|
+
}
|
|
1041
|
+
|
|
1042
|
+
/**
|
|
1043
|
+
* @brief Saves the index to a file.
|
|
1044
|
+
* @param[in] path The path to the file.
|
|
1045
|
+
* @param[in] config Configuration parameters for exports.
|
|
1046
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
1047
|
+
*/
|
|
1048
|
+
template <typename progress_at = dummy_progress_t>
|
|
1049
|
+
serialization_result_t save(output_file_t file, serialization_config_t config = {},
|
|
1050
|
+
progress_at&& progress = {}) const {
|
|
1051
|
+
|
|
1052
|
+
serialization_result_t io_result = file.open_if_not();
|
|
1053
|
+
if (!io_result)
|
|
1054
|
+
return io_result;
|
|
1055
|
+
|
|
1056
|
+
serialization_result_t stream_result = save_to_stream(
|
|
1057
|
+
[&](void const* buffer, std::size_t length) {
|
|
1058
|
+
io_result = file.write(buffer, length);
|
|
1059
|
+
return !!io_result;
|
|
1060
|
+
},
|
|
1061
|
+
config, std::forward<progress_at>(progress));
|
|
1062
|
+
|
|
1063
|
+
if (!stream_result) {
|
|
1064
|
+
io_result.error.release();
|
|
1065
|
+
return stream_result;
|
|
1066
|
+
}
|
|
1067
|
+
return io_result;
|
|
1068
|
+
}
|
|
1069
|
+
|
|
1070
|
+
/**
|
|
1071
|
+
* @brief Memory-maps the serialized binary index representation from disk,
|
|
1072
|
+
* @b without copying data into RAM, and fetching it on-demand.
|
|
1073
|
+
*/
|
|
1074
|
+
template <typename progress_at = dummy_progress_t>
|
|
1075
|
+
serialization_result_t save(memory_mapped_file_t file, //
|
|
1076
|
+
std::size_t offset = 0, //
|
|
1077
|
+
serialization_config_t config = {}, //
|
|
1078
|
+
progress_at&& progress = {}) const {
|
|
1079
|
+
|
|
1080
|
+
serialization_result_t io_result = file.open_if_not();
|
|
1081
|
+
if (!io_result)
|
|
1082
|
+
return io_result;
|
|
1083
|
+
|
|
1084
|
+
serialization_result_t stream_result = save_to_stream(
|
|
1085
|
+
[&](void const* buffer, std::size_t length) {
|
|
1086
|
+
if (offset + length > file.size())
|
|
1087
|
+
return false;
|
|
1088
|
+
std::memcpy(file.data() + offset, buffer, length);
|
|
1089
|
+
offset += length;
|
|
1090
|
+
return true;
|
|
1091
|
+
},
|
|
1092
|
+
config, std::forward<progress_at>(progress));
|
|
1093
|
+
|
|
1094
|
+
return stream_result;
|
|
1095
|
+
}
|
|
1096
|
+
|
|
1097
|
+
/**
|
|
1098
|
+
* @brief Parses the index from file to RAM.
|
|
1099
|
+
* @param[in] path The path to the file.
|
|
1100
|
+
* @param[in] config Configuration parameters for imports.
|
|
1101
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
1102
|
+
*/
|
|
1103
|
+
template <typename progress_at = dummy_progress_t>
|
|
1104
|
+
serialization_result_t load(input_file_t file, serialization_config_t config = {}, progress_at&& progress = {}) {
|
|
1105
|
+
|
|
1106
|
+
serialization_result_t io_result = file.open_if_not();
|
|
1107
|
+
if (!io_result)
|
|
1108
|
+
return io_result;
|
|
1109
|
+
|
|
1110
|
+
serialization_result_t stream_result = load_from_stream(
|
|
1111
|
+
[&](void* buffer, std::size_t length) {
|
|
1112
|
+
io_result = file.read(buffer, length);
|
|
1113
|
+
return !!io_result;
|
|
1114
|
+
},
|
|
1115
|
+
config, std::forward<progress_at>(progress));
|
|
1116
|
+
|
|
1117
|
+
if (!stream_result) {
|
|
1118
|
+
io_result.error.release();
|
|
1119
|
+
return stream_result;
|
|
1120
|
+
}
|
|
1121
|
+
return io_result;
|
|
1122
|
+
}
|
|
1123
|
+
|
|
1124
|
+
/**
|
|
1125
|
+
* @brief Memory-maps the serialized binary index representation from disk,
|
|
1126
|
+
* @b without copying data into RAM, and fetching it on-demand.
|
|
1127
|
+
*/
|
|
1128
|
+
template <typename progress_at = dummy_progress_t>
|
|
1129
|
+
serialization_result_t load(memory_mapped_file_t file, //
|
|
1130
|
+
std::size_t offset = 0, //
|
|
1131
|
+
serialization_config_t config = {}, //
|
|
1132
|
+
progress_at&& progress = {}) {
|
|
1133
|
+
|
|
1134
|
+
serialization_result_t io_result = file.open_if_not();
|
|
1135
|
+
if (!io_result)
|
|
1136
|
+
return io_result;
|
|
1137
|
+
|
|
1138
|
+
serialization_result_t stream_result = load_from_stream(
|
|
1139
|
+
[&](void* buffer, std::size_t length) {
|
|
1140
|
+
if (offset + length > file.size())
|
|
1141
|
+
return false;
|
|
1142
|
+
std::memcpy(buffer, file.data() + offset, length);
|
|
1143
|
+
offset += length;
|
|
1144
|
+
return true;
|
|
1145
|
+
},
|
|
1146
|
+
config, std::forward<progress_at>(progress));
|
|
1147
|
+
|
|
1148
|
+
return stream_result;
|
|
1149
|
+
}
|
|
1150
|
+
|
|
1151
|
+
template <typename progress_at = dummy_progress_t>
|
|
1152
|
+
serialization_result_t save(char const* file_path, //
|
|
1153
|
+
serialization_config_t config = {}, //
|
|
1154
|
+
progress_at&& progress = {}) const {
|
|
1155
|
+
return save(output_file_t(file_path), config, std::forward<progress_at>(progress));
|
|
1156
|
+
}
|
|
1157
|
+
|
|
1158
|
+
template <typename progress_at = dummy_progress_t>
|
|
1159
|
+
serialization_result_t load(char const* file_path, //
|
|
1160
|
+
serialization_config_t config = {}, //
|
|
1161
|
+
progress_at&& progress = {}) {
|
|
1162
|
+
return load(input_file_t(file_path), config, std::forward<progress_at>(progress));
|
|
1163
|
+
}
|
|
1164
|
+
|
|
1165
|
+
/**
|
|
1166
|
+
* @brief Checks if a vector with specified key is present.
|
|
1167
|
+
* @return `true` if the key is present in the index, `false` otherwise.
|
|
1168
|
+
*/
|
|
1169
|
+
bool contains(vector_key_t key) const {
|
|
1170
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1171
|
+
return slot_lookup_.contains(key_and_slot_t::any_slot(key));
|
|
1172
|
+
}
|
|
1173
|
+
|
|
1174
|
+
/**
|
|
1175
|
+
* @brief Count the number of vectors with specified key present.
|
|
1176
|
+
* @return Zero if nothing is found, a positive integer otherwise.
|
|
1177
|
+
*/
|
|
1178
|
+
std::size_t count(vector_key_t key) const {
|
|
1179
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1180
|
+
return slot_lookup_.count(key_and_slot_t::any_slot(key));
|
|
1181
|
+
}
|
|
1182
|
+
|
|
1183
|
+
struct labeling_result_t {
|
|
1184
|
+
error_t error{};
|
|
1185
|
+
std::size_t completed{};
|
|
1186
|
+
|
|
1187
|
+
explicit operator bool() const noexcept { return !error; }
|
|
1188
|
+
labeling_result_t failed(error_t message) noexcept {
|
|
1189
|
+
error = std::move(message);
|
|
1190
|
+
return std::move(*this);
|
|
1191
|
+
}
|
|
1192
|
+
};
|
|
1193
|
+
|
|
1194
|
+
/**
|
|
1195
|
+
* @brief Removes an entry with the specified key from the index.
|
|
1196
|
+
* @param[in] key The key of the entry to remove.
|
|
1197
|
+
* @return The ::labeling_result_t indicating the result of the removal operation.
|
|
1198
|
+
* If the removal was successful, `result.completed` will be `true`.
|
|
1199
|
+
* If the key was not found in the index, `result.completed` will be `false`.
|
|
1200
|
+
* If an error occurred during the removal operation, `result.error` will contain an error message.
|
|
1201
|
+
*/
|
|
1202
|
+
labeling_result_t remove(vector_key_t key) {
|
|
1203
|
+
labeling_result_t result;
|
|
1204
|
+
|
|
1205
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
1206
|
+
auto matching_slots = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
|
|
1207
|
+
if (matching_slots.first == matching_slots.second)
|
|
1208
|
+
return result;
|
|
1209
|
+
|
|
1210
|
+
// Grow the removed entries ring, if needed
|
|
1211
|
+
std::size_t matching_count = std::distance(matching_slots.first, matching_slots.second);
|
|
1212
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
1213
|
+
if (!free_keys_.reserve(free_keys_.size() + matching_count))
|
|
1214
|
+
return result.failed("Can't allocate memory for a free-list");
|
|
1215
|
+
|
|
1216
|
+
// A removed entry would be:
|
|
1217
|
+
// - present in `free_keys_`
|
|
1218
|
+
// - missing in the `slot_lookup_`
|
|
1219
|
+
// - marked in the `typed_` index with a `free_key_`
|
|
1220
|
+
for (auto slots_it = matching_slots.first; slots_it != matching_slots.second; ++slots_it) {
|
|
1221
|
+
compressed_slot_t slot = (*slots_it).slot;
|
|
1222
|
+
free_keys_.push(slot);
|
|
1223
|
+
typed_->at(slot).key = free_key_;
|
|
1224
|
+
}
|
|
1225
|
+
slot_lookup_.erase(key);
|
|
1226
|
+
result.completed = matching_count;
|
|
1227
|
+
|
|
1228
|
+
return result;
|
|
1229
|
+
}
|
|
1230
|
+
|
|
1231
|
+
/**
|
|
1232
|
+
* @brief Removes multiple entries with the specified keys from the index.
|
|
1233
|
+
* @param[in] keys_begin The beginning of the keys range.
|
|
1234
|
+
* @param[in] keys_end The ending of the keys range.
|
|
1235
|
+
* @return The ::labeling_result_t indicating the result of the removal operation.
|
|
1236
|
+
* `result.completed` will contain the number of keys that were successfully removed.
|
|
1237
|
+
* `result.error` will contain an error message if an error occurred during the removal operation.
|
|
1238
|
+
*/
|
|
1239
|
+
template <typename keys_iterator_at>
|
|
1240
|
+
labeling_result_t remove(keys_iterator_at keys_begin, keys_iterator_at keys_end) {
|
|
1241
|
+
|
|
1242
|
+
labeling_result_t result;
|
|
1243
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
1244
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
1245
|
+
// Grow the removed entries ring, if needed
|
|
1246
|
+
std::size_t matching_count = 0;
|
|
1247
|
+
for (auto keys_it = keys_begin; keys_it != keys_end; ++keys_it)
|
|
1248
|
+
matching_count += slot_lookup_.count(key_and_slot_t::any_slot(*keys_it));
|
|
1249
|
+
|
|
1250
|
+
if (!free_keys_.reserve(free_keys_.size() + matching_count))
|
|
1251
|
+
return result.failed("Can't allocate memory for a free-list");
|
|
1252
|
+
|
|
1253
|
+
// Remove them one-by-one
|
|
1254
|
+
for (auto keys_it = keys_begin; keys_it != keys_end; ++keys_it) {
|
|
1255
|
+
vector_key_t key = *keys_it;
|
|
1256
|
+
auto matching_slots = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
|
|
1257
|
+
// A removed entry would be:
|
|
1258
|
+
// - present in `free_keys_`
|
|
1259
|
+
// - missing in the `slot_lookup_`
|
|
1260
|
+
// - marked in the `typed_` index with a `free_key_`
|
|
1261
|
+
matching_count = 0;
|
|
1262
|
+
for (auto slots_it = matching_slots.first; slots_it != matching_slots.second; ++slots_it) {
|
|
1263
|
+
compressed_slot_t slot = (*slots_it).slot;
|
|
1264
|
+
free_keys_.push(slot);
|
|
1265
|
+
typed_->at(slot).key = free_key_;
|
|
1266
|
+
++matching_count;
|
|
1267
|
+
}
|
|
1268
|
+
|
|
1269
|
+
slot_lookup_.erase(key);
|
|
1270
|
+
result.completed += matching_count;
|
|
1271
|
+
}
|
|
1272
|
+
|
|
1273
|
+
return result;
|
|
1274
|
+
}
|
|
1275
|
+
|
|
1276
|
+
/**
|
|
1277
|
+
* @brief Renames an entry with the specified key to a new key.
|
|
1278
|
+
* @param[in] from The current key of the entry to rename.
|
|
1279
|
+
* @param[in] to The new key to assign to the entry.
|
|
1280
|
+
* @return The ::labeling_result_t indicating the result of the rename operation.
|
|
1281
|
+
* If the rename was successful, `result.completed` will be `true`.
|
|
1282
|
+
* If the entry with the current key was not found, `result.completed` will be `false`.
|
|
1283
|
+
*/
|
|
1284
|
+
labeling_result_t rename(vector_key_t from, vector_key_t to) {
|
|
1285
|
+
labeling_result_t result;
|
|
1286
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
1287
|
+
|
|
1288
|
+
if (!multi() && slot_lookup_.contains(key_and_slot_t::any_slot(to)))
|
|
1289
|
+
return result.failed("Renaming impossible, the key is already in use");
|
|
1290
|
+
|
|
1291
|
+
// The `from` may map to multiple entries
|
|
1292
|
+
while (true) {
|
|
1293
|
+
key_and_slot_t key_and_slot_removed;
|
|
1294
|
+
if (!slot_lookup_.pop_first(key_and_slot_t::any_slot(from), key_and_slot_removed))
|
|
1295
|
+
break;
|
|
1296
|
+
|
|
1297
|
+
key_and_slot_t key_and_slot_replacing{to, key_and_slot_removed.slot};
|
|
1298
|
+
slot_lookup_.try_emplace(key_and_slot_replacing); // This can't fail
|
|
1299
|
+
typed_->at(key_and_slot_removed.slot).key = to;
|
|
1300
|
+
++result.completed;
|
|
1301
|
+
}
|
|
1302
|
+
|
|
1303
|
+
return result;
|
|
1304
|
+
}
|
|
1305
|
+
|
|
1306
|
+
/**
|
|
1307
|
+
* @brief Exports a range of keys for the vectors present in the index.
|
|
1308
|
+
* @param[out] keys Pointer to the array where the keys will be exported.
|
|
1309
|
+
* @param[in] offset The number of keys to skip. Useful for pagination.
|
|
1310
|
+
* @param[in] limit The maximum number of keys to export, that can fit in ::keys.
|
|
1311
|
+
*/
|
|
1312
|
+
void export_keys(vector_key_t* keys, std::size_t offset, std::size_t limit) const {
|
|
1313
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1314
|
+
offset = (std::min)(offset, slot_lookup_.size());
|
|
1315
|
+
slot_lookup_.for_each([&](key_and_slot_t const& key_and_slot) {
|
|
1316
|
+
if (offset)
|
|
1317
|
+
// Skip the first `offset` entries
|
|
1318
|
+
--offset;
|
|
1319
|
+
else if (limit) {
|
|
1320
|
+
*keys = key_and_slot.key;
|
|
1321
|
+
++keys;
|
|
1322
|
+
--limit;
|
|
1323
|
+
}
|
|
1324
|
+
});
|
|
1325
|
+
}
|
|
1326
|
+
|
|
1327
|
+
struct copy_result_t {
|
|
1328
|
+
index_dense_gt index;
|
|
1329
|
+
error_t error;
|
|
1330
|
+
|
|
1331
|
+
explicit operator bool() const noexcept { return !error; }
|
|
1332
|
+
copy_result_t failed(error_t message) noexcept {
|
|
1333
|
+
error = std::move(message);
|
|
1334
|
+
return std::move(*this);
|
|
1335
|
+
}
|
|
1336
|
+
};
|
|
1337
|
+
|
|
1338
|
+
/**
|
|
1339
|
+
* @brief Copies the ::index_dense_gt @b with all the data in it.
|
|
1340
|
+
* @param config The copy configuration (optional).
|
|
1341
|
+
* @return A copy of the ::index_dense_gt instance.
|
|
1342
|
+
*/
|
|
1343
|
+
copy_result_t copy(index_dense_copy_config_t config = {}) const {
|
|
1344
|
+
copy_result_t result = fork();
|
|
1345
|
+
if (!result)
|
|
1346
|
+
return result;
|
|
1347
|
+
|
|
1348
|
+
auto typed_result = typed_->copy(config);
|
|
1349
|
+
if (!typed_result)
|
|
1350
|
+
return result.failed(std::move(typed_result.error));
|
|
1351
|
+
|
|
1352
|
+
// Export the free (removed) slot numbers
|
|
1353
|
+
index_dense_gt& copy = result.index;
|
|
1354
|
+
if (!copy.free_keys_.reserve(free_keys_.size()))
|
|
1355
|
+
return result.failed(std::move(typed_result.error));
|
|
1356
|
+
for (std::size_t i = 0; i != free_keys_.size(); ++i)
|
|
1357
|
+
copy.free_keys_.push(free_keys_[i]);
|
|
1358
|
+
|
|
1359
|
+
// Allocate buffers and move the vectors themselves
|
|
1360
|
+
if (!config.force_vector_copy && copy.config_.exclude_vectors)
|
|
1361
|
+
copy.vectors_lookup_ = vectors_lookup_;
|
|
1362
|
+
else {
|
|
1363
|
+
copy.vectors_lookup_.resize(vectors_lookup_.size());
|
|
1364
|
+
for (std::size_t slot = 0; slot != vectors_lookup_.size(); ++slot)
|
|
1365
|
+
copy.vectors_lookup_[slot] = copy.vectors_tape_allocator_.allocate(copy.metric_.bytes_per_vector());
|
|
1366
|
+
if (std::count(copy.vectors_lookup_.begin(), copy.vectors_lookup_.end(), nullptr))
|
|
1367
|
+
return result.failed("Out of memory!");
|
|
1368
|
+
for (std::size_t slot = 0; slot != vectors_lookup_.size(); ++slot)
|
|
1369
|
+
std::memcpy(copy.vectors_lookup_[slot], vectors_lookup_[slot], metric_.bytes_per_vector());
|
|
1370
|
+
}
|
|
1371
|
+
|
|
1372
|
+
copy.slot_lookup_ = slot_lookup_;
|
|
1373
|
+
*copy.typed_ = std::move(typed_result.index);
|
|
1374
|
+
return result;
|
|
1375
|
+
}
|
|
1376
|
+
|
|
1377
|
+
/**
|
|
1378
|
+
* @brief Copies the ::index_dense_gt model @b without any data.
|
|
1379
|
+
* @return A similarly configured ::index_dense_gt instance.
|
|
1380
|
+
*/
|
|
1381
|
+
copy_result_t fork() const {
|
|
1382
|
+
copy_result_t result;
|
|
1383
|
+
index_dense_gt& other = result.index;
|
|
1384
|
+
|
|
1385
|
+
other.config_ = config_;
|
|
1386
|
+
other.cast_buffer_ = cast_buffer_;
|
|
1387
|
+
other.casts_ = casts_;
|
|
1388
|
+
|
|
1389
|
+
other.metric_ = metric_;
|
|
1390
|
+
other.available_threads_ = available_threads_;
|
|
1391
|
+
other.free_key_ = free_key_;
|
|
1392
|
+
|
|
1393
|
+
index_t* raw = index_allocator_t{}.allocate(1);
|
|
1394
|
+
if (!raw)
|
|
1395
|
+
return result.failed("Can't allocate the index");
|
|
1396
|
+
|
|
1397
|
+
new (raw) index_t(config());
|
|
1398
|
+
other.typed_ = raw;
|
|
1399
|
+
return result;
|
|
1400
|
+
}
|
|
1401
|
+
|
|
1402
|
+
struct compaction_result_t {
|
|
1403
|
+
error_t error{};
|
|
1404
|
+
std::size_t pruned_edges{};
|
|
1405
|
+
|
|
1406
|
+
explicit operator bool() const noexcept { return !error; }
|
|
1407
|
+
compaction_result_t failed(error_t message) noexcept {
|
|
1408
|
+
error = std::move(message);
|
|
1409
|
+
return std::move(*this);
|
|
1410
|
+
}
|
|
1411
|
+
};
|
|
1412
|
+
|
|
1413
|
+
/**
|
|
1414
|
+
* @brief Performs compaction on the index, pruning links to removed entries.
|
|
1415
|
+
* @param executor The executor parallel processing. Default ::dummy_executor_t single-threaded.
|
|
1416
|
+
* @param progress The progress tracker instance to use. Default ::dummy_progress_t reports nothing.
|
|
1417
|
+
* @return The ::compaction_result_t indicating the result of the compaction operation.
|
|
1418
|
+
* `result.pruned_edges` will contain the number of edges that were removed.
|
|
1419
|
+
* `result.error` will contain an error message if an error occurred during the compaction operation.
|
|
1420
|
+
*/
|
|
1421
|
+
template <typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
|
|
1422
|
+
compaction_result_t isolate(executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
|
|
1423
|
+
compaction_result_t result;
|
|
1424
|
+
std::atomic<std::size_t> pruned_edges;
|
|
1425
|
+
auto disallow = [&](member_cref_t const& member) noexcept {
|
|
1426
|
+
bool freed = member.key == free_key_;
|
|
1427
|
+
pruned_edges += freed;
|
|
1428
|
+
return freed;
|
|
1429
|
+
};
|
|
1430
|
+
typed_->isolate(disallow, std::forward<executor_at>(executor), std::forward<progress_at>(progress));
|
|
1431
|
+
result.pruned_edges = pruned_edges;
|
|
1432
|
+
return result;
|
|
1433
|
+
}
|
|
1434
|
+
|
|
1435
|
+
class values_proxy_t {
|
|
1436
|
+
index_dense_gt const* index_;
|
|
1437
|
+
|
|
1438
|
+
public:
|
|
1439
|
+
values_proxy_t(index_dense_gt const& index) noexcept : index_(&index) {}
|
|
1440
|
+
byte_t const* operator[](compressed_slot_t slot) const noexcept { return index_->vectors_lookup_[slot]; }
|
|
1441
|
+
byte_t const* operator[](member_citerator_t it) const noexcept { return index_->vectors_lookup_[get_slot(it)]; }
|
|
1442
|
+
};
|
|
1443
|
+
|
|
1444
|
+
/**
|
|
1445
|
+
* @brief Performs compaction on the index, pruning links to removed entries.
|
|
1446
|
+
* @param executor The executor parallel processing. Default ::dummy_executor_t single-threaded.
|
|
1447
|
+
* @param progress The progress tracker instance to use. Default ::dummy_progress_t reports nothing.
|
|
1448
|
+
* @return The ::compaction_result_t indicating the result of the compaction operation.
|
|
1449
|
+
* `result.pruned_edges` will contain the number of edges that were removed.
|
|
1450
|
+
* `result.error` will contain an error message if an error occurred during the compaction operation.
|
|
1451
|
+
*/
|
|
1452
|
+
template <typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
|
|
1453
|
+
compaction_result_t compact(executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
|
|
1454
|
+
compaction_result_t result;
|
|
1455
|
+
|
|
1456
|
+
std::vector<byte_t*> new_vectors_lookup(vectors_lookup_.size());
|
|
1457
|
+
vectors_tape_allocator_t new_vectors_allocator;
|
|
1458
|
+
|
|
1459
|
+
auto track_slot_change = [&](vector_key_t, compressed_slot_t old_slot, compressed_slot_t new_slot) {
|
|
1460
|
+
byte_t* new_vector = new_vectors_allocator.allocate(metric_.bytes_per_vector());
|
|
1461
|
+
byte_t* old_vector = vectors_lookup_[old_slot];
|
|
1462
|
+
std::memcpy(new_vector, old_vector, metric_.bytes_per_vector());
|
|
1463
|
+
new_vectors_lookup[new_slot] = new_vector;
|
|
1464
|
+
};
|
|
1465
|
+
typed_->compact(values_proxy_t{*this}, metric_proxy_t{*this}, track_slot_change,
|
|
1466
|
+
std::forward<executor_at>(executor), std::forward<progress_at>(progress));
|
|
1467
|
+
vectors_lookup_ = std::move(new_vectors_lookup);
|
|
1468
|
+
vectors_tape_allocator_ = std::move(new_vectors_allocator);
|
|
1469
|
+
return result;
|
|
1470
|
+
}
|
|
1471
|
+
|
|
1472
|
+
template < //
|
|
1473
|
+
typename man_to_woman_at = dummy_key_to_key_mapping_t, //
|
|
1474
|
+
typename woman_to_man_at = dummy_key_to_key_mapping_t, //
|
|
1475
|
+
typename executor_at = dummy_executor_t, //
|
|
1476
|
+
typename progress_at = dummy_progress_t //
|
|
1477
|
+
>
|
|
1478
|
+
join_result_t join( //
|
|
1479
|
+
index_dense_gt const& women, //
|
|
1480
|
+
index_join_config_t config = {}, //
|
|
1481
|
+
man_to_woman_at&& man_to_woman = man_to_woman_at{}, //
|
|
1482
|
+
woman_to_man_at&& woman_to_man = woman_to_man_at{}, //
|
|
1483
|
+
executor_at&& executor = executor_at{}, //
|
|
1484
|
+
progress_at&& progress = progress_at{}) const {
|
|
1485
|
+
|
|
1486
|
+
index_dense_gt const& men = *this;
|
|
1487
|
+
return unum::usearch::join( //
|
|
1488
|
+
*men.typed_, *women.typed_, //
|
|
1489
|
+
values_proxy_t{men}, values_proxy_t{women}, //
|
|
1490
|
+
metric_proxy_t{men}, metric_proxy_t{women}, //
|
|
1491
|
+
config, //
|
|
1492
|
+
std::forward<man_to_woman_at>(man_to_woman), //
|
|
1493
|
+
std::forward<woman_to_man_at>(woman_to_man), //
|
|
1494
|
+
std::forward<executor_at>(executor), //
|
|
1495
|
+
std::forward<progress_at>(progress));
|
|
1496
|
+
}
|
|
1497
|
+
|
|
1498
|
+
struct clustering_result_t {
|
|
1499
|
+
error_t error{};
|
|
1500
|
+
std::size_t clusters{};
|
|
1501
|
+
std::size_t visited_members{};
|
|
1502
|
+
std::size_t computed_distances{};
|
|
1503
|
+
|
|
1504
|
+
explicit operator bool() const noexcept { return !error; }
|
|
1505
|
+
clustering_result_t failed(error_t message) noexcept {
|
|
1506
|
+
error = std::move(message);
|
|
1507
|
+
return std::move(*this);
|
|
1508
|
+
}
|
|
1509
|
+
};
|
|
1510
|
+
|
|
1511
|
+
/**
|
|
1512
|
+
* @brief Implements clustering, classifying the given objects (vectors of member keys)
|
|
1513
|
+
* into a given number of clusters.
|
|
1514
|
+
*
|
|
1515
|
+
* @param[in] queries_begin Iterator pointing to the first query.
|
|
1516
|
+
* @param[in] queries_end Iterator pointing to the last query.
|
|
1517
|
+
* @param[in] executor Thread-pool to execute the job in parallel.
|
|
1518
|
+
* @param[in] progress Callback to report the execution progress.
|
|
1519
|
+
* @param[in] config Configuration parameters for clustering.
|
|
1520
|
+
*
|
|
1521
|
+
* @param[out] cluster_keys Pointer to the array where the cluster keys will be exported.
|
|
1522
|
+
* @param[out] cluster_distances Pointer to the array where the distances to those centroids will be exported.
|
|
1523
|
+
*/
|
|
1524
|
+
template < //
|
|
1525
|
+
typename queries_iterator_at, //
|
|
1526
|
+
typename executor_at = dummy_executor_t, //
|
|
1527
|
+
typename progress_at = dummy_progress_t //
|
|
1528
|
+
>
|
|
1529
|
+
clustering_result_t cluster( //
|
|
1530
|
+
queries_iterator_at queries_begin, //
|
|
1531
|
+
queries_iterator_at queries_end, //
|
|
1532
|
+
index_dense_clustering_config_t config, //
|
|
1533
|
+
vector_key_t* cluster_keys, //
|
|
1534
|
+
distance_t* cluster_distances, //
|
|
1535
|
+
executor_at&& executor = executor_at{}, //
|
|
1536
|
+
progress_at&& progress = progress_at{}) {
|
|
1537
|
+
|
|
1538
|
+
std::size_t const queries_count = queries_end - queries_begin;
|
|
1539
|
+
|
|
1540
|
+
// Find the first level (top -> down) that has enough nodes to exceed `config.min_clusters`.
|
|
1541
|
+
std::size_t level = max_level();
|
|
1542
|
+
if (config.min_clusters) {
|
|
1543
|
+
for (; level > 1; --level) {
|
|
1544
|
+
if (stats(level).nodes > config.min_clusters)
|
|
1545
|
+
break;
|
|
1546
|
+
}
|
|
1547
|
+
} else
|
|
1548
|
+
level = 1, config.max_clusters = stats(1).nodes, config.min_clusters = 2;
|
|
1549
|
+
|
|
1550
|
+
clustering_result_t result;
|
|
1551
|
+
if (max_level() < 2)
|
|
1552
|
+
return result.failed("Index too small to cluster!");
|
|
1553
|
+
|
|
1554
|
+
// A structure used to track the popularity of a specific cluster
|
|
1555
|
+
struct cluster_t {
|
|
1556
|
+
vector_key_t centroid;
|
|
1557
|
+
vector_key_t merged_into;
|
|
1558
|
+
std::size_t popularity;
|
|
1559
|
+
byte_t* vector;
|
|
1560
|
+
};
|
|
1561
|
+
|
|
1562
|
+
auto centroid_id = [](cluster_t const& a, cluster_t const& b) { return a.centroid < b.centroid; };
|
|
1563
|
+
auto higher_popularity = [](cluster_t const& a, cluster_t const& b) { return a.popularity > b.popularity; };
|
|
1564
|
+
|
|
1565
|
+
std::atomic<std::size_t> visited_members(0);
|
|
1566
|
+
std::atomic<std::size_t> computed_distances(0);
|
|
1567
|
+
std::atomic<char const*> atomic_error{nullptr};
|
|
1568
|
+
|
|
1569
|
+
using dynamic_allocator_traits_t = std::allocator_traits<dynamic_allocator_t>;
|
|
1570
|
+
using clusters_allocator_t = typename dynamic_allocator_traits_t::template rebind_alloc<cluster_t>;
|
|
1571
|
+
buffer_gt<cluster_t, clusters_allocator_t> clusters(queries_count);
|
|
1572
|
+
if (!clusters)
|
|
1573
|
+
return result.failed("Out of memory!");
|
|
1574
|
+
|
|
1575
|
+
map_to_clusters:
|
|
1576
|
+
// Concurrently perform search until a certain depth
|
|
1577
|
+
executor.dynamic(queries_count, [&](std::size_t thread_idx, std::size_t query_idx) {
|
|
1578
|
+
auto result = cluster(queries_begin[query_idx], level, thread_idx);
|
|
1579
|
+
if (!result) {
|
|
1580
|
+
atomic_error = result.error.release();
|
|
1581
|
+
return false;
|
|
1582
|
+
}
|
|
1583
|
+
|
|
1584
|
+
cluster_keys[query_idx] = result.cluster.member.key;
|
|
1585
|
+
cluster_distances[query_idx] = result.cluster.distance;
|
|
1586
|
+
|
|
1587
|
+
// Export in case we need to refine afterwards
|
|
1588
|
+
clusters[query_idx].centroid = result.cluster.member.key;
|
|
1589
|
+
clusters[query_idx].vector = vectors_lookup_[result.cluster.member.slot];
|
|
1590
|
+
clusters[query_idx].merged_into = free_key();
|
|
1591
|
+
clusters[query_idx].popularity = 1;
|
|
1592
|
+
|
|
1593
|
+
visited_members += result.visited_members;
|
|
1594
|
+
computed_distances += result.computed_distances;
|
|
1595
|
+
return true;
|
|
1596
|
+
});
|
|
1597
|
+
|
|
1598
|
+
if (atomic_error)
|
|
1599
|
+
return result.failed(atomic_error.load());
|
|
1600
|
+
|
|
1601
|
+
// Now once we have identified the closest clusters,
|
|
1602
|
+
// we can try reducing their quantity, refining
|
|
1603
|
+
std::sort(clusters.begin(), clusters.end(), centroid_id);
|
|
1604
|
+
|
|
1605
|
+
// Transform into run-length encoding, computing the number of unique clusters
|
|
1606
|
+
std::size_t unique_clusters = 0;
|
|
1607
|
+
{
|
|
1608
|
+
std::size_t last_idx = 0;
|
|
1609
|
+
for (std::size_t current_idx = 1; current_idx != clusters.size(); ++current_idx) {
|
|
1610
|
+
if (clusters[last_idx].centroid == clusters[current_idx].centroid) {
|
|
1611
|
+
clusters[last_idx].popularity++;
|
|
1612
|
+
} else {
|
|
1613
|
+
last_idx++;
|
|
1614
|
+
clusters[last_idx] = clusters[current_idx];
|
|
1615
|
+
}
|
|
1616
|
+
}
|
|
1617
|
+
unique_clusters = last_idx + 1;
|
|
1618
|
+
}
|
|
1619
|
+
|
|
1620
|
+
// In some cases the queries may be co-located, all mapping into the same cluster on that
|
|
1621
|
+
// level. In that case we refine the granularity and dive deeper into clusters:
|
|
1622
|
+
if (unique_clusters < config.min_clusters && level > 1) {
|
|
1623
|
+
level--;
|
|
1624
|
+
goto map_to_clusters;
|
|
1625
|
+
}
|
|
1626
|
+
|
|
1627
|
+
std::sort(clusters.data(), clusters.data() + unique_clusters, higher_popularity);
|
|
1628
|
+
|
|
1629
|
+
// If clusters are too numerous, merge the ones that are too close to each other.
|
|
1630
|
+
std::size_t merge_cycles = 0;
|
|
1631
|
+
merge_nearby_clusters:
|
|
1632
|
+
if (unique_clusters > config.max_clusters) {
|
|
1633
|
+
|
|
1634
|
+
cluster_t& merge_source = clusters[unique_clusters - 1];
|
|
1635
|
+
std::size_t merge_target_idx = 0;
|
|
1636
|
+
distance_t merge_distance = std::numeric_limits<distance_t>::max();
|
|
1637
|
+
|
|
1638
|
+
for (std::size_t candidate_idx = 0; candidate_idx + 1 < unique_clusters; ++candidate_idx) {
|
|
1639
|
+
distance_t distance = metric_(merge_source.vector, clusters[candidate_idx].vector);
|
|
1640
|
+
if (distance < merge_distance) {
|
|
1641
|
+
merge_distance = distance;
|
|
1642
|
+
merge_target_idx = candidate_idx;
|
|
1643
|
+
}
|
|
1644
|
+
}
|
|
1645
|
+
|
|
1646
|
+
merge_source.merged_into = clusters[merge_target_idx].centroid;
|
|
1647
|
+
clusters[merge_target_idx].popularity += exchange(merge_source.popularity, 0);
|
|
1648
|
+
|
|
1649
|
+
// The target object may have to be swapped a few times to get to optimal position.
|
|
1650
|
+
while (merge_target_idx &&
|
|
1651
|
+
clusters[merge_target_idx - 1].popularity < clusters[merge_target_idx].popularity)
|
|
1652
|
+
std::swap(clusters[merge_target_idx - 1], clusters[merge_target_idx]), --merge_target_idx;
|
|
1653
|
+
|
|
1654
|
+
unique_clusters--;
|
|
1655
|
+
merge_cycles++;
|
|
1656
|
+
goto merge_nearby_clusters;
|
|
1657
|
+
}
|
|
1658
|
+
|
|
1659
|
+
// Replace evicted clusters
|
|
1660
|
+
if (merge_cycles) {
|
|
1661
|
+
// Sort dropped clusters by name to accelerate future lookups
|
|
1662
|
+
auto clusters_end = clusters.data() + config.max_clusters + merge_cycles;
|
|
1663
|
+
std::sort(clusters.data(), clusters_end, centroid_id);
|
|
1664
|
+
|
|
1665
|
+
executor.dynamic(queries_count, [&](std::size_t thread_idx, std::size_t query_idx) {
|
|
1666
|
+
vector_key_t& cluster_key = cluster_keys[query_idx];
|
|
1667
|
+
distance_t& cluster_distance = cluster_distances[query_idx];
|
|
1668
|
+
|
|
1669
|
+
// Recursively trace replacements of that cluster
|
|
1670
|
+
while (true) {
|
|
1671
|
+
// To avoid implementing heterogeneous comparisons, lets wrap the `cluster_key`
|
|
1672
|
+
cluster_t updated_cluster;
|
|
1673
|
+
updated_cluster.centroid = cluster_key;
|
|
1674
|
+
updated_cluster = *std::lower_bound(clusters.data(), clusters_end, updated_cluster, centroid_id);
|
|
1675
|
+
if (updated_cluster.merged_into == free_key())
|
|
1676
|
+
break;
|
|
1677
|
+
cluster_key = updated_cluster.merged_into;
|
|
1678
|
+
}
|
|
1679
|
+
|
|
1680
|
+
cluster_distance = distance_between(cluster_key, queries_begin[query_idx], thread_idx).mean;
|
|
1681
|
+
return true;
|
|
1682
|
+
});
|
|
1683
|
+
}
|
|
1684
|
+
|
|
1685
|
+
result.computed_distances = computed_distances;
|
|
1686
|
+
result.visited_members = visited_members;
|
|
1687
|
+
|
|
1688
|
+
(void)progress;
|
|
1689
|
+
return result;
|
|
1690
|
+
}
|
|
1691
|
+
|
|
1692
|
+
private:
|
|
1693
|
+
struct thread_lock_t {
|
|
1694
|
+
index_dense_gt const& parent;
|
|
1695
|
+
std::size_t thread_id;
|
|
1696
|
+
bool engaged;
|
|
1697
|
+
|
|
1698
|
+
~thread_lock_t() {
|
|
1699
|
+
if (engaged)
|
|
1700
|
+
parent.thread_unlock_(thread_id);
|
|
1701
|
+
}
|
|
1702
|
+
};
|
|
1703
|
+
|
|
1704
|
+
thread_lock_t thread_lock_(std::size_t thread_id) const {
|
|
1705
|
+
if (thread_id != any_thread())
|
|
1706
|
+
return {*this, thread_id, false};
|
|
1707
|
+
|
|
1708
|
+
available_threads_mutex_.lock();
|
|
1709
|
+
thread_id = available_threads_.back();
|
|
1710
|
+
available_threads_.pop_back();
|
|
1711
|
+
available_threads_mutex_.unlock();
|
|
1712
|
+
return {*this, thread_id, true};
|
|
1713
|
+
}
|
|
1714
|
+
|
|
1715
|
+
void thread_unlock_(std::size_t thread_id) const {
|
|
1716
|
+
available_threads_mutex_.lock();
|
|
1717
|
+
available_threads_.push_back(thread_id);
|
|
1718
|
+
available_threads_mutex_.unlock();
|
|
1719
|
+
}
|
|
1720
|
+
|
|
1721
|
+
template <typename scalar_at>
|
|
1722
|
+
add_result_t add_( //
|
|
1723
|
+
vector_key_t key, scalar_at const* vector, //
|
|
1724
|
+
std::size_t thread, bool force_vector_copy, cast_t const& cast) {
|
|
1725
|
+
|
|
1726
|
+
if (!multi() && contains(key))
|
|
1727
|
+
return add_result_t{}.failed("Duplicate keys not allowed in high-level wrappers");
|
|
1728
|
+
|
|
1729
|
+
// Cast the vector, if needed for compatibility with `metric_`
|
|
1730
|
+
thread_lock_t lock = thread_lock_(thread);
|
|
1731
|
+
bool copy_vector = !config_.exclude_vectors || force_vector_copy;
|
|
1732
|
+
byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
|
|
1733
|
+
{
|
|
1734
|
+
byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
|
|
1735
|
+
bool casted = cast(vector_data, dimensions(), casted_data);
|
|
1736
|
+
if (casted)
|
|
1737
|
+
vector_data = casted_data, copy_vector = true;
|
|
1738
|
+
}
|
|
1739
|
+
|
|
1740
|
+
// Check if there are some removed entries, whose nodes we can reuse
|
|
1741
|
+
compressed_slot_t free_slot = default_free_value<compressed_slot_t>();
|
|
1742
|
+
{
|
|
1743
|
+
std::unique_lock<std::mutex> lock(free_keys_mutex_);
|
|
1744
|
+
free_keys_.try_pop(free_slot);
|
|
1745
|
+
}
|
|
1746
|
+
|
|
1747
|
+
// Perform the insertion or the update
|
|
1748
|
+
bool reuse_node = free_slot != default_free_value<compressed_slot_t>();
|
|
1749
|
+
auto on_success = [&](member_ref_t member) {
|
|
1750
|
+
unique_lock_t slot_lock(slot_lookup_mutex_);
|
|
1751
|
+
slot_lookup_.try_emplace(key_and_slot_t{key, static_cast<compressed_slot_t>(member.slot)});
|
|
1752
|
+
if (copy_vector) {
|
|
1753
|
+
if (!reuse_node)
|
|
1754
|
+
vectors_lookup_[member.slot] = vectors_tape_allocator_.allocate(metric_.bytes_per_vector());
|
|
1755
|
+
std::memcpy(vectors_lookup_[member.slot], vector_data, metric_.bytes_per_vector());
|
|
1756
|
+
} else
|
|
1757
|
+
vectors_lookup_[member.slot] = (byte_t*)vector_data;
|
|
1758
|
+
};
|
|
1759
|
+
|
|
1760
|
+
index_update_config_t update_config;
|
|
1761
|
+
update_config.thread = lock.thread_id;
|
|
1762
|
+
update_config.expansion = config_.expansion_add;
|
|
1763
|
+
|
|
1764
|
+
metric_proxy_t metric{*this};
|
|
1765
|
+
return reuse_node //
|
|
1766
|
+
? typed_->update(typed_->iterator_at(free_slot), key, vector_data, metric, update_config, on_success)
|
|
1767
|
+
: typed_->add(key, vector_data, metric, update_config, on_success);
|
|
1768
|
+
}
|
|
1769
|
+
|
|
1770
|
+
template <typename scalar_at>
|
|
1771
|
+
search_result_t search_( //
|
|
1772
|
+
scalar_at const* vector, std::size_t wanted, //
|
|
1773
|
+
std::size_t thread, bool exact, cast_t const& cast) const {
|
|
1774
|
+
|
|
1775
|
+
// Cast the vector, if needed for compatibility with `metric_`
|
|
1776
|
+
thread_lock_t lock = thread_lock_(thread);
|
|
1777
|
+
byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
|
|
1778
|
+
{
|
|
1779
|
+
byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
|
|
1780
|
+
bool casted = cast(vector_data, dimensions(), casted_data);
|
|
1781
|
+
if (casted)
|
|
1782
|
+
vector_data = casted_data;
|
|
1783
|
+
}
|
|
1784
|
+
|
|
1785
|
+
index_search_config_t search_config;
|
|
1786
|
+
search_config.thread = lock.thread_id;
|
|
1787
|
+
search_config.expansion = config_.expansion_search;
|
|
1788
|
+
search_config.exact = exact;
|
|
1789
|
+
|
|
1790
|
+
auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
|
|
1791
|
+
return typed_->search(vector_data, wanted, metric_proxy_t{*this}, search_config, allow);
|
|
1792
|
+
}
|
|
1793
|
+
|
|
1794
|
+
template <typename scalar_at>
|
|
1795
|
+
cluster_result_t cluster_( //
|
|
1796
|
+
scalar_at const* vector, std::size_t level, //
|
|
1797
|
+
std::size_t thread, cast_t const& cast) const {
|
|
1798
|
+
|
|
1799
|
+
// Cast the vector, if needed for compatibility with `metric_`
|
|
1800
|
+
thread_lock_t lock = thread_lock_(thread);
|
|
1801
|
+
byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
|
|
1802
|
+
{
|
|
1803
|
+
byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
|
|
1804
|
+
bool casted = cast(vector_data, dimensions(), casted_data);
|
|
1805
|
+
if (casted)
|
|
1806
|
+
vector_data = casted_data;
|
|
1807
|
+
}
|
|
1808
|
+
|
|
1809
|
+
index_cluster_config_t cluster_config;
|
|
1810
|
+
cluster_config.thread = lock.thread_id;
|
|
1811
|
+
cluster_config.expansion = config_.expansion_search;
|
|
1812
|
+
|
|
1813
|
+
auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
|
|
1814
|
+
return typed_->cluster(vector_data, level, metric_proxy_t{*this}, cluster_config, allow);
|
|
1815
|
+
}
|
|
1816
|
+
|
|
1817
|
+
template <typename scalar_at>
|
|
1818
|
+
aggregated_distances_t distance_between_( //
|
|
1819
|
+
vector_key_t key, scalar_at const* vector, //
|
|
1820
|
+
std::size_t thread, cast_t const& cast) const {
|
|
1821
|
+
|
|
1822
|
+
// Cast the vector, if needed for compatibility with `metric_`
|
|
1823
|
+
thread_lock_t lock = thread_lock_(thread);
|
|
1824
|
+
byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
|
|
1825
|
+
{
|
|
1826
|
+
byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
|
|
1827
|
+
bool casted = cast(vector_data, dimensions(), casted_data);
|
|
1828
|
+
if (casted)
|
|
1829
|
+
vector_data = casted_data;
|
|
1830
|
+
}
|
|
1831
|
+
|
|
1832
|
+
// Check if such `key` is even present.
|
|
1833
|
+
shared_lock_t slots_lock(slot_lookup_mutex_);
|
|
1834
|
+
auto key_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
|
|
1835
|
+
aggregated_distances_t result;
|
|
1836
|
+
if (key_range.first == key_range.second)
|
|
1837
|
+
return result;
|
|
1838
|
+
|
|
1839
|
+
result.min = std::numeric_limits<distance_t>::max();
|
|
1840
|
+
result.max = std::numeric_limits<distance_t>::min();
|
|
1841
|
+
result.mean = 0;
|
|
1842
|
+
result.count = 0;
|
|
1843
|
+
|
|
1844
|
+
while (key_range.first != key_range.second) {
|
|
1845
|
+
key_and_slot_t key_and_slot = *key_range.first;
|
|
1846
|
+
byte_t const* a_vector = vectors_lookup_[key_and_slot.slot];
|
|
1847
|
+
byte_t const* b_vector = vector_data;
|
|
1848
|
+
distance_t a_b_distance = metric_(a_vector, b_vector);
|
|
1849
|
+
|
|
1850
|
+
result.mean += a_b_distance;
|
|
1851
|
+
result.min = (std::min)(result.min, a_b_distance);
|
|
1852
|
+
result.max = (std::max)(result.max, a_b_distance);
|
|
1853
|
+
result.count++;
|
|
1854
|
+
|
|
1855
|
+
//
|
|
1856
|
+
++key_range.first;
|
|
1857
|
+
}
|
|
1858
|
+
|
|
1859
|
+
result.mean /= result.count;
|
|
1860
|
+
return result;
|
|
1861
|
+
}
|
|
1862
|
+
|
|
1863
|
+
void reindex_keys_() {
|
|
1864
|
+
|
|
1865
|
+
// Estimate number of entries first
|
|
1866
|
+
std::size_t count_total = typed_->size();
|
|
1867
|
+
std::size_t count_removed = 0;
|
|
1868
|
+
for (std::size_t i = 0; i != count_total; ++i) {
|
|
1869
|
+
member_cref_t member = typed_->at(i);
|
|
1870
|
+
count_removed += member.key == free_key_;
|
|
1871
|
+
}
|
|
1872
|
+
|
|
1873
|
+
if (!count_removed && !config_.enable_key_lookups)
|
|
1874
|
+
return;
|
|
1875
|
+
|
|
1876
|
+
// Pull entries from the underlying `typed_` into either
|
|
1877
|
+
// into `slot_lookup_`, or `free_keys_` if they are unused.
|
|
1878
|
+
unique_lock_t lock(slot_lookup_mutex_);
|
|
1879
|
+
slot_lookup_.clear();
|
|
1880
|
+
if (config_.enable_key_lookups)
|
|
1881
|
+
slot_lookup_.reserve(count_total - count_removed);
|
|
1882
|
+
free_keys_.clear();
|
|
1883
|
+
free_keys_.reserve(count_removed);
|
|
1884
|
+
for (std::size_t i = 0; i != typed_->size(); ++i) {
|
|
1885
|
+
member_cref_t member = typed_->at(i);
|
|
1886
|
+
if (member.key == free_key_)
|
|
1887
|
+
free_keys_.push(static_cast<compressed_slot_t>(i));
|
|
1888
|
+
else if (config_.enable_key_lookups)
|
|
1889
|
+
slot_lookup_.try_emplace(key_and_slot_t{vector_key_t(member.key), static_cast<compressed_slot_t>(i)});
|
|
1890
|
+
}
|
|
1891
|
+
}
|
|
1892
|
+
|
|
1893
|
+
template <typename scalar_at>
|
|
1894
|
+
std::size_t get_(vector_key_t key, scalar_at* reconstructed, std::size_t vectors_limit, cast_t const& cast) const {
|
|
1895
|
+
|
|
1896
|
+
if (!multi()) {
|
|
1897
|
+
compressed_slot_t slot;
|
|
1898
|
+
// Find the matching ID
|
|
1899
|
+
{
|
|
1900
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1901
|
+
auto it = slot_lookup_.find(key_and_slot_t::any_slot(key));
|
|
1902
|
+
if (it == slot_lookup_.end())
|
|
1903
|
+
return false;
|
|
1904
|
+
slot = (*it).slot;
|
|
1905
|
+
}
|
|
1906
|
+
// Export the entry
|
|
1907
|
+
byte_t const* punned_vector = reinterpret_cast<byte_t const*>(vectors_lookup_[slot]);
|
|
1908
|
+
bool casted = cast(punned_vector, dimensions(), (byte_t*)reconstructed);
|
|
1909
|
+
if (!casted)
|
|
1910
|
+
std::memcpy(reconstructed, punned_vector, metric_.bytes_per_vector());
|
|
1911
|
+
return true;
|
|
1912
|
+
} else {
|
|
1913
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1914
|
+
auto equal_range_pair = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
|
|
1915
|
+
std::size_t count_exported = 0;
|
|
1916
|
+
for (auto begin = equal_range_pair.first;
|
|
1917
|
+
begin != equal_range_pair.second && count_exported != vectors_limit; ++begin, ++count_exported) {
|
|
1918
|
+
//
|
|
1919
|
+
compressed_slot_t slot = (*begin).slot;
|
|
1920
|
+
byte_t const* punned_vector = reinterpret_cast<byte_t const*>(vectors_lookup_[slot]);
|
|
1921
|
+
byte_t* reconstructed_vector = (byte_t*)reconstructed + metric_.bytes_per_vector() * count_exported;
|
|
1922
|
+
bool casted = cast(punned_vector, dimensions(), reconstructed_vector);
|
|
1923
|
+
if (!casted)
|
|
1924
|
+
std::memcpy(reconstructed_vector, punned_vector, metric_.bytes_per_vector());
|
|
1925
|
+
}
|
|
1926
|
+
return count_exported;
|
|
1927
|
+
}
|
|
1928
|
+
}
|
|
1929
|
+
|
|
1930
|
+
template <typename to_scalar_at> static casts_t make_casts_() {
|
|
1931
|
+
casts_t result;
|
|
1932
|
+
|
|
1933
|
+
result.from_b1x8 = cast_gt<b1x8_t, to_scalar_at>{};
|
|
1934
|
+
result.from_i8 = cast_gt<i8_t, to_scalar_at>{};
|
|
1935
|
+
result.from_f16 = cast_gt<f16_t, to_scalar_at>{};
|
|
1936
|
+
result.from_f32 = cast_gt<f32_t, to_scalar_at>{};
|
|
1937
|
+
result.from_f64 = cast_gt<f64_t, to_scalar_at>{};
|
|
1938
|
+
|
|
1939
|
+
result.to_b1x8 = cast_gt<to_scalar_at, b1x8_t>{};
|
|
1940
|
+
result.to_i8 = cast_gt<to_scalar_at, i8_t>{};
|
|
1941
|
+
result.to_f16 = cast_gt<to_scalar_at, f16_t>{};
|
|
1942
|
+
result.to_f32 = cast_gt<to_scalar_at, f32_t>{};
|
|
1943
|
+
result.to_f64 = cast_gt<to_scalar_at, f64_t>{};
|
|
1944
|
+
|
|
1945
|
+
return result;
|
|
1946
|
+
}
|
|
1947
|
+
|
|
1948
|
+
static casts_t make_casts_(scalar_kind_t scalar_kind) {
|
|
1949
|
+
switch (scalar_kind) {
|
|
1950
|
+
case scalar_kind_t::f64_k: return make_casts_<f64_t>();
|
|
1951
|
+
case scalar_kind_t::f32_k: return make_casts_<f32_t>();
|
|
1952
|
+
case scalar_kind_t::f16_k: return make_casts_<f16_t>();
|
|
1953
|
+
case scalar_kind_t::i8_k: return make_casts_<i8_t>();
|
|
1954
|
+
case scalar_kind_t::b1x8_k: return make_casts_<b1x8_t>();
|
|
1955
|
+
default: return {};
|
|
1956
|
+
}
|
|
1957
|
+
}
|
|
1958
|
+
};
|
|
1959
|
+
|
|
1960
|
+
using index_dense_t = index_dense_gt<>;
|
|
1961
|
+
using index_dense_big_t = index_dense_gt<uuid_t, uint40_t>;
|
|
1962
|
+
|
|
1963
|
+
/**
|
|
1964
|
+
* @brief Adapts the Male-Optimal Stable Marriage algorithm for unequal sets
|
|
1965
|
+
* to perform fast one-to-one matching between two large collections
|
|
1966
|
+
* of vectors, using approximate nearest neighbors search.
|
|
1967
|
+
*
|
|
1968
|
+
* @param[inout] man_to_woman Container to map ::first keys to ::second.
|
|
1969
|
+
* @param[inout] woman_to_man Container to map ::second keys to ::first.
|
|
1970
|
+
* @param[in] executor Thread-pool to execute the job in parallel.
|
|
1971
|
+
* @param[in] progress Callback to report the execution progress.
|
|
1972
|
+
*/
|
|
1973
|
+
template < //
|
|
1974
|
+
|
|
1975
|
+
typename men_key_at, //
|
|
1976
|
+
typename women_key_at, //
|
|
1977
|
+
typename men_slot_at, //
|
|
1978
|
+
typename women_slot_at, //
|
|
1979
|
+
|
|
1980
|
+
typename man_to_woman_at = dummy_key_to_key_mapping_t, //
|
|
1981
|
+
typename woman_to_man_at = dummy_key_to_key_mapping_t, //
|
|
1982
|
+
typename executor_at = dummy_executor_t, //
|
|
1983
|
+
typename progress_at = dummy_progress_t //
|
|
1984
|
+
>
|
|
1985
|
+
static join_result_t join( //
|
|
1986
|
+
index_dense_gt<men_key_at, men_slot_at> const& men, //
|
|
1987
|
+
index_dense_gt<women_key_at, women_slot_at> const& women, //
|
|
1988
|
+
|
|
1989
|
+
index_join_config_t config = {}, //
|
|
1990
|
+
man_to_woman_at&& man_to_woman = man_to_woman_at{}, //
|
|
1991
|
+
woman_to_man_at&& woman_to_man = woman_to_man_at{}, //
|
|
1992
|
+
executor_at&& executor = executor_at{}, //
|
|
1993
|
+
progress_at&& progress = progress_at{}) noexcept {
|
|
1994
|
+
|
|
1995
|
+
return men.join( //
|
|
1996
|
+
women, config, //
|
|
1997
|
+
std::forward<woman_to_man_at>(woman_to_man), //
|
|
1998
|
+
std::forward<man_to_woman_at>(man_to_woman), //
|
|
1999
|
+
std::forward<executor_at>(executor), //
|
|
2000
|
+
std::forward<progress_at>(progress));
|
|
2001
|
+
}
|
|
2002
|
+
|
|
2003
|
+
} // namespace usearch
|
|
2004
|
+
} // namespace unum
|