usearch 0.20.0 → 0.21.0

package/include/usearch/index_punned_dense.hpp

@@ -170,6 +170,7 @@ class index_punned_dense_gt {
     using search_result_t = typename index_t::search_result_t;
     using add_result_t = typename index_t::add_result_t;
     using serialization_result_t = typename index_t::serialization_result_t;
+    using join_result_t = typename index_t::join_result_t;
     using stats_t = typename index_t::stats_t;
 
     index_punned_dense_gt() = default;
@@ -184,12 +185,15 @@ class index_punned_dense_gt {
     void swap(index_punned_dense_gt& other) {
         std::swap(dimensions_, other.dimensions_);
         std::swap(scalar_words_, other.scalar_words_);
+        std::swap(expansion_add_, other.expansion_add_);
+        std::swap(expansion_search_, other.expansion_search_);
         std::swap(casted_vector_bytes_, other.casted_vector_bytes_);
         std::swap(typed_, other.typed_);
         std::swap(cast_buffer_, other.cast_buffer_);
         std::swap(casts_, other.casts_);
         std::swap(root_metric_, other.root_metric_);
         std::swap(available_threads_, other.available_threads_);
+        std::swap(lookup_table_, other.lookup_table_);
     }
 
     static index_config_t optimize(index_config_t config) { return index_t::optimize(config); }
@@ -199,6 +203,7 @@ class index_punned_dense_gt {
     std::size_t connectivity() const { return typed_->connectivity(); }
     std::size_t size() const { return typed_->size(); }
     std::size_t capacity() const { return typed_->capacity(); }
+    std::size_t max_level() const noexcept { return typed_->max_level(); }
     index_config_t const& config() const { return typed_->config(); }
     index_limits_t const& limits() const { return typed_->limits(); }
     void clear() { return typed_->clear(); }
@@ -323,23 +328,76 @@ class index_punned_dense_gt {
             metric_t(metric), make_casts_(accuracy));
     }
 
-    index_punned_dense_gt fork() const {
-        index_punned_dense_gt result;
+    struct copy_result_t {
+        index_punned_dense_gt index;
+        error_t error;
+
+        explicit operator bool() const noexcept { return !error; }
+        copy_result_t failed(error_t message) noexcept {
+            error = std::move(message);
+            return std::move(*this);
+        }
+    };
+
+    copy_result_t copy(copy_config_t config = {}) const {
+        copy_result_t result = fork();
+        if (!result)
+            return result;
+        auto typed_result = typed_->copy(config);
+        if (!typed_result)
+            return result.failed(std::move(typed_result.error));
+        *result.index.typed_ = std::move(typed_result.index);
+        result.index.lookup_table_ = lookup_table_;
+        return result;
+    }
 
-        result.dimensions_ = dimensions_;
-        result.scalar_words_ = scalar_words_;
-        result.casted_vector_bytes_ = casted_vector_bytes_;
-        result.cast_buffer_ = cast_buffer_;
-        result.casts_ = casts_;
+    copy_result_t fork() const {
+        copy_result_t result;
+        index_punned_dense_gt& other = result.index;
+
+        other.dimensions_ = dimensions_;
+        other.scalar_words_ = scalar_words_;
+        other.expansion_add_ = expansion_add_;
+        other.expansion_search_ = expansion_search_;
+        other.casted_vector_bytes_ = casted_vector_bytes_;
+        other.cast_buffer_ = cast_buffer_;
+        other.casts_ = casts_;
+
+        other.root_metric_ = root_metric_;
+        other.available_threads_ = available_threads_;
 
-        result.root_metric_ = root_metric_;
         index_t* raw = index_allocator_t{}.allocate(1);
-        new (raw) index_t(config(), root_metric_);
-        result.typed_ = raw;
+        if (!raw)
+            return result.failed("Can't allocate the index");
 
+        new (raw) index_t(config(), root_metric_);
+        other.typed_ = raw;
         return result;
     }
 
+    template <                                                        //
+        typename first_to_second_at = dummy_label_to_label_mapping_t, //
+        typename second_to_first_at = dummy_label_to_label_mapping_t, //
+        typename executor_at = dummy_executor_t,                      //
+        typename progress_at = dummy_progress_t                       //
+        >
+    static join_result_t join(                                       //
+        index_punned_dense_gt const& first,                          //
+        index_punned_dense_gt const& second,                         //
+        join_config_t config = {},                                   //
+        first_to_second_at&& first_to_second = first_to_second_at{}, //
+        second_to_first_at&& second_to_first = second_to_first_at{}, //
+        executor_at&& executor = executor_at{},                      //
+        progress_at&& progress = progress_at{}) noexcept {
+
+        return index_t::join(                                  //
+            *first.typed_, *second.typed_, config,             //
+            std::forward<first_to_second_at>(first_to_second), //
+            std::forward<second_to_first_at>(second_to_first), //
+            std::forward<executor_at>(executor),               //
+            std::forward<progress_at>(progress));
+    }
+
   private:
     struct thread_lock_t {
         index_punned_dense_gt const& parent;

package/include/usearch/index_punned_helpers.hpp

@@ -313,9 +313,11 @@ class executor_stl_t {
     std::size_t threads_count_{};
 
   public:
-    executor_stl_t(std::size_t threads_count) noexcept
+    executor_stl_t(std::size_t threads_count = 0) noexcept
         : threads_count_(threads_count ? threads_count : std::thread::hardware_concurrency()) {}
 
+    std::size_t size() const noexcept { return threads_count_; }
+
     template <typename thread_aware_function_at>
     void execute_bulk(std::size_t tasks, thread_aware_function_at&& thread_aware_function) noexcept(false) {
         std::vector<std::thread> threads_pool;
@@ -330,22 +332,39 @@ class executor_stl_t {
         for (std::size_t thread_idx = 0; thread_idx != threads_count_; ++thread_idx)
             threads_pool[thread_idx].join();
     }
+
+    template <typename thread_aware_function_at>
+    void execute_bulk(thread_aware_function_at&& thread_aware_function) noexcept(false) {
+        std::vector<std::thread> threads_pool;
+        for (std::size_t thread_idx = 0; thread_idx != threads_count_; ++thread_idx)
+            threads_pool.emplace_back([=]() { thread_aware_function(thread_idx); });
+        for (std::size_t thread_idx = 0; thread_idx != threads_count_; ++thread_idx)
+            threads_pool[thread_idx].join();
+    }
 };
 
 #if USEARCH_USE_OPENMP
 
 class executor_openmp_t {
   public:
-    executor_openmp_t(std::size_t threads_count) noexcept {
+    executor_openmp_t(std::size_t threads_count = 0) noexcept {
         omp_set_num_threads(threads_count ? threads_count : std::thread::hardware_concurrency());
     }
 
+    std::size_t size() const noexcept { return omp_get_num_threads(); }
+
     template <typename thread_aware_function_at>
     void execute_bulk(std::size_t tasks, thread_aware_function_at&& thread_aware_function) noexcept(false) {
 #pragma omp parallel for schedule(dynamic)
-        for (std::size_t i = 0; i < tasks; ++i)
+        for (std::size_t i = 0; i != tasks; ++i)
             thread_aware_function(omp_get_thread_num(), i);
     }
+
+    template <typename thread_aware_function_at>
+    void execute_bulk(thread_aware_function_at&& thread_aware_function) noexcept(false) {
+#pragma omp parallel
+        { thread_aware_function(omp_get_thread_num()); }
+    }
 };
 
 using executor_default_t = executor_openmp_t;
@@ -451,6 +470,12 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
         return *this;
     }
 
+    memory_mapping_allocator_gt(memory_mapping_allocator_gt const&) noexcept {}
+    memory_mapping_allocator_gt& operator=(memory_mapping_allocator_gt const&) noexcept {
+        reset();
+        return *this;
+    }
+
     ~memory_mapping_allocator_gt() noexcept { reset(); }
 
     inline byte_t* allocate(std::size_t count_bytes) noexcept {

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "usearch",
-    "version": "0.20.0",
+    "version": "0.21.0",
     "description": "Smaller & Faster Single-File Vector Search Engine from Unum",
     "author": "Ashot Vardanian",
     "license": "Apache 2.0",

package/python/lib.cpp

@@ -18,6 +18,7 @@
 #define PY_SSIZE_T_CLEAN
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include <usearch/index_punned_dense.hpp>
 
@@ -356,7 +357,6 @@ static py::tuple search_many_in_index( //
 
     Py_ssize_t vectors_count = vectors_info.shape[0];
     Py_ssize_t vectors_dimensions = vectors_info.shape[1];
-    char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
     if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
         throw std::invalid_argument("The number of vector dimensions doesn't match!");
 
@@ -381,10 +381,48 @@ static py::tuple search_many_in_index( //
     return results;
 }
 
+static std::unordered_map<label_t, label_t> join_index(   //
+    dense_index_py_t const& a, dense_index_py_t const& b, //
+    std::size_t max_proposals, bool exact) {
+
+    std::unordered_map<label_t, label_t> a_to_b;
+    a_to_b.reserve((std::min)(a.size(), b.size()));
+
+    using join_result_t = typename dense_index_py_t::join_result_t;
+    join_config_t config;
+
+    config.max_proposals = max_proposals;
+    config.exact = exact;
+    config.expansion = (std::max)(a.expansion_search(), b.expansion_search());
+    std::size_t threads = (std::min)(a.limits().threads(), b.limits().threads());
+    executor_default_t executor{threads};
+    join_result_t result = dense_index_py_t::join( //
+        a, b, config,                              //
+        a_to_b,                                    //
+        dummy_label_to_label_mapping_t{},          //
+        executor);
+    result.error.raise();
+    return a_to_b;
+}
+
+static dense_index_py_t copy_index(dense_index_py_t const& index) {
+
+    using copy_result_t = typename dense_index_py_t::copy_result_t;
+    copy_config_t config;
+
+    copy_result_t result = index.copy();
+    result.error.raise();
+    return std::move(result.index);
+}
+
 // clang-format off
 template <typename index_at> void save_index(index_at const& index, std::string const& path) { index.save(path.c_str()).error.raise(); }
 template <typename index_at> void load_index(index_at& index, std::string const& path) { index.load(path.c_str()).error.raise(); }
 template <typename index_at> void view_index(index_at& index, std::string const& path) { index.view(path.c_str()).error.raise(); }
+template <typename index_at> void clear_index(index_at& index) { index.clear(); }
+template <typename index_at> std::size_t max_level(index_at const &index) { return index.max_level(); }
+template <typename index_at> typename index_at::stats_t compute_stats(index_at const &index) { return index.stats(); }
+template <typename index_at> typename index_at::stats_t compute_level_stats(index_at const &index, std::size_t level) { return index.stats(level); }
 // clang-format on
 
 template <typename internal_at, typename external_at = internal_at, typename index_at = void>
@@ -650,7 +688,21 @@ PYBIND11_MODULE(compiled, m) {
     i.def("save", &save_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
     i.def("load", &load_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
     i.def("view", &view_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
-    i.def("clear", &dense_index_py_t::clear, py::call_guard<py::gil_scoped_release>());
+    i.def("clear", &clear_index<dense_index_py_t>, py::call_guard<py::gil_scoped_release>());
+    i.def("copy", &copy_index, py::call_guard<py::gil_scoped_release>());
+    i.def("join", &join_index, py::arg("other"), py::arg("max_proposals") = 0, py::arg("exact") = false,
+          py::call_guard<py::gil_scoped_release>());
+
+    using punned_index_stats_t = typename dense_index_py_t::stats_t;
+    auto i_stats = py::class_<punned_index_stats_t>(m, "IndexStats");
+    i_stats.def_readonly("nodes", &punned_index_stats_t::nodes);
+    i_stats.def_readonly("edges", &punned_index_stats_t::edges);
+    i_stats.def_readonly("max_edges", &punned_index_stats_t::max_edges);
+    i_stats.def_readonly("allocated_bytes", &punned_index_stats_t::allocated_bytes);
+
+    i.def_property_readonly("max_level", &max_level<dense_index_py_t>);
+    i.def_property_readonly("levels_stats", &compute_stats<dense_index_py_t>);
+    i.def("level_stats", &compute_level_stats<dense_index_py_t>, py::arg("level"));
 
     auto si = py::class_<sparse_index_py_t>(m, "SparseIndex");
 
@@ -704,8 +756,8 @@ PYBIND11_MODULE(compiled, m) {
     si.def_property_readonly("connectivity", &sparse_index_py_t::connectivity);
     si.def_property_readonly("capacity", &sparse_index_py_t::capacity);
 
-    si.def("save", &sparse_index_py_t::save, py::arg("path"));
-    si.def("load", &sparse_index_py_t::load, py::arg("path"));
-    si.def("view", &sparse_index_py_t::view, py::arg("path"));
+    si.def("save", &save_index<sparse_index_py_t>, py::arg("path"));
+    si.def("load", &load_index<sparse_index_py_t>, py::arg("path"));
+    si.def("view", &view_index<sparse_index_py_t>, py::arg("path"));
     si.def("clear", &sparse_index_py_t::clear);
 }

package/python/scripts/join.py

@@ -0,0 +1,151 @@
+"""
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clip_images.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clip_texts.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/images.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/texts.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/uform_english_images.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/uform_english_texts.fbin -P datasets/cc_3M/
+rm -rf datasets/cc_3M/*.usearch
+python python/scripts/join.py
+"""
+from numpy import dot
+from numpy.linalg import norm
+from tqdm import tqdm
+from simsimd import cos_f32x4_neon, to_int
+
+from usearch.index import Index, MetricKind, CompiledMetric, MetricSignature
+from usearch.io import load_matrix
+from usearch.eval import measure_seconds
+
+k = 10
+exact = False
+batch_size = 1024 * 4
+max_elements = 1000000
+
+a_name = "cc_3M/texts"
+b_name = "cc_3M/images"
+
+a_mat = load_matrix(f"datasets/{a_name}.fbin", view=True)
+b_mat = load_matrix(f"datasets/{b_name}.fbin", view=True)
+
+a_mat = a_mat[:max_elements]
+b_mat = b_mat[:max_elements]
+
+print(f"Loaded two datasets of shape: {a_mat.shape}, {b_mat.shape}")
+print("--------------------------------------")
+print("---------------Indexing---------------")
+print("--------------------------------------")
+
+metric = CompiledMetric(
+    pointer=to_int(cos_f32x4_neon),
+    kind=MetricKind.Cos,
+    signature=MetricSignature.ArrayArraySize,
+)
+
+a = Index(
+    a_mat.shape[1],
+    metric=metric,
+    path=f"datasets/{a_name}.f32.usearch",
+    dtype="f32",
+)
+b = Index(
+    b_mat.shape[1],
+    metric=metric,
+    path=f"datasets/{b_name}.f32.usearch",
+    dtype="f32",
+)
+
+if len(a) != a_mat.shape[0]:
+    a.clear()
+    a.add(None, a_mat, log=True, batch_size=batch_size)
+    a.save()
+
+if len(b) != b_mat.shape[0]:
+    b.clear()
+    b.add(None, b_mat, log=True, batch_size=batch_size)
+    b.save()
+
+
+print(
+    f"Loaded two indexes of size: {len(a):,} for {a_name} and {len(b):,} for {b_name}"
+)
+min_elements = min(len(a), len(b))
+
+run_diagnostics = input("Would you like to run diagnostics? [Y/n]: ")
+if len(run_diagnostics) == 0 or run_diagnostics.lower() == "y":
+    print("--------------------------------------")
+    print("-------------Diagnostics--------------")
+    print("--------------------------------------")
+
+    mean_similarity = 0.0
+    mean_recovered_similarity = 0.0
+
+    for i in tqdm(range(min_elements), desc="Pairwise Similarity"):
+        a_vec = a_mat[i]
+        b_vec = b_mat[i]
+        cos_similarity = dot(a_vec, b_vec) / (norm(a_vec) * norm(b_vec))
+        mean_similarity += cos_similarity
+
+        a_vec = a[i]
+        b_vec = b[i]
+        cos_similarity = dot(a_vec, b_vec) / (norm(a_vec) * norm(b_vec))
+        mean_recovered_similarity += cos_similarity
+
+    mean_similarity /= min_elements
+    mean_recovered_similarity /= min_elements
+    print(
+        f"Average vector similarity is {mean_similarity:.4f} in original dataset, "
+        f"and {mean_recovered_similarity:.4f} in recovered state in index"
+    )
+
+    dt = measure_seconds
+    args = dict(
+        k=k,
+        batch_size=batch_size,
+        log=True,
+        exact=exact,
+    )
+
+    secs, a_self_recall = dt(lambda: a.search(a.vectors, **args).recall(a.labels))
+    print(
+        "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
+            k, a_name, a_self_recall * 100, secs
+        )
+    )
+
+    secs, b_self_recall = dt(lambda: b.search(b.vectors, **args).recall(b.labels))
+    print(
+        "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
+            k, b_name, b_self_recall * 100, secs
+        )
+    )
+
+    secs, ab_recall = dt(lambda: b.search(a.vectors, **args).recall(b.labels))
+    print(
+        "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
+            k, a_name, b_name, ab_recall * 100, secs
+        )
+    )
+
+    secs, ba_recall = dt(lambda: a.search(b.vectors, **args).recall(a.labels))
+    print(
+        "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
+            k, b_name, a_name, ba_recall * 100, secs
+        )
+    )
+
+
+print("--------------------------------------")
+print("-----------------Join-----------------")
+print("--------------------------------------")
+
+secs, bimapping = measure_seconds(lambda: a.join(b, max_proposals=100))
+mapping_size = len(bimapping)
+recall = 0
+for i, j in bimapping.items():
+    recall += i == j
+
+recall *= 100.0 / min_elements
+print(
+    f"Took {secs:.2f}s to find {mapping_size:,} pairings with {recall:.2f}% being exact"
+)

package/python/scripts/test.py

@@ -115,6 +115,11 @@ def test_index(
     assert matches[0] == 42
     assert distances[0] == pytest.approx(0, abs=1e-3)
 
+    assert index.max_level >= 0
+    assert index.levels_stats.nodes >= 1
+    assert index.level_stats(0).nodes == 1
+    assert str(index).startswith("usearch.index.Index")
+
     index.save("tmp.usearch")
     index.clear()
     assert len(index) == 0
@@ -127,6 +132,10 @@ def test_index(
     assert len(index) == 1
     assert len(index[42]) == ndim
 
+    index_copy = index.copy()
+    assert len(index_copy) == 1
+    assert len(index_copy[42]) == ndim
+
     # Cleanup
     os.remove("tmp.usearch")
 
@@ -157,6 +166,10 @@ def test_index_batch(
     assert matches.counts.shape[0] == batch_size
     assert np.all(np.sort(index.labels) == np.sort(labels))
 
+    assert index.max_level >= 0  # TODO: This should be 1
+    assert index.levels_stats.nodes >= batch_size
+    assert index.level_stats(0).nodes == batch_size
+
     index.save("tmp.usearch")
     index.clear()
     assert len(index) == 0
@@ -205,6 +218,12 @@ def test_exact_recall(
     for i in range(batch_size):
         assert found_labels[i, 0] == i
 
+    # Match entries aginst themselves
+    index_copy: Index = index.copy()
+    mapping: dict = index.join(index_copy, exact=True)
+    for man, woman in mapping.items():
+        assert man == woman, "Stable marriage failed"
+
 
 @pytest.mark.parametrize("ndim", dimensions)
 @pytest.mark.parametrize("batch_size", batch_sizes)

package/python/usearch/eval.py

@@ -3,6 +3,7 @@ from time import time_ns
 from typing import Tuple, Any, Callable, Union, Optional, List
 from dataclasses import dataclass, asdict
 from collections import defaultdict
+from math import ceil
 
 import numpy as np
 
@@ -47,7 +48,7 @@ def random_vectors(
         return x
 
 
-def recall_members(index: Index, *args, **kwargs) -> float:
+def recall_members(index: Index, sample: float = 1, **kwargs) -> float:
     """Simplest benchmark for a quality of search, which queries every
     existing member of the index, to make sure approximate search finds
     the point itself.
@@ -59,9 +60,16 @@ def recall_members(index: Index, *args, **kwargs) -> float:
     """
     if len(index) == 0:
         return 0
+    if "k" not in kwargs:
+        kwargs["k"] = 1
 
-    matches: Matches = index.search(index.vectors, 1, *args, **kwargs)
-    return matches.recall_first(index.labels)
+    labels = index.labels
+    if sample != 1:
+        labels = np.random.choice(labels, int(ceil(len(labels) * sample)))
+
+    queries = index.get_vectors(labels, index.dtype)
+    matches: Matches = index.search(queries, **kwargs)
+    return matches.recall_first(labels)
 
 
 def measure_seconds(f: Callable) -> Tuple[float, Any]: