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.
Files changed (71) hide show
  1. package/.github/workflows/prerelease.yml +3 -3
  2. package/.github/workflows/release.yml +3 -3
  3. package/.github/workflows/update_version.sh +1 -1
  4. package/.gitmodules +1 -4
  5. package/Cargo.toml +1 -1
  6. package/Package.swift +2 -1
  7. package/README.md +129 -46
  8. package/binding.gyp +16 -12
  9. package/build.gradle +1 -1
  10. package/build.rs +0 -1
  11. package/c/CMakeLists.txt +11 -7
  12. package/c/lib.cpp +36 -24
  13. package/c/usearch.h +42 -12
  14. package/cluster.ipynb +167 -0
  15. package/conanfile.py +1 -1
  16. package/cpp/CMakeLists.txt +8 -3
  17. package/cpp/bench.cpp +2 -2
  18. package/cpp/test.cpp +42 -25
  19. package/docs/benchmarks.md +8 -0
  20. package/docs/compilation.md +7 -7
  21. package/docs/conf.py +27 -24
  22. package/docs/javascript/reference.rst +3 -0
  23. package/docs/wolfram/index.rst +6 -0
  24. package/golang/README.md +31 -7
  25. package/golang/lib.go +1 -1
  26. package/include/usearch/index.hpp +348 -69
  27. package/include/usearch/index_dense.hpp +587 -179
  28. package/include/usearch/index_plugins.hpp +194 -16
  29. package/java/cloud/unum/usearch/Index.java +13 -0
  30. package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +40 -9
  31. package/java/cloud/unum/usearch/cloud_unum_usearch_Index.h +16 -0
  32. package/javascript/docs.js +111 -0
  33. package/javascript/lib.cpp +29 -9
  34. package/javascript/usearch.d.ts +9 -1
  35. package/objc/USearchObjective.mm +11 -11
  36. package/package.json +2 -2
  37. package/pyproject.toml +1 -1
  38. package/python/lib.cpp +492 -159
  39. package/python/scripts/index_faiss.py +3 -3
  40. package/python/scripts/join.py +6 -6
  41. package/python/scripts/test_index.py +219 -0
  42. package/python/scripts/test_sparse.py +21 -21
  43. package/python/scripts/test_tooling.py +111 -0
  44. package/python/usearch/eval.py +9 -5
  45. package/python/usearch/index.py +491 -216
  46. package/rust/lib.cpp +6 -24
  47. package/rust/lib.hpp +1 -3
  48. package/rust/lib.rs +89 -0
  49. package/setup.py +5 -6
  50. package/wasm/CMakeLists.txt +0 -1
  51. package/wasmer.toml +1 -1
  52. package/python/scripts/test.py +0 -364
  53. package/robin-map/.codecov.yml +0 -5
  54. package/robin-map/.github/workflows/ci.yml +0 -136
  55. package/robin-map/CMakeLists.txt +0 -86
  56. package/robin-map/LICENSE +0 -21
  57. package/robin-map/README.md +0 -483
  58. package/robin-map/cmake/tsl-robin-mapConfig.cmake.in +0 -9
  59. package/robin-map/doxygen.conf +0 -2483
  60. package/robin-map/include/tsl/robin_growth_policy.h +0 -406
  61. package/robin-map/include/tsl/robin_hash.h +0 -1639
  62. package/robin-map/include/tsl/robin_map.h +0 -807
  63. package/robin-map/include/tsl/robin_set.h +0 -660
  64. package/robin-map/tests/CMakeLists.txt +0 -26
  65. package/robin-map/tests/custom_allocator_tests.cpp +0 -138
  66. package/robin-map/tests/main.cpp +0 -26
  67. package/robin-map/tests/policy_tests.cpp +0 -97
  68. package/robin-map/tests/robin_map_tests.cpp +0 -1451
  69. package/robin-map/tests/robin_set_tests.cpp +0 -164
  70. package/robin-map/tests/utils.h +0 -443
  71. 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
- #define key_t typename index_dense_t::key_t
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 add(std::shared_ptr<dense_index_py_t> shard) { shards_.push_back(shard); }
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 add_paths(std::vector<std::string> const& paths, bool view = true, std::size_t threads = 0) {
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}.execute_bulk(paths.size(), [&](std::size_t, std::size_t task_idx) {
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 typed_udf( //
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 udf( //
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: return typed_udf<b1x8_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
138
- case scalar_kind_t::i8_k: return typed_udf<i8_bits_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
139
- case scalar_kind_t::f16_k: return typed_udf<f16_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
140
- case scalar_kind_t::f32_k: return typed_udf<f32_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
141
- case scalar_kind_t::f64_k: return typed_udf<f64_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
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
- ? udf(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
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}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
191
- index_dense_update_config_t config;
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, config);
197
- result.error.raise();
198
- if (PyErr_CheckSignals() != 0)
199
- throw py::error_already_set();
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(key_t))
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("Labels must be placed in a single-dimensional array!");
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<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
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
- executor_default_t{threads}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
266
- index_search_config_t config;
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, config);
271
- result.error.raise();
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
- if (PyErr_CheckSignals() != 0)
278
- throw py::error_already_set();
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<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
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
- std::vector<std::mutex> query_mutexes(static_cast<std::size_t>(vectors_count));
302
- executor_default_t{threads}.execute_bulk(indexes.shards_.size(), [&](std::size_t, std::size_t task_idx) {
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
- throw std::bad_alloc();
311
-
312
- index_search_config_t config;
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, config);
319
- result.error.raise();
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
- std::unique_lock<std::mutex> lock(query_mutexes[vector_idx]);
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
- if (PyErr_CheckSignals() != 0)
332
- throw py::error_already_set();
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<key_t> keys_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
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<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py) {
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
- visits_bitset_t query_mutexes(static_cast<std::size_t>(queries_count));
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}.execute_bulk(tasks_count, [&](std::size_t, std::size_t task_idx) {
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
- key_t* keys = &keys_py2d(query_idx, 0);
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(key_t));
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<key_t>(dataset_idx);
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
- if (PyErr_CheckSignals() != 0)
447
- throw py::error_already_set();
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<key_t> keys_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
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
- ? udf(metric_kind, metric_signature, metric_uintptr, queries_kind, dimensions)
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
- static std::unordered_map<key_t, key_t> join_index( //
509
- dense_index_py_t const& a, dense_index_py_t const& b, //
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<key_t, key_t> a_to_b;
513
- dummy_label_to_label_mapping_t b_to_a;
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.error.raise();
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
- index_copy_config_t config;
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.error.raise();
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::object get_typed_member(index_at const& index, key_t key) {
560
- py::array_t<external_at> result_py(static_cast<Py_ssize_t>(index.scalar_words()));
561
- auto result_py1d = result_py.template mutable_unchecked<1>();
562
- if (!index.get(key, (internal_at*)&result_py1d(0)))
563
- return py::none();
564
- return py::object(result_py);
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::object get_member(index_at const& index, key_t key, scalar_kind_t scalar_kind) {
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 get_typed_member<f32_t>(index, key);
804
+ return get_typed_vectors_for_keys<f32_t>(index, keys);
570
805
  else if (scalar_kind == scalar_kind_t::f64_k)
571
- return get_typed_member<f64_t>(index, key);
806
+ return get_typed_vectors_for_keys<f64_t>(index, keys);
572
807
  else if (scalar_kind == scalar_kind_t::f16_k)
573
- return get_typed_member<f16_t, std::uint16_t>(index, key);
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 get_typed_member<i8_bits_t, std::int8_t>(index, key);
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 get_typed_member<b1x8_t, std::uint8_t>(index, key);
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::cos_k)
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::cos_k);
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.error.raise();
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
- "add", &add_many_to_index<dense_index_py_t>, //
709
- py::arg("keys"), //
710
- py::arg("vectors"), //
711
- py::kw_only(), //
712
- py::arg("copy") = true, //
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
- "search", &search_many_in_index<dense_index_py_t>, //
718
- py::arg("query"), //
719
- py::arg("count") = 10, //
720
- py::arg("exact") = false, //
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
- "rename",
726
- [](dense_index_py_t& index, key_t from, key_t to) -> bool {
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.error.raise();
957
+ forward_error(result);
729
958
  return result.completed;
730
959
  },
731
- py::arg("from"), py::arg("to"));
960
+ py::arg("from_"), py::arg("to"));
732
961
 
733
962
  i.def(
734
- "remove",
735
- [](dense_index_py_t& index, key_t key, bool compact, std::size_t threads) -> bool {
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.error.raise();
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
- "remove",
753
- [](dense_index_py_t& index, std::vector<key_t> const& keys, bool compact, std::size_t threads) -> std::size_t {
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.error.raise();
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
- ? udf(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
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("keys", &get_all_keys<dense_index_py_t>);
802
- i.def("get_keys", &get_keys<dense_index_py_t>, py::arg("offset") = 0,
803
- py::arg("limit") = std::numeric_limits<std::size_t>::max());
804
- i.def("__contains__", &dense_index_py_t::contains);
805
- i.def("__getitem__", &get_member<dense_index_py_t>, py::arg("key"), py::arg("dtype") = scalar_kind_t::f32_k);
806
-
807
- i.def("save", &save_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
808
- i.def("load", &load_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
809
- i.def("view", &view_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
810
- i.def("reset", &reset_index<dense_index_py_t>, py::call_guard<py::gil_scoped_release>());
811
- i.def("clear", &clear_index<dense_index_py_t>, py::call_guard<py::gil_scoped_release>());
812
- i.def("copy", &copy_index, py::call_guard<py::gil_scoped_release>());
813
- i.def("compact", &compact_index, py::call_guard<py::gil_scoped_release>());
814
- i.def("join", &join_index, py::arg("other"), py::arg("max_proposals") = 0, py::arg("exact") = false,
815
- py::call_guard<py::gil_scoped_release>());
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", &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("add", &dense_indexes_py_t::add);
832
- is.def("add_paths", &dense_indexes_py_t::add_paths, py::arg("paths"), py::arg("view") = true,
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
- "search", &search_many_in_index<dense_indexes_py_t>, //
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
  }