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
@@ -1,28 +1,21 @@
1
1
  #pragma once
2
2
  #include <stdlib.h> // `aligned_alloc`
3
3
 
4
- #include <functional> // `std::function`
5
- #include <numeric> // `std::iota`
6
- #include <shared_mutex> // `std::shared_mutex`
7
- #include <thread> // `std::thread`
8
- #include <vector> // `std::vector`
4
+ #include <functional> // `std::function`
5
+ #include <numeric> // `std::iota`
6
+ #include <shared_mutex> // `std::shared_mutex`
7
+ #include <thread> // `std::thread`
8
+ #include <unordered_set> // `std::unordered_multiset`
9
+ #include <vector> // `std::vector`
9
10
 
10
11
  #include <usearch/index.hpp>
11
12
  #include <usearch/index_plugins.hpp>
12
13
 
13
- #include <tsl/robin_set.h>
14
-
15
14
  namespace unum {
16
15
  namespace usearch {
17
16
 
18
17
  template <typename, typename> class index_dense_gt;
19
18
 
20
- // clang-format off
21
- template <typename key_at, typename std::enable_if<std::is_integral<key_at>::value>::type* = nullptr> key_at default_free_value() { return std::numeric_limits<key_at>::max(); }
22
- template <typename key_at, typename std::enable_if<std::is_same<key_at, uint40_t>::value>::type* = nullptr> uint40_t default_free_value() { return uint40_t::max(); }
23
- template <typename key_at, typename std::enable_if<!std::is_integral<key_at>::value && !std::is_same<key_at, uint40_t>::value>::type* = nullptr> key_at default_free_value() { return key_at(); }
24
- // clang-format on
25
-
26
19
  /**
27
20
  * @brief The "magic" sequence helps infer the type of the file.
28
21
  * USearch indexes start with the "usearch" string.
@@ -64,19 +57,21 @@ struct index_dense_head_t {
64
57
  misaligned_ref_gt<std::uint64_t> count_present;
65
58
  misaligned_ref_gt<std::uint64_t> count_deleted;
66
59
  misaligned_ref_gt<std::uint64_t> dimensions;
60
+ misaligned_ref_gt<bool> multi;
67
61
 
68
62
  index_dense_head_t(byte_t* ptr) noexcept
69
- : magic((char const*)exchange(ptr, ptr + sizeof(magic_t))), //
70
- version_major(exchange(ptr, ptr + sizeof(version_t))), //
71
- version_minor(exchange(ptr, ptr + sizeof(version_t))), //
72
- version_patch(exchange(ptr, ptr + sizeof(version_t))), //
73
- kind_metric(exchange(ptr, ptr + sizeof(metric_kind_t))), //
74
- kind_scalar(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
75
- kind_key(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
76
- kind_compressed_slot(exchange(ptr, ptr + sizeof(scalar_kind_t))),
77
- count_present(exchange(ptr, ptr + sizeof(std::uint64_t))),
78
- count_deleted(exchange(ptr, ptr + sizeof(std::uint64_t))),
79
- dimensions(exchange(ptr, ptr + sizeof(std::uint64_t))) {}
63
+ : magic((char const*)exchange(ptr, ptr + sizeof(magic_t))), //
64
+ version_major(exchange(ptr, ptr + sizeof(version_t))), //
65
+ version_minor(exchange(ptr, ptr + sizeof(version_t))), //
66
+ version_patch(exchange(ptr, ptr + sizeof(version_t))), //
67
+ kind_metric(exchange(ptr, ptr + sizeof(metric_kind_t))), //
68
+ kind_scalar(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
69
+ kind_key(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
70
+ kind_compressed_slot(exchange(ptr, ptr + sizeof(scalar_kind_t))), //
71
+ count_present(exchange(ptr, ptr + sizeof(std::uint64_t))), //
72
+ count_deleted(exchange(ptr, ptr + sizeof(std::uint64_t))), //
73
+ dimensions(exchange(ptr, ptr + sizeof(std::uint64_t))), //
74
+ multi(exchange(ptr, ptr + sizeof(bool))) {}
80
75
  };
81
76
 
82
77
  struct index_dense_head_result_t {
@@ -96,7 +91,7 @@ struct index_dense_config_t : public index_config_t {
96
91
  std::size_t expansion_add = default_expansion_add();
97
92
  std::size_t expansion_search = default_expansion_search();
98
93
  bool exclude_vectors = false;
99
- bool allow_key_collisions = false;
94
+ bool multi = false;
100
95
 
101
96
  index_dense_config_t(index_config_t base) noexcept : index_config_t(base) {}
102
97
 
@@ -109,18 +104,20 @@ struct index_dense_config_t : public index_config_t {
109
104
  expansion_search(expansion_search ? expansion_search : default_expansion_search()) {}
110
105
  };
111
106
 
107
+ struct index_dense_clustering_config_t {
108
+ std::size_t min_clusters = 0;
109
+ std::size_t max_clusters = 0;
110
+ enum mode_t {
111
+ merge_smallest_k,
112
+ merge_closest_k,
113
+ } mode = merge_smallest_k;
114
+ };
115
+
112
116
  struct index_dense_serialization_config_t {
113
117
  bool exclude_vectors = false;
114
118
  bool use_64_bit_dimensions = false;
115
119
  };
116
120
 
117
- struct index_dense_update_config_t : public index_update_config_t {
118
- bool force_vector_copy = false;
119
-
120
- index_dense_update_config_t() = default;
121
- index_dense_update_config_t(index_update_config_t base) noexcept : index_update_config_t(base) {}
122
- };
123
-
124
121
  struct index_dense_copy_config_t : public index_copy_config_t {
125
122
  bool force_vector_copy = true;
126
123
 
@@ -255,13 +252,16 @@ class index_dense_gt {
255
252
 
256
253
  using serialization_config_t = index_dense_serialization_config_t;
257
254
 
255
+ using dynamic_allocator_t = aligned_allocator_gt<byte_t, 64>;
256
+ using tape_allocator_t = memory_mapping_allocator_gt<64>;
257
+
258
258
  private:
259
259
  /// @brief Schema: input buffer, bytes in input buffer, output buffer.
260
260
  using cast_t = std::function<bool(byte_t const*, std::size_t, byte_t*)>;
261
261
  /// @brief Punned index.
262
262
  using index_t = index_gt< //
263
263
  distance_t, key_t, compressed_slot_t, //
264
- aligned_allocator_gt<byte_t, 64>, memory_mapping_allocator_gt<64>>;
264
+ dynamic_allocator_t, tape_allocator_t>;
265
265
  using index_allocator_t = aligned_allocator_gt<index_t, 64>;
266
266
 
267
267
  using member_iterator_t = typename index_t::member_iterator_t;
@@ -323,13 +323,16 @@ class index_dense_gt {
323
323
  /// @brief Mutex, controlling concurrent access to `available_threads_`.
324
324
  mutable std::mutex available_threads_mutex_;
325
325
 
326
- using shared_mutex_t = std::mutex; // TODO: Find an OS-compatible solution
327
- using shared_lock_t = std::unique_lock<shared_mutex_t>;
326
+ using shared_mutex_t = unfair_shared_mutex_t;
327
+ using shared_lock_t = shared_lock_gt<shared_mutex_t>;
328
328
  using unique_lock_t = std::unique_lock<shared_mutex_t>;
329
329
 
330
330
  struct key_and_slot_t {
331
331
  key_t key;
332
332
  compressed_slot_t slot;
333
+
334
+ bool any_slot() const { return slot == default_free_value<compressed_slot_t>(); }
335
+ static key_and_slot_t any_slot(key_t key) { return {key, default_free_value<compressed_slot_t>()}; }
333
336
  };
334
337
 
335
338
  struct lookup_key_hash_t {
@@ -342,13 +345,11 @@ class index_dense_gt {
342
345
  using is_transparent = void;
343
346
  bool operator()(key_and_slot_t const& a, key_t const& b) const noexcept { return a.key == b; }
344
347
  bool operator()(key_t const& a, key_and_slot_t const& b) const noexcept { return a == b.key; }
345
- bool operator()(key_and_slot_t const& a, key_and_slot_t const& b) const noexcept {
346
- return a.key == b.key && a.slot == b.slot;
347
- }
348
+ bool operator()(key_and_slot_t const& a, key_and_slot_t const& b) const noexcept { return a.key == b.key; }
348
349
  };
349
350
 
350
351
  /// @brief Multi-Map from keys to IDs, and allocated vectors.
351
- tsl::robin_set<key_and_slot_t, lookup_key_hash_t, lookup_key_same_t> slot_lookup_;
352
+ std::unordered_multiset<key_and_slot_t, lookup_key_hash_t, lookup_key_same_t> slot_lookup_;
352
353
 
353
354
  /// @brief Mutex, controlling concurrent access to `slot_lookup_`.
354
355
  mutable shared_mutex_t slot_lookup_mutex_;
@@ -364,6 +365,7 @@ class index_dense_gt {
364
365
 
365
366
  public:
366
367
  using search_result_t = typename index_t::search_result_t;
368
+ using cluster_result_t = typename index_t::cluster_result_t;
367
369
  using add_result_t = typename index_t::add_result_t;
368
370
  using stats_t = typename index_t::stats_t;
369
371
  using match_t = typename index_t::match_t;
@@ -422,13 +424,13 @@ class index_dense_gt {
422
424
  * @brief Constructs an instance of ::index_dense_gt.
423
425
  * @param[in] metric One of the provided or an @b ad-hoc metric, type-punned.
424
426
  * @param[in] config The index configuration (optional).
425
- * @param[in] free_label The key used for freed vectors (optional).
427
+ * @param[in] free_key The key used for freed vectors (optional).
426
428
  * @return An instance of ::index_dense_gt.
427
429
  */
428
430
  static index_dense_gt make( //
429
431
  metric_t metric, //
430
432
  index_dense_config_t config = {}, //
431
- key_t free_label = default_free_value<key_t>()) {
433
+ key_t free_key = default_free_value<key_t>()) {
432
434
 
433
435
  scalar_kind_t scalar_kind = metric.scalar_kind();
434
436
  std::size_t hardware_threads = std::thread::hardware_concurrency();
@@ -438,7 +440,7 @@ class index_dense_gt {
438
440
  result.cast_buffer_.resize(hardware_threads * metric.bytes_per_vector());
439
441
  result.casts_ = make_casts_(scalar_kind);
440
442
  result.metric_ = metric;
441
- result.free_key_ = free_label;
443
+ result.free_key_ = free_key;
442
444
 
443
445
  // Fill the thread IDs.
444
446
  result.available_threads_.resize(hardware_threads);
@@ -473,6 +475,7 @@ class index_dense_gt {
473
475
  std::size_t max_level() const noexcept { return typed_->max_level(); }
474
476
  index_dense_config_t const& config() const { return config_; }
475
477
  index_limits_t const& limits() const { return typed_->limits(); }
478
+ bool multi() const { return config_.multi; }
476
479
 
477
480
  // The metric and its properties
478
481
  metric_t const& metric() const { return metric_; }
@@ -499,6 +502,9 @@ class index_dense_gt {
499
502
  stats_t stats() const { return typed_->stats(); }
500
503
  stats_t stats(std::size_t level) const { return typed_->stats(level); }
501
504
 
505
+ dynamic_allocator_t const& allocator() const { return typed_->dynamic_allocator(); }
506
+ key_t const& free_key() const { return free_key_; }
507
+
502
508
  /**
503
509
  * @brief A relatively accurate lower bound on the amount of memory consumed by the system.
504
510
  * In practice it's error will be below 10%.
@@ -513,40 +519,146 @@ class index_dense_gt {
513
519
  vectors_tape_allocator_.total_allocated();
514
520
  }
515
521
 
522
+ static constexpr std::size_t any_thread() { return std::numeric_limits<std::size_t>::max(); }
523
+ static constexpr distance_t infinite_distance() { return std::numeric_limits<distance_t>::max(); }
524
+
525
+ struct aggregated_distances_t {
526
+ std::size_t count = 0;
527
+ distance_t mean = infinite_distance();
528
+ distance_t min = infinite_distance();
529
+ distance_t max = infinite_distance();
530
+ };
531
+
516
532
  // clang-format off
517
- add_result_t add(key_t key, b1x8_t const* vector) { return add_(key, vector, casts_.from_b1x8); }
518
- add_result_t add(key_t key, i8_bits_t const* vector) { return add_(key, vector, casts_.from_i8); }
519
- add_result_t add(key_t key, f16_t const* vector) { return add_(key, vector, casts_.from_f16); }
520
- add_result_t add(key_t key, f32_t const* vector) { return add_(key, vector, casts_.from_f32); }
521
- add_result_t add(key_t key, f64_t const* vector) { return add_(key, vector, casts_.from_f64); }
522
-
523
- add_result_t add(key_t key, b1x8_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_b1x8); }
524
- add_result_t add(key_t key, i8_bits_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_i8); }
525
- add_result_t add(key_t key, f16_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_f16); }
526
- add_result_t add(key_t key, f32_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_f32); }
527
- add_result_t add(key_t key, f64_t const* vector, index_dense_update_config_t config) { return add_(key, vector, config, casts_.from_f64); }
528
-
529
- search_result_t search(b1x8_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_b1x8); }
530
- search_result_t search(i8_bits_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_i8); }
531
- search_result_t search(f16_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_f16); }
532
- search_result_t search(f32_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_f32); }
533
- search_result_t search(f64_t const* vector, std::size_t wanted) const { return search_(vector, wanted, casts_.from_f64); }
534
-
535
- search_result_t search(b1x8_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_b1x8); }
536
- search_result_t search(i8_bits_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_i8); }
537
- search_result_t search(f16_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_f16); }
538
- search_result_t search(f32_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_f32); }
539
- search_result_t search(f64_t const* vector, std::size_t wanted, index_search_config_t config) const { return search_(vector, wanted, config, casts_.from_f64); }
540
-
541
- bool get(key_t key, b1x8_t* vector) const { return get_(key, vector, casts_.to_b1x8); }
542
- bool get(key_t key, i8_bits_t* vector) const { return get_(key, vector, casts_.to_i8); }
543
- bool get(key_t key, f16_t* vector) const { return get_(key, vector, casts_.to_f16); }
544
- bool get(key_t key, f32_t* vector) const { return get_(key, vector, casts_.to_f32); }
545
- bool get(key_t key, f64_t* vector) const { return get_(key, vector, casts_.to_f64); }
533
+ add_result_t add(key_t key, b1x8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_b1x8); }
534
+ add_result_t add(key_t key, i8_bits_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_i8); }
535
+ add_result_t add(key_t key, f16_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f16); }
536
+ add_result_t add(key_t key, f32_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f32); }
537
+ add_result_t add(key_t key, f64_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f64); }
538
+
539
+ search_result_t search(b1x8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_b1x8); }
540
+ search_result_t search(i8_bits_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_i8); }
541
+ search_result_t search(f16_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f16); }
542
+ search_result_t search(f32_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f32); }
543
+ search_result_t search(f64_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f64); }
544
+
545
+ bool get(key_t key, b1x8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_b1x8); }
546
+ bool get(key_t key, i8_bits_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_i8); }
547
+ bool get(key_t key, f16_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f16); }
548
+ bool get(key_t key, f32_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f32); }
549
+ bool get(key_t key, f64_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f64); }
550
+
551
+ cluster_result_t cluster(b1x8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_b1x8); }
552
+ cluster_result_t cluster(i8_bits_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_i8); }
553
+ cluster_result_t cluster(f16_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f16); }
554
+ cluster_result_t cluster(f32_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f32); }
555
+ cluster_result_t cluster(f64_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f64); }
556
+
557
+ aggregated_distances_t distance_between(key_t key, b1x8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_b1x8); }
558
+ aggregated_distances_t distance_between(key_t key, i8_bits_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_i8); }
559
+ aggregated_distances_t distance_between(key_t key, f16_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f16); }
560
+ aggregated_distances_t distance_between(key_t key, f32_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f32); }
561
+ aggregated_distances_t distance_between(key_t key, f64_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f64); }
546
562
  // clang-format on
547
563
 
548
564
  /**
549
- * @brief Reserves memory for the index and the labeled lookup.
565
+ * @brief Computes the distance between two managed entities.
566
+ * If either key maps into more than one vector, will aggregate results
567
+ * exporting the mean, maximum, and minimum values.
568
+ */
569
+ aggregated_distances_t distance_between(key_t a, key_t b, std::size_t = any_thread()) const {
570
+ shared_lock_t lock(slot_lookup_mutex_);
571
+ aggregated_distances_t result;
572
+ if (!multi()) {
573
+ auto a_it = slot_lookup_.find(key_and_slot_t::any_slot(a));
574
+ auto b_it = slot_lookup_.find(key_and_slot_t::any_slot(b));
575
+ bool a_missing = a_it == slot_lookup_.end();
576
+ bool b_missing = b_it == slot_lookup_.end();
577
+ if (a_missing || b_missing)
578
+ return result;
579
+
580
+ key_and_slot_t a_key_and_slot = *a_it;
581
+ byte_t const* a_vector = vectors_lookup_[a_key_and_slot.slot];
582
+ key_and_slot_t b_key_and_slot = *b_it;
583
+ byte_t const* b_vector = vectors_lookup_[b_key_and_slot.slot];
584
+ distance_t a_b_distance = metric_(a_vector, b_vector);
585
+
586
+ result.mean = result.min = result.max = a_b_distance;
587
+ result.count = 1;
588
+ return result;
589
+ }
590
+
591
+ auto a_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(a));
592
+ auto b_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(b));
593
+ bool a_missing = a_range.first == a_range.second;
594
+ bool b_missing = b_range.first == b_range.second;
595
+ if (a_missing || b_missing)
596
+ return result;
597
+
598
+ result.min = std::numeric_limits<distance_t>::max();
599
+ result.max = std::numeric_limits<distance_t>::min();
600
+ result.mean = 0;
601
+ result.count = 0;
602
+
603
+ while (a_range.first != a_range.second) {
604
+ key_and_slot_t a_key_and_slot = *a_range.first;
605
+ byte_t const* a_vector = vectors_lookup_[a_key_and_slot.slot];
606
+ while (b_range.first != b_range.second) {
607
+ key_and_slot_t b_key_and_slot = *b_range.first;
608
+ byte_t const* b_vector = vectors_lookup_[b_key_and_slot.slot];
609
+ distance_t a_b_distance = metric_(a_vector, b_vector);
610
+
611
+ result.mean += a_b_distance;
612
+ result.min = (std::min)(result.min, a_b_distance);
613
+ result.max = (std::max)(result.max, a_b_distance);
614
+ result.count++;
615
+
616
+ //
617
+ ++b_range.first;
618
+ }
619
+ ++a_range.first;
620
+ }
621
+
622
+ result.mean /= result.count;
623
+ return result;
624
+ }
625
+
626
+ /**
627
+ * @brief Identifies a node in a given `level`, that is the closest to the `key`.
628
+ */
629
+ cluster_result_t cluster(key_t key, std::size_t level, std::size_t thread = any_thread()) const {
630
+
631
+ // Check if such `key` is even present.
632
+ shared_lock_t slots_lock(slot_lookup_mutex_);
633
+ auto key_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
634
+ cluster_result_t result;
635
+ if (key_range.first == key_range.second)
636
+ return result.failed("Key missing!");
637
+
638
+ index_cluster_config_t cluster_config;
639
+ thread_lock_t lock = thread_lock_(thread);
640
+ cluster_config.thread = lock.thread_id;
641
+ cluster_config.expansion = config_.expansion_search;
642
+ metric_proxy_t metric{*this};
643
+ auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
644
+
645
+ // Find the closest cluster for any vector under that key.
646
+ while (key_range.first != key_range.second) {
647
+ key_and_slot_t key_and_slot = *key_range.first;
648
+ byte_t const* vector_data = vectors_lookup_[key_and_slot.slot];
649
+ cluster_result_t new_result = typed_->cluster(vector_data, level, metric, cluster_config, allow);
650
+ if (!new_result)
651
+ return new_result;
652
+ if (new_result.cluster.distance < result.cluster.distance)
653
+ result = std::move(new_result);
654
+
655
+ ++key_range.first;
656
+ }
657
+ return result;
658
+ }
659
+
660
+ /**
661
+ * @brief Reserves memory for the index and the keyed lookup.
550
662
  * @return `true` if the memory reservation was successful, `false` otherwise.
551
663
  */
552
664
  bool reserve(index_limits_t limits) {
@@ -592,7 +704,10 @@ class index_dense_gt {
592
704
  vectors_lookup_.clear();
593
705
  free_keys_.clear();
594
706
  vectors_tape_allocator_.reset();
595
- available_threads_.clear();
707
+
708
+ // Reset the thread IDs.
709
+ available_threads_.resize(std::thread::hardware_concurrency());
710
+ std::iota(available_threads_.begin(), available_threads_.end(), 0ul);
596
711
  }
597
712
 
598
713
  /**
@@ -617,7 +732,7 @@ class index_dense_gt {
617
732
  * @brief Saves serialized binary index representation to a stream.
618
733
  */
619
734
  template <typename output_callback_at, typename progress_at = dummy_progress_t>
620
- serialization_result_t stream(output_callback_at&& callback, serialization_config_t config = {}) const noexcept {
735
+ serialization_result_t stream(output_callback_at&& callback, serialization_config_t config = {}) const {
621
736
 
622
737
  serialization_result_t result;
623
738
  std::uint64_t matrix_rows = 0;
@@ -674,6 +789,7 @@ class index_dense_gt {
674
789
  head.count_present = size();
675
790
  head.count_deleted = typed_->size() - size();
676
791
  head.dimensions = dimensions();
792
+ head.multi = multi();
677
793
 
678
794
  if (!callback(&buffer, sizeof(buffer)))
679
795
  return result.failed("Failed to serialize into stream");
@@ -762,6 +878,7 @@ class index_dense_gt {
762
878
  return result.failed("Slot type doesn't match, consider rebuilding");
763
879
 
764
880
  metric_ = metric_t(head.dimensions, head.kind_metric, head.kind_scalar);
881
+ config_.multi = head.multi;
765
882
  }
766
883
 
767
884
  // Pull the actual proximity graph
@@ -771,7 +888,7 @@ class index_dense_gt {
771
888
  if (typed_->size() != static_cast<std::size_t>(matrix_rows))
772
889
  return result.failed("Index size and the number of vectors doesn't match");
773
890
 
774
- reindex_labels_();
891
+ reindex_keys_();
775
892
  return result;
776
893
  }
777
894
 
@@ -838,6 +955,7 @@ class index_dense_gt {
838
955
  return result.failed("Slot type doesn't match, consider rebuilding");
839
956
 
840
957
  metric_ = metric_t(head.dimensions, head.kind_metric, head.kind_scalar);
958
+ config_.multi = head.multi;
841
959
  offset += sizeof(buffer);
842
960
  }
843
961
 
@@ -854,7 +972,7 @@ class index_dense_gt {
854
972
  for (std::uint64_t slot = 0; slot != matrix_rows; ++slot)
855
973
  vectors_lookup_[slot] = (byte_t*)vectors_buffer.data() + matrix_cols * slot;
856
974
 
857
- reindex_labels_();
975
+ reindex_keys_();
858
976
  return result;
859
977
  }
860
978
 
@@ -864,7 +982,7 @@ class index_dense_gt {
864
982
  */
865
983
  bool contains(key_t key) const {
866
984
  shared_lock_t lock(slot_lookup_mutex_);
867
- return slot_lookup_.contains(key);
985
+ return slot_lookup_.find(key_and_slot_t::any_slot(key)) != slot_lookup_.end();
868
986
  }
869
987
 
870
988
  /**
@@ -873,7 +991,7 @@ class index_dense_gt {
873
991
  */
874
992
  std::size_t count(key_t key) const {
875
993
  shared_lock_t lock(slot_lookup_mutex_);
876
- return slot_lookup_.count(key);
994
+ return slot_lookup_.count(key_and_slot_t::any_slot(key));
877
995
  }
878
996
 
879
997
  struct labeling_result_t {
@@ -899,65 +1017,70 @@ class index_dense_gt {
899
1017
  labeling_result_t result;
900
1018
 
901
1019
  unique_lock_t lookup_lock(slot_lookup_mutex_);
902
- auto labeled_iterator = slot_lookup_.find(key);
903
- if (labeled_iterator == slot_lookup_.end())
1020
+ auto matching_slots = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
1021
+ if (matching_slots.first == matching_slots.second)
904
1022
  return result;
905
1023
 
906
1024
  // Grow the removed entries ring, if needed
1025
+ std::size_t matching_count = std::distance(matching_slots.first, matching_slots.second);
907
1026
  std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
908
- if (free_keys_.size() == free_keys_.capacity())
909
- if (!free_keys_.reserve((std::max<std::size_t>)(free_keys_.capacity() * 2, 64ul)))
910
- return result.failed("Can't allocate memory for a free-list");
1027
+ if (!free_keys_.reserve(free_keys_.size() + matching_count))
1028
+ return result.failed("Can't allocate memory for a free-list");
911
1029
 
912
1030
  // A removed entry would be:
913
1031
  // - present in `free_keys_`
914
1032
  // - missing in the `slot_lookup_`
915
1033
  // - marked in the `typed_` index with a `free_key_`
916
- compressed_slot_t slot = (*labeled_iterator).slot;
917
- free_keys_.push(slot);
918
- slot_lookup_.erase(labeled_iterator);
919
- typed_->at(slot).key = free_key_;
920
- result.completed = true;
1034
+ for (auto slots_it = matching_slots.first; slots_it != matching_slots.second; ++slots_it) {
1035
+ compressed_slot_t slot = (*slots_it).slot;
1036
+ free_keys_.push(slot);
1037
+ typed_->at(slot).key = free_key_;
1038
+ }
1039
+ slot_lookup_.erase(matching_slots.first, matching_slots.second);
1040
+ result.completed = matching_count;
921
1041
 
922
1042
  return result;
923
1043
  }
924
1044
 
925
1045
  /**
926
1046
  * @brief Removes multiple entries with the specified keys from the index.
927
- * @param[in] labels_begin The beginning of the keys range.
928
- * @param[in] labels_end The ending of the keys range.
1047
+ * @param[in] keys_begin The beginning of the keys range.
1048
+ * @param[in] keys_end The ending of the keys range.
929
1049
  * @return The ::labeling_result_t indicating the result of the removal operation.
930
1050
  * `result.completed` will contain the number of keys that were successfully removed.
931
1051
  * `result.error` will contain an error message if an error occurred during the removal operation.
932
1052
  */
933
- template <typename labels_iterator_at>
934
- labeling_result_t remove(labels_iterator_at&& labels_begin, labels_iterator_at&& labels_end) {
1053
+ template <typename keys_iterator_at>
1054
+ labeling_result_t remove(keys_iterator_at keys_begin, keys_iterator_at keys_end) {
935
1055
 
936
1056
  labeling_result_t result;
937
1057
  unique_lock_t lookup_lock(slot_lookup_mutex_);
938
1058
  std::unique_lock<std::mutex> free_lock(free_keys_mutex_);
939
-
940
1059
  // Grow the removed entries ring, if needed
941
- std::size_t count_requests = std::distance(labels_begin, labels_end);
942
- if (!free_keys_.reserve(free_keys_.size() + count_requests))
1060
+ std::size_t matching_count = 0;
1061
+ for (auto keys_it = keys_begin; keys_it != keys_end; ++keys_it)
1062
+ matching_count += slot_lookup_.count(key_and_slot_t::any_slot(*keys_it));
1063
+
1064
+ if (!free_keys_.reserve(free_keys_.size() + matching_count))
943
1065
  return result.failed("Can't allocate memory for a free-list");
944
1066
 
945
1067
  // Remove them one-by-one
946
- for (auto label_it = labels_begin; label_it != labels_end; ++label_it) {
947
- key_t key = *label_it;
948
- auto labeled_iterator = slot_lookup_.find(key);
949
- if (labeled_iterator == slot_lookup_.end())
950
- continue;
951
-
1068
+ for (auto keys_it = keys_begin; keys_it != keys_end; ++keys_it) {
1069
+ key_t key = *keys_it;
1070
+ auto matching_slots = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
952
1071
  // A removed entry would be:
953
1072
  // - present in `free_keys_`
954
1073
  // - missing in the `slot_lookup_`
955
1074
  // - marked in the `typed_` index with a `free_key_`
956
- compressed_slot_t slot = (*labeled_iterator).slot;
957
- free_keys_.push(slot);
958
- slot_lookup_.erase(labeled_iterator);
959
- typed_->at(slot).key = free_key_;
960
- result.completed += 1;
1075
+ for (auto slots_it = matching_slots.first; slots_it != matching_slots.second; ++slots_it) {
1076
+ compressed_slot_t slot = (*slots_it).slot;
1077
+ free_keys_.push(slot);
1078
+ typed_->at(slot).key = free_key_;
1079
+ }
1080
+
1081
+ matching_count = std::distance(matching_slots.first, matching_slots.second);
1082
+ slot_lookup_.erase(matching_slots.first, matching_slots.second);
1083
+ result.completed += matching_count;
961
1084
  }
962
1085
 
963
1086
  return result;
@@ -974,16 +1097,24 @@ class index_dense_gt {
974
1097
  labeling_result_t rename(key_t from, key_t to) {
975
1098
  labeling_result_t result;
976
1099
  unique_lock_t lookup_lock(slot_lookup_mutex_);
977
- auto labeled_iterator = slot_lookup_.find(from);
978
- if (labeled_iterator == slot_lookup_.end())
979
- return result;
980
1100
 
981
- compressed_slot_t slot = (*labeled_iterator).slot;
982
- key_and_slot_t key_and_slot{to, slot};
983
- slot_lookup_.erase(labeled_iterator);
984
- slot_lookup_.insert(key_and_slot);
985
- typed_->at(slot).key = to;
986
- result.completed = true;
1101
+ if (!multi() && slot_lookup_.count(key_and_slot_t::any_slot(to)))
1102
+ return result.failed("Renaming impossible, the key is already in use");
1103
+
1104
+ // The `from` may map to multiple entries
1105
+ while (true) {
1106
+ auto slots_it = slot_lookup_.find(key_and_slot_t::any_slot(from));
1107
+ if (slots_it == slot_lookup_.end())
1108
+ break;
1109
+
1110
+ compressed_slot_t slot = (*slots_it).slot;
1111
+ key_and_slot_t key_and_slot{to, slot};
1112
+ slot_lookup_.erase(slots_it);
1113
+ slot_lookup_.insert(key_and_slot);
1114
+ typed_->at(slot).key = to;
1115
+ ++result.completed;
1116
+ }
1117
+
987
1118
  return result;
988
1119
  }
989
1120
 
@@ -1147,11 +1278,11 @@ class index_dense_gt {
1147
1278
  return result;
1148
1279
  }
1149
1280
 
1150
- template < //
1151
- typename man_to_woman_at = dummy_label_to_label_mapping_t, //
1152
- typename woman_to_man_at = dummy_label_to_label_mapping_t, //
1153
- typename executor_at = dummy_executor_t, //
1154
- typename progress_at = dummy_progress_t //
1281
+ template < //
1282
+ typename man_to_woman_at = dummy_key_to_key_mapping_t, //
1283
+ typename woman_to_man_at = dummy_key_to_key_mapping_t, //
1284
+ typename executor_at = dummy_executor_t, //
1285
+ typename progress_at = dummy_progress_t //
1155
1286
  >
1156
1287
  join_result_t join( //
1157
1288
  index_dense_gt const& women, //
@@ -1173,20 +1304,215 @@ class index_dense_gt {
1173
1304
  std::forward<progress_at>(progress));
1174
1305
  }
1175
1306
 
1307
+ struct clustering_result_t {
1308
+ error_t error{};
1309
+ std::size_t clusters{};
1310
+ std::size_t visited_members{};
1311
+ std::size_t computed_distances{};
1312
+
1313
+ explicit operator bool() const noexcept { return !error; }
1314
+ clustering_result_t failed(error_t message) noexcept {
1315
+ error = std::move(message);
1316
+ return std::move(*this);
1317
+ }
1318
+ };
1319
+
1320
+ /**
1321
+ * @brief Implements clustering, classifying the given objects (vectors of member keys)
1322
+ * into a given number of clusters.
1323
+ *
1324
+ * @param[in] queries_begin Iterator targeting the fiest query.
1325
+ * @param[in] queries_end
1326
+ * @param[in] executor Thread-pool to execute the job in parallel.
1327
+ * @param[in] progress Callback to report the execution progress.
1328
+ */
1329
+ template < //
1330
+ typename queries_iterator_at, //
1331
+ typename executor_at = dummy_executor_t, //
1332
+ typename progress_at = dummy_progress_t //
1333
+ >
1334
+ clustering_result_t cluster( //
1335
+ queries_iterator_at queries_begin, //
1336
+ queries_iterator_at queries_end, //
1337
+ index_dense_clustering_config_t config, //
1338
+ key_t* cluster_keys, //
1339
+ distance_t* cluster_distances, //
1340
+ executor_at&& executor = executor_at{}, //
1341
+ progress_at&& progress = progress_at{}) {
1342
+
1343
+ std::size_t const queries_count = queries_end - queries_begin;
1344
+
1345
+ // Find the first level (top -> down) that has enough nodes to exceed `config.min_clusters`.
1346
+ std::size_t level = max_level();
1347
+ if (config.min_clusters) {
1348
+ for (; level > 1; --level) {
1349
+ if (stats(level).nodes > config.min_clusters)
1350
+ break;
1351
+ }
1352
+ } else
1353
+ level = 1, config.max_clusters = stats(1).nodes, config.min_clusters = 2;
1354
+
1355
+ clustering_result_t result;
1356
+ if (max_level() < 2)
1357
+ return result.failed("Index too small to cluster!");
1358
+
1359
+ // A structure used to track the popularity of a specific cluster
1360
+ struct cluster_t {
1361
+ key_t centroid;
1362
+ key_t merged_into;
1363
+ std::size_t popularity;
1364
+ byte_t* vector;
1365
+ };
1366
+
1367
+ auto centroid_id = [](cluster_t const& a, cluster_t const& b) { return a.centroid < b.centroid; };
1368
+ auto higher_popularity = [](cluster_t const& a, cluster_t const& b) { return a.popularity > b.popularity; };
1369
+
1370
+ std::atomic<std::size_t> visited_members(0);
1371
+ std::atomic<std::size_t> computed_distances(0);
1372
+ std::atomic<char const*> atomic_error{nullptr};
1373
+
1374
+ using dynamic_allocator_traits_t = std::allocator_traits<dynamic_allocator_t>;
1375
+ using clusters_allocator_t = typename dynamic_allocator_traits_t::template rebind_alloc<cluster_t>;
1376
+ buffer_gt<cluster_t, clusters_allocator_t> clusters(queries_count);
1377
+ if (!clusters)
1378
+ return result.failed("Out of memory!");
1379
+
1380
+ map_to_clusters:
1381
+ // Concurrently perform search until a certain depth
1382
+ executor.dynamic(queries_count, [&](std::size_t thread_idx, std::size_t query_idx) {
1383
+ auto result = cluster(queries_begin[query_idx], level, thread_idx);
1384
+ if (!result) {
1385
+ atomic_error = result.error.release();
1386
+ return false;
1387
+ }
1388
+
1389
+ cluster_keys[query_idx] = result.cluster.member.key;
1390
+ cluster_distances[query_idx] = result.cluster.distance;
1391
+
1392
+ // Export in case we need to refine afterwards
1393
+ clusters[query_idx].centroid = result.cluster.member.key;
1394
+ clusters[query_idx].vector = vectors_lookup_[result.cluster.member.slot];
1395
+ clusters[query_idx].merged_into = free_key();
1396
+ clusters[query_idx].popularity = 1;
1397
+
1398
+ visited_members += result.visited_members;
1399
+ computed_distances += result.computed_distances;
1400
+ return true;
1401
+ });
1402
+
1403
+ if (atomic_error)
1404
+ return result.failed(atomic_error.load());
1405
+
1406
+ // Now once we have identified the closest clusters,
1407
+ // we can try reducing their quantity, refining
1408
+ std::sort(clusters.begin(), clusters.end(), centroid_id);
1409
+
1410
+ // Transform into run-length encoding, computing the number of unique clusters
1411
+ std::size_t unique_clusters = 0;
1412
+ {
1413
+ std::size_t last_idx = 0;
1414
+ for (std::size_t current_idx = 1; current_idx != clusters.size(); ++current_idx) {
1415
+ if (clusters[last_idx].centroid == clusters[current_idx].centroid) {
1416
+ clusters[last_idx].popularity++;
1417
+ } else {
1418
+ last_idx++;
1419
+ clusters[last_idx] = clusters[current_idx];
1420
+ }
1421
+ }
1422
+ unique_clusters = last_idx + 1;
1423
+ }
1424
+
1425
+ // In some cases the queries may be co-located, all mapping into the same cluster on that
1426
+ // level. In that case we refine the granularity and dive deeper into clusters:
1427
+ if (unique_clusters < config.min_clusters && level > 1) {
1428
+ level--;
1429
+ goto map_to_clusters;
1430
+ }
1431
+
1432
+ std::sort(clusters.data(), clusters.data() + unique_clusters, higher_popularity);
1433
+
1434
+ // If clusters are too numerous, merge the ones that are too close to each other.
1435
+ std::size_t merge_cycles = 0;
1436
+ merge_nearby_clusters:
1437
+ if (unique_clusters > config.max_clusters) {
1438
+
1439
+ cluster_t& merge_source = clusters[unique_clusters - 1];
1440
+ std::size_t merge_target_idx = 0;
1441
+ distance_t merge_distance = std::numeric_limits<distance_t>::max();
1442
+
1443
+ for (std::size_t candidate_idx = 0; candidate_idx + 1 < unique_clusters; ++candidate_idx) {
1444
+ distance_t distance = metric_(merge_source.vector, clusters[candidate_idx].vector);
1445
+ if (distance < merge_distance) {
1446
+ merge_distance = distance;
1447
+ merge_target_idx = candidate_idx;
1448
+ }
1449
+ }
1450
+
1451
+ merge_source.merged_into = clusters[merge_target_idx].centroid;
1452
+ clusters[merge_target_idx].popularity += exchange(merge_source.popularity, 0);
1453
+
1454
+ // The target object may have to be swapped a few times to get to optimal position.
1455
+ while (merge_target_idx &&
1456
+ clusters[merge_target_idx - 1].popularity < clusters[merge_target_idx].popularity)
1457
+ std::swap(clusters[merge_target_idx - 1], clusters[merge_target_idx]), --merge_target_idx;
1458
+
1459
+ unique_clusters--;
1460
+ merge_cycles++;
1461
+ goto merge_nearby_clusters;
1462
+ }
1463
+
1464
+ // Replace evicted clusters
1465
+ if (merge_cycles) {
1466
+ // Sort dropped clusters by name to accelerate future lookups
1467
+ auto clusters_end = clusters.data() + config.max_clusters + merge_cycles;
1468
+ std::sort(clusters.data(), clusters_end, centroid_id);
1469
+
1470
+ executor.dynamic(queries_count, [&](std::size_t thread_idx, std::size_t query_idx) {
1471
+ key_t& cluster_key = cluster_keys[query_idx];
1472
+ distance_t& cluster_distance = cluster_distances[query_idx];
1473
+
1474
+ // Recursively trace replacements of that cluster
1475
+ while (true) {
1476
+ // To avoid implementing heterogeneous comparisons, lets wrap the `cluster_key`
1477
+ cluster_t updated_cluster;
1478
+ updated_cluster.centroid = cluster_key;
1479
+ updated_cluster = *std::lower_bound(clusters.data(), clusters_end, updated_cluster, centroid_id);
1480
+ if (updated_cluster.merged_into == free_key())
1481
+ break;
1482
+ cluster_key = updated_cluster.merged_into;
1483
+ }
1484
+
1485
+ cluster_distance = distance_between(cluster_key, queries_begin[query_idx], thread_idx).mean;
1486
+ return true;
1487
+ });
1488
+ }
1489
+
1490
+ result.computed_distances = computed_distances;
1491
+ result.visited_members = visited_members;
1492
+ return result;
1493
+ }
1494
+
1176
1495
  private:
1177
1496
  struct thread_lock_t {
1178
1497
  index_dense_gt const& parent;
1179
1498
  std::size_t thread_id;
1499
+ bool engaged;
1180
1500
 
1181
- ~thread_lock_t() { parent.thread_unlock_(thread_id); }
1501
+ ~thread_lock_t() {
1502
+ if (engaged)
1503
+ parent.thread_unlock_(thread_id);
1504
+ }
1182
1505
  };
1183
1506
 
1184
- thread_lock_t thread_lock_() const {
1507
+ thread_lock_t thread_lock_(std::size_t thread_id) const {
1508
+ if (thread_id != any_thread())
1509
+ return {*this, thread_id, false};
1510
+
1185
1511
  available_threads_mutex_.lock();
1186
- std::size_t thread_id = available_threads_.back();
1512
+ thread_id = available_threads_.back();
1187
1513
  available_threads_.pop_back();
1188
1514
  available_threads_mutex_.unlock();
1189
- return {*this, thread_id};
1515
+ return {*this, thread_id, true};
1190
1516
  }
1191
1517
 
1192
1518
  void thread_unlock_(std::size_t thread_id) const {
@@ -1196,16 +1522,19 @@ class index_dense_gt {
1196
1522
  }
1197
1523
 
1198
1524
  template <typename scalar_at>
1199
- add_result_t add_(key_t key, scalar_at const* vector, index_dense_update_config_t config, cast_t const& cast) {
1525
+ add_result_t add_( //
1526
+ key_t key, scalar_at const* vector, //
1527
+ std::size_t thread, bool force_vector_copy, cast_t const& cast) {
1200
1528
 
1201
- if (!config_.allow_key_collisions && contains(key))
1529
+ if (!multi() && contains(key))
1202
1530
  return add_result_t{}.failed("Duplicate keys not allowed in high-level wrappers");
1203
1531
 
1204
1532
  // Cast the vector, if needed for compatibility with `metric_`
1205
- bool copy_vector = !config_.exclude_vectors || config.force_vector_copy;
1533
+ thread_lock_t lock = thread_lock_(thread);
1534
+ bool copy_vector = !config_.exclude_vectors || force_vector_copy;
1206
1535
  byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
1207
1536
  {
1208
- byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * config.thread;
1537
+ byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
1209
1538
  bool casted = cast(vector_data, dimensions(), casted_data);
1210
1539
  if (casted)
1211
1540
  vector_data = casted_data, copy_vector = true;
@@ -1222,7 +1551,7 @@ class index_dense_gt {
1222
1551
  bool reuse_node = free_slot != default_free_value<compressed_slot_t>();
1223
1552
  auto on_success = [&](member_ref_t member) {
1224
1553
  unique_lock_t slot_lock(slot_lookup_mutex_);
1225
- slot_lookup_.insert({key, static_cast<compressed_slot_t>(member.slot)});
1554
+ slot_lookup_.insert(key_and_slot_t{key, static_cast<compressed_slot_t>(member.slot)});
1226
1555
  if (copy_vector) {
1227
1556
  if (!reuse_node)
1228
1557
  vectors_lookup_[member.slot] = vectors_tape_allocator_.allocate(metric_.bytes_per_vector());
@@ -1231,28 +1560,107 @@ class index_dense_gt {
1231
1560
  vectors_lookup_[member.slot] = (byte_t*)vector_data;
1232
1561
  };
1233
1562
 
1563
+ index_update_config_t update_config;
1564
+ update_config.thread = lock.thread_id;
1565
+ update_config.expansion = config_.expansion_add;
1566
+
1234
1567
  metric_proxy_t metric{*this};
1235
1568
  return reuse_node //
1236
- ? typed_->update(typed_->iterator_at(free_slot), key, vector_data, metric, config, on_success)
1237
- : typed_->add(key, vector_data, metric, config, on_success);
1569
+ ? typed_->update(typed_->iterator_at(free_slot), key, vector_data, metric, update_config, on_success)
1570
+ : typed_->add(key, vector_data, metric, update_config, on_success);
1238
1571
  }
1239
1572
 
1240
1573
  template <typename scalar_at>
1241
1574
  search_result_t search_( //
1242
1575
  scalar_at const* vector, std::size_t wanted, //
1243
- index_search_config_t config, cast_t const& cast) const {
1576
+ std::size_t thread, bool exact, cast_t const& cast) const {
1244
1577
 
1245
1578
  // Cast the vector, if needed for compatibility with `metric_`
1579
+ thread_lock_t lock = thread_lock_(thread);
1246
1580
  byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
1247
1581
  {
1248
- byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * config.thread;
1582
+ byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
1249
1583
  bool casted = cast(vector_data, dimensions(), casted_data);
1250
1584
  if (casted)
1251
1585
  vector_data = casted_data;
1252
1586
  }
1253
1587
 
1588
+ index_search_config_t search_config;
1589
+ search_config.thread = lock.thread_id;
1590
+ search_config.expansion = config_.expansion_search;
1591
+ search_config.exact = exact;
1592
+
1254
1593
  auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
1255
- return typed_->search(vector_data, wanted, metric_proxy_t{*this}, config, allow);
1594
+ return typed_->search(vector_data, wanted, metric_proxy_t{*this}, search_config, allow);
1595
+ }
1596
+
1597
+ template <typename scalar_at>
1598
+ cluster_result_t cluster_( //
1599
+ scalar_at const* vector, std::size_t level, //
1600
+ std::size_t thread, cast_t const& cast) const {
1601
+
1602
+ // Cast the vector, if needed for compatibility with `metric_`
1603
+ thread_lock_t lock = thread_lock_(thread);
1604
+ byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
1605
+ {
1606
+ byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
1607
+ bool casted = cast(vector_data, dimensions(), casted_data);
1608
+ if (casted)
1609
+ vector_data = casted_data;
1610
+ }
1611
+
1612
+ index_cluster_config_t cluster_config;
1613
+ cluster_config.thread = lock.thread_id;
1614
+ cluster_config.expansion = config_.expansion_search;
1615
+
1616
+ auto allow = [=](member_cref_t const& member) noexcept { return member.key != free_key_; };
1617
+ return typed_->cluster(vector_data, level, metric_proxy_t{*this}, cluster_config, allow);
1618
+ }
1619
+
1620
+ template <typename scalar_at>
1621
+ aggregated_distances_t distance_between_( //
1622
+ key_t key, scalar_at const* vector, //
1623
+ std::size_t thread, cast_t const& cast) const {
1624
+
1625
+ // Cast the vector, if needed for compatibility with `metric_`
1626
+ thread_lock_t lock = thread_lock_(thread);
1627
+ byte_t const* vector_data = reinterpret_cast<byte_t const*>(vector);
1628
+ {
1629
+ byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id;
1630
+ bool casted = cast(vector_data, dimensions(), casted_data);
1631
+ if (casted)
1632
+ vector_data = casted_data;
1633
+ }
1634
+
1635
+ // Check if such `key` is even present.
1636
+ shared_lock_t slots_lock(slot_lookup_mutex_);
1637
+ auto key_range = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
1638
+ aggregated_distances_t result;
1639
+ if (key_range.first == key_range.second)
1640
+ return result;
1641
+
1642
+ result.min = std::numeric_limits<distance_t>::max();
1643
+ result.max = std::numeric_limits<distance_t>::min();
1644
+ result.mean = 0;
1645
+ result.count = 0;
1646
+
1647
+ while (key_range.first != key_range.second) {
1648
+ key_and_slot_t key_and_slot = *key_range.first;
1649
+ byte_t const* a_vector = vectors_lookup_[key_and_slot.slot];
1650
+ byte_t const* b_vector = vector_data;
1651
+ distance_t a_b_distance = metric_(a_vector, b_vector);
1652
+
1653
+ result.mean += a_b_distance;
1654
+ result.min = (std::min)(result.min, a_b_distance);
1655
+ result.max = (std::max)(result.max, a_b_distance);
1656
+ result.count++;
1657
+
1658
+ //
1659
+ ++key_range.first;
1660
+ }
1661
+
1662
+ result.mean /= result.count;
1663
+ return result;
1256
1664
  }
1257
1665
 
1258
1666
  compressed_slot_t lookup_id_(key_t key) const {
@@ -1260,7 +1668,7 @@ class index_dense_gt {
1260
1668
  return slot_lookup_.at(key);
1261
1669
  }
1262
1670
 
1263
- void reindex_labels_() {
1671
+ void reindex_keys_() {
1264
1672
 
1265
1673
  // Estimate number of entries first
1266
1674
  std::size_t count_total = typed_->size();
@@ -1286,41 +1694,41 @@ class index_dense_gt {
1286
1694
  }
1287
1695
  }
1288
1696
 
1289
- template <typename scalar_at> bool get_(key_t key, scalar_at* reconstructed, cast_t const& cast) const {
1290
- compressed_slot_t slot;
1291
- // Find the matching ID
1292
- {
1697
+ template <typename scalar_at>
1698
+ std::size_t get_(key_t key, scalar_at* reconstructed, std::size_t vectors_limit, cast_t const& cast) const {
1699
+
1700
+ if (!multi()) {
1701
+ compressed_slot_t slot;
1702
+ // Find the matching ID
1703
+ {
1704
+ shared_lock_t lock(slot_lookup_mutex_);
1705
+ auto it = slot_lookup_.find(key_and_slot_t::any_slot(key));
1706
+ if (it == slot_lookup_.end())
1707
+ return false;
1708
+ slot = (*it).slot;
1709
+ }
1710
+ // Export the entry
1711
+ byte_t const* punned_vector = reinterpret_cast<byte_t const*>(vectors_lookup_[slot]);
1712
+ bool casted = cast(punned_vector, dimensions(), (byte_t*)reconstructed);
1713
+ if (!casted)
1714
+ std::memcpy(reconstructed, punned_vector, metric_.bytes_per_vector());
1715
+ return true;
1716
+ } else {
1293
1717
  shared_lock_t lock(slot_lookup_mutex_);
1294
- auto it = slot_lookup_.find(key);
1295
- if (it == slot_lookup_.end())
1296
- return false;
1297
- slot = (*it).slot;
1718
+ auto equal_range_pair = slot_lookup_.equal_range(key_and_slot_t::any_slot(key));
1719
+ std::size_t count_exported = 0;
1720
+ for (auto begin = equal_range_pair.first;
1721
+ begin != equal_range_pair.second && count_exported != vectors_limit; ++begin, ++count_exported) {
1722
+ //
1723
+ compressed_slot_t slot = (*begin).slot;
1724
+ byte_t const* punned_vector = reinterpret_cast<byte_t const*>(vectors_lookup_[slot]);
1725
+ byte_t* reconstructed_vector = (byte_t*)reconstructed + metric_.bytes_per_vector() * count_exported;
1726
+ bool casted = cast(punned_vector, dimensions(), reconstructed_vector);
1727
+ if (!casted)
1728
+ std::memcpy(reconstructed_vector, punned_vector, metric_.bytes_per_vector());
1729
+ }
1730
+ return count_exported;
1298
1731
  }
1299
- // Export the entry
1300
- byte_t const* punned_vector = reinterpret_cast<byte_t const*>(vectors_lookup_[slot]);
1301
- bool casted = cast(punned_vector, dimensions(), (byte_t*)reconstructed);
1302
- if (!casted)
1303
- std::memcpy(reconstructed, punned_vector, metric_.bytes_per_vector());
1304
- return true;
1305
- }
1306
-
1307
- template <typename scalar_at> add_result_t add_(key_t key, scalar_at const* vector, cast_t const& cast) {
1308
- thread_lock_t lock = thread_lock_();
1309
- index_dense_update_config_t update_config;
1310
- update_config.thread = lock.thread_id;
1311
- update_config.expansion = config_.expansion_add;
1312
- return add_(key, vector, update_config, cast);
1313
- }
1314
-
1315
- template <typename scalar_at>
1316
- search_result_t search_( //
1317
- scalar_at const* vector, std::size_t wanted, //
1318
- cast_t const& cast) const {
1319
- thread_lock_t lock = thread_lock_();
1320
- index_search_config_t search_config;
1321
- search_config.thread = lock.thread_id;
1322
- search_config.expansion = config_.expansion_search;
1323
- return search_(vector, wanted, search_config, cast);
1324
1732
  }
1325
1733
 
1326
1734
  template <typename to_scalar_at> static casts_t make_casts_() {
@@ -1368,19 +1776,19 @@ using index_dense_big_t = index_dense_gt<uuid_t, uint40_t>;
1368
1776
  */
1369
1777
  template < //
1370
1778
 
1371
- typename men_label_at, //
1372
- typename women_label_at, //
1373
- typename men_slot_at, //
1374
- typename women_slot_at, //
1779
+ typename men_key_at, //
1780
+ typename women_key_at, //
1781
+ typename men_slot_at, //
1782
+ typename women_slot_at, //
1375
1783
 
1376
- typename man_to_woman_at = dummy_label_to_label_mapping_t, //
1377
- typename woman_to_man_at = dummy_label_to_label_mapping_t, //
1378
- typename executor_at = dummy_executor_t, //
1379
- typename progress_at = dummy_progress_t //
1784
+ typename man_to_woman_at = dummy_key_to_key_mapping_t, //
1785
+ typename woman_to_man_at = dummy_key_to_key_mapping_t, //
1786
+ typename executor_at = dummy_executor_t, //
1787
+ typename progress_at = dummy_progress_t //
1380
1788
  >
1381
- static join_result_t join( //
1382
- index_dense_gt<men_label_at, men_slot_at> const& men, //
1383
- index_dense_gt<women_label_at, women_slot_at> const& women, //
1789
+ static join_result_t join( //
1790
+ index_dense_gt<men_key_at, men_slot_at> const& men, //
1791
+ index_dense_gt<women_key_at, women_slot_at> const& women, //
1384
1792
 
1385
1793
  index_join_config_t config = {}, //
1386
1794
  man_to_woman_at&& man_to_woman = man_to_woman_at{}, //