usearch 1.1.0 → 1.2.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/.github/workflows/release.yml +3 -3
- package/.github/workflows/update_version.sh +1 -1
- package/.gitmodules +1 -4
- package/Cargo.toml +1 -1
- package/Package.swift +2 -1
- package/README.md +129 -46
- package/binding.gyp +16 -12
- package/build.gradle +1 -1
- package/build.rs +0 -1
- package/c/CMakeLists.txt +11 -7
- package/c/lib.cpp +36 -24
- package/c/usearch.h +42 -12
- package/cluster.ipynb +167 -0
- package/conanfile.py +1 -1
- package/cpp/CMakeLists.txt +8 -3
- package/cpp/bench.cpp +2 -2
- package/cpp/test.cpp +42 -25
- package/docs/benchmarks.md +8 -0
- package/docs/compilation.md +7 -7
- package/docs/conf.py +27 -24
- package/docs/javascript/reference.rst +3 -0
- package/docs/wolfram/index.rst +6 -0
- package/golang/README.md +31 -7
- package/golang/lib.go +1 -1
- package/include/usearch/index.hpp +348 -69
- package/include/usearch/index_dense.hpp +587 -179
- package/include/usearch/index_plugins.hpp +194 -16
- package/java/cloud/unum/usearch/Index.java +13 -0
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +40 -9
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.h +16 -0
- package/javascript/docs.js +111 -0
- package/javascript/lib.cpp +29 -9
- package/javascript/usearch.d.ts +9 -1
- package/objc/USearchObjective.mm +11 -11
- package/package.json +2 -2
- package/pyproject.toml +1 -1
- package/python/lib.cpp +492 -159
- package/python/scripts/index_faiss.py +3 -3
- package/python/scripts/join.py +6 -6
- package/python/scripts/test_index.py +219 -0
- package/python/scripts/test_sparse.py +21 -21
- package/python/scripts/test_tooling.py +111 -0
- package/python/usearch/eval.py +9 -5
- package/python/usearch/index.py +491 -216
- package/rust/lib.cpp +6 -24
- package/rust/lib.hpp +1 -3
- package/rust/lib.rs +89 -0
- package/setup.py +5 -6
- package/wasm/CMakeLists.txt +0 -1
- package/wasmer.toml +1 -1
- package/python/scripts/test.py +0 -364
- package/robin-map/.codecov.yml +0 -5
- package/robin-map/.github/workflows/ci.yml +0 -136
- package/robin-map/CMakeLists.txt +0 -86
- package/robin-map/LICENSE +0 -21
- package/robin-map/README.md +0 -483
- package/robin-map/cmake/tsl-robin-mapConfig.cmake.in +0 -9
- package/robin-map/doxygen.conf +0 -2483
- package/robin-map/include/tsl/robin_growth_policy.h +0 -406
- package/robin-map/include/tsl/robin_hash.h +0 -1639
- package/robin-map/include/tsl/robin_map.h +0 -807
- package/robin-map/include/tsl/robin_set.h +0 -660
- package/robin-map/tests/CMakeLists.txt +0 -26
- package/robin-map/tests/custom_allocator_tests.cpp +0 -138
- package/robin-map/tests/main.cpp +0 -26
- package/robin-map/tests/policy_tests.cpp +0 -97
- package/robin-map/tests/robin_map_tests.cpp +0 -1451
- package/robin-map/tests/robin_set_tests.cpp +0 -164
- package/robin-map/tests/utils.h +0 -443
- package/robin-map/tsl-robin-map.natvis +0 -78
package/python/lib.cpp
CHANGED
|
@@ -12,7 +12,10 @@
|
|
|
12
12
|
*
|
|
13
13
|
* @copyright Copyright (c) 2023
|
|
14
14
|
*/
|
|
15
|
+
#if !defined(__cpp_exceptions)
|
|
15
16
|
#define __cpp_exceptions 1
|
|
17
|
+
#endif
|
|
18
|
+
|
|
16
19
|
#include <limits> // `std::numeric_limits`
|
|
17
20
|
#include <thread> // `std::thread`
|
|
18
21
|
|
|
@@ -41,13 +44,16 @@ enum class metric_signature_t {
|
|
|
41
44
|
|
|
42
45
|
namespace py = pybind11;
|
|
43
46
|
using py_shape_t = py::array::ShapeContainer;
|
|
44
|
-
|
|
47
|
+
|
|
45
48
|
using metric_t = metric_punned_t;
|
|
46
49
|
using distance_t = distance_punned_t;
|
|
47
50
|
|
|
51
|
+
using dense_key_t = typename index_dense_t::key_t;
|
|
48
52
|
using dense_add_result_t = typename index_dense_t::add_result_t;
|
|
49
53
|
using dense_search_result_t = typename index_dense_t::search_result_t;
|
|
50
54
|
using dense_labeling_result_t = typename index_dense_t::labeling_result_t;
|
|
55
|
+
using dense_cluster_result_t = typename index_dense_t::cluster_result_t;
|
|
56
|
+
using dense_clustering_result_t = typename index_dense_t::clustering_result_t;
|
|
51
57
|
|
|
52
58
|
struct dense_index_py_t : public index_dense_t {
|
|
53
59
|
using native_t = index_dense_t;
|
|
@@ -63,26 +69,27 @@ struct dense_index_py_t : public index_dense_t {
|
|
|
63
69
|
struct dense_indexes_py_t {
|
|
64
70
|
std::vector<std::shared_ptr<dense_index_py_t>> shards_;
|
|
65
71
|
|
|
66
|
-
void
|
|
72
|
+
void merge(std::shared_ptr<dense_index_py_t> shard) { shards_.push_back(shard); }
|
|
67
73
|
std::size_t bytes_per_vector() const noexcept { return shards_.empty() ? 0 : shards_[0]->bytes_per_vector(); }
|
|
68
74
|
std::size_t scalar_words() const noexcept { return shards_.empty() ? 0 : shards_[0]->scalar_words(); }
|
|
69
75
|
index_limits_t limits() const noexcept { return {size(), std::numeric_limits<std::size_t>::max()}; }
|
|
70
76
|
|
|
71
|
-
void
|
|
77
|
+
void merge_paths(std::vector<std::string> const& paths, bool view = true, std::size_t threads = 0) {
|
|
72
78
|
if (!threads)
|
|
73
79
|
threads = std::thread::hardware_concurrency();
|
|
74
80
|
|
|
75
81
|
shards_.reserve(shards_.size() + paths.size());
|
|
76
82
|
std::mutex shards_mutex;
|
|
77
|
-
executor_default_t{threads}.
|
|
83
|
+
executor_default_t{threads}.dynamic(paths.size(), [&](std::size_t, std::size_t task_idx) {
|
|
78
84
|
index_dense_t index = index_dense_t::make(paths[task_idx].c_str(), view);
|
|
79
85
|
if (!index)
|
|
80
|
-
return;
|
|
86
|
+
return false;
|
|
81
87
|
auto shared_index = std::make_shared<dense_index_py_t>(std::move(index));
|
|
82
88
|
std::unique_lock<std::mutex> lock(shards_mutex);
|
|
83
89
|
shards_.push_back(shared_index);
|
|
84
90
|
if (PyErr_CheckSignals() != 0)
|
|
85
91
|
throw py::error_already_set();
|
|
92
|
+
return true;
|
|
86
93
|
});
|
|
87
94
|
}
|
|
88
95
|
|
|
@@ -95,7 +102,7 @@ struct dense_indexes_py_t {
|
|
|
95
102
|
};
|
|
96
103
|
|
|
97
104
|
template <typename scalar_at>
|
|
98
|
-
metric_t
|
|
105
|
+
metric_t wrap_typed_user_defined_metric( //
|
|
99
106
|
metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, //
|
|
100
107
|
scalar_kind_t scalar_kind, std::size_t dimensions) {
|
|
101
108
|
//
|
|
@@ -129,16 +136,21 @@ metric_t typed_udf(
|
|
|
129
136
|
return metric_t(stl_function, dimensions, kind, scalar_kind);
|
|
130
137
|
}
|
|
131
138
|
|
|
132
|
-
metric_t
|
|
139
|
+
metric_t wrap_user_defined_metric( //
|
|
133
140
|
metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, //
|
|
134
141
|
scalar_kind_t scalar_kind, std::size_t dimensions) {
|
|
135
142
|
|
|
136
143
|
switch (scalar_kind) {
|
|
137
|
-
case scalar_kind_t::b1x8_k:
|
|
138
|
-
|
|
139
|
-
case scalar_kind_t::
|
|
140
|
-
|
|
141
|
-
case scalar_kind_t::
|
|
144
|
+
case scalar_kind_t::b1x8_k:
|
|
145
|
+
return wrap_typed_user_defined_metric<b1x8_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
146
|
+
case scalar_kind_t::i8_k:
|
|
147
|
+
return wrap_typed_user_defined_metric<i8_bits_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
148
|
+
case scalar_kind_t::f16_k:
|
|
149
|
+
return wrap_typed_user_defined_metric<f16_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
150
|
+
case scalar_kind_t::f32_k:
|
|
151
|
+
return wrap_typed_user_defined_metric<f32_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
152
|
+
case scalar_kind_t::f64_k:
|
|
153
|
+
return wrap_typed_user_defined_metric<f64_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
142
154
|
default: return {};
|
|
143
155
|
}
|
|
144
156
|
}
|
|
@@ -151,12 +163,15 @@ static dense_index_py_t make_index( //
|
|
|
151
163
|
std::size_t expansion_search, //
|
|
152
164
|
metric_kind_t metric_kind, //
|
|
153
165
|
metric_signature_t metric_signature, //
|
|
154
|
-
std::uintptr_t metric_uintptr
|
|
166
|
+
std::uintptr_t metric_uintptr, //
|
|
167
|
+
bool multi) {
|
|
155
168
|
|
|
156
169
|
index_dense_config_t config(connectivity, expansion_add, expansion_search);
|
|
170
|
+
config.multi = multi;
|
|
171
|
+
|
|
157
172
|
metric_t metric = //
|
|
158
173
|
metric_uintptr //
|
|
159
|
-
?
|
|
174
|
+
? wrap_user_defined_metric(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
|
|
160
175
|
: metric_t(dimensions, metric_kind, scalar_kind);
|
|
161
176
|
return index_dense_t::make(metric, config);
|
|
162
177
|
}
|
|
@@ -177,6 +192,18 @@ scalar_kind_t numpy_string_to_kind(std::string const& name) {
|
|
|
177
192
|
return scalar_kind_t::unknown_k;
|
|
178
193
|
}
|
|
179
194
|
|
|
195
|
+
template <typename result_at> void forward_error(result_at&& result) {
|
|
196
|
+
|
|
197
|
+
if (!result)
|
|
198
|
+
throw std::invalid_argument(result.error.release());
|
|
199
|
+
|
|
200
|
+
int signals = PyErr_CheckSignals();
|
|
201
|
+
if (signals != 0)
|
|
202
|
+
throw py::error_already_set();
|
|
203
|
+
}
|
|
204
|
+
|
|
205
|
+
using atomic_error_t = std::atomic<char const*>;
|
|
206
|
+
|
|
180
207
|
template <typename scalar_at>
|
|
181
208
|
static void add_typed_to_index( //
|
|
182
209
|
dense_index_py_t& index, //
|
|
@@ -186,18 +213,30 @@ static void add_typed_to_index( //
|
|
|
186
213
|
Py_ssize_t vectors_count = vectors_info.shape[0];
|
|
187
214
|
byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);
|
|
188
215
|
byte_t const* keys_data = reinterpret_cast<byte_t const*>(keys_info.ptr);
|
|
216
|
+
atomic_error_t atomic_error{nullptr};
|
|
189
217
|
|
|
190
|
-
executor_default_t{threads}.
|
|
191
|
-
|
|
192
|
-
config.force_vector_copy = force_copy;
|
|
193
|
-
config.thread = thread_idx;
|
|
194
|
-
key_t key = *reinterpret_cast<key_t const*>(keys_data + task_idx * keys_info.strides[0]);
|
|
218
|
+
executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
219
|
+
dense_key_t key = *reinterpret_cast<dense_key_t const*>(keys_data + task_idx * keys_info.strides[0]);
|
|
195
220
|
scalar_at const* vector = reinterpret_cast<scalar_at const*>(vectors_data + task_idx * vectors_info.strides[0]);
|
|
196
|
-
dense_add_result_t result = index.add(key, vector,
|
|
197
|
-
result
|
|
198
|
-
|
|
199
|
-
|
|
221
|
+
dense_add_result_t result = index.add(key, vector, thread_idx, force_copy);
|
|
222
|
+
if (!result) {
|
|
223
|
+
atomic_error = result.error.release();
|
|
224
|
+
return false;
|
|
225
|
+
}
|
|
226
|
+
|
|
227
|
+
// We don't want to check for signals from multiple threads
|
|
228
|
+
if (thread_idx == 0)
|
|
229
|
+
if (PyErr_CheckSignals() != 0)
|
|
230
|
+
return false;
|
|
231
|
+
return true;
|
|
200
232
|
});
|
|
233
|
+
|
|
234
|
+
// Raise the error from a single thread
|
|
235
|
+
auto error = atomic_error.load();
|
|
236
|
+
if (error) {
|
|
237
|
+
PyErr_SetString(PyExc_RuntimeError, error);
|
|
238
|
+
throw py::error_already_set();
|
|
239
|
+
}
|
|
201
240
|
}
|
|
202
241
|
|
|
203
242
|
template <typename index_at>
|
|
@@ -208,11 +247,11 @@ static void add_many_to_index( //
|
|
|
208
247
|
py::buffer_info keys_info = keys.request();
|
|
209
248
|
py::buffer_info vectors_info = vectors.request();
|
|
210
249
|
|
|
211
|
-
if (keys_info.itemsize != sizeof(
|
|
250
|
+
if (keys_info.itemsize != sizeof(dense_key_t))
|
|
212
251
|
throw std::invalid_argument("Incompatible key type!");
|
|
213
252
|
|
|
214
253
|
if (keys_info.ndim != 1)
|
|
215
|
-
throw std::invalid_argument("
|
|
254
|
+
throw std::invalid_argument("Keys must be placed in a single-dimensional array!");
|
|
216
255
|
|
|
217
256
|
if (vectors_info.ndim != 2)
|
|
218
257
|
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
@@ -247,7 +286,7 @@ template <typename scalar_at>
|
|
|
247
286
|
static void search_typed( //
|
|
248
287
|
dense_index_py_t& index, py::buffer_info& vectors_info, //
|
|
249
288
|
std::size_t wanted, bool exact, std::size_t threads, //
|
|
250
|
-
py::array_t<
|
|
289
|
+
py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
|
|
251
290
|
std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances) {
|
|
252
291
|
|
|
253
292
|
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
@@ -262,28 +301,41 @@ static void search_typed( //
|
|
|
262
301
|
if (!index.reserve(index_limits_t(index.size(), threads)))
|
|
263
302
|
throw std::invalid_argument("Out of memory!");
|
|
264
303
|
|
|
265
|
-
|
|
266
|
-
|
|
267
|
-
config.thread = thread_idx;
|
|
268
|
-
config.exact = exact;
|
|
304
|
+
atomic_error_t atomic_error{nullptr};
|
|
305
|
+
executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
269
306
|
scalar_at const* vector = (scalar_at const*)(vectors_data + task_idx * vectors_info.strides[0]);
|
|
270
|
-
dense_search_result_t result = index.search(vector, wanted,
|
|
271
|
-
result
|
|
307
|
+
dense_search_result_t result = index.search(vector, wanted, thread_idx, exact);
|
|
308
|
+
if (!result) {
|
|
309
|
+
atomic_error = result.error.release();
|
|
310
|
+
return false;
|
|
311
|
+
}
|
|
312
|
+
|
|
272
313
|
counts_py1d(task_idx) =
|
|
273
314
|
static_cast<Py_ssize_t>(result.dump_to(&keys_py2d(task_idx, 0), &distances_py2d(task_idx, 0)));
|
|
274
315
|
|
|
275
316
|
stats_visited_members += result.visited_members;
|
|
276
317
|
stats_computed_distances += result.computed_distances;
|
|
277
|
-
|
|
278
|
-
|
|
318
|
+
|
|
319
|
+
// We don't want to check for signals from multiple threads
|
|
320
|
+
if (thread_idx == 0)
|
|
321
|
+
if (PyErr_CheckSignals() != 0)
|
|
322
|
+
return false;
|
|
323
|
+
return true;
|
|
279
324
|
});
|
|
325
|
+
|
|
326
|
+
// Raise the error from a single thread
|
|
327
|
+
auto error = atomic_error.load();
|
|
328
|
+
if (error) {
|
|
329
|
+
PyErr_SetString(PyExc_RuntimeError, error);
|
|
330
|
+
throw py::error_already_set();
|
|
331
|
+
}
|
|
280
332
|
}
|
|
281
333
|
|
|
282
334
|
template <typename scalar_at>
|
|
283
335
|
static void search_typed( //
|
|
284
336
|
dense_indexes_py_t& indexes, py::buffer_info& vectors_info, //
|
|
285
337
|
std::size_t wanted, bool exact, std::size_t threads, //
|
|
286
|
-
py::array_t<
|
|
338
|
+
py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
|
|
287
339
|
std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances) {
|
|
288
340
|
|
|
289
341
|
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
@@ -298,27 +350,33 @@ static void search_typed( //
|
|
|
298
350
|
if (!threads)
|
|
299
351
|
threads = std::thread::hardware_concurrency();
|
|
300
352
|
|
|
301
|
-
|
|
302
|
-
|
|
353
|
+
bitset_t query_mutexes(static_cast<std::size_t>(vectors_count));
|
|
354
|
+
if (!query_mutexes)
|
|
355
|
+
throw std::bad_alloc();
|
|
356
|
+
|
|
357
|
+
atomic_error_t atomic_error{nullptr};
|
|
358
|
+
executor_default_t{threads}.dynamic(indexes.shards_.size(), [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
303
359
|
dense_index_py_t& index = *indexes.shards_[task_idx].get();
|
|
304
360
|
|
|
305
361
|
index_limits_t limits;
|
|
306
362
|
limits.members = index.size();
|
|
307
363
|
limits.threads_add = 0;
|
|
308
364
|
limits.threads_search = 1;
|
|
309
|
-
if (!index.reserve(limits))
|
|
310
|
-
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
config.thread = 0;
|
|
314
|
-
config.exact = exact;
|
|
365
|
+
if (!index.reserve(limits)) {
|
|
366
|
+
atomic_error = "Out of memory!";
|
|
367
|
+
return false;
|
|
368
|
+
}
|
|
315
369
|
|
|
316
370
|
for (std::size_t vector_idx = 0; vector_idx != static_cast<std::size_t>(vectors_count); ++vector_idx) {
|
|
317
371
|
scalar_at const* vector = (scalar_at const*)(vectors_data + vector_idx * vectors_info.strides[0]);
|
|
318
|
-
dense_search_result_t result = index.search(vector, wanted,
|
|
319
|
-
result
|
|
372
|
+
dense_search_result_t result = index.search(vector, wanted, 0, exact);
|
|
373
|
+
if (!result) {
|
|
374
|
+
atomic_error = result.error.release();
|
|
375
|
+
return false;
|
|
376
|
+
}
|
|
377
|
+
|
|
320
378
|
{
|
|
321
|
-
|
|
379
|
+
auto lock = query_mutexes.lock(vector_idx);
|
|
322
380
|
counts_py1d(vector_idx) = static_cast<Py_ssize_t>(result.merge_into( //
|
|
323
381
|
&keys_py2d(vector_idx, 0), //
|
|
324
382
|
&distances_py2d(vector_idx, 0), //
|
|
@@ -328,10 +386,21 @@ static void search_typed( //
|
|
|
328
386
|
|
|
329
387
|
stats_visited_members += result.visited_members;
|
|
330
388
|
stats_computed_distances += result.computed_distances;
|
|
331
|
-
|
|
332
|
-
|
|
389
|
+
|
|
390
|
+
// We don't want to check for signals from multiple threads
|
|
391
|
+
if (thread_idx == 0)
|
|
392
|
+
if (PyErr_CheckSignals() != 0)
|
|
393
|
+
return false;
|
|
394
|
+
return true;
|
|
333
395
|
}
|
|
334
396
|
});
|
|
397
|
+
|
|
398
|
+
// Raise the error from a single thread
|
|
399
|
+
auto error = atomic_error.load();
|
|
400
|
+
if (error) {
|
|
401
|
+
PyErr_SetString(PyExc_RuntimeError, error);
|
|
402
|
+
throw py::error_already_set();
|
|
403
|
+
}
|
|
335
404
|
}
|
|
336
405
|
|
|
337
406
|
/**
|
|
@@ -341,7 +410,9 @@ static void search_typed( //
|
|
|
341
410
|
* @return Tuple with:
|
|
342
411
|
* 1. matrix of neighbors,
|
|
343
412
|
* 2. matrix of distances,
|
|
344
|
-
* 3. array with match counts
|
|
413
|
+
* 3. array with match counts,
|
|
414
|
+
* 4. number of visited nodes,
|
|
415
|
+
* 4. number of computed pairwise distances.
|
|
345
416
|
*/
|
|
346
417
|
template <typename index_at>
|
|
347
418
|
static py::tuple search_many_in_index( //
|
|
@@ -362,7 +433,7 @@ static py::tuple search_many_in_index( //
|
|
|
362
433
|
if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
|
|
363
434
|
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
364
435
|
|
|
365
|
-
py::array_t<
|
|
436
|
+
py::array_t<dense_key_t> keys_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
|
|
366
437
|
py::array_t<distance_t> distances_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
|
|
367
438
|
py::array_t<Py_ssize_t> counts_py(vectors_count);
|
|
368
439
|
std::atomic<std::size_t> stats_visited_members(0);
|
|
@@ -392,7 +463,7 @@ template <typename scalar_at>
|
|
|
392
463
|
static void search_typed_brute_force( //
|
|
393
464
|
py::buffer_info& dataset_info, py::buffer_info& queries_info, //
|
|
394
465
|
std::size_t wanted, std::size_t threads, metric_t const& metric, //
|
|
395
|
-
py::array_t<
|
|
466
|
+
py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py) {
|
|
396
467
|
|
|
397
468
|
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
398
469
|
auto distances_py2d = distances_py.template mutable_unchecked<2>();
|
|
@@ -400,7 +471,6 @@ static void search_typed_brute_force( //
|
|
|
400
471
|
|
|
401
472
|
std::size_t dataset_count = static_cast<std::size_t>(dataset_info.shape[0]);
|
|
402
473
|
std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);
|
|
403
|
-
std::size_t dimensions = static_cast<std::size_t>(dataset_info.shape[1]);
|
|
404
474
|
|
|
405
475
|
byte_t const* dataset_data = reinterpret_cast<byte_t const*>(dataset_info.ptr);
|
|
406
476
|
byte_t const* queries_data = reinterpret_cast<byte_t const*>(queries_info.ptr);
|
|
@@ -411,11 +481,11 @@ static void search_typed_brute_force( //
|
|
|
411
481
|
threads = std::thread::hardware_concurrency();
|
|
412
482
|
|
|
413
483
|
std::size_t tasks_count = static_cast<std::size_t>(dataset_count * queries_count);
|
|
414
|
-
|
|
484
|
+
bitset_t query_mutexes(static_cast<std::size_t>(queries_count));
|
|
415
485
|
if (!query_mutexes)
|
|
416
486
|
throw std::bad_alloc();
|
|
417
487
|
|
|
418
|
-
executor_default_t{threads}.
|
|
488
|
+
executor_default_t{threads}.dynamic(tasks_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
419
489
|
//
|
|
420
490
|
std::size_t dataset_idx = task_idx / queries_count;
|
|
421
491
|
std::size_t query_idx = task_idx % queries_count;
|
|
@@ -426,25 +496,28 @@ static void search_typed_brute_force( //
|
|
|
426
496
|
|
|
427
497
|
{
|
|
428
498
|
auto lock = query_mutexes.lock(query_idx);
|
|
429
|
-
|
|
499
|
+
dense_key_t* keys = &keys_py2d(query_idx, 0);
|
|
430
500
|
distance_t* distances = &distances_py2d(query_idx, 0);
|
|
431
501
|
std::size_t& matches = reinterpret_cast<std::size_t&>(counts_py1d(query_idx));
|
|
432
502
|
if (matches == wanted)
|
|
433
503
|
if (distances[wanted - 1] <= distance)
|
|
434
|
-
return;
|
|
504
|
+
return true;
|
|
435
505
|
|
|
436
506
|
std::size_t offset = std::lower_bound(distances, distances + matches, distance) - distances;
|
|
437
507
|
|
|
438
508
|
std::size_t count_worse = matches - offset - (wanted == matches);
|
|
439
|
-
std::memmove(keys + offset + 1, keys + offset, count_worse * sizeof(
|
|
509
|
+
std::memmove(keys + offset + 1, keys + offset, count_worse * sizeof(dense_key_t));
|
|
440
510
|
std::memmove(distances + offset + 1, distances + offset, count_worse * sizeof(distance_t));
|
|
441
|
-
keys[offset] = static_cast<
|
|
511
|
+
keys[offset] = static_cast<dense_key_t>(dataset_idx);
|
|
442
512
|
distances[offset] = distance;
|
|
443
513
|
matches += matches != wanted;
|
|
444
514
|
}
|
|
445
515
|
|
|
446
|
-
|
|
447
|
-
|
|
516
|
+
// We don't want to check for signals from multiple threads
|
|
517
|
+
if (thread_idx == 0)
|
|
518
|
+
if (PyErr_CheckSignals() != 0)
|
|
519
|
+
return false;
|
|
520
|
+
return true;
|
|
448
521
|
});
|
|
449
522
|
}
|
|
450
523
|
|
|
@@ -475,14 +548,14 @@ static py::tuple search_many_brute_force( //
|
|
|
475
548
|
if (dataset_kind != queries_kind)
|
|
476
549
|
throw std::invalid_argument("The types of vectors don't match!");
|
|
477
550
|
|
|
478
|
-
py::array_t<
|
|
551
|
+
py::array_t<dense_key_t> keys_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
|
|
479
552
|
py::array_t<distance_t> distances_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
|
|
480
553
|
py::array_t<Py_ssize_t> counts_py(dataset_count);
|
|
481
554
|
|
|
482
555
|
std::size_t dimensions = static_cast<std::size_t>(queries_dimensions);
|
|
483
556
|
metric_t metric = //
|
|
484
557
|
metric_uintptr //
|
|
485
|
-
?
|
|
558
|
+
? wrap_user_defined_metric(metric_kind, metric_signature, metric_uintptr, queries_kind, dimensions)
|
|
486
559
|
: metric_t(dimensions, metric_kind, queries_kind);
|
|
487
560
|
|
|
488
561
|
// clang-format off
|
|
@@ -505,12 +578,153 @@ static py::tuple search_many_brute_force( //
|
|
|
505
578
|
return results;
|
|
506
579
|
}
|
|
507
580
|
|
|
508
|
-
|
|
509
|
-
|
|
581
|
+
template <typename scalar_at> struct rows_lookup_gt {
|
|
582
|
+
byte_t* data_;
|
|
583
|
+
std::size_t stride_;
|
|
584
|
+
|
|
585
|
+
rows_lookup_gt(void* data, std::size_t stride) noexcept : data_((byte_t*)data), stride_(stride) {}
|
|
586
|
+
scalar_at* operator[](std::size_t i) const noexcept { return reinterpret_cast<scalar_at*>(data_ + i * stride_); }
|
|
587
|
+
std::ptrdiff_t operator-(rows_lookup_gt const& other) const noexcept { return (data_ - other.data_) / stride_; }
|
|
588
|
+
rows_lookup_gt operator+(std::size_t n) const noexcept { return {data_ + stride_ * n, stride_}; }
|
|
589
|
+
template <typename other_scalar_at> rows_lookup_gt<other_scalar_at> as() const noexcept { return {data_, stride_}; }
|
|
590
|
+
};
|
|
591
|
+
|
|
592
|
+
/**
|
|
593
|
+
* @param queries Matrix of vectors to search for.
|
|
594
|
+
* @param count Number of clusters to produce.
|
|
595
|
+
*
|
|
596
|
+
* @return Tuple with:
|
|
597
|
+
* 1. vector of cluster IDs,
|
|
598
|
+
* 2. vector of distances to those clusters,
|
|
599
|
+
* 3. array with match counts, set to all ones,
|
|
600
|
+
* 4. number of visited nodes,
|
|
601
|
+
* 4. number of computed pairwise distances.
|
|
602
|
+
*/
|
|
603
|
+
template <typename index_at>
|
|
604
|
+
static py::tuple cluster_vectors( //
|
|
605
|
+
index_at& index, py::buffer queries, //
|
|
606
|
+
std::size_t min_count, std::size_t max_count, std::size_t threads) {
|
|
607
|
+
|
|
608
|
+
if (index.limits().threads_search < threads)
|
|
609
|
+
throw std::invalid_argument("Can't use that many threads!");
|
|
610
|
+
|
|
611
|
+
py::buffer_info queries_info = queries.request();
|
|
612
|
+
if (queries_info.ndim != 2)
|
|
613
|
+
throw std::invalid_argument("Expects a matrix of queries to add!");
|
|
614
|
+
|
|
615
|
+
std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);
|
|
616
|
+
std::size_t queries_stride = static_cast<std::size_t>(queries_info.strides[0]);
|
|
617
|
+
std::size_t queries_dimensions = static_cast<std::size_t>(queries_info.shape[1]);
|
|
618
|
+
if (queries_dimensions != index.scalar_words())
|
|
619
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
620
|
+
|
|
621
|
+
py::array_t<dense_key_t> keys_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
|
|
622
|
+
py::array_t<distance_t> distances_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
|
|
623
|
+
dense_clustering_result_t cluster_result;
|
|
624
|
+
executor_default_t executor{threads};
|
|
625
|
+
|
|
626
|
+
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
627
|
+
auto distances_py2d = distances_py.template mutable_unchecked<2>();
|
|
628
|
+
dense_key_t* keys_ptr = reinterpret_cast<dense_key_t*>(&keys_py2d(0, 0));
|
|
629
|
+
distance_t* distances_ptr = reinterpret_cast<distance_t*>(&distances_py2d(0, 0));
|
|
630
|
+
|
|
631
|
+
index_dense_clustering_config_t config;
|
|
632
|
+
config.min_clusters = min_count;
|
|
633
|
+
config.max_clusters = max_count;
|
|
634
|
+
|
|
635
|
+
rows_lookup_gt<byte_t const> queries_begin(queries_info.ptr, queries_stride);
|
|
636
|
+
rows_lookup_gt<byte_t const> queries_end = queries_begin + queries_count;
|
|
637
|
+
|
|
638
|
+
// clang-format off
|
|
639
|
+
switch (numpy_string_to_kind(queries_info.format)) {
|
|
640
|
+
case scalar_kind_t::b1x8_k: cluster_result = index.cluster(queries_begin.as<b1x8_t const>(), queries_end.as<b1x8_t const>(), config, keys_ptr, distances_ptr, executor); break;
|
|
641
|
+
case scalar_kind_t::i8_k: cluster_result = index.cluster(queries_begin.as<i8_bits_t const>(), queries_end.as<i8_bits_t const>(), config, keys_ptr, distances_ptr, executor); break;
|
|
642
|
+
case scalar_kind_t::f16_k: cluster_result = index.cluster(queries_begin.as<f16_t const>(), queries_end.as<f16_t const>(), config, keys_ptr, distances_ptr, executor); break;
|
|
643
|
+
case scalar_kind_t::f32_k: cluster_result = index.cluster(queries_begin.as<f32_t const>(), queries_end.as<f32_t const>(), config, keys_ptr, distances_ptr, executor); break;
|
|
644
|
+
case scalar_kind_t::f64_k: cluster_result = index.cluster(queries_begin.as<f64_t const>(), queries_end.as<f64_t const>(), config, keys_ptr, distances_ptr, executor); break;
|
|
645
|
+
default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + queries_info.format);
|
|
646
|
+
}
|
|
647
|
+
// clang-format on
|
|
648
|
+
|
|
649
|
+
cluster_result.error.raise();
|
|
650
|
+
|
|
651
|
+
// Those would be set to 1 for all entries, in case of success
|
|
652
|
+
py::array_t<Py_ssize_t> counts_py(queries_count);
|
|
653
|
+
auto counts_py1d = counts_py.template mutable_unchecked<1>();
|
|
654
|
+
for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)
|
|
655
|
+
counts_py1d(static_cast<Py_ssize_t>(query_idx)) = 1;
|
|
656
|
+
|
|
657
|
+
py::tuple results(5);
|
|
658
|
+
results[0] = keys_py;
|
|
659
|
+
results[1] = distances_py;
|
|
660
|
+
results[2] = counts_py;
|
|
661
|
+
results[3] = cluster_result.visited_members;
|
|
662
|
+
results[4] = cluster_result.computed_distances;
|
|
663
|
+
return results;
|
|
664
|
+
}
|
|
665
|
+
|
|
666
|
+
/**
|
|
667
|
+
* @param queries Array of keys to cluster.
|
|
668
|
+
* @param count Number of clusters to produce.
|
|
669
|
+
*
|
|
670
|
+
* @return Tuple with:
|
|
671
|
+
* 1. vector of cluster IDs,
|
|
672
|
+
* 2. vector of distances to those clusters,
|
|
673
|
+
* 3. array with match counts, set to all ones,
|
|
674
|
+
* 4. number of visited nodes,
|
|
675
|
+
* 4. number of computed pairwise distances.
|
|
676
|
+
*/
|
|
677
|
+
template <typename index_at>
|
|
678
|
+
static py::tuple cluster_keys( //
|
|
679
|
+
index_at& index, py::array_t<dense_key_t> queries_py, //
|
|
680
|
+
std::size_t min_count, std::size_t max_count, std::size_t threads) {
|
|
681
|
+
|
|
682
|
+
if (index.limits().threads_search < threads)
|
|
683
|
+
throw std::invalid_argument("Can't use that many threads!");
|
|
684
|
+
|
|
685
|
+
std::size_t queries_count = static_cast<std::size_t>(queries_py.size());
|
|
686
|
+
auto queries_py1d = queries_py.template unchecked<1>();
|
|
687
|
+
dense_key_t const* queries_begin = &queries_py1d(0);
|
|
688
|
+
dense_key_t const* queries_end = queries_begin + queries_count;
|
|
689
|
+
|
|
690
|
+
py::array_t<dense_key_t> keys_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
|
|
691
|
+
py::array_t<distance_t> distances_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
|
|
692
|
+
executor_default_t executor{threads};
|
|
693
|
+
|
|
694
|
+
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
695
|
+
auto distances_py2d = distances_py.template mutable_unchecked<2>();
|
|
696
|
+
dense_key_t* keys_ptr = reinterpret_cast<dense_key_t*>(&keys_py2d(0, 0));
|
|
697
|
+
distance_t* distances_ptr = reinterpret_cast<distance_t*>(&distances_py2d(0, 0));
|
|
698
|
+
|
|
699
|
+
index_dense_clustering_config_t config;
|
|
700
|
+
config.min_clusters = min_count;
|
|
701
|
+
config.max_clusters = max_count;
|
|
702
|
+
|
|
703
|
+
dense_clustering_result_t cluster_result =
|
|
704
|
+
index.cluster(queries_begin, queries_end, config, keys_ptr, distances_ptr, executor);
|
|
705
|
+
cluster_result.error.raise();
|
|
706
|
+
|
|
707
|
+
// Those would be set to 1 for all entries, in case of success
|
|
708
|
+
py::array_t<Py_ssize_t> counts_py(queries_count);
|
|
709
|
+
auto counts_py1d = counts_py.template mutable_unchecked<1>();
|
|
710
|
+
for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)
|
|
711
|
+
counts_py1d(static_cast<Py_ssize_t>(query_idx)) = 1;
|
|
712
|
+
|
|
713
|
+
py::tuple results(5);
|
|
714
|
+
results[0] = keys_py;
|
|
715
|
+
results[1] = distances_py;
|
|
716
|
+
results[2] = counts_py;
|
|
717
|
+
results[3] = cluster_result.visited_members;
|
|
718
|
+
results[4] = cluster_result.computed_distances;
|
|
719
|
+
return results;
|
|
720
|
+
}
|
|
721
|
+
|
|
722
|
+
static std::unordered_map<dense_key_t, dense_key_t> join_index( //
|
|
723
|
+
dense_index_py_t const& a, dense_index_py_t const& b, //
|
|
510
724
|
std::size_t max_proposals, bool exact) {
|
|
511
725
|
|
|
512
|
-
std::unordered_map<
|
|
513
|
-
|
|
726
|
+
std::unordered_map<dense_key_t, dense_key_t> a_to_b;
|
|
727
|
+
dummy_key_to_key_mapping_t b_to_a;
|
|
514
728
|
a_to_b.reserve((std::min)(a.size(), b.size()));
|
|
515
729
|
|
|
516
730
|
index_join_config_t config;
|
|
@@ -520,17 +734,18 @@ static std::unordered_map<key_t, key_t> join_index( //
|
|
|
520
734
|
std::size_t threads = (std::min)(a.limits().threads(), b.limits().threads());
|
|
521
735
|
executor_default_t executor{threads};
|
|
522
736
|
join_result_t result = a.join(b, config, a_to_b, b_to_a, executor);
|
|
523
|
-
result
|
|
737
|
+
forward_error(result);
|
|
524
738
|
|
|
525
739
|
return a_to_b;
|
|
526
740
|
}
|
|
527
741
|
|
|
528
|
-
static dense_index_py_t copy_index(dense_index_py_t const& index) {
|
|
742
|
+
static dense_index_py_t copy_index(dense_index_py_t const& index, bool force_copy) {
|
|
529
743
|
|
|
530
744
|
using copy_result_t = typename dense_index_py_t::copy_result_t;
|
|
531
|
-
|
|
745
|
+
index_dense_copy_config_t config;
|
|
746
|
+
config.force_vector_copy = force_copy;
|
|
532
747
|
copy_result_t result = index.copy(config);
|
|
533
|
-
result
|
|
748
|
+
forward_error(result);
|
|
534
749
|
return std::move(result.index);
|
|
535
750
|
}
|
|
536
751
|
|
|
@@ -556,53 +771,49 @@ template <typename index_at> typename index_at::stats_t compute_level_stats(inde
|
|
|
556
771
|
// clang-format on
|
|
557
772
|
|
|
558
773
|
template <typename internal_at, typename external_at = internal_at, typename index_at = void>
|
|
559
|
-
py::
|
|
560
|
-
|
|
561
|
-
|
|
562
|
-
if (
|
|
563
|
-
|
|
564
|
-
|
|
774
|
+
static py::tuple get_typed_vectors_for_keys(index_at const& index, py::buffer keys) {
|
|
775
|
+
|
|
776
|
+
py::buffer_info keys_info = keys.request();
|
|
777
|
+
if (keys_info.ndim != 1)
|
|
778
|
+
throw std::invalid_argument("Keys must be placed in a single-dimensional array!");
|
|
779
|
+
|
|
780
|
+
Py_ssize_t keys_count = keys_info.shape[0];
|
|
781
|
+
byte_t const* keys_data = reinterpret_cast<byte_t const*>(keys_info.ptr);
|
|
782
|
+
py::tuple results(keys_count);
|
|
783
|
+
|
|
784
|
+
for (Py_ssize_t task_idx = 0; task_idx != keys_count; ++task_idx) {
|
|
785
|
+
dense_key_t key = *reinterpret_cast<dense_key_t const*>(keys_data + task_idx * keys_info.strides[0]);
|
|
786
|
+
std::size_t vectors_count = index.count(key);
|
|
787
|
+
if (!vectors_count) {
|
|
788
|
+
results[task_idx] = py::none();
|
|
789
|
+
continue;
|
|
790
|
+
}
|
|
791
|
+
|
|
792
|
+
py::array_t<external_at> result_py({static_cast<Py_ssize_t>(vectors_count), //
|
|
793
|
+
static_cast<Py_ssize_t>(index.scalar_words())});
|
|
794
|
+
auto result_py2d = result_py.template mutable_unchecked<2>();
|
|
795
|
+
index.get(key, (internal_at*)&result_py2d(0, 0), vectors_count);
|
|
796
|
+
results[task_idx] = result_py;
|
|
797
|
+
}
|
|
798
|
+
|
|
799
|
+
return results;
|
|
565
800
|
}
|
|
566
801
|
|
|
567
|
-
template <typename index_at> py::
|
|
802
|
+
template <typename index_at> py::tuple get_many(index_at const& index, py::buffer keys, scalar_kind_t scalar_kind) {
|
|
568
803
|
if (scalar_kind == scalar_kind_t::f32_k)
|
|
569
|
-
return
|
|
804
|
+
return get_typed_vectors_for_keys<f32_t>(index, keys);
|
|
570
805
|
else if (scalar_kind == scalar_kind_t::f64_k)
|
|
571
|
-
return
|
|
806
|
+
return get_typed_vectors_for_keys<f64_t>(index, keys);
|
|
572
807
|
else if (scalar_kind == scalar_kind_t::f16_k)
|
|
573
|
-
return
|
|
808
|
+
return get_typed_vectors_for_keys<f16_t, std::uint16_t>(index, keys);
|
|
574
809
|
else if (scalar_kind == scalar_kind_t::i8_k)
|
|
575
|
-
return
|
|
810
|
+
return get_typed_vectors_for_keys<i8_bits_t, std::int8_t>(index, keys);
|
|
576
811
|
else if (scalar_kind == scalar_kind_t::b1x8_k)
|
|
577
|
-
return
|
|
812
|
+
return get_typed_vectors_for_keys<b1x8_t, std::uint8_t>(index, keys);
|
|
578
813
|
else
|
|
579
814
|
throw std::invalid_argument("Incompatible scalars in the query matrix!");
|
|
580
815
|
}
|
|
581
816
|
|
|
582
|
-
template <typename index_at> py::array_t<key_t> get_keys(index_at const& index, std::size_t offset, std::size_t limit) {
|
|
583
|
-
limit = std::min(index.size(), limit);
|
|
584
|
-
py::array_t<key_t> result_py(static_cast<Py_ssize_t>(limit));
|
|
585
|
-
auto result_py1d = result_py.template mutable_unchecked<1>();
|
|
586
|
-
index.export_keys(&result_py1d(0), offset, limit);
|
|
587
|
-
return result_py;
|
|
588
|
-
}
|
|
589
|
-
|
|
590
|
-
template <typename index_at> py::array_t<key_t> get_all_keys(index_at const& index) {
|
|
591
|
-
return get_keys(index, 0, index.size());
|
|
592
|
-
}
|
|
593
|
-
|
|
594
|
-
template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
|
|
595
|
-
if (begin == end)
|
|
596
|
-
return false;
|
|
597
|
-
element_at const* last = begin;
|
|
598
|
-
begin++;
|
|
599
|
-
for (; begin != end; ++begin, ++last) {
|
|
600
|
-
if (*begin == *last)
|
|
601
|
-
return true;
|
|
602
|
-
}
|
|
603
|
-
return false;
|
|
604
|
-
}
|
|
605
|
-
|
|
606
817
|
PYBIND11_MODULE(compiled, m) {
|
|
607
818
|
m.doc() = "Smaller & Faster Single-File Vector Search Engine from Unum";
|
|
608
819
|
|
|
@@ -622,7 +833,7 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
622
833
|
py::enum_<metric_kind_t>(m, "MetricKind")
|
|
623
834
|
.value("Unknown", metric_kind_t::unknown_k)
|
|
624
835
|
|
|
625
|
-
.value("IP", metric_kind_t::
|
|
836
|
+
.value("IP", metric_kind_t::ip_k)
|
|
626
837
|
.value("Cos", metric_kind_t::cos_k)
|
|
627
838
|
.value("L2sq", metric_kind_t::l2sq_k)
|
|
628
839
|
|
|
@@ -634,7 +845,7 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
634
845
|
.value("Sorensen", metric_kind_t::sorensen_k)
|
|
635
846
|
|
|
636
847
|
.value("Cosine", metric_kind_t::cos_k)
|
|
637
|
-
.value("InnerProduct", metric_kind_t::
|
|
848
|
+
.value("InnerProduct", metric_kind_t::ip_k);
|
|
638
849
|
|
|
639
850
|
py::enum_<scalar_kind_t>(m, "ScalarKind")
|
|
640
851
|
.value("Unknown", scalar_kind_t::unknown_k)
|
|
@@ -656,7 +867,8 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
656
867
|
|
|
657
868
|
m.def("index_dense_metadata", [](std::string const& path) -> py::dict {
|
|
658
869
|
index_dense_metadata_result_t meta = index_dense_metadata(path.c_str());
|
|
659
|
-
meta
|
|
870
|
+
forward_error(meta);
|
|
871
|
+
|
|
660
872
|
index_dense_head_t const& head = meta.head;
|
|
661
873
|
|
|
662
874
|
py::dict result;
|
|
@@ -701,40 +913,87 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
701
913
|
py::arg("expansion_search") = default_expansion_search(), //
|
|
702
914
|
py::arg("metric_kind") = metric_kind_t::cos_k, //
|
|
703
915
|
py::arg("metric_signature") = metric_signature_t::array_array_k, //
|
|
704
|
-
py::arg("metric_pointer") = 0
|
|
916
|
+
py::arg("metric_pointer") = 0, //
|
|
917
|
+
py::arg("multi") = false //
|
|
918
|
+
);
|
|
919
|
+
|
|
920
|
+
i.def( //
|
|
921
|
+
"add_many", &add_many_to_index<dense_index_py_t>, //
|
|
922
|
+
py::arg("keys"), //
|
|
923
|
+
py::arg("vectors"), //
|
|
924
|
+
py::kw_only(), //
|
|
925
|
+
py::arg("copy") = true, //
|
|
926
|
+
py::arg("threads") = 0 //
|
|
927
|
+
);
|
|
928
|
+
|
|
929
|
+
i.def( //
|
|
930
|
+
"search_many", &search_many_in_index<dense_index_py_t>, //
|
|
931
|
+
py::arg("queries"), //
|
|
932
|
+
py::arg("count") = 10, //
|
|
933
|
+
py::arg("exact") = false, //
|
|
934
|
+
py::arg("threads") = 0 //
|
|
705
935
|
);
|
|
706
936
|
|
|
707
|
-
i.def(
|
|
708
|
-
"
|
|
709
|
-
py::arg("
|
|
710
|
-
py::arg("
|
|
711
|
-
py::
|
|
712
|
-
py::arg("
|
|
713
|
-
py::arg("threads") = 0 //
|
|
937
|
+
i.def( //
|
|
938
|
+
"cluster_vectors", &cluster_vectors<dense_index_py_t>, //
|
|
939
|
+
py::arg("queries"), //
|
|
940
|
+
py::arg("min_count") = 0, //
|
|
941
|
+
py::arg("max_count") = 0, //
|
|
942
|
+
py::arg("threads") = 0 //
|
|
714
943
|
);
|
|
715
944
|
|
|
716
|
-
i.def(
|
|
717
|
-
"
|
|
718
|
-
py::arg("
|
|
719
|
-
py::arg("
|
|
720
|
-
py::arg("
|
|
721
|
-
py::arg("threads") = 0
|
|
945
|
+
i.def( //
|
|
946
|
+
"cluster_keys", &cluster_keys<dense_index_py_t>, //
|
|
947
|
+
py::arg("queries"), //
|
|
948
|
+
py::arg("min_count") = 0, //
|
|
949
|
+
py::arg("max_count") = 0, //
|
|
950
|
+
py::arg("threads") = 0 //
|
|
722
951
|
);
|
|
723
952
|
|
|
724
953
|
i.def(
|
|
725
|
-
"
|
|
726
|
-
[](dense_index_py_t& index,
|
|
954
|
+
"rename_one_to_one",
|
|
955
|
+
[](dense_index_py_t& index, dense_key_t from, dense_key_t to) -> bool {
|
|
727
956
|
dense_labeling_result_t result = index.rename(from, to);
|
|
728
|
-
result
|
|
957
|
+
forward_error(result);
|
|
729
958
|
return result.completed;
|
|
730
959
|
},
|
|
731
|
-
py::arg("
|
|
960
|
+
py::arg("from_"), py::arg("to"));
|
|
732
961
|
|
|
733
962
|
i.def(
|
|
734
|
-
"
|
|
735
|
-
[](dense_index_py_t& index,
|
|
963
|
+
"rename_many_to_many",
|
|
964
|
+
[](dense_index_py_t& index, std::vector<dense_key_t> const& from,
|
|
965
|
+
std::vector<dense_key_t> const& to) -> std::vector<bool> {
|
|
966
|
+
if (from.size() != to.size())
|
|
967
|
+
throw std::invalid_argument("Sizes of `from` and `to` arrays don't match!");
|
|
968
|
+
|
|
969
|
+
std::vector<bool> results(from.size(), false);
|
|
970
|
+
for (std::size_t i = 0; i != from.size(); ++i) {
|
|
971
|
+
dense_labeling_result_t result = index.rename(from[i], to[i]);
|
|
972
|
+
results[i] = result.completed;
|
|
973
|
+
forward_error(result);
|
|
974
|
+
}
|
|
975
|
+
return results;
|
|
976
|
+
},
|
|
977
|
+
py::arg("from_"), py::arg("to"));
|
|
978
|
+
|
|
979
|
+
i.def(
|
|
980
|
+
"rename_many_to_one",
|
|
981
|
+
[](dense_index_py_t& index, std::vector<dense_key_t> const& from, dense_key_t to) -> std::vector<bool> {
|
|
982
|
+
std::vector<bool> results(from.size(), false);
|
|
983
|
+
for (std::size_t i = 0; i != from.size(); ++i) {
|
|
984
|
+
dense_labeling_result_t result = index.rename(from[i], to);
|
|
985
|
+
results[i] = result.completed;
|
|
986
|
+
forward_error(result);
|
|
987
|
+
}
|
|
988
|
+
return results;
|
|
989
|
+
},
|
|
990
|
+
py::arg("from_"), py::arg("to"));
|
|
991
|
+
|
|
992
|
+
i.def(
|
|
993
|
+
"remove_one",
|
|
994
|
+
[](dense_index_py_t& index, dense_key_t key, bool compact, std::size_t threads) -> bool {
|
|
736
995
|
dense_labeling_result_t result = index.remove(key);
|
|
737
|
-
result
|
|
996
|
+
forward_error(result);
|
|
738
997
|
if (!compact)
|
|
739
998
|
return result.completed;
|
|
740
999
|
|
|
@@ -749,10 +1008,11 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
749
1008
|
py::arg("key"), py::arg("compact"), py::arg("threads"));
|
|
750
1009
|
|
|
751
1010
|
i.def(
|
|
752
|
-
"
|
|
753
|
-
[](dense_index_py_t& index, std::vector<
|
|
1011
|
+
"remove_many",
|
|
1012
|
+
[](dense_index_py_t& index, std::vector<dense_key_t> const& keys, bool compact,
|
|
1013
|
+
std::size_t threads) -> std::size_t {
|
|
754
1014
|
dense_labeling_result_t result = index.remove(keys.begin(), keys.end());
|
|
755
|
-
result
|
|
1015
|
+
forward_error(result);
|
|
756
1016
|
if (!compact)
|
|
757
1017
|
return result.completed;
|
|
758
1018
|
|
|
@@ -775,8 +1035,7 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
775
1035
|
i.def_property_readonly( //
|
|
776
1036
|
"dtype", [](dense_index_py_t const& index) -> scalar_kind_t { return index.scalar_kind(); });
|
|
777
1037
|
i.def_property_readonly( //
|
|
778
|
-
"memory_usage", [](dense_index_py_t const& index) -> std::size_t { return index.memory_usage(); }
|
|
779
|
-
py::call_guard<py::gil_scoped_release>());
|
|
1038
|
+
"memory_usage", [](dense_index_py_t const& index) -> std::size_t { return index.memory_usage(); });
|
|
780
1039
|
|
|
781
1040
|
i.def_property("expansion_add", &dense_index_py_t::expansion_add, &dense_index_py_t::change_expansion_add);
|
|
782
1041
|
i.def_property("expansion_search", &dense_index_py_t::expansion_search, &dense_index_py_t::change_expansion_search);
|
|
@@ -789,7 +1048,7 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
789
1048
|
std::size_t dimensions = index.dimensions();
|
|
790
1049
|
metric_t metric = //
|
|
791
1050
|
metric_uintptr //
|
|
792
|
-
?
|
|
1051
|
+
? wrap_user_defined_metric(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
|
|
793
1052
|
: metric_t(dimensions, metric_kind, scalar_kind);
|
|
794
1053
|
index.change_metric(std::move(metric));
|
|
795
1054
|
},
|
|
@@ -798,21 +1057,95 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
798
1057
|
py::arg("metric_pointer") = 0 //
|
|
799
1058
|
);
|
|
800
1059
|
|
|
801
|
-
i.def_property_readonly("
|
|
802
|
-
|
|
803
|
-
|
|
804
|
-
|
|
805
|
-
i.def("
|
|
806
|
-
|
|
807
|
-
|
|
808
|
-
i.def(
|
|
809
|
-
|
|
810
|
-
|
|
811
|
-
|
|
812
|
-
|
|
813
|
-
|
|
814
|
-
|
|
815
|
-
|
|
1060
|
+
i.def_property_readonly("hardware_acceleration", [](dense_index_py_t const& index) -> py::str {
|
|
1061
|
+
return isa_name(index.metric().isa_kind());
|
|
1062
|
+
});
|
|
1063
|
+
|
|
1064
|
+
i.def("contains_one", &dense_index_py_t::contains);
|
|
1065
|
+
i.def("count_one", &dense_index_py_t::count);
|
|
1066
|
+
|
|
1067
|
+
i.def( //
|
|
1068
|
+
"contains_many",
|
|
1069
|
+
[](dense_index_py_t const& index, py::array_t<dense_key_t> const& keys_py) -> py::array_t<bool> {
|
|
1070
|
+
py::array_t<bool> results_py(keys_py.size());
|
|
1071
|
+
auto results_py1d = results_py.template mutable_unchecked<1>();
|
|
1072
|
+
auto keys_py1d = keys_py.template unchecked<1>();
|
|
1073
|
+
for (Py_ssize_t task_idx = 0; task_idx != keys_py.size(); ++task_idx)
|
|
1074
|
+
results_py1d(task_idx) = index.contains(keys_py1d(task_idx));
|
|
1075
|
+
return results_py;
|
|
1076
|
+
});
|
|
1077
|
+
|
|
1078
|
+
i.def( //
|
|
1079
|
+
"count_many",
|
|
1080
|
+
[](dense_index_py_t const& index, py::array_t<dense_key_t> const& keys_py) -> py::array_t<std::size_t> {
|
|
1081
|
+
py::array_t<std::size_t> results_py(keys_py.size());
|
|
1082
|
+
auto results_py1d = results_py.template mutable_unchecked<1>();
|
|
1083
|
+
auto keys_py1d = keys_py.template unchecked<1>();
|
|
1084
|
+
for (Py_ssize_t task_idx = 0; task_idx != keys_py.size(); ++task_idx)
|
|
1085
|
+
results_py1d(task_idx) = index.count(keys_py1d(task_idx));
|
|
1086
|
+
return results_py;
|
|
1087
|
+
});
|
|
1088
|
+
|
|
1089
|
+
i.def( //
|
|
1090
|
+
"pairwise_distances",
|
|
1091
|
+
[](dense_index_py_t const& index, py::array_t<dense_key_t> const& left_py,
|
|
1092
|
+
py::array_t<dense_key_t> const& right_py) -> py::array_t<distance_t> {
|
|
1093
|
+
py::array_t<distance_t> results_py(left_py.size());
|
|
1094
|
+
auto results_py1d = results_py.template mutable_unchecked<1>();
|
|
1095
|
+
auto left_py1d = left_py.template unchecked<1>();
|
|
1096
|
+
auto right_py1d = right_py.template unchecked<1>();
|
|
1097
|
+
for (Py_ssize_t task_idx = 0; task_idx != left_py.size(); ++task_idx)
|
|
1098
|
+
results_py1d(task_idx) = index.distance_between(left_py1d(task_idx), right_py1d(task_idx)).min;
|
|
1099
|
+
return results_py;
|
|
1100
|
+
});
|
|
1101
|
+
|
|
1102
|
+
i.def( //
|
|
1103
|
+
"pairwise_distance", [](dense_index_py_t const& index, dense_key_t left, dense_key_t right) -> distance_t {
|
|
1104
|
+
return index.distance_between(left, right).min;
|
|
1105
|
+
});
|
|
1106
|
+
|
|
1107
|
+
i.def("get_many", &get_many<dense_index_py_t>, py::arg("keys"), py::arg("dtype") = scalar_kind_t::f32_k);
|
|
1108
|
+
|
|
1109
|
+
i.def(
|
|
1110
|
+
"get_keys_in_slice",
|
|
1111
|
+
[](dense_index_py_t const& index, std::size_t offset, std::size_t limit) -> py::array_t<dense_key_t> {
|
|
1112
|
+
limit = std::min(index.size(), limit);
|
|
1113
|
+
py::array_t<dense_key_t> result_py(static_cast<Py_ssize_t>(limit));
|
|
1114
|
+
auto result_py1d = result_py.template mutable_unchecked<1>();
|
|
1115
|
+
index.export_keys(&result_py1d(0), offset, limit);
|
|
1116
|
+
return result_py;
|
|
1117
|
+
},
|
|
1118
|
+
py::arg("offset") = 0, py::arg("limit") = std::numeric_limits<std::size_t>::max());
|
|
1119
|
+
|
|
1120
|
+
i.def(
|
|
1121
|
+
"get_keys_at_offsets",
|
|
1122
|
+
[](dense_index_py_t const& index, py::array_t<Py_ssize_t> const& offsets_py) -> py::array_t<dense_key_t> {
|
|
1123
|
+
py::array_t<dense_key_t> result_py(offsets_py.size());
|
|
1124
|
+
auto result_py1d = result_py.template mutable_unchecked<1>();
|
|
1125
|
+
auto offsets_py1d = offsets_py.template unchecked<1>();
|
|
1126
|
+
for (Py_ssize_t task_idx = 0; task_idx != offsets_py.size(); ++task_idx)
|
|
1127
|
+
index.export_keys(&result_py1d(task_idx), offsets_py1d(task_idx), 1);
|
|
1128
|
+
return result_py;
|
|
1129
|
+
},
|
|
1130
|
+
py::arg("offsets"));
|
|
1131
|
+
|
|
1132
|
+
i.def(
|
|
1133
|
+
"get_key_at_offset",
|
|
1134
|
+
[](dense_index_py_t const& index, std::size_t offset) -> dense_key_t {
|
|
1135
|
+
dense_key_t result;
|
|
1136
|
+
index.export_keys(&result, offset, 1);
|
|
1137
|
+
return result;
|
|
1138
|
+
},
|
|
1139
|
+
py::arg("offset"));
|
|
1140
|
+
|
|
1141
|
+
i.def("save", &save_index<dense_index_py_t>, py::arg("path"));
|
|
1142
|
+
i.def("load", &load_index<dense_index_py_t>, py::arg("path"));
|
|
1143
|
+
i.def("view", &view_index<dense_index_py_t>, py::arg("path"));
|
|
1144
|
+
i.def("reset", &reset_index<dense_index_py_t>);
|
|
1145
|
+
i.def("clear", &clear_index<dense_index_py_t>);
|
|
1146
|
+
i.def("copy", ©_index, py::kw_only(), py::arg("copy") = true);
|
|
1147
|
+
i.def("compact", &compact_index);
|
|
1148
|
+
i.def("join", &join_index, py::arg("other"), py::arg("max_proposals") = 0, py::arg("exact") = false);
|
|
816
1149
|
|
|
817
1150
|
using punned_index_stats_t = typename dense_index_py_t::stats_t;
|
|
818
1151
|
auto i_stats = py::class_<punned_index_stats_t>(m, "IndexStats");
|
|
@@ -828,14 +1161,14 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
828
1161
|
auto is = py::class_<dense_indexes_py_t>(m, "Indexes");
|
|
829
1162
|
is.def(py::init());
|
|
830
1163
|
is.def("__len__", &dense_indexes_py_t::size);
|
|
831
|
-
is.def("
|
|
832
|
-
is.def("
|
|
1164
|
+
is.def("merge", &dense_indexes_py_t::merge);
|
|
1165
|
+
is.def("merge_paths", &dense_indexes_py_t::merge_paths, py::arg("paths"), py::arg("view") = true,
|
|
833
1166
|
py::arg("threads") = 0);
|
|
834
|
-
is.def(
|
|
835
|
-
"
|
|
836
|
-
py::arg("query"),
|
|
837
|
-
py::arg("count") = 10,
|
|
838
|
-
py::arg("exact") = false,
|
|
839
|
-
py::arg("threads") = 0
|
|
1167
|
+
is.def( //
|
|
1168
|
+
"search_many", &search_many_in_index<dense_indexes_py_t>, //
|
|
1169
|
+
py::arg("query"), //
|
|
1170
|
+
py::arg("count") = 10, //
|
|
1171
|
+
py::arg("exact") = false, //
|
|
1172
|
+
py::arg("threads") = 0 //
|
|
840
1173
|
);
|
|
841
1174
|
}
|