usearch 0.22.2 → 1.0.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.github/workflows/prerelease.yml +3 -3
- package/.releaserc +9 -0
- package/.vscode/launch.json +15 -11
- package/.vscode/settings.json +12 -2
- package/.vscode/tasks.json +4 -4
- package/CMakeLists.txt +6 -1
- package/Cargo.toml +2 -2
- package/README.md +39 -46
- package/build.rs +2 -2
- package/c/CMakeLists.txt +48 -0
- package/c/README.md +34 -0
- package/c/lib.cpp +83 -72
- package/c/test.c +20 -20
- package/c/usearch.h +132 -32
- package/conanfile.py +1 -1
- package/cpp/CMakeLists.txt +3 -1
- package/cpp/README.md +16 -16
- package/cpp/bench.cpp +50 -126
- package/cpp/test.cpp +164 -77
- package/docs/benchmarks.md +26 -5
- package/docs/c/reference.rst +3 -3
- package/docs/compilation.md +49 -45
- package/docs/cpp/index.rst +2 -2
- package/docs/cpp/reference.rst +13 -3
- package/docs/index.rst +8 -8
- package/docs/java/reference.rst +1 -1
- package/golang/lib.go +275 -141
- package/golang/lib_test.go +105 -0
- package/include/usearch/index.hpp +1733 -1659
- package/include/usearch/index_dense.hpp +1400 -0
- package/include/usearch/index_plugins.hpp +1263 -0
- package/java/README.md +1 -1
- package/java/cloud/unum/usearch/Index.java +10 -10
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +34 -37
- package/java/test/IndexTest.java +1 -1
- package/javascript/README.md +7 -7
- package/javascript/lib.cpp +94 -55
- package/javascript/test.js +25 -11
- package/javascript/usearch.d.ts +87 -0
- package/javascript/usearch.js +2 -118
- package/objc/USearchObjective.mm +57 -31
- package/objc/include/USearchObjective.h +32 -22
- package/package.json +5 -3
- package/pyproject.toml +3 -3
- package/python/README.md +23 -14
- package/python/lib.cpp +472 -444
- package/python/scripts/bench.ipynb +3 -2
- package/python/scripts/bench.py +4 -4
- package/python/scripts/index_faiss.py +5 -5
- package/python/scripts/join.py +4 -4
- package/python/scripts/test.py +152 -61
- package/python/scripts/test_jit.py +10 -12
- package/python/scripts/test_sparse.py +32 -0
- package/python/usearch/client.py +26 -28
- package/python/usearch/eval.py +73 -23
- package/python/usearch/index.py +520 -228
- package/python/usearch/io.py +57 -25
- package/python/usearch/numba.py +4 -4
- package/python/usearch/server.py +10 -10
- package/rust/lib.cpp +22 -27
- package/rust/lib.hpp +11 -12
- package/rust/lib.rs +3 -3
- package/setup.py +5 -2
- package/simsimd/.github/workflows/prerelease.yml +73 -0
- package/simsimd/.github/workflows/release.yml +110 -0
- package/simsimd/.github/workflows/update_version.sh +3 -0
- package/simsimd/.releaserc +77 -0
- package/simsimd/.vscode/launch.json +27 -1
- package/simsimd/.vscode/settings.json +22 -1
- package/simsimd/.vscode/tasks.json +6 -0
- package/simsimd/CMakeLists.txt +3 -2
- package/simsimd/README.md +82 -46
- package/simsimd/bench.cxx +63 -24
- package/simsimd/include/simsimd/simsimd.h +88 -66
- package/simsimd/include/simsimd/simsimd_chem.h +126 -0
- package/simsimd/package.json +8 -0
- package/simsimd/pyproject.toml +42 -0
- package/simsimd/python/python.c +81 -0
- package/simsimd/python/run.sh +1 -0
- package/simsimd/python/test.c +109 -0
- package/simsimd/python/test.py +23 -0
- package/simsimd/setup.py +44 -0
- package/swift/Index+Sugar.swift +75 -23
- package/swift/README.md +26 -2
- package/swift/Test.swift +19 -7
- package/swift/USearch.swift +2 -2
- package/tsconfig.json +8 -0
- package/wasm/CMakeLists.txt +1 -0
- package/wasmer.toml +5 -1
- package/wolfram/CMakeLists.txt +1 -0
- package/wolfram/lib.cpp +11 -13
- package/golang/usearch.h +0 -9
- package/golang/usearch_par_test.go +0 -40
- package/golang/usearch_test.go +0 -82
- package/include/usearch/index_punned_dense.hpp +0 -871
- package/include/usearch/index_punned_helpers.hpp +0 -566
|
@@ -0,0 +1,1400 @@
|
|
|
1
|
+
#pragma once
|
|
2
|
+
#include <stdlib.h> // `aligned_alloc`
|
|
3
|
+
|
|
4
|
+
#include <functional> // `std::function`
|
|
5
|
+
#include <numeric> // `std::iota`
|
|
6
|
+
#include <shared_mutex> // `std::shared_mutex`
|
|
7
|
+
#include <thread> // `std::thread`
|
|
8
|
+
#include <vector> // `std::vector`
|
|
9
|
+
|
|
10
|
+
#include <usearch/index.hpp>
|
|
11
|
+
#include <usearch/index_plugins.hpp>
|
|
12
|
+
|
|
13
|
+
#include <tsl/robin_set.h>
|
|
14
|
+
|
|
15
|
+
namespace unum {
|
|
16
|
+
namespace usearch {
|
|
17
|
+
|
|
18
|
+
template <typename, typename> class index_dense_gt;
|
|
19
|
+
|
|
20
|
+
// clang-format off
|
|
21
|
+
template <typename key_at, typename std::enable_if<std::is_integral<key_at>::value>::type* = nullptr> key_at default_free_value() { return std::numeric_limits<key_at>::max(); }
|
|
22
|
+
template <typename key_at, typename std::enable_if<std::is_same<key_at, uint40_t>::value>::type* = nullptr> uint40_t default_free_value() { return uint40_t::max(); }
|
|
23
|
+
template <typename key_at, typename std::enable_if<!std::is_integral<key_at>::value && !std::is_same<key_at, uint40_t>::value>::type* = nullptr> key_at default_free_value() { return key_at(); }
|
|
24
|
+
// clang-format on
|
|
25
|
+
|
|
26
|
+
/**
|
|
27
|
+
* @brief The "magic" sequence helps infer the type of the file.
|
|
28
|
+
* USearch indexes start with the "usearch" string.
|
|
29
|
+
*/
|
|
30
|
+
constexpr char const* default_magic() { return "usearch"; }
|
|
31
|
+
|
|
32
|
+
using index_dense_head_buffer_t = byte_t[64];
|
|
33
|
+
|
|
34
|
+
static_assert(sizeof(index_dense_head_buffer_t) == 64, "File header should be exactly 64 bytes");
|
|
35
|
+
|
|
36
|
+
/**
|
|
37
|
+
* @brief Serialized binary representations of the USearch index start with metadata.
|
|
38
|
+
* Metadata is parsed into a `index_dense_head_t`, containing the USearch package version,
|
|
39
|
+
* and the properties of the index.
|
|
40
|
+
*
|
|
41
|
+
* It uses: 13 bytes for file versioning, 22 bytes for structural information = 35 bytes.
|
|
42
|
+
* The following 24 bytes contain binary size of the graph, of the vectors, and the checksum,
|
|
43
|
+
* leaving 5 bytes at the end vacant.
|
|
44
|
+
*/
|
|
45
|
+
struct index_dense_head_t {
|
|
46
|
+
|
|
47
|
+
// Versioning:
|
|
48
|
+
using magic_t = char[7];
|
|
49
|
+
using version_t = std::uint16_t;
|
|
50
|
+
|
|
51
|
+
// Versioning: 7 + 2 * 3 = 13 bytes
|
|
52
|
+
char const* magic;
|
|
53
|
+
misaligned_ref_gt<version_t> version_major;
|
|
54
|
+
misaligned_ref_gt<version_t> version_minor;
|
|
55
|
+
misaligned_ref_gt<version_t> version_patch;
|
|
56
|
+
|
|
57
|
+
// Structural: 4 * 3 = 12 bytes
|
|
58
|
+
misaligned_ref_gt<metric_kind_t> kind_metric;
|
|
59
|
+
misaligned_ref_gt<scalar_kind_t> kind_scalar;
|
|
60
|
+
misaligned_ref_gt<scalar_kind_t> kind_key;
|
|
61
|
+
misaligned_ref_gt<scalar_kind_t> kind_compressed_slot;
|
|
62
|
+
|
|
63
|
+
// Population: 8 * 3 = 24 bytes
|
|
64
|
+
misaligned_ref_gt<std::uint64_t> count_present;
|
|
65
|
+
misaligned_ref_gt<std::uint64_t> count_deleted;
|
|
66
|
+
misaligned_ref_gt<std::uint64_t> dimensions;
|
|
67
|
+
|
|
68
|
+
index_dense_head_t(byte_t* ptr) noexcept
|
|
69
|
+
: magic((char const*)exchange(ptr, ptr + sizeof(magic_t))), //
|
|
70
|
+
version_major(exchange(ptr, ptr + sizeof(version_t))), //
|
|
71
|
+
version_minor(exchange(ptr, ptr + sizeof(version_t))), //
|
|
72
|
+
version_patch(exchange(ptr, ptr + sizeof(version_t))), //
|
|
73
|
+
kind_metric(exchange(ptr, ptr + sizeof(metric_kind_t))), //
|
|
74
|
+
kind_scalar(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
|
|
75
|
+
kind_key(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
|
|
76
|
+
kind_compressed_slot(exchange(ptr, ptr + sizeof(scalar_kind_t))),
|
|
77
|
+
count_present(exchange(ptr, ptr + sizeof(std::uint64_t))),
|
|
78
|
+
count_deleted(exchange(ptr, ptr + sizeof(std::uint64_t))),
|
|
79
|
+
dimensions(exchange(ptr, ptr + sizeof(std::uint64_t))) {}
|
|
80
|
+
};
|
|
81
|
+
|
|
82
|
+
struct index_dense_head_result_t {
|
|
83
|
+
|
|
84
|
+
index_dense_head_buffer_t buffer;
|
|
85
|
+
index_dense_head_t head;
|
|
86
|
+
error_t error;
|
|
87
|
+
|
|
88
|
+
explicit operator bool() const noexcept { return !error; }
|
|
89
|
+
index_dense_head_result_t failed(error_t message) noexcept {
|
|
90
|
+
error = std::move(message);
|
|
91
|
+
return std::move(*this);
|
|
92
|
+
}
|
|
93
|
+
};
|
|
94
|
+
|
|
95
|
+
struct index_dense_config_t : public index_config_t {
|
|
96
|
+
std::size_t expansion_add = default_expansion_add();
|
|
97
|
+
std::size_t expansion_search = default_expansion_search();
|
|
98
|
+
bool exclude_vectors = false;
|
|
99
|
+
bool allow_key_collisions = false;
|
|
100
|
+
|
|
101
|
+
index_dense_config_t(index_config_t base) noexcept : index_config_t(base) {}
|
|
102
|
+
|
|
103
|
+
index_dense_config_t( //
|
|
104
|
+
std::size_t connectivity = default_connectivity(), //
|
|
105
|
+
std::size_t expansion_add = default_expansion_add(),
|
|
106
|
+
std::size_t expansion_search = default_expansion_search()) noexcept
|
|
107
|
+
: index_config_t(connectivity), //
|
|
108
|
+
expansion_add(expansion_add ? expansion_add : default_expansion_add()),
|
|
109
|
+
expansion_search(expansion_search ? expansion_search : default_expansion_search()) {}
|
|
110
|
+
};
|
|
111
|
+
|
|
112
|
+
struct index_dense_serialization_config_t {
|
|
113
|
+
bool exclude_vectors = false;
|
|
114
|
+
bool use_64_bit_dimensions = false;
|
|
115
|
+
};
|
|
116
|
+
|
|
117
|
+
struct index_dense_update_config_t : public index_update_config_t {
|
|
118
|
+
bool force_vector_copy = false;
|
|
119
|
+
|
|
120
|
+
index_dense_update_config_t() = default;
|
|
121
|
+
index_dense_update_config_t(index_update_config_t base) noexcept : index_update_config_t(base) {}
|
|
122
|
+
};
|
|
123
|
+
|
|
124
|
+
struct index_dense_copy_config_t : public index_copy_config_t {
|
|
125
|
+
bool force_vector_copy = true;
|
|
126
|
+
|
|
127
|
+
index_dense_copy_config_t() = default;
|
|
128
|
+
index_dense_copy_config_t(index_copy_config_t base) noexcept : index_copy_config_t(base) {}
|
|
129
|
+
};
|
|
130
|
+
|
|
131
|
+
struct index_dense_metadata_result_t {
|
|
132
|
+
index_dense_serialization_config_t config;
|
|
133
|
+
index_dense_head_buffer_t head_buffer;
|
|
134
|
+
index_dense_head_t head;
|
|
135
|
+
error_t error;
|
|
136
|
+
|
|
137
|
+
explicit operator bool() const noexcept { return !error; }
|
|
138
|
+
index_dense_metadata_result_t failed(error_t message) noexcept {
|
|
139
|
+
error = std::move(message);
|
|
140
|
+
return std::move(*this);
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
index_dense_metadata_result_t() noexcept : config(), head_buffer(), head(head_buffer), error() {}
|
|
144
|
+
|
|
145
|
+
index_dense_metadata_result_t(index_dense_metadata_result_t&& other) noexcept
|
|
146
|
+
: config(), head_buffer(), head(head_buffer), error(std::move(other.error)) {
|
|
147
|
+
std::memcpy(&config, &other.config, sizeof(other.config));
|
|
148
|
+
std::memcpy(&head_buffer, &other.head_buffer, sizeof(other.head_buffer));
|
|
149
|
+
}
|
|
150
|
+
|
|
151
|
+
index_dense_metadata_result_t& operator=(index_dense_metadata_result_t&& other) noexcept {
|
|
152
|
+
std::memcpy(&config, &other.config, sizeof(other.config));
|
|
153
|
+
std::memcpy(&head_buffer, &other.head_buffer, sizeof(other.head_buffer));
|
|
154
|
+
error = std::move(other.error);
|
|
155
|
+
return *this;
|
|
156
|
+
}
|
|
157
|
+
};
|
|
158
|
+
|
|
159
|
+
/**
|
|
160
|
+
* @brief Extracts metadata from pre-constructed index on disk,
|
|
161
|
+
* without loading it or mapping the whole binary file.
|
|
162
|
+
*/
|
|
163
|
+
inline index_dense_metadata_result_t index_dense_metadata(char const* file_path) noexcept {
|
|
164
|
+
index_dense_metadata_result_t result;
|
|
165
|
+
std::unique_ptr<std::FILE, int (*)(std::FILE*)> file(std::fopen(file_path, "rb"), &std::fclose);
|
|
166
|
+
if (!file)
|
|
167
|
+
return result.failed(std::strerror(errno));
|
|
168
|
+
|
|
169
|
+
// Read the header
|
|
170
|
+
std::size_t read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get());
|
|
171
|
+
if (!read)
|
|
172
|
+
return result.failed(std::feof(file.get()) ? "End of file reached!" : std::strerror(errno));
|
|
173
|
+
|
|
174
|
+
// Check if the file immediately starts with the index, instead of vectors
|
|
175
|
+
result.config.exclude_vectors = true;
|
|
176
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
177
|
+
return result;
|
|
178
|
+
|
|
179
|
+
if (std::fseek(file.get(), 0L, SEEK_END) != 0)
|
|
180
|
+
return result.failed("Can't infer file size");
|
|
181
|
+
|
|
182
|
+
// Check if it starts with 32-bit
|
|
183
|
+
std::size_t const file_size = std::ftell(file.get());
|
|
184
|
+
|
|
185
|
+
std::uint32_t dimensions_u32[2]{0};
|
|
186
|
+
std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32));
|
|
187
|
+
std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32);
|
|
188
|
+
|
|
189
|
+
std::uint64_t dimensions_u64[2]{0};
|
|
190
|
+
std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64));
|
|
191
|
+
std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64);
|
|
192
|
+
|
|
193
|
+
// Check if it starts with 32-bit
|
|
194
|
+
if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) {
|
|
195
|
+
if (std::fseek(file.get(), offset_if_u32, SEEK_SET) != 0)
|
|
196
|
+
return result.failed(std::strerror(errno));
|
|
197
|
+
read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get());
|
|
198
|
+
if (!read)
|
|
199
|
+
return result.failed(std::feof(file.get()) ? "End of file reached!" : std::strerror(errno));
|
|
200
|
+
|
|
201
|
+
result.config.exclude_vectors = false;
|
|
202
|
+
result.config.use_64_bit_dimensions = false;
|
|
203
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
204
|
+
return result;
|
|
205
|
+
}
|
|
206
|
+
|
|
207
|
+
// Check if it starts with 64-bit
|
|
208
|
+
if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) {
|
|
209
|
+
if (std::fseek(file.get(), offset_if_u64, SEEK_SET) != 0)
|
|
210
|
+
return result.failed(std::strerror(errno));
|
|
211
|
+
read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get());
|
|
212
|
+
if (!read)
|
|
213
|
+
return result.failed(std::feof(file.get()) ? "End of file reached!" : std::strerror(errno));
|
|
214
|
+
|
|
215
|
+
// Check if it starts with 64-bit
|
|
216
|
+
result.config.exclude_vectors = false;
|
|
217
|
+
result.config.use_64_bit_dimensions = true;
|
|
218
|
+
if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
|
|
219
|
+
return result;
|
|
220
|
+
}
|
|
221
|
+
|
|
222
|
+
return result.failed("Not a dense USearch index!");
|
|
223
|
+
}
|
|
224
|
+
|
|
225
|
+
/**
|
|
226
|
+
* @brief Oversimplified type-punned index for equidimensional vectors
|
|
227
|
+
* with automatic @b down-casting, hardware-specific @b SIMD metrics,
|
|
228
|
+
* and ability to @b remove existing vectors, common in Semantic Caching
|
|
229
|
+
* applications.
|
|
230
|
+
*
|
|
231
|
+
* @section Serialization
|
|
232
|
+
*
|
|
233
|
+
* The serialized binary form of `index_dense_gt` is made up of three parts:
|
|
234
|
+
* 1. Binary matrix, aka the `.bbin` part,
|
|
235
|
+
* 2. Metadata about used metrics, number of used vs free slots,
|
|
236
|
+
* 3. The HNSW index in a binary form.
|
|
237
|
+
* The first (1.) generally starts with 2 integers - number of rows (vectors) and @b single-byte columns.
|
|
238
|
+
* The second (2.) starts with @b "usearch"-magic-string, used to infer the file type on open.
|
|
239
|
+
* The third (3.) is implemented by the underlying `index_gt` class.
|
|
240
|
+
*/
|
|
241
|
+
template <typename key_at = default_key_t, typename compressed_slot_at = default_slot_t> //
|
|
242
|
+
class index_dense_gt {
|
|
243
|
+
public:
|
|
244
|
+
using key_t = key_at;
|
|
245
|
+
using compressed_slot_t = compressed_slot_at;
|
|
246
|
+
using distance_t = distance_punned_t;
|
|
247
|
+
using metric_t = metric_punned_t;
|
|
248
|
+
|
|
249
|
+
using member_ref_t = member_ref_gt<key_t>;
|
|
250
|
+
using member_cref_t = member_cref_gt<key_t>;
|
|
251
|
+
|
|
252
|
+
using head_t = index_dense_head_t;
|
|
253
|
+
using head_buffer_t = index_dense_head_buffer_t;
|
|
254
|
+
using head_result_t = index_dense_head_result_t;
|
|
255
|
+
|
|
256
|
+
using serialization_config_t = index_dense_serialization_config_t;
|
|
257
|
+
|
|
258
|
+
private:
|
|
259
|
+
/// @brief Schema: input buffer, bytes in input buffer, output buffer.
|
|
260
|
+
using cast_t = std::function<bool(byte_t const*, std::size_t, byte_t*)>;
|
|
261
|
+
/// @brief Punned index.
|
|
262
|
+
using index_t = index_gt< //
|
|
263
|
+
distance_t, key_t, compressed_slot_t, //
|
|
264
|
+
aligned_allocator_gt<byte_t, 64>, memory_mapping_allocator_gt<64>>;
|
|
265
|
+
using index_allocator_t = aligned_allocator_gt<index_t, 64>;
|
|
266
|
+
|
|
267
|
+
using member_iterator_t = typename index_t::member_iterator_t;
|
|
268
|
+
using member_citerator_t = typename index_t::member_citerator_t;
|
|
269
|
+
|
|
270
|
+
/// @brief Punned metric object.
|
|
271
|
+
class metric_proxy_t {
|
|
272
|
+
index_dense_gt const* index_ = nullptr;
|
|
273
|
+
|
|
274
|
+
public:
|
|
275
|
+
metric_proxy_t(index_dense_gt const& index) noexcept : index_(&index) {}
|
|
276
|
+
|
|
277
|
+
inline distance_t operator()(byte_t const* a, member_cref_t b) const noexcept { return f(a, v(b)); }
|
|
278
|
+
inline distance_t operator()(member_cref_t a, member_cref_t b) const noexcept { return f(v(a), v(b)); }
|
|
279
|
+
|
|
280
|
+
inline distance_t operator()(byte_t const* a, member_citerator_t b) const noexcept { return f(a, v(b)); }
|
|
281
|
+
inline distance_t operator()(member_citerator_t a, member_citerator_t b) const noexcept {
|
|
282
|
+
return f(v(a), v(b));
|
|
283
|
+
}
|
|
284
|
+
|
|
285
|
+
inline distance_t operator()(byte_t const* a, byte_t const* b) const noexcept { return f(a, b); }
|
|
286
|
+
|
|
287
|
+
inline byte_t const* v(member_cref_t m) const noexcept { return index_->vectors_lookup_[get_slot(m)]; }
|
|
288
|
+
inline byte_t const* v(member_citerator_t m) const noexcept { return index_->vectors_lookup_[get_slot(m)]; }
|
|
289
|
+
inline distance_t f(byte_t const* a, byte_t const* b) const noexcept { return index_->metric_(a, b); }
|
|
290
|
+
};
|
|
291
|
+
|
|
292
|
+
index_dense_config_t config_;
|
|
293
|
+
index_t* typed_ = nullptr;
|
|
294
|
+
|
|
295
|
+
mutable std::vector<byte_t> cast_buffer_;
|
|
296
|
+
struct casts_t {
|
|
297
|
+
cast_t from_b1x8;
|
|
298
|
+
cast_t from_i8;
|
|
299
|
+
cast_t from_f16;
|
|
300
|
+
cast_t from_f32;
|
|
301
|
+
cast_t from_f64;
|
|
302
|
+
|
|
303
|
+
cast_t to_b1x8;
|
|
304
|
+
cast_t to_i8;
|
|
305
|
+
cast_t to_f16;
|
|
306
|
+
cast_t to_f32;
|
|
307
|
+
cast_t to_f64;
|
|
308
|
+
} casts_;
|
|
309
|
+
|
|
310
|
+
/// @brief An instance of a potentially stateful `metric_t` used to initialize copies and forks.
|
|
311
|
+
metric_t metric_;
|
|
312
|
+
|
|
313
|
+
using vectors_tape_allocator_t = memory_mapping_allocator_gt<8>;
|
|
314
|
+
/// @brief Allocator for the copied vectors, aligned to widest double-precision scalars.
|
|
315
|
+
vectors_tape_allocator_t vectors_tape_allocator_;
|
|
316
|
+
|
|
317
|
+
/// @brief For every managed `compressed_slot_t` stores a pointer to the allocated vector copy.
|
|
318
|
+
mutable std::vector<byte_t*> vectors_lookup_;
|
|
319
|
+
|
|
320
|
+
/// @brief Originally forms and array of integers [0, threads], marking all
|
|
321
|
+
mutable std::vector<std::size_t> available_threads_;
|
|
322
|
+
|
|
323
|
+
/// @brief Mutex, controlling concurrent access to `available_threads_`.
|
|
324
|
+
mutable std::mutex available_threads_mutex_;
|
|
325
|
+
|
|
326
|
+
using shared_mutex_t = std::mutex; // TODO: Find an OS-compatible solution
|
|
327
|
+
using shared_lock_t = std::unique_lock<shared_mutex_t>;
|
|
328
|
+
using unique_lock_t = std::unique_lock<shared_mutex_t>;
|
|
329
|
+
|
|
330
|
+
struct key_and_slot_t {
|
|
331
|
+
key_t key;
|
|
332
|
+
compressed_slot_t slot;
|
|
333
|
+
};
|
|
334
|
+
|
|
335
|
+
struct lookup_key_hash_t {
|
|
336
|
+
using is_transparent = void;
|
|
337
|
+
std::size_t operator()(key_and_slot_t const& k) const noexcept { return std::hash<key_t>{}(k.key); }
|
|
338
|
+
std::size_t operator()(key_t const& k) const noexcept { return std::hash<key_t>{}(k); }
|
|
339
|
+
};
|
|
340
|
+
|
|
341
|
+
struct lookup_key_same_t {
|
|
342
|
+
using is_transparent = void;
|
|
343
|
+
bool operator()(key_and_slot_t const& a, key_t const& b) const noexcept { return a.key == b; }
|
|
344
|
+
bool operator()(key_t const& a, key_and_slot_t const& b) const noexcept { return a == b.key; }
|
|
345
|
+
bool operator()(key_and_slot_t const& a, key_and_slot_t const& b) const noexcept {
|
|
346
|
+
return a.key == b.key && a.slot == b.slot;
|
|
347
|
+
}
|
|
348
|
+
};
|
|
349
|
+
|
|
350
|
+
/// @brief Multi-Map from keys to IDs, and allocated vectors.
|
|
351
|
+
tsl::robin_set<key_and_slot_t, lookup_key_hash_t, lookup_key_same_t> slot_lookup_;
|
|
352
|
+
|
|
353
|
+
/// @brief Mutex, controlling concurrent access to `slot_lookup_`.
|
|
354
|
+
mutable shared_mutex_t slot_lookup_mutex_;
|
|
355
|
+
|
|
356
|
+
/// @brief Ring-shaped queue of deleted entries, to be reused on future insertions.
|
|
357
|
+
ring_gt<compressed_slot_t> free_keys_;
|
|
358
|
+
|
|
359
|
+
/// @brief Mutex, controlling concurrent access to `free_keys_`.
|
|
360
|
+
mutable std::mutex free_keys_mutex_;
|
|
361
|
+
|
|
362
|
+
/// @brief A constant for the reserved key value, used to mark deleted entries.
|
|
363
|
+
key_t free_key_ = default_free_value<key_t>();
|
|
364
|
+
|
|
365
|
+
public:
|
|
366
|
+
using search_result_t = typename index_t::search_result_t;
|
|
367
|
+
using add_result_t = typename index_t::add_result_t;
|
|
368
|
+
using stats_t = typename index_t::stats_t;
|
|
369
|
+
using match_t = typename index_t::match_t;
|
|
370
|
+
|
|
371
|
+
index_dense_gt() = default;
|
|
372
|
+
index_dense_gt(index_dense_gt&& other)
|
|
373
|
+
: config_(std::move(other.config_)),
|
|
374
|
+
|
|
375
|
+
typed_(exchange(other.typed_, nullptr)), //
|
|
376
|
+
cast_buffer_(std::move(other.cast_buffer_)), //
|
|
377
|
+
casts_(std::move(other.casts_)), //
|
|
378
|
+
metric_(std::move(other.metric_)), //
|
|
379
|
+
|
|
380
|
+
vectors_tape_allocator_(std::move(other.vectors_tape_allocator_)), //
|
|
381
|
+
vectors_lookup_(std::move(other.vectors_lookup_)), //
|
|
382
|
+
|
|
383
|
+
available_threads_(std::move(other.available_threads_)), //
|
|
384
|
+
slot_lookup_(std::move(other.slot_lookup_)), //
|
|
385
|
+
free_keys_(std::move(other.free_keys_)), //
|
|
386
|
+
free_key_(std::move(other.free_key_)) {} //
|
|
387
|
+
|
|
388
|
+
index_dense_gt& operator=(index_dense_gt&& other) {
|
|
389
|
+
swap(other);
|
|
390
|
+
return *this;
|
|
391
|
+
}
|
|
392
|
+
|
|
393
|
+
/**
|
|
394
|
+
* @brief Swaps the contents of this index with another index.
|
|
395
|
+
* @param other The other index to swap with.
|
|
396
|
+
*/
|
|
397
|
+
void swap(index_dense_gt& other) {
|
|
398
|
+
std::swap(config_, other.config_);
|
|
399
|
+
|
|
400
|
+
std::swap(typed_, other.typed_);
|
|
401
|
+
std::swap(cast_buffer_, other.cast_buffer_);
|
|
402
|
+
std::swap(casts_, other.casts_);
|
|
403
|
+
std::swap(metric_, other.metric_);
|
|
404
|
+
|
|
405
|
+
std::swap(vectors_tape_allocator_, other.vectors_tape_allocator_);
|
|
406
|
+
std::swap(vectors_lookup_, other.vectors_lookup_);
|
|
407
|
+
|
|
408
|
+
std::swap(available_threads_, other.available_threads_);
|
|
409
|
+
std::swap(slot_lookup_, other.slot_lookup_);
|
|
410
|
+
std::swap(free_keys_, other.free_keys_);
|
|
411
|
+
std::swap(free_key_, other.free_key_);
|
|
412
|
+
}
|
|
413
|
+
|
|
414
|
+
~index_dense_gt() {
|
|
415
|
+
if (typed_)
|
|
416
|
+
typed_->~index_t();
|
|
417
|
+
index_allocator_t{}.deallocate(typed_, 1);
|
|
418
|
+
typed_ = nullptr;
|
|
419
|
+
}
|
|
420
|
+
|
|
421
|
+
/**
|
|
422
|
+
* @brief Constructs an instance of ::index_dense_gt.
|
|
423
|
+
* @param[in] metric One of the provided or an @b ad-hoc metric, type-punned.
|
|
424
|
+
* @param[in] config The index configuration (optional).
|
|
425
|
+
* @param[in] free_label The key used for freed vectors (optional).
|
|
426
|
+
* @return An instance of ::index_dense_gt.
|
|
427
|
+
*/
|
|
428
|
+
static index_dense_gt make( //
|
|
429
|
+
metric_t metric, //
|
|
430
|
+
index_dense_config_t config = {}, //
|
|
431
|
+
key_t free_label = default_free_value<key_t>()) {
|
|
432
|
+
|
|
433
|
+
scalar_kind_t scalar_kind = metric.scalar_kind();
|
|
434
|
+
std::size_t hardware_threads = std::thread::hardware_concurrency();
|
|
435
|
+
|
|
436
|
+
index_dense_gt result;
|
|
437
|
+
result.config_ = config;
|
|
438
|
+
result.cast_buffer_.resize(hardware_threads * metric.bytes_per_vector());
|
|
439
|
+
result.casts_ = make_casts_(scalar_kind);
|
|
440
|
+
result.metric_ = metric;
|
|
441
|
+
result.free_key_ = free_label;
|
|
442
|
+
|
|
443
|
+
// Fill the thread IDs.
|
|
444
|
+
result.available_threads_.resize(hardware_threads);
|
|
445
|
+
std::iota(result.available_threads_.begin(), result.available_threads_.end(), 0ul);
|
|
446
|
+
|
|
447
|
+
// Available since C11, but only C++17, so we use the C version.
|
|
448
|
+
index_t* raw = index_allocator_t{}.allocate(1);
|
|
449
|
+
new (raw) index_t(config);
|
|
450
|
+
result.typed_ = raw;
|
|
451
|
+
return result;
|
|
452
|
+
}
|
|
453
|
+
|
|
454
|
+
static index_dense_gt make(char const* path, bool view = false) {
|
|
455
|
+
index_dense_metadata_result_t meta = index_dense_metadata(path);
|
|
456
|
+
if (!meta)
|
|
457
|
+
return {};
|
|
458
|
+
metric_punned_t metric(meta.head.dimensions, meta.head.kind_metric, meta.head.kind_scalar);
|
|
459
|
+
index_dense_gt result = make(metric);
|
|
460
|
+
if (!result)
|
|
461
|
+
return result;
|
|
462
|
+
if (view)
|
|
463
|
+
result.view(path);
|
|
464
|
+
else
|
|
465
|
+
result.load(path);
|
|
466
|
+
return result;
|
|
467
|
+
}
|
|
468
|
+
|
|
469
|
+
explicit operator bool() const { return typed_; }
|
|
470
|
+
std::size_t connectivity() const { return typed_->connectivity(); }
|
|
471
|
+
std::size_t size() const { return typed_->size() - free_keys_.size(); }
|
|
472
|
+
std::size_t capacity() const { return typed_->capacity(); }
|
|
473
|
+
std::size_t max_level() const noexcept { return typed_->max_level(); }
|
|
474
|
+
index_dense_config_t const& config() const { return config_; }
|
|
475
|
+
index_limits_t const& limits() const { return typed_->limits(); }
|
|
476
|
+
|
|
477
|
+
// The metric and its properties
|
|
478
|
+
metric_t const& metric() const { return metric_; }
|
|
479
|
+
void change_metric(metric_t metric) { metric_ = std::move(metric); }
|
|
480
|
+
|
|
481
|
+
scalar_kind_t scalar_kind() const noexcept { return metric_.scalar_kind(); }
|
|
482
|
+
std::size_t bytes_per_vector() const noexcept { return metric_.bytes_per_vector(); }
|
|
483
|
+
std::size_t scalar_words() const noexcept { return metric_.scalar_words(); }
|
|
484
|
+
std::size_t dimensions() const noexcept { return metric_.dimensions(); }
|
|
485
|
+
|
|
486
|
+
// Fetching and changing search criteria
|
|
487
|
+
std::size_t expansion_add() const { return config_.expansion_add; }
|
|
488
|
+
std::size_t expansion_search() const { return config_.expansion_search; }
|
|
489
|
+
void change_expansion_add(std::size_t n) { config_.expansion_add = n; }
|
|
490
|
+
void change_expansion_search(std::size_t n) { config_.expansion_search = n; }
|
|
491
|
+
|
|
492
|
+
member_citerator_t cbegin() const { return typed_->cbegin(); }
|
|
493
|
+
member_citerator_t cend() const { return typed_->cend(); }
|
|
494
|
+
member_citerator_t begin() const { return typed_->begin(); }
|
|
495
|
+
member_citerator_t end() const { return typed_->end(); }
|
|
496
|
+
member_iterator_t begin() { return typed_->begin(); }
|
|
497
|
+
member_iterator_t end() { return typed_->end(); }
|
|
498
|
+
|
|
499
|
+
stats_t stats() const { return typed_->stats(); }
|
|
500
|
+
stats_t stats(std::size_t level) const { return typed_->stats(level); }
|
|
501
|
+
|
|
502
|
+
/**
|
|
503
|
+
* @brief A relatively accurate lower bound on the amount of memory consumed by the system.
|
|
504
|
+
* In practice it's error will be below 10%.
|
|
505
|
+
*
|
|
506
|
+
* @see `stream_length` for the length of the binary serialized representation.
|
|
507
|
+
*/
|
|
508
|
+
std::size_t memory_usage() const {
|
|
509
|
+
return //
|
|
510
|
+
typed_->memory_usage(0) + //
|
|
511
|
+
typed_->tape_allocator().total_wasted() + //
|
|
512
|
+
typed_->tape_allocator().total_reserved() + //
|
|
513
|
+
vectors_tape_allocator_.total_allocated();
|
|
514
|
+
}
|
|
515
|
+
|
|
516
|
+
// clang-format off
|
|
517
|
+
add_result_t add(key_t key, b1x8_t const* vector) { return add_(key, vector, casts_.from_b1x8); }
|
|
518
|
+
add_result_t add(key_t key, i8_bits_t const* vector) { return add_(key, vector, casts_.from_i8); }
|
|
519
|
+
add_result_t add(key_t key, f16_t const* vector) { return add_(key, vector, casts_.from_f16); }
|
|
520
|
+
add_result_t add(key_t key, f32_t const* vector) { return add_(key, vector, casts_.from_f32); }
|
|
521
|
+
add_result_t add(key_t key, f64_t const* vector) { return add_(key, vector, casts_.from_f64); }
|
|
522
|
+
|
|
523
|
+
add_result_t add(key_t key, b1x8_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_b1x8); }
|
|
524
|
+
add_result_t add(key_t key, i8_bits_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_i8); }
|
|
525
|
+
add_result_t add(key_t key, f16_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_f16); }
|
|
526
|
+
add_result_t add(key_t key, f32_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_f32); }
|
|
527
|
+
add_result_t add(key_t key, f64_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_f64); }
|
|
528
|
+
|
|
529
|
+
search_result_t search(b1x8_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_b1x8); }
|
|
530
|
+
search_result_t search(i8_bits_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_i8); }
|
|
531
|
+
search_result_t search(f16_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_f16); }
|
|
532
|
+
search_result_t search(f32_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_f32); }
|
|
533
|
+
search_result_t search(f64_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_f64); }
|
|
534
|
+
|
|
535
|
+
search_result_t search(b1x8_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_b1x8); }
|
|
536
|
+
search_result_t search(i8_bits_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_i8); }
|
|
537
|
+
search_result_t search(f16_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_f16); }
|
|
538
|
+
search_result_t search(f32_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_f32); }
|
|
539
|
+
search_result_t search(f64_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_f64); }
|
|
540
|
+
|
|
541
|
+
bool get(key_t key, b1x8_t* vector) const { return get_(key, vector, casts_.to_b1x8); }
|
|
542
|
+
bool get(key_t key, i8_bits_t* vector) const { return get_(key, vector, casts_.to_i8); }
|
|
543
|
+
bool get(key_t key, f16_t* vector) const { return get_(key, vector, casts_.to_f16); }
|
|
544
|
+
bool get(key_t key, f32_t* vector) const { return get_(key, vector, casts_.to_f32); }
|
|
545
|
+
bool get(key_t key, f64_t* vector) const { return get_(key, vector, casts_.to_f64); }
|
|
546
|
+
// clang-format on
|
|
547
|
+
|
|
548
|
+
/**
|
|
549
|
+
* @brief Reserves memory for the index and the labeled lookup.
|
|
550
|
+
* @return `true` if the memory reservation was successful, `false` otherwise.
|
|
551
|
+
*/
|
|
552
|
+
bool reserve(index_limits_t limits) {
|
|
553
|
+
{
|
|
554
|
+
unique_lock_t lock(slot_lookup_mutex_);
|
|
555
|
+
slot_lookup_.reserve(limits.members);
|
|
556
|
+
vectors_lookup_.resize(limits.members);
|
|
557
|
+
}
|
|
558
|
+
return typed_->reserve(limits);
|
|
559
|
+
}
|
|
560
|
+
|
|
561
|
+
/**
|
|
562
|
+
* @brief Erases all the vectors from the index.
|
|
563
|
+
*
|
|
564
|
+
* Will change `size()` to zero, but will keep the same `capacity()`.
|
|
565
|
+
* Will keep the number of available threads/contexts the same as it was.
|
|
566
|
+
*/
|
|
567
|
+
void clear() {
|
|
568
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
569
|
+
|
|
570
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
571
|
+
typed_->clear();
|
|
572
|
+
slot_lookup_.clear();
|
|
573
|
+
vectors_lookup_.clear();
|
|
574
|
+
free_keys_.clear();
|
|
575
|
+
vectors_tape_allocator_.reset();
|
|
576
|
+
}
|
|
577
|
+
|
|
578
|
+
/**
|
|
579
|
+
* @brief Erases all members from index, closing files, and returning RAM to OS.
|
|
580
|
+
*
|
|
581
|
+
* Will change both `size()` and `capacity()` to zero.
|
|
582
|
+
* Will deallocate all threads/contexts.
|
|
583
|
+
* If the index is memory-mapped - releases the mapping and the descriptor.
|
|
584
|
+
*/
|
|
585
|
+
void reset() {
|
|
586
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
587
|
+
|
|
588
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
589
|
+
std::unique_lock<std::mutex> available_threads_lock(available_threads_mutex_);
|
|
590
|
+
typed_->reset();
|
|
591
|
+
slot_lookup_.clear();
|
|
592
|
+
vectors_lookup_.clear();
|
|
593
|
+
free_keys_.clear();
|
|
594
|
+
vectors_tape_allocator_.reset();
|
|
595
|
+
available_threads_.clear();
|
|
596
|
+
}
|
|
597
|
+
|
|
598
|
+
/**
|
|
599
|
+
* @brief Saves the index to a file.
|
|
600
|
+
* @param[in] path The path to the file.
|
|
601
|
+
* @param[in] config Configuration parameters for exports.
|
|
602
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
603
|
+
*/
|
|
604
|
+
serialization_result_t save(output_file_t file, serialization_config_t config = {}) const {
|
|
605
|
+
serialization_result_t result = file.open_if_not();
|
|
606
|
+
if (result)
|
|
607
|
+
stream(
|
|
608
|
+
[&](void* buffer, std::size_t length) {
|
|
609
|
+
result = file.write(buffer, length);
|
|
610
|
+
return !!result;
|
|
611
|
+
},
|
|
612
|
+
config);
|
|
613
|
+
return result;
|
|
614
|
+
}
|
|
615
|
+
|
|
616
|
+
/**
|
|
617
|
+
* @brief Saves serialized binary index representation to a stream.
|
|
618
|
+
*/
|
|
619
|
+
template <typename output_callback_at, typename progress_at = dummy_progress_t>
|
|
620
|
+
serialization_result_t stream(output_callback_at&& callback, serialization_config_t config = {}) const noexcept {
|
|
621
|
+
|
|
622
|
+
serialization_result_t result;
|
|
623
|
+
std::uint64_t matrix_rows = 0;
|
|
624
|
+
std::uint64_t matrix_cols = 0;
|
|
625
|
+
|
|
626
|
+
// We may not want to put the vectors into the same file
|
|
627
|
+
if (!config.exclude_vectors) {
|
|
628
|
+
// Save the matrix size
|
|
629
|
+
if (!config.use_64_bit_dimensions) {
|
|
630
|
+
std::uint32_t dimensions[2];
|
|
631
|
+
dimensions[0] = static_cast<std::uint32_t>(typed_->size());
|
|
632
|
+
dimensions[1] = static_cast<std::uint32_t>(metric_.bytes_per_vector());
|
|
633
|
+
if (!callback(&dimensions, sizeof(dimensions)))
|
|
634
|
+
return result.failed("Failed to serialize into stream");
|
|
635
|
+
matrix_rows = dimensions[0];
|
|
636
|
+
matrix_cols = dimensions[1];
|
|
637
|
+
} else {
|
|
638
|
+
std::uint64_t dimensions[2];
|
|
639
|
+
dimensions[0] = static_cast<std::uint64_t>(typed_->size());
|
|
640
|
+
dimensions[1] = static_cast<std::uint64_t>(metric_.bytes_per_vector());
|
|
641
|
+
if (!callback(&dimensions, sizeof(dimensions)))
|
|
642
|
+
return result.failed("Failed to serialize into stream");
|
|
643
|
+
matrix_rows = dimensions[0];
|
|
644
|
+
matrix_cols = dimensions[1];
|
|
645
|
+
}
|
|
646
|
+
|
|
647
|
+
// Dump the vectors one after another
|
|
648
|
+
for (std::uint64_t i = 0; i != matrix_rows; ++i) {
|
|
649
|
+
byte_t* vector = vectors_lookup_[i];
|
|
650
|
+
if (!callback(vector, matrix_cols))
|
|
651
|
+
return result.failed("Failed to serialize into stream");
|
|
652
|
+
}
|
|
653
|
+
}
|
|
654
|
+
|
|
655
|
+
// Augment metadata
|
|
656
|
+
{
|
|
657
|
+
index_dense_head_buffer_t buffer;
|
|
658
|
+
std::memset(buffer, 0, sizeof(buffer));
|
|
659
|
+
index_dense_head_t head{buffer};
|
|
660
|
+
std::memcpy(buffer, default_magic(), std::strlen(default_magic()));
|
|
661
|
+
|
|
662
|
+
// Describe software version
|
|
663
|
+
using version_t = index_dense_head_t::version_t;
|
|
664
|
+
head.version_major = static_cast<version_t>(USEARCH_VERSION_MAJOR);
|
|
665
|
+
head.version_minor = static_cast<version_t>(USEARCH_VERSION_MINOR);
|
|
666
|
+
head.version_patch = static_cast<version_t>(USEARCH_VERSION_PATCH);
|
|
667
|
+
|
|
668
|
+
// Describes types used
|
|
669
|
+
head.kind_metric = metric_.metric_kind();
|
|
670
|
+
head.kind_scalar = metric_.scalar_kind();
|
|
671
|
+
head.kind_key = unum::usearch::scalar_kind<key_t>();
|
|
672
|
+
head.kind_compressed_slot = unum::usearch::scalar_kind<compressed_slot_t>();
|
|
673
|
+
|
|
674
|
+
head.count_present = size();
|
|
675
|
+
head.count_deleted = typed_->size() - size();
|
|
676
|
+
head.dimensions = dimensions();
|
|
677
|
+
|
|
678
|
+
if (!callback(&buffer, sizeof(buffer)))
|
|
679
|
+
return result.failed("Failed to serialize into stream");
|
|
680
|
+
}
|
|
681
|
+
|
|
682
|
+
// Save the actual proximity graph
|
|
683
|
+
return typed_->stream(std::forward<output_callback_at>(callback));
|
|
684
|
+
}
|
|
685
|
+
|
|
686
|
+
/**
|
|
687
|
+
* @brief Estimate the binary length (in bytes) of the serialized index.
|
|
688
|
+
*/
|
|
689
|
+
std::size_t stream_length(serialization_config_t config = {}) const noexcept {
|
|
690
|
+
std::size_t dimensions_length = 0;
|
|
691
|
+
std::size_t matrix_length = 0;
|
|
692
|
+
if (!config.exclude_vectors) {
|
|
693
|
+
dimensions_length = config.use_64_bit_dimensions ? sizeof(std::uint64_t) * 2 : sizeof(std::uint32_t) * 2;
|
|
694
|
+
matrix_length = typed_->size() * metric_.bytes_per_vector();
|
|
695
|
+
}
|
|
696
|
+
return dimensions_length + matrix_length + sizeof(index_dense_head_buffer_t) + typed_->stream_length();
|
|
697
|
+
}
|
|
698
|
+
|
|
699
|
+
/**
|
|
700
|
+
* @brief Parses the index from file to RAM.
|
|
701
|
+
* @param[in] path The path to the file.
|
|
702
|
+
* @param[in] config Configuration parameters for imports.
|
|
703
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
704
|
+
*/
|
|
705
|
+
serialization_result_t load(input_file_t file, serialization_config_t config = {}) {
|
|
706
|
+
|
|
707
|
+
serialization_result_t result = file.open_if_not();
|
|
708
|
+
if (!result)
|
|
709
|
+
return result;
|
|
710
|
+
|
|
711
|
+
std::uint64_t matrix_rows = 0;
|
|
712
|
+
std::uint64_t matrix_cols = 0;
|
|
713
|
+
|
|
714
|
+
// We may not want to load the vectors from the same file, or allow attaching them afterwards
|
|
715
|
+
if (!config.exclude_vectors) {
|
|
716
|
+
// Save the matrix size
|
|
717
|
+
if (!config.use_64_bit_dimensions) {
|
|
718
|
+
std::uint32_t dimensions[2];
|
|
719
|
+
result = file.read(&dimensions, sizeof(dimensions));
|
|
720
|
+
if (!result)
|
|
721
|
+
return result;
|
|
722
|
+
matrix_rows = dimensions[0];
|
|
723
|
+
matrix_cols = dimensions[1];
|
|
724
|
+
} else {
|
|
725
|
+
std::uint64_t dimensions[2];
|
|
726
|
+
result = file.read(&dimensions, sizeof(dimensions));
|
|
727
|
+
if (!result)
|
|
728
|
+
return result;
|
|
729
|
+
matrix_rows = dimensions[0];
|
|
730
|
+
matrix_cols = dimensions[1];
|
|
731
|
+
}
|
|
732
|
+
// Load the vectors one after another
|
|
733
|
+
vectors_lookup_.resize(matrix_rows);
|
|
734
|
+
for (std::uint64_t slot = 0; slot != matrix_rows; ++slot) {
|
|
735
|
+
byte_t* vector = vectors_tape_allocator_.allocate(matrix_cols);
|
|
736
|
+
result = file.read(vector, matrix_cols);
|
|
737
|
+
if (!result)
|
|
738
|
+
return result;
|
|
739
|
+
vectors_lookup_[slot] = vector;
|
|
740
|
+
}
|
|
741
|
+
}
|
|
742
|
+
|
|
743
|
+
// Load metadata and choose the right metric
|
|
744
|
+
{
|
|
745
|
+
index_dense_head_buffer_t buffer;
|
|
746
|
+
result = file.read(buffer, sizeof(buffer));
|
|
747
|
+
if (!result)
|
|
748
|
+
return result;
|
|
749
|
+
|
|
750
|
+
index_dense_head_t head{buffer};
|
|
751
|
+
if (std::memcmp(buffer, default_magic(), std::strlen(default_magic())) != 0)
|
|
752
|
+
return result.failed("Magic header mismatch - the file isn't an index");
|
|
753
|
+
|
|
754
|
+
// Validate the software version
|
|
755
|
+
if (head.version_major != USEARCH_VERSION_MAJOR)
|
|
756
|
+
return result.failed("File format may be different, please rebuild");
|
|
757
|
+
|
|
758
|
+
// Check the types used
|
|
759
|
+
if (head.kind_key != unum::usearch::scalar_kind<key_t>())
|
|
760
|
+
return result.failed("Key type doesn't match, consider rebuilding");
|
|
761
|
+
if (head.kind_compressed_slot != unum::usearch::scalar_kind<compressed_slot_t>())
|
|
762
|
+
return result.failed("Slot type doesn't match, consider rebuilding");
|
|
763
|
+
|
|
764
|
+
metric_ = metric_t(head.dimensions, head.kind_metric, head.kind_scalar);
|
|
765
|
+
}
|
|
766
|
+
|
|
767
|
+
// Pull the actual proximity graph
|
|
768
|
+
result = typed_->load(std::move(file));
|
|
769
|
+
if (!result)
|
|
770
|
+
return result;
|
|
771
|
+
if (typed_->size() != static_cast<std::size_t>(matrix_rows))
|
|
772
|
+
return result.failed("Index size and the number of vectors doesn't match");
|
|
773
|
+
|
|
774
|
+
reindex_labels_();
|
|
775
|
+
return result;
|
|
776
|
+
}
|
|
777
|
+
|
|
778
|
+
/**
|
|
779
|
+
* @brief Parses the index from file, without loading it into RAM.
|
|
780
|
+
* @param[in] path The path to the file.
|
|
781
|
+
* @param[in] config Configuration parameters for imports.
|
|
782
|
+
* @return Outcome descriptor explicitly convertible to boolean.
|
|
783
|
+
*/
|
|
784
|
+
serialization_result_t view(memory_mapped_file_t file, std::size_t offset = 0, serialization_config_t config = {}) {
|
|
785
|
+
|
|
786
|
+
serialization_result_t result = file.open_if_not();
|
|
787
|
+
if (!result)
|
|
788
|
+
return result;
|
|
789
|
+
|
|
790
|
+
std::uint64_t matrix_rows = 0;
|
|
791
|
+
std::uint64_t matrix_cols = 0;
|
|
792
|
+
span_punned_t vectors_buffer;
|
|
793
|
+
|
|
794
|
+
// We may not want to fetch the vectors from the same file, or allow attaching them afterwards
|
|
795
|
+
if (!config.exclude_vectors) {
|
|
796
|
+
// Save the matrix size
|
|
797
|
+
if (!config.use_64_bit_dimensions) {
|
|
798
|
+
std::uint32_t dimensions[2];
|
|
799
|
+
if (file.size() - offset < sizeof(dimensions))
|
|
800
|
+
return result.failed("File is corrupted and lacks matrix dimensions");
|
|
801
|
+
std::memcpy(&dimensions, file.data() + offset, sizeof(dimensions));
|
|
802
|
+
matrix_rows = dimensions[0];
|
|
803
|
+
matrix_cols = dimensions[1];
|
|
804
|
+
offset += sizeof(dimensions);
|
|
805
|
+
} else {
|
|
806
|
+
std::uint64_t dimensions[2];
|
|
807
|
+
if (file.size() - offset < sizeof(dimensions))
|
|
808
|
+
return result.failed("File is corrupted and lacks matrix dimensions");
|
|
809
|
+
std::memcpy(&dimensions, file.data() + offset, sizeof(dimensions));
|
|
810
|
+
matrix_rows = dimensions[0];
|
|
811
|
+
matrix_cols = dimensions[1];
|
|
812
|
+
offset += sizeof(dimensions);
|
|
813
|
+
}
|
|
814
|
+
vectors_buffer = {file.data() + offset, static_cast<std::size_t>(matrix_rows * matrix_cols)};
|
|
815
|
+
offset += vectors_buffer.size();
|
|
816
|
+
}
|
|
817
|
+
|
|
818
|
+
// Load metadata and choose the right metric
|
|
819
|
+
{
|
|
820
|
+
index_dense_head_buffer_t buffer;
|
|
821
|
+
if (file.size() - offset < sizeof(buffer))
|
|
822
|
+
return result.failed("File is corrupted and lacks a header");
|
|
823
|
+
|
|
824
|
+
std::memcpy(buffer, file.data() + offset, sizeof(buffer));
|
|
825
|
+
|
|
826
|
+
index_dense_head_t head{buffer};
|
|
827
|
+
if (std::memcmp(buffer, default_magic(), std::strlen(default_magic())) != 0)
|
|
828
|
+
return result.failed("Magic header mismatch - the file isn't an index");
|
|
829
|
+
|
|
830
|
+
// Validate the software version
|
|
831
|
+
if (head.version_major != USEARCH_VERSION_MAJOR)
|
|
832
|
+
return result.failed("File format may be different, please rebuild");
|
|
833
|
+
|
|
834
|
+
// Check the types used
|
|
835
|
+
if (head.kind_key != unum::usearch::scalar_kind<key_t>())
|
|
836
|
+
return result.failed("Key type doesn't match, consider rebuilding");
|
|
837
|
+
if (head.kind_compressed_slot != unum::usearch::scalar_kind<compressed_slot_t>())
|
|
838
|
+
return result.failed("Slot type doesn't match, consider rebuilding");
|
|
839
|
+
|
|
840
|
+
metric_ = metric_t(head.dimensions, head.kind_metric, head.kind_scalar);
|
|
841
|
+
offset += sizeof(buffer);
|
|
842
|
+
}
|
|
843
|
+
|
|
844
|
+
// Pull the actual proximity graph
|
|
845
|
+
result = typed_->view(std::move(file), offset);
|
|
846
|
+
if (!result)
|
|
847
|
+
return result;
|
|
848
|
+
if (typed_->size() != static_cast<std::size_t>(matrix_rows))
|
|
849
|
+
return result.failed("Index size and the number of vectors doesn't match");
|
|
850
|
+
|
|
851
|
+
// Address the vectors
|
|
852
|
+
vectors_lookup_.resize(matrix_rows);
|
|
853
|
+
if (!config.exclude_vectors)
|
|
854
|
+
for (std::uint64_t slot = 0; slot != matrix_rows; ++slot)
|
|
855
|
+
vectors_lookup_[slot] = (byte_t*)vectors_buffer.data() + matrix_cols * slot;
|
|
856
|
+
|
|
857
|
+
reindex_labels_();
|
|
858
|
+
return result;
|
|
859
|
+
}
|
|
860
|
+
|
|
861
|
+
/**
|
|
862
|
+
* @brief Checks if a vector with specified key is present.
|
|
863
|
+
* @return `true` if the key is present in the index, `false` otherwise.
|
|
864
|
+
*/
|
|
865
|
+
bool contains(key_t key) const {
|
|
866
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
867
|
+
return slot_lookup_.contains(key);
|
|
868
|
+
}
|
|
869
|
+
|
|
870
|
+
/**
|
|
871
|
+
* @brief Count the number of vectors with specified key present.
|
|
872
|
+
* @return Zero if nothing is found, a positive integer otherwise.
|
|
873
|
+
*/
|
|
874
|
+
std::size_t count(key_t key) const {
|
|
875
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
876
|
+
return slot_lookup_.count(key);
|
|
877
|
+
}
|
|
878
|
+
|
|
879
|
+
struct labeling_result_t {
|
|
880
|
+
error_t error{};
|
|
881
|
+
std::size_t completed{};
|
|
882
|
+
|
|
883
|
+
explicit operator bool() const noexcept { return !error; }
|
|
884
|
+
labeling_result_t failed(error_t message) noexcept {
|
|
885
|
+
error = std::move(message);
|
|
886
|
+
return std::move(*this);
|
|
887
|
+
}
|
|
888
|
+
};
|
|
889
|
+
|
|
890
|
+
/**
|
|
891
|
+
* @brief Removes an entry with the specified key from the index.
|
|
892
|
+
* @param[in] key The key of the entry to remove.
|
|
893
|
+
* @return The ::labeling_result_t indicating the result of the removal operation.
|
|
894
|
+
* If the removal was successful, `result.completed` will be `true`.
|
|
895
|
+
* If the key was not found in the index, `result.completed` will be `false`.
|
|
896
|
+
* If an error occurred during the removal operation, `result.error` will contain an error message.
|
|
897
|
+
*/
|
|
898
|
+
labeling_result_t remove(key_t key) {
|
|
899
|
+
labeling_result_t result;
|
|
900
|
+
|
|
901
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
902
|
+
auto labeled_iterator = slot_lookup_.find(key);
|
|
903
|
+
if (labeled_iterator == slot_lookup_.end())
|
|
904
|
+
return result;
|
|
905
|
+
|
|
906
|
+
// Grow the removed entries ring, if needed
|
|
907
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
908
|
+
if (free_keys_.size() == free_keys_.capacity())
|
|
909
|
+
if (!free_keys_.reserve((std::max<std::size_t>)(free_keys_.capacity() * 2, 64ul)))
|
|
910
|
+
return result.failed("Can't allocate memory for a free-list");
|
|
911
|
+
|
|
912
|
+
// A removed entry would be:
|
|
913
|
+
// - present in `free_keys_`
|
|
914
|
+
// - missing in the `slot_lookup_`
|
|
915
|
+
// - marked in the `typed_` index with a `free_key_`
|
|
916
|
+
compressed_slot_t slot = (*labeled_iterator).slot;
|
|
917
|
+
free_keys_.push(slot);
|
|
918
|
+
slot_lookup_.erase(labeled_iterator);
|
|
919
|
+
typed_->at(slot).key = free_key_;
|
|
920
|
+
result.completed = true;
|
|
921
|
+
|
|
922
|
+
return result;
|
|
923
|
+
}
|
|
924
|
+
|
|
925
|
+
/**
|
|
926
|
+
* @brief Removes multiple entries with the specified keys from the index.
|
|
927
|
+
* @param[in] labels_begin The beginning of the keys range.
|
|
928
|
+
* @param[in] labels_end The ending of the keys range.
|
|
929
|
+
* @return The ::labeling_result_t indicating the result of the removal operation.
|
|
930
|
+
* `result.completed` will contain the number of keys that were successfully removed.
|
|
931
|
+
* `result.error` will contain an error message if an error occurred during the removal operation.
|
|
932
|
+
*/
|
|
933
|
+
template <typename labels_iterator_at>
|
|
934
|
+
labeling_result_t remove(labels_iterator_at&& labels_begin, labels_iterator_at&& labels_end) {
|
|
935
|
+
|
|
936
|
+
labeling_result_t result;
|
|
937
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
938
|
+
std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
|
|
939
|
+
|
|
940
|
+
// Grow the removed entries ring, if needed
|
|
941
|
+
std::size_t count_requests = std::distance(labels_begin, labels_end);
|
|
942
|
+
if (!free_keys_.reserve(free_keys_.size() + count_requests))
|
|
943
|
+
return result.failed("Can't allocate memory for a free-list");
|
|
944
|
+
|
|
945
|
+
// Remove them one-by-one
|
|
946
|
+
for (auto label_it = labels_begin; label_it != labels_end; ++label_it) {
|
|
947
|
+
key_t key = *label_it;
|
|
948
|
+
auto labeled_iterator = slot_lookup_.find(key);
|
|
949
|
+
if (labeled_iterator == slot_lookup_.end())
|
|
950
|
+
continue;
|
|
951
|
+
|
|
952
|
+
// A removed entry would be:
|
|
953
|
+
// - present in `free_keys_`
|
|
954
|
+
// - missing in the `slot_lookup_`
|
|
955
|
+
// - marked in the `typed_` index with a `free_key_`
|
|
956
|
+
compressed_slot_t slot = (*labeled_iterator).slot;
|
|
957
|
+
free_keys_.push(slot);
|
|
958
|
+
slot_lookup_.erase(labeled_iterator);
|
|
959
|
+
typed_->at(slot).key = free_key_;
|
|
960
|
+
result.completed += 1;
|
|
961
|
+
}
|
|
962
|
+
|
|
963
|
+
return result;
|
|
964
|
+
}
|
|
965
|
+
|
|
966
|
+
/**
|
|
967
|
+
* @brief Renames an entry with the specified key to a new key.
|
|
968
|
+
* @param[in] from The current key of the entry to rename.
|
|
969
|
+
* @param[in] to The new key to assign to the entry.
|
|
970
|
+
* @return The ::labeling_result_t indicating the result of the rename operation.
|
|
971
|
+
* If the rename was successful, `result.completed` will be `true`.
|
|
972
|
+
* If the entry with the current key was not found, `result.completed` will be `false`.
|
|
973
|
+
*/
|
|
974
|
+
labeling_result_t rename(key_t from, key_t to) {
|
|
975
|
+
labeling_result_t result;
|
|
976
|
+
unique_lock_t lookup_lock(slot_lookup_mutex_);
|
|
977
|
+
auto labeled_iterator = slot_lookup_.find(from);
|
|
978
|
+
if (labeled_iterator == slot_lookup_.end())
|
|
979
|
+
return result;
|
|
980
|
+
|
|
981
|
+
compressed_slot_t slot = (*labeled_iterator).slot;
|
|
982
|
+
key_and_slot_t key_and_slot{to, slot};
|
|
983
|
+
slot_lookup_.erase(labeled_iterator);
|
|
984
|
+
slot_lookup_.insert(key_and_slot);
|
|
985
|
+
typed_->at(slot).key = to;
|
|
986
|
+
result.completed = true;
|
|
987
|
+
return result;
|
|
988
|
+
}
|
|
989
|
+
|
|
990
|
+
/**
|
|
991
|
+
* @brief Exports a range of keys for the vectors present in the index.
|
|
992
|
+
* @param[out] keys Pointer to the array where the keys will be exported.
|
|
993
|
+
* @param[in] offset The number of keys to skip. Useful for pagination.
|
|
994
|
+
* @param[in] limit The maximum number of keys to export, that can fit in ::keys.
|
|
995
|
+
*/
|
|
996
|
+
void export_keys(key_t* keys, std::size_t offset, std::size_t limit) const {
|
|
997
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
998
|
+
auto it = slot_lookup_.begin();
|
|
999
|
+
offset = (std::min)(offset, slot_lookup_.size());
|
|
1000
|
+
std::advance(it, offset);
|
|
1001
|
+
for (; it != slot_lookup_.end() && limit; ++it, ++keys, --limit)
|
|
1002
|
+
*keys = (*it).key;
|
|
1003
|
+
}
|
|
1004
|
+
|
|
1005
|
+
struct copy_result_t {
|
|
1006
|
+
index_dense_gt index;
|
|
1007
|
+
error_t error;
|
|
1008
|
+
|
|
1009
|
+
explicit operator bool() const noexcept { return !error; }
|
|
1010
|
+
copy_result_t failed(error_t message) noexcept {
|
|
1011
|
+
error = std::move(message);
|
|
1012
|
+
return std::move(*this);
|
|
1013
|
+
}
|
|
1014
|
+
};
|
|
1015
|
+
|
|
1016
|
+
/**
|
|
1017
|
+
* @brief Copies the ::index_dense_gt @b with all the data in it.
|
|
1018
|
+
* @param config The copy configuration (optional).
|
|
1019
|
+
* @return A copy of the ::index_dense_gt instance.
|
|
1020
|
+
*/
|
|
1021
|
+
copy_result_t copy(index_dense_copy_config_t config = {}) const {
|
|
1022
|
+
copy_result_t result = fork();
|
|
1023
|
+
if (!result)
|
|
1024
|
+
return result;
|
|
1025
|
+
|
|
1026
|
+
auto typed_result = typed_->copy(config);
|
|
1027
|
+
if (!typed_result)
|
|
1028
|
+
return result.failed(std::move(typed_result.error));
|
|
1029
|
+
|
|
1030
|
+
// Export the free (removed) slot numbers
|
|
1031
|
+
index_dense_gt& copy = result.index;
|
|
1032
|
+
if (!copy.free_keys_.reserve(free_keys_.size()))
|
|
1033
|
+
return result.failed(std::move(typed_result.error));
|
|
1034
|
+
for (std::size_t i = 0; i != free_keys_.size(); ++i)
|
|
1035
|
+
copy.free_keys_.push(free_keys_[i]);
|
|
1036
|
+
|
|
1037
|
+
// Allocate buffers and move the vectors themselves
|
|
1038
|
+
if (!config.force_vector_copy && copy.config_.exclude_vectors)
|
|
1039
|
+
copy.vectors_lookup_ = vectors_lookup_;
|
|
1040
|
+
else {
|
|
1041
|
+
copy.vectors_lookup_.resize(vectors_lookup_.size());
|
|
1042
|
+
for (std::size_t slot = 0; slot != vectors_lookup_.size(); ++slot)
|
|
1043
|
+
copy.vectors_lookup_[slot] = copy.vectors_tape_allocator_.allocate(copy.metric_.bytes_per_vector());
|
|
1044
|
+
if (std::count(copy.vectors_lookup_.begin(), copy.vectors_lookup_.end(), nullptr))
|
|
1045
|
+
return result.failed("Out of memory!");
|
|
1046
|
+
for (std::size_t slot = 0; slot != vectors_lookup_.size(); ++slot)
|
|
1047
|
+
std::memcpy(copy.vectors_lookup_[slot], vectors_lookup_[slot], metric_.bytes_per_vector());
|
|
1048
|
+
}
|
|
1049
|
+
|
|
1050
|
+
copy.slot_lookup_ = slot_lookup_;
|
|
1051
|
+
*copy.typed_ = std::move(typed_result.index);
|
|
1052
|
+
return result;
|
|
1053
|
+
}
|
|
1054
|
+
|
|
1055
|
+
/**
|
|
1056
|
+
* @brief Copies the ::index_dense_gt model @b without any data.
|
|
1057
|
+
* @return A similarly configured ::index_dense_gt instance.
|
|
1058
|
+
*/
|
|
1059
|
+
copy_result_t fork() const {
|
|
1060
|
+
copy_result_t result;
|
|
1061
|
+
index_dense_gt& other = result.index;
|
|
1062
|
+
|
|
1063
|
+
other.config_ = config_;
|
|
1064
|
+
other.cast_buffer_ = cast_buffer_;
|
|
1065
|
+
other.casts_ = casts_;
|
|
1066
|
+
|
|
1067
|
+
other.metric_ = metric_;
|
|
1068
|
+
other.available_threads_ = available_threads_;
|
|
1069
|
+
other.free_key_ = free_key_;
|
|
1070
|
+
|
|
1071
|
+
index_t* raw = index_allocator_t{}.allocate(1);
|
|
1072
|
+
if (!raw)
|
|
1073
|
+
return result.failed("Can't allocate the index");
|
|
1074
|
+
|
|
1075
|
+
new (raw) index_t(config());
|
|
1076
|
+
other.typed_ = raw;
|
|
1077
|
+
return result;
|
|
1078
|
+
}
|
|
1079
|
+
|
|
1080
|
+
struct compaction_result_t {
|
|
1081
|
+
error_t error{};
|
|
1082
|
+
std::size_t pruned_edges{};
|
|
1083
|
+
|
|
1084
|
+
explicit operator bool() const noexcept { return !error; }
|
|
1085
|
+
compaction_result_t failed(error_t message) noexcept {
|
|
1086
|
+
error = std::move(message);
|
|
1087
|
+
return std::move(*this);
|
|
1088
|
+
}
|
|
1089
|
+
};
|
|
1090
|
+
|
|
1091
|
+
/**
|
|
1092
|
+
* @brief Performs compaction on the index, pruning links to removed entries.
|
|
1093
|
+
* @param executor The executor parallel processing. Default ::dummy_executor_t single-threaded.
|
|
1094
|
+
* @param progress The progress tracker instance to use. Default ::dummy_progress_t reports nothing.
|
|
1095
|
+
* @return The ::compaction_result_t indicating the result of the compaction operation.
|
|
1096
|
+
* `result.pruned_edges` will contain the number of edges that were removed.
|
|
1097
|
+
* `result.error` will contain an error message if an error occurred during the compaction operation.
|
|
1098
|
+
*/
|
|
1099
|
+
template <typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
|
|
1100
|
+
compaction_result_t isolate(executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
|
|
1101
|
+
compaction_result_t result;
|
|
1102
|
+
std::atomic<std::size_t> pruned_edges;
|
|
1103
|
+
auto disallow = [&](member_cref_t const& member) noexcept {
|
|
1104
|
+
bool freed = member.key == free_key_;
|
|
1105
|
+
pruned_edges += freed;
|
|
1106
|
+
return freed;
|
|
1107
|
+
};
|
|
1108
|
+
typed_->isolate(disallow, std::forward<executor_at>(executor), std::forward<progress_at>(progress));
|
|
1109
|
+
result.pruned_edges = pruned_edges;
|
|
1110
|
+
return result;
|
|
1111
|
+
}
|
|
1112
|
+
|
|
1113
|
+
class values_proxy_t {
|
|
1114
|
+
index_dense_gt const* index_;
|
|
1115
|
+
|
|
1116
|
+
public:
|
|
1117
|
+
values_proxy_t(index_dense_gt const& index) noexcept : index_(&index) {}
|
|
1118
|
+
byte_t const* operator[](compressed_slot_t slot) const noexcept { return index_->vectors_lookup_[slot]; }
|
|
1119
|
+
byte_t const* operator[](member_citerator_t it) const noexcept { return index_->vectors_lookup_[get_slot(it)]; }
|
|
1120
|
+
};
|
|
1121
|
+
|
|
1122
|
+
/**
|
|
1123
|
+
* @brief Performs compaction on the index, pruning links to removed entries.
|
|
1124
|
+
* @param executor The executor parallel processing. Default ::dummy_executor_t single-threaded.
|
|
1125
|
+
* @param progress The progress tracker instance to use. Default ::dummy_progress_t reports nothing.
|
|
1126
|
+
* @return The ::compaction_result_t indicating the result of the compaction operation.
|
|
1127
|
+
* `result.pruned_edges` will contain the number of edges that were removed.
|
|
1128
|
+
* `result.error` will contain an error message if an error occurred during the compaction operation.
|
|
1129
|
+
*/
|
|
1130
|
+
template <typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
|
|
1131
|
+
compaction_result_t compact(executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
|
|
1132
|
+
compaction_result_t result;
|
|
1133
|
+
|
|
1134
|
+
std::vector<byte_t*> new_vectors_lookup(vectors_lookup_.size());
|
|
1135
|
+
vectors_tape_allocator_t new_vectors_allocator;
|
|
1136
|
+
|
|
1137
|
+
auto track_slot_change = [&](key_t, compressed_slot_t old_slot, compressed_slot_t new_slot) {
|
|
1138
|
+
byte_t* new_vector = new_vectors_allocator.allocate(metric_.bytes_per_vector());
|
|
1139
|
+
byte_t* old_vector = vectors_lookup_[old_slot];
|
|
1140
|
+
std::memcpy(new_vector, old_vector, metric_.bytes_per_vector());
|
|
1141
|
+
new_vectors_lookup[new_slot] = new_vector;
|
|
1142
|
+
};
|
|
1143
|
+
typed_->compact(values_proxy_t{*this}, metric_proxy_t{*this}, track_slot_change,
|
|
1144
|
+
std::forward<executor_at>(executor), std::forward<progress_at>(progress));
|
|
1145
|
+
vectors_lookup_ = std::move(new_vectors_lookup);
|
|
1146
|
+
vectors_tape_allocator_ = std::move(new_vectors_allocator);
|
|
1147
|
+
return result;
|
|
1148
|
+
}
|
|
1149
|
+
|
|
1150
|
+
template < //
|
|
1151
|
+
typename man_to_woman_at = dummy_label_to_label_mapping_t, //
|
|
1152
|
+
typename woman_to_man_at = dummy_label_to_label_mapping_t, //
|
|
1153
|
+
typename executor_at = dummy_executor_t, //
|
|
1154
|
+
typename progress_at = dummy_progress_t //
|
|
1155
|
+
>
|
|
1156
|
+
join_result_t join( //
|
|
1157
|
+
index_dense_gt const& women, //
|
|
1158
|
+
index_join_config_t config = {}, //
|
|
1159
|
+
man_to_woman_at&& man_to_woman = man_to_woman_at{}, //
|
|
1160
|
+
woman_to_man_at&& woman_to_man = woman_to_man_at{}, //
|
|
1161
|
+
executor_at&& executor = executor_at{}, //
|
|
1162
|
+
progress_at&& progress = progress_at{}) const {
|
|
1163
|
+
|
|
1164
|
+
index_dense_gt const& men = *this;
|
|
1165
|
+
return unum::usearch::join( //
|
|
1166
|
+
*men.typed_, *women.typed_, //
|
|
1167
|
+
values_proxy_t{men}, values_proxy_t{women}, //
|
|
1168
|
+
metric_proxy_t{men}, metric_proxy_t{women}, //
|
|
1169
|
+
config, //
|
|
1170
|
+
std::forward<man_to_woman_at>(man_to_woman), //
|
|
1171
|
+
std::forward<woman_to_man_at>(woman_to_man), //
|
|
1172
|
+
std::forward<executor_at>(executor), //
|
|
1173
|
+
std::forward<progress_at>(progress));
|
|
1174
|
+
}
|
|
1175
|
+
|
|
1176
|
+
private:
|
|
1177
|
+
struct thread_lock_t {
|
|
1178
|
+
index_dense_gt const& parent;
|
|
1179
|
+
std::size_t thread_id;
|
|
1180
|
+
|
|
1181
|
+
~thread_lock_t() { parent.thread_unlock_(thread_id); }
|
|
1182
|
+
};
|
|
1183
|
+
|
|
1184
|
+
thread_lock_t thread_lock_() const {
|
|
1185
|
+
available_threads_mutex_.lock();
|
|
1186
|
+
std::size_t thread_id = available_threads_.back();
|
|
1187
|
+
available_threads_.pop_back();
|
|
1188
|
+
available_threads_mutex_.unlock();
|
|
1189
|
+
return {*this, thread_id};
|
|
1190
|
+
}
|
|
1191
|
+
|
|
1192
|
+
void thread_unlock_(std::size_t thread_id) const {
|
|
1193
|
+
available_threads_mutex_.lock();
|
|
1194
|
+
available_threads_.push_back(thread_id);
|
|
1195
|
+
available_threads_mutex_.unlock();
|
|
1196
|
+
}
|
|
1197
|
+
|
|
1198
|
+
template <typename scalar_at>
|
|
1199
|
+
add_result_t add_(key_t key, scalar_at const* vector, index_dense_update_config_t config, cast_t const& cast) {
|
|
1200
|
+
|
|
1201
|
+
if (!config_.allow_key_collisions && contains(key))
|
|
1202
|
+
return add_result_t{}.failed("Duplicate keys not allowed in high-level wrappers");
|
|
1203
|
+
|
|
1204
|
+
// Cast the vector, if needed for compatibility with `metric_`
|
|
1205
|
+
bool copy_vector = !config_.exclude_vectors || config.force_vector_copy;
|
|
1206
|
+
byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
|
|
1207
|
+
{
|
|
1208
|
+
byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * config.thread;
|
|
1209
|
+
bool casted = cast(vector_data, dimensions(), casted_data);
|
|
1210
|
+
if (casted)
|
|
1211
|
+
vector_data = casted_data, copy_vector = true;
|
|
1212
|
+
}
|
|
1213
|
+
|
|
1214
|
+
// Check if there are some removed entries, whose nodes we can reuse
|
|
1215
|
+
compressed_slot_t free_slot = default_free_value<compressed_slot_t>();
|
|
1216
|
+
{
|
|
1217
|
+
std::unique_lock<std::mutex> lock(free_keys_mutex_);
|
|
1218
|
+
free_keys_.try_pop(free_slot);
|
|
1219
|
+
}
|
|
1220
|
+
|
|
1221
|
+
// Perform the insertion or the update
|
|
1222
|
+
bool reuse_node = free_slot != default_free_value<compressed_slot_t>();
|
|
1223
|
+
auto on_success = [&](member_ref_t member) {
|
|
1224
|
+
unique_lock_t slot_lock(slot_lookup_mutex_);
|
|
1225
|
+
slot_lookup_.insert({key, static_cast<compressed_slot_t>(member.slot)});
|
|
1226
|
+
if (copy_vector) {
|
|
1227
|
+
if (!reuse_node)
|
|
1228
|
+
vectors_lookup_[member.slot] = vectors_tape_allocator_.allocate(metric_.bytes_per_vector());
|
|
1229
|
+
std::memcpy(vectors_lookup_[member.slot], vector_data, metric_.bytes_per_vector());
|
|
1230
|
+
} else
|
|
1231
|
+
vectors_lookup_[member.slot] = (byte_t*)vector_data;
|
|
1232
|
+
};
|
|
1233
|
+
|
|
1234
|
+
metric_proxy_t metric{*this};
|
|
1235
|
+
return reuse_node //
|
|
1236
|
+
? typed_->update(typed_->iterator_at(free_slot), key, vector_data, metric, config, on_success)
|
|
1237
|
+
: typed_->add(key, vector_data, metric, config, on_success);
|
|
1238
|
+
}
|
|
1239
|
+
|
|
1240
|
+
template <typename scalar_at>
|
|
1241
|
+
search_result_t search_( //
|
|
1242
|
+
scalar_at const* vector, std::size_t wanted, //
|
|
1243
|
+
index_search_config_t config, cast_t const& cast) const {
|
|
1244
|
+
|
|
1245
|
+
// Cast the vector, if needed for compatibility with `metric_`
|
|
1246
|
+
byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
|
|
1247
|
+
{
|
|
1248
|
+
byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * config.thread;
|
|
1249
|
+
bool casted = cast(vector_data, dimensions(), casted_data);
|
|
1250
|
+
if (casted)
|
|
1251
|
+
vector_data = casted_data;
|
|
1252
|
+
}
|
|
1253
|
+
|
|
1254
|
+
auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
|
|
1255
|
+
return typed_->search(vector_data, wanted, metric_proxy_t{*this}, config, allow);
|
|
1256
|
+
}
|
|
1257
|
+
|
|
1258
|
+
compressed_slot_t lookup_id_(key_t key) const {
|
|
1259
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1260
|
+
return slot_lookup_.at(key);
|
|
1261
|
+
}
|
|
1262
|
+
|
|
1263
|
+
void reindex_labels_() {
|
|
1264
|
+
|
|
1265
|
+
// Estimate number of entries first
|
|
1266
|
+
std::size_t count_total = typed_->size();
|
|
1267
|
+
std::size_t count_removed = 0;
|
|
1268
|
+
for (std::size_t i = 0; i != count_total; ++i) {
|
|
1269
|
+
member_cref_t member = typed_->at(i);
|
|
1270
|
+
count_removed += member.key == free_key_;
|
|
1271
|
+
}
|
|
1272
|
+
|
|
1273
|
+
// Pull entries from the underlying `typed_` into either
|
|
1274
|
+
// into `slot_lookup_`, or `free_keys_` if they are unused.
|
|
1275
|
+
unique_lock_t lock(slot_lookup_mutex_);
|
|
1276
|
+
slot_lookup_.clear();
|
|
1277
|
+
slot_lookup_.reserve(count_total - count_removed);
|
|
1278
|
+
free_keys_.clear();
|
|
1279
|
+
free_keys_.reserve(count_removed);
|
|
1280
|
+
for (std::size_t i = 0; i != typed_->size(); ++i) {
|
|
1281
|
+
member_cref_t member = typed_->at(i);
|
|
1282
|
+
if (member.key == free_key_)
|
|
1283
|
+
free_keys_.push(static_cast<compressed_slot_t>(i));
|
|
1284
|
+
else
|
|
1285
|
+
slot_lookup_.insert(key_and_slot_t{key_t(member.key), static_cast<compressed_slot_t>(i)});
|
|
1286
|
+
}
|
|
1287
|
+
}
|
|
1288
|
+
|
|
1289
|
+
template <typename scalar_at> bool get_(key_t key, scalar_at* reconstructed, cast_t const& cast) const {
|
|
1290
|
+
compressed_slot_t slot;
|
|
1291
|
+
// Find the matching ID
|
|
1292
|
+
{
|
|
1293
|
+
shared_lock_t lock(slot_lookup_mutex_);
|
|
1294
|
+
auto it = slot_lookup_.find(key);
|
|
1295
|
+
if (it == slot_lookup_.end())
|
|
1296
|
+
return false;
|
|
1297
|
+
slot = (*it).slot;
|
|
1298
|
+
}
|
|
1299
|
+
// Export the entry
|
|
1300
|
+
byte_t const* punned_vector = reinterpret_cast<byte_t const*>(vectors_lookup_[slot]);
|
|
1301
|
+
bool casted = cast(punned_vector, dimensions(), (byte_t*)reconstructed);
|
|
1302
|
+
if (!casted)
|
|
1303
|
+
std::memcpy(reconstructed, punned_vector, metric_.bytes_per_vector());
|
|
1304
|
+
return true;
|
|
1305
|
+
}
|
|
1306
|
+
|
|
1307
|
+
template <typename scalar_at> add_result_t add_(key_t key, scalar_at const* vector, cast_t const& cast) {
|
|
1308
|
+
thread_lock_t lock = thread_lock_();
|
|
1309
|
+
index_dense_update_config_t update_config;
|
|
1310
|
+
update_config.thread = lock.thread_id;
|
|
1311
|
+
update_config.expansion = config_.expansion_add;
|
|
1312
|
+
return add_(key, vector, update_config, cast);
|
|
1313
|
+
}
|
|
1314
|
+
|
|
1315
|
+
template <typename scalar_at>
|
|
1316
|
+
search_result_t search_( //
|
|
1317
|
+
scalar_at const* vector, std::size_t wanted, //
|
|
1318
|
+
cast_t const& cast) const {
|
|
1319
|
+
thread_lock_t lock = thread_lock_();
|
|
1320
|
+
index_search_config_t search_config;
|
|
1321
|
+
search_config.thread = lock.thread_id;
|
|
1322
|
+
search_config.expansion = config_.expansion_search;
|
|
1323
|
+
return search_(vector, wanted, search_config, cast);
|
|
1324
|
+
}
|
|
1325
|
+
|
|
1326
|
+
template <typename to_scalar_at> static casts_t make_casts_() {
|
|
1327
|
+
casts_t result;
|
|
1328
|
+
|
|
1329
|
+
result.from_b1x8 = cast_gt<b1x8_t, to_scalar_at>{};
|
|
1330
|
+
result.from_i8 = cast_gt<i8_bits_t, to_scalar_at>{};
|
|
1331
|
+
result.from_f16 = cast_gt<f16_t, to_scalar_at>{};
|
|
1332
|
+
result.from_f32 = cast_gt<f32_t, to_scalar_at>{};
|
|
1333
|
+
result.from_f64 = cast_gt<f64_t, to_scalar_at>{};
|
|
1334
|
+
|
|
1335
|
+
result.to_b1x8 = cast_gt<to_scalar_at, b1x8_t>{};
|
|
1336
|
+
result.to_i8 = cast_gt<to_scalar_at, i8_bits_t>{};
|
|
1337
|
+
result.to_f16 = cast_gt<to_scalar_at, f16_t>{};
|
|
1338
|
+
result.to_f32 = cast_gt<to_scalar_at, f32_t>{};
|
|
1339
|
+
result.to_f64 = cast_gt<to_scalar_at, f64_t>{};
|
|
1340
|
+
|
|
1341
|
+
return result;
|
|
1342
|
+
}
|
|
1343
|
+
|
|
1344
|
+
static casts_t make_casts_(scalar_kind_t scalar_kind) {
|
|
1345
|
+
switch (scalar_kind) {
|
|
1346
|
+
case scalar_kind_t::f64_k: return make_casts_<f64_t>();
|
|
1347
|
+
case scalar_kind_t::f32_k: return make_casts_<f32_t>();
|
|
1348
|
+
case scalar_kind_t::f16_k: return make_casts_<f16_t>();
|
|
1349
|
+
case scalar_kind_t::i8_k: return make_casts_<i8_bits_t>();
|
|
1350
|
+
case scalar_kind_t::b1x8_k: return make_casts_<b1x8_t>();
|
|
1351
|
+
default: return {};
|
|
1352
|
+
}
|
|
1353
|
+
}
|
|
1354
|
+
};
|
|
1355
|
+
|
|
1356
|
+
using index_dense_t = index_dense_gt<>;
|
|
1357
|
+
using index_dense_big_t = index_dense_gt<uuid_t, uint40_t>;
|
|
1358
|
+
|
|
1359
|
+
/**
|
|
1360
|
+
* @brief Adapts the Male-Optimal Stable Marriage algorithm for unequal sets
|
|
1361
|
+
* to perform fast one-to-one matching between two large collections
|
|
1362
|
+
* of vectors, using approximate nearest neighbors search.
|
|
1363
|
+
*
|
|
1364
|
+
* @param[inout] man_to_woman Container to map ::first keys to ::second.
|
|
1365
|
+
* @param[inout] woman_to_man Container to map ::second keys to ::first.
|
|
1366
|
+
* @param[in] executor Thread-pool to execute the job in parallel.
|
|
1367
|
+
* @param[in] progress Callback to report the execution progress.
|
|
1368
|
+
*/
|
|
1369
|
+
template < //
|
|
1370
|
+
|
|
1371
|
+
typename men_label_at, //
|
|
1372
|
+
typename women_label_at, //
|
|
1373
|
+
typename men_slot_at, //
|
|
1374
|
+
typename women_slot_at, //
|
|
1375
|
+
|
|
1376
|
+
typename man_to_woman_at = dummy_label_to_label_mapping_t, //
|
|
1377
|
+
typename woman_to_man_at = dummy_label_to_label_mapping_t, //
|
|
1378
|
+
typename executor_at = dummy_executor_t, //
|
|
1379
|
+
typename progress_at = dummy_progress_t //
|
|
1380
|
+
>
|
|
1381
|
+
static join_result_t join( //
|
|
1382
|
+
index_dense_gt<men_label_at, men_slot_at> const& men, //
|
|
1383
|
+
index_dense_gt<women_label_at, women_slot_at> const& women, //
|
|
1384
|
+
|
|
1385
|
+
index_join_config_t config = {}, //
|
|
1386
|
+
man_to_woman_at&& man_to_woman = man_to_woman_at{}, //
|
|
1387
|
+
woman_to_man_at&& woman_to_man = woman_to_man_at{}, //
|
|
1388
|
+
executor_at&& executor = executor_at{}, //
|
|
1389
|
+
progress_at&& progress = progress_at{}) noexcept {
|
|
1390
|
+
|
|
1391
|
+
return men.join( //
|
|
1392
|
+
women, config, //
|
|
1393
|
+
std::forward<woman_to_man_at>(woman_to_man), //
|
|
1394
|
+
std::forward<man_to_woman_at>(man_to_woman), //
|
|
1395
|
+
std::forward<executor_at>(executor), //
|
|
1396
|
+
std::forward<progress_at>(progress));
|
|
1397
|
+
}
|
|
1398
|
+
|
|
1399
|
+
} // namespace usearch
|
|
1400
|
+
} // namespace unum
|