usearch 1.1.1 → 1.2.1

package/python/lib.cpp

@@ -12,7 +12,10 @@
  *
  *  @copyright Copyright (c) 2023
  */
+#if !defined(__cpp_exceptions)
 #define __cpp_exceptions 1
+#endif
+
 #include <limits> // `std::numeric_limits`
 #include <thread> // `std::thread`
 
@@ -41,13 +44,16 @@ enum class metric_signature_t {
 
 namespace py = pybind11;
 using py_shape_t = py::array::ShapeContainer;
-#define key_t typename index_dense_t::key_t
+
 using metric_t = metric_punned_t;
 using distance_t = distance_punned_t;
 
+using dense_key_t = typename index_dense_t::key_t;
 using dense_add_result_t = typename index_dense_t::add_result_t;
 using dense_search_result_t = typename index_dense_t::search_result_t;
 using dense_labeling_result_t = typename index_dense_t::labeling_result_t;
+using dense_cluster_result_t = typename index_dense_t::cluster_result_t;
+using dense_clustering_result_t = typename index_dense_t::clustering_result_t;
 
 struct dense_index_py_t : public index_dense_t {
     using native_t = index_dense_t;
@@ -63,26 +69,27 @@ struct dense_index_py_t : public index_dense_t {
 struct dense_indexes_py_t {
     std::vector<std::shared_ptr<dense_index_py_t>> shards_;
 
-    void add(std::shared_ptr<dense_index_py_t> shard) { shards_.push_back(shard); }
+    void merge(std::shared_ptr<dense_index_py_t> shard) { shards_.push_back(shard); }
     std::size_t bytes_per_vector() const noexcept { return shards_.empty() ? 0 : shards_[0]->bytes_per_vector(); }
     std::size_t scalar_words() const noexcept { return shards_.empty() ? 0 : shards_[0]->scalar_words(); }
     index_limits_t limits() const noexcept { return {size(), std::numeric_limits<std::size_t>::max()}; }
 
-    void add_paths(std::vector<std::string> const& paths, bool view = true, std::size_t threads = 0) {
+    void merge_paths(std::vector<std::string> const& paths, bool view = true, std::size_t threads = 0) {
         if (!threads)
             threads = std::thread::hardware_concurrency();
 
         shards_.reserve(shards_.size() + paths.size());
         std::mutex shards_mutex;
-        executor_default_t{threads}.execute_bulk(paths.size(), [&](std::size_t, std::size_t task_idx) {
+        executor_default_t{threads}.dynamic(paths.size(), [&](std::size_t, std::size_t task_idx) {
             index_dense_t index = index_dense_t::make(paths[task_idx].c_str(), view);
             if (!index)
-                return;
+                return false;
             auto shared_index = std::make_shared<dense_index_py_t>(std::move(index));
             std::unique_lock<std::mutex> lock(shards_mutex);
             shards_.push_back(shared_index);
             if (PyErr_CheckSignals() != 0)
                 throw py::error_already_set();
+            return true;
         });
     }
 
@@ -95,7 +102,7 @@ struct dense_indexes_py_t {
 };
 
 template <typename scalar_at>
-metric_t typed_udf(                                                                  //
+metric_t wrap_typed_user_defined_metric(                                             //
     metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, //
     scalar_kind_t scalar_kind, std::size_t dimensions) {
     //
@@ -129,16 +136,21 @@ metric_t typed_udf(
     return metric_t(stl_function, dimensions, kind, scalar_kind);
 }
 
-metric_t udf(                                                                        //
+metric_t wrap_user_defined_metric(                                                   //
     metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, //
     scalar_kind_t scalar_kind, std::size_t dimensions) {
 
     switch (scalar_kind) {
-    case scalar_kind_t::b1x8_k: return typed_udf<b1x8_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
-    case scalar_kind_t::i8_k: return typed_udf<i8_bits_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
-    case scalar_kind_t::f16_k: return typed_udf<f16_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
-    case scalar_kind_t::f32_k: return typed_udf<f32_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
-    case scalar_kind_t::f64_k: return typed_udf<f64_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
+    case scalar_kind_t::b1x8_k:
+        return wrap_typed_user_defined_metric<b1x8_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
+    case scalar_kind_t::i8_k:
+        return wrap_typed_user_defined_metric<i8_bits_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
+    case scalar_kind_t::f16_k:
+        return wrap_typed_user_defined_metric<f16_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
+    case scalar_kind_t::f32_k:
+        return wrap_typed_user_defined_metric<f32_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
+    case scalar_kind_t::f64_k:
+        return wrap_typed_user_defined_metric<f64_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
     default: return {};
     }
 }
@@ -151,12 +163,15 @@ static dense_index_py_t make_index(      //
     std::size_t expansion_search,        //
     metric_kind_t metric_kind,           //
     metric_signature_t metric_signature, //
-    std::uintptr_t metric_uintptr) {
+    std::uintptr_t metric_uintptr,       //
+    bool multi) {
 
     index_dense_config_t config(connectivity, expansion_add, expansion_search);
+    config.multi = multi;
+
     metric_t metric =  //
         metric_uintptr //
-            ? udf(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
+            ? wrap_user_defined_metric(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
             : metric_t(dimensions, metric_kind, scalar_kind);
     return index_dense_t::make(metric, config);
 }
@@ -177,6 +192,18 @@ scalar_kind_t numpy_string_to_kind(std::string const& name) {
         return scalar_kind_t::unknown_k;
 }
 
+template <typename result_at> void forward_error(result_at&& result) {
+
+    if (!result)
+        throw std::invalid_argument(result.error.release());
+
+    int signals = PyErr_CheckSignals();
+    if (signals != 0)
+        throw py::error_already_set();
+}
+
+using atomic_error_t = std::atomic<char const*>;
+
 template <typename scalar_at>
 static void add_typed_to_index(                                            //
     dense_index_py_t& index,                                               //
@@ -186,18 +213,30 @@ static void add_typed_to_index(                                            //
     Py_ssize_t vectors_count = vectors_info.shape[0];
     byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);
     byte_t const* keys_data = reinterpret_cast<byte_t const*>(keys_info.ptr);
+    atomic_error_t atomic_error{nullptr};
 
-    executor_default_t{threads}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
-        index_dense_update_config_t config;
-        config.force_vector_copy = force_copy;
-        config.thread = thread_idx;
-        key_t key = *reinterpret_cast<key_t const*>(keys_data + task_idx * keys_info.strides[0]);
+    executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+        dense_key_t key = *reinterpret_cast<dense_key_t const*>(keys_data + task_idx * keys_info.strides[0]);
         scalar_at const* vector = reinterpret_cast<scalar_at const*>(vectors_data + task_idx * vectors_info.strides[0]);
-        dense_add_result_t result = index.add(key, vector, config);
-        result.error.raise();
-        if (PyErr_CheckSignals() != 0)
-            throw py::error_already_set();
+        dense_add_result_t result = index.add(key, vector, thread_idx, force_copy);
+        if (!result) {
+            atomic_error = result.error.release();
+            return false;
+        }
+
+        // We don't want to check for signals from multiple threads
+        if (thread_idx == 0)
+            if (PyErr_CheckSignals() != 0)
+                return false;
+        return true;
     });
+
+    // Raise the error from a single thread
+    auto error = atomic_error.load();
+    if (error) {
+        PyErr_SetString(PyExc_RuntimeError, error);
+        throw py::error_already_set();
+    }
 }
 
 template <typename index_at>
@@ -208,11 +247,11 @@ static void add_many_to_index(                            //
     py::buffer_info keys_info = keys.request();
     py::buffer_info vectors_info = vectors.request();
 
-    if (keys_info.itemsize != sizeof(key_t))
+    if (keys_info.itemsize != sizeof(dense_key_t))
         throw std::invalid_argument("Incompatible key type!");
 
     if (keys_info.ndim != 1)
-        throw std::invalid_argument("Labels must be placed in a single-dimensional array!");
+        throw std::invalid_argument("Keys must be placed in a single-dimensional array!");
 
     if (vectors_info.ndim != 2)
         throw std::invalid_argument("Expects a matrix of vectors to add!");
@@ -247,7 +286,7 @@ template <typename scalar_at>
 static void search_typed(                                   //
     dense_index_py_t& index, py::buffer_info& vectors_info, //
     std::size_t wanted, bool exact, std::size_t threads,    //
-    py::array_t<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
+    py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
     std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances) {
 
     auto keys_py2d = keys_py.template mutable_unchecked<2>();
@@ -262,28 +301,41 @@ static void search_typed(                                   //
     if (!index.reserve(index_limits_t(index.size(), threads)))
         throw std::invalid_argument("Out of memory!");
 
-    executor_default_t{threads}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
-        index_search_config_t config;
-        config.thread = thread_idx;
-        config.exact = exact;
+    atomic_error_t atomic_error{nullptr};
+    executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
         scalar_at const* vector = (scalar_at const*)(vectors_data + task_idx * vectors_info.strides[0]);
-        dense_search_result_t result = index.search(vector, wanted, config);
-        result.error.raise();
+        dense_search_result_t result = index.search(vector, wanted, thread_idx, exact);
+        if (!result) {
+            atomic_error = result.error.release();
+            return false;
+        }
+
         counts_py1d(task_idx) =
             static_cast<Py_ssize_t>(result.dump_to(&keys_py2d(task_idx, 0), &distances_py2d(task_idx, 0)));
 
         stats_visited_members += result.visited_members;
         stats_computed_distances += result.computed_distances;
-        if (PyErr_CheckSignals() != 0)
-            throw py::error_already_set();
+
+        // We don't want to check for signals from multiple threads
+        if (thread_idx == 0)
+            if (PyErr_CheckSignals() != 0)
+                return false;
+        return true;
     });
+
+    // Raise the error from a single thread
+    auto error = atomic_error.load();
+    if (error) {
+        PyErr_SetString(PyExc_RuntimeError, error);
+        throw py::error_already_set();
+    }
 }
 
 template <typename scalar_at>
 static void search_typed(                                       //
     dense_indexes_py_t& indexes, py::buffer_info& vectors_info, //
     std::size_t wanted, bool exact, std::size_t threads,        //
-    py::array_t<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
+    py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
     std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances) {
 
     auto keys_py2d = keys_py.template mutable_unchecked<2>();
@@ -298,28 +350,31 @@ static void search_typed(                                       //
     if (!threads)
         threads = std::thread::hardware_concurrency();
 
-    visits_bitset_t query_mutexes(static_cast<std::size_t>(vectors_count));
+    bitset_t query_mutexes(static_cast<std::size_t>(vectors_count));
     if (!query_mutexes)
         throw std::bad_alloc();
 
-    executor_default_t{threads}.execute_bulk(indexes.shards_.size(), [&](std::size_t, std::size_t task_idx) {
+    atomic_error_t atomic_error{nullptr};
+    executor_default_t{threads}.dynamic(indexes.shards_.size(), [&](std::size_t thread_idx, std::size_t task_idx) {
         dense_index_py_t& index = *indexes.shards_[task_idx].get();
 
         index_limits_t limits;
         limits.members = index.size();
         limits.threads_add = 0;
         limits.threads_search = 1;
-        if (!index.reserve(limits))
-            throw std::bad_alloc();
-
-        index_search_config_t config;
-        config.thread = 0;
-        config.exact = exact;
+        if (!index.reserve(limits)) {
+            atomic_error = "Out of memory!";
+            return false;
+        }
 
         for (std::size_t vector_idx = 0; vector_idx != static_cast<std::size_t>(vectors_count); ++vector_idx) {
             scalar_at const* vector = (scalar_at const*)(vectors_data + vector_idx * vectors_info.strides[0]);
-            dense_search_result_t result = index.search(vector, wanted, config);
-            result.error.raise();
+            dense_search_result_t result = index.search(vector, wanted, 0, exact);
+            if (!result) {
+                atomic_error = result.error.release();
+                return false;
+            }
+
             {
                 auto lock = query_mutexes.lock(vector_idx);
                 counts_py1d(vector_idx) = static_cast<Py_ssize_t>(result.merge_into( //
@@ -331,10 +386,21 @@ static void search_typed(                                       //
 
             stats_visited_members += result.visited_members;
             stats_computed_distances += result.computed_distances;
-            if (PyErr_CheckSignals() != 0)
-                throw py::error_already_set();
+
+            // We don't want to check for signals from multiple threads
+            if (thread_idx == 0)
+                if (PyErr_CheckSignals() != 0)
+                    return false;
+            return true;
         }
     });
+
+    // Raise the error from a single thread
+    auto error = atomic_error.load();
+    if (error) {
+        PyErr_SetString(PyExc_RuntimeError, error);
+        throw py::error_already_set();
+    }
 }
 
 /**
@@ -344,7 +410,9 @@ static void search_typed(                                       //
  *  @return Tuple with:
  *      1. matrix of neighbors,
  *      2. matrix of distances,
- *      3. array with match counts.
+ *      3. array with match counts,
+ *      4. number of visited nodes,
+ *      4. number of computed pairwise distances.
  */
 template <typename index_at>
 static py::tuple search_many_in_index( //
@@ -365,7 +433,7 @@ static py::tuple search_many_in_index( //
     if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
         throw std::invalid_argument("The number of vector dimensions doesn't match!");
 
-    py::array_t<key_t> keys_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
+    py::array_t<dense_key_t> keys_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
     py::array_t<distance_t> distances_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
     py::array_t<Py_ssize_t> counts_py(vectors_count);
     std::atomic<std::size_t> stats_visited_members(0);
@@ -395,7 +463,7 @@ template <typename scalar_at>
 static void search_typed_brute_force(                                //
     py::buffer_info& dataset_info, py::buffer_info& queries_info,    //
     std::size_t wanted, std::size_t threads, metric_t const& metric, //
-    py::array_t<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py) {
+    py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py) {
 
     auto keys_py2d = keys_py.template mutable_unchecked<2>();
     auto distances_py2d = distances_py.template mutable_unchecked<2>();
@@ -403,7 +471,6 @@ static void search_typed_brute_force(                                //
 
     std::size_t dataset_count = static_cast<std::size_t>(dataset_info.shape[0]);
     std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);
-    std::size_t dimensions = static_cast<std::size_t>(dataset_info.shape[1]);
 
     byte_t const* dataset_data = reinterpret_cast<byte_t const*>(dataset_info.ptr);
     byte_t const* queries_data = reinterpret_cast<byte_t const*>(queries_info.ptr);
@@ -414,11 +481,11 @@ static void search_typed_brute_force(                                //
         threads = std::thread::hardware_concurrency();
 
     std::size_t tasks_count = static_cast<std::size_t>(dataset_count * queries_count);
-    visits_bitset_t query_mutexes(static_cast<std::size_t>(queries_count));
+    bitset_t query_mutexes(static_cast<std::size_t>(queries_count));
     if (!query_mutexes)
         throw std::bad_alloc();
 
-    executor_default_t{threads}.execute_bulk(tasks_count, [&](std::size_t, std::size_t task_idx) {
+    executor_default_t{threads}.dynamic(tasks_count, [&](std::size_t thread_idx, std::size_t task_idx) {
         //
         std::size_t dataset_idx = task_idx / queries_count;
         std::size_t query_idx = task_idx % queries_count;
@@ -429,25 +496,28 @@ static void search_typed_brute_force(                                //
 
         {
             auto lock = query_mutexes.lock(query_idx);
-            key_t* keys = &keys_py2d(query_idx, 0);
+            dense_key_t* keys = &keys_py2d(query_idx, 0);
             distance_t* distances = &distances_py2d(query_idx, 0);
             std::size_t& matches = reinterpret_cast<std::size_t&>(counts_py1d(query_idx));
             if (matches == wanted)
                 if (distances[wanted - 1] <= distance)
-                    return;
+                    return true;
 
             std::size_t offset = std::lower_bound(distances, distances + matches, distance) - distances;
 
             std::size_t count_worse = matches - offset - (wanted == matches);
-            std::memmove(keys + offset + 1, keys + offset, count_worse * sizeof(key_t));
+            std::memmove(keys + offset + 1, keys + offset, count_worse * sizeof(dense_key_t));
             std::memmove(distances + offset + 1, distances + offset, count_worse * sizeof(distance_t));
-            keys[offset] = static_cast<key_t>(dataset_idx);
+            keys[offset] = static_cast<dense_key_t>(dataset_idx);
             distances[offset] = distance;
             matches += matches != wanted;
         }
 
-        if (PyErr_CheckSignals() != 0)
-            throw py::error_already_set();
+        // We don't want to check for signals from multiple threads
+        if (thread_idx == 0)
+            if (PyErr_CheckSignals() != 0)
+                return false;
+        return true;
     });
 }
 
@@ -478,14 +548,14 @@ static py::tuple search_many_brute_force(    //
     if (dataset_kind != queries_kind)
         throw std::invalid_argument("The types of vectors don't match!");
 
-    py::array_t<key_t> keys_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
+    py::array_t<dense_key_t> keys_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
     py::array_t<distance_t> distances_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
     py::array_t<Py_ssize_t> counts_py(dataset_count);
 
     std::size_t dimensions = static_cast<std::size_t>(queries_dimensions);
     metric_t metric =  //
         metric_uintptr //
-            ? udf(metric_kind, metric_signature, metric_uintptr, queries_kind, dimensions)
+            ? wrap_user_defined_metric(metric_kind, metric_signature, metric_uintptr, queries_kind, dimensions)
             : metric_t(dimensions, metric_kind, queries_kind);
 
     // clang-format off
@@ -508,12 +578,153 @@ static py::tuple search_many_brute_force(    //
     return results;
 }
 
-static std::unordered_map<key_t, key_t> join_index(       //
-    dense_index_py_t const& a, dense_index_py_t const& b, //
+template <typename scalar_at> struct rows_lookup_gt {
+    byte_t* data_;
+    std::size_t stride_;
+
+    rows_lookup_gt(void* data, std::size_t stride) noexcept : data_((byte_t*)data), stride_(stride) {}
+    scalar_at* operator[](std::size_t i) const noexcept { return reinterpret_cast<scalar_at*>(data_ + i * stride_); }
+    std::ptrdiff_t operator-(rows_lookup_gt const& other) const noexcept { return (data_ - other.data_) / stride_; }
+    rows_lookup_gt operator+(std::size_t n) const noexcept { return {data_ + stride_ * n, stride_}; }
+    template <typename other_scalar_at> rows_lookup_gt<other_scalar_at> as() const noexcept { return {data_, stride_}; }
+};
+
+/**
+ *  @param queries Matrix of vectors to search for.
+ *  @param count Number of clusters to produce.
+ *
+ *  @return Tuple with:
+ *      1. vector of cluster IDs,
+ *      2. vector of distances to those clusters,
+ *      3. array with match counts, set to all ones,
+ *      4. number of visited nodes,
+ *      4. number of computed pairwise distances.
+ */
+template <typename index_at>
+static py::tuple cluster_vectors(        //
+    index_at& index, py::buffer queries, //
+    std::size_t min_count, std::size_t max_count, std::size_t threads) {
+
+    if (index.limits().threads_search < threads)
+        throw std::invalid_argument("Can't use that many threads!");
+
+    py::buffer_info queries_info = queries.request();
+    if (queries_info.ndim != 2)
+        throw std::invalid_argument("Expects a matrix of queries to add!");
+
+    std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);
+    std::size_t queries_stride = static_cast<std::size_t>(queries_info.strides[0]);
+    std::size_t queries_dimensions = static_cast<std::size_t>(queries_info.shape[1]);
+    if (queries_dimensions != index.scalar_words())
+        throw std::invalid_argument("The number of vector dimensions doesn't match!");
+
+    py::array_t<dense_key_t> keys_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
+    py::array_t<distance_t> distances_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
+    dense_clustering_result_t cluster_result;
+    executor_default_t executor{threads};
+
+    auto keys_py2d = keys_py.template mutable_unchecked<2>();
+    auto distances_py2d = distances_py.template mutable_unchecked<2>();
+    dense_key_t* keys_ptr = reinterpret_cast<dense_key_t*>(&keys_py2d(0, 0));
+    distance_t* distances_ptr = reinterpret_cast<distance_t*>(&distances_py2d(0, 0));
+
+    index_dense_clustering_config_t config;
+    config.min_clusters = min_count;
+    config.max_clusters = max_count;
+
+    rows_lookup_gt<byte_t const> queries_begin(queries_info.ptr, queries_stride);
+    rows_lookup_gt<byte_t const> queries_end = queries_begin + queries_count;
+
+    // clang-format off
+    switch (numpy_string_to_kind(queries_info.format)) {
+    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;
+    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;
+    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;
+    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;
+    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;
+    default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + queries_info.format);
+    }
+    // clang-format on
+
+    cluster_result.error.raise();
+
+    // Those would be set to 1 for all entries, in case of success
+    py::array_t<Py_ssize_t> counts_py(queries_count);
+    auto counts_py1d = counts_py.template mutable_unchecked<1>();
+    for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)
+        counts_py1d(static_cast<Py_ssize_t>(query_idx)) = 1;
+
+    py::tuple results(5);
+    results[0] = keys_py;
+    results[1] = distances_py;
+    results[2] = counts_py;
+    results[3] = cluster_result.visited_members;
+    results[4] = cluster_result.computed_distances;
+    return results;
+}
+
+/**
+ *  @param queries Array of keys to cluster.
+ *  @param count Number of clusters to produce.
+ *
+ *  @return Tuple with:
+ *      1. vector of cluster IDs,
+ *      2. vector of distances to those clusters,
+ *      3. array with match counts, set to all ones,
+ *      4. number of visited nodes,
+ *      4. number of computed pairwise distances.
+ */
+template <typename index_at>
+static py::tuple cluster_keys(                            //
+    index_at& index, py::array_t<dense_key_t> queries_py, //
+    std::size_t min_count, std::size_t max_count, std::size_t threads) {
+
+    if (index.limits().threads_search < threads)
+        throw std::invalid_argument("Can't use that many threads!");
+
+    std::size_t queries_count = static_cast<std::size_t>(queries_py.size());
+    auto queries_py1d = queries_py.template unchecked<1>();
+    dense_key_t const* queries_begin = &queries_py1d(0);
+    dense_key_t const* queries_end = queries_begin + queries_count;
+
+    py::array_t<dense_key_t> keys_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
+    py::array_t<distance_t> distances_py({Py_ssize_t(queries_count), Py_ssize_t(1)});
+    executor_default_t executor{threads};
+
+    auto keys_py2d = keys_py.template mutable_unchecked<2>();
+    auto distances_py2d = distances_py.template mutable_unchecked<2>();
+    dense_key_t* keys_ptr = reinterpret_cast<dense_key_t*>(&keys_py2d(0, 0));
+    distance_t* distances_ptr = reinterpret_cast<distance_t*>(&distances_py2d(0, 0));
+
+    index_dense_clustering_config_t config;
+    config.min_clusters = min_count;
+    config.max_clusters = max_count;
+
+    dense_clustering_result_t cluster_result =
+        index.cluster(queries_begin, queries_end, config, keys_ptr, distances_ptr, executor);
+    cluster_result.error.raise();
+
+    // Those would be set to 1 for all entries, in case of success
+    py::array_t<Py_ssize_t> counts_py(queries_count);
+    auto counts_py1d = counts_py.template mutable_unchecked<1>();
+    for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)
+        counts_py1d(static_cast<Py_ssize_t>(query_idx)) = 1;
+
+    py::tuple results(5);
+    results[0] = keys_py;
+    results[1] = distances_py;
+    results[2] = counts_py;
+    results[3] = cluster_result.visited_members;
+    results[4] = cluster_result.computed_distances;
+    return results;
+}
+
+static std::unordered_map<dense_key_t, dense_key_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<key_t, key_t> a_to_b;
-    dummy_label_to_label_mapping_t b_to_a;
+    std::unordered_map<dense_key_t, dense_key_t> a_to_b;
+    dummy_key_to_key_mapping_t b_to_a;
     a_to_b.reserve((std::min)(a.size(), b.size()));
 
     index_join_config_t config;
@@ -523,17 +734,18 @@ static std::unordered_map<key_t, key_t> join_index(       //
     std::size_t threads = (std::min)(a.limits().threads(), b.limits().threads());
     executor_default_t executor{threads};
     join_result_t result = a.join(b, config, a_to_b, b_to_a, executor);
-    result.error.raise();
+    forward_error(result);
 
     return a_to_b;
 }
 
-static dense_index_py_t copy_index(dense_index_py_t const& index) {
+static dense_index_py_t copy_index(dense_index_py_t const& index, bool force_copy) {
 
     using copy_result_t = typename dense_index_py_t::copy_result_t;
-    index_copy_config_t config;
+    index_dense_copy_config_t config;
+    config.force_vector_copy = force_copy;
     copy_result_t result = index.copy(config);
-    result.error.raise();
+    forward_error(result);
     return std::move(result.index);
 }
 
@@ -559,53 +771,49 @@ template <typename index_at> typename index_at::stats_t compute_level_stats(inde
 // clang-format on
 
 template <typename internal_at, typename external_at = internal_at, typename index_at = void>
-py::object get_typed_member(index_at const& index, key_t key) {
-    py::array_t<external_at> result_py(static_cast<Py_ssize_t>(index.scalar_words()));
-    auto result_py1d = result_py.template mutable_unchecked<1>();
-    if (!index.get(key, (internal_at*)&result_py1d(0)))
-        return py::none();
-    return py::object(result_py);
+static py::tuple get_typed_vectors_for_keys(index_at const& index, py::buffer keys) {
+
+    py::buffer_info keys_info = keys.request();
+    if (keys_info.ndim != 1)
+        throw std::invalid_argument("Keys must be placed in a single-dimensional array!");
+
+    Py_ssize_t keys_count = keys_info.shape[0];
+    byte_t const* keys_data = reinterpret_cast<byte_t const*>(keys_info.ptr);
+    py::tuple results(keys_count);
+
+    for (Py_ssize_t task_idx = 0; task_idx != keys_count; ++task_idx) {
+        dense_key_t key = *reinterpret_cast<dense_key_t const*>(keys_data + task_idx * keys_info.strides[0]);
+        std::size_t vectors_count = index.count(key);
+        if (!vectors_count) {
+            results[task_idx] = py::none();
+            continue;
+        }
+
+        py::array_t<external_at> result_py({static_cast<Py_ssize_t>(vectors_count), //
+                                            static_cast<Py_ssize_t>(index.scalar_words())});
+        auto result_py2d = result_py.template mutable_unchecked<2>();
+        index.get(key, (internal_at*)&result_py2d(0, 0), vectors_count);
+        results[task_idx] = result_py;
+    }
+
+    return results;
 }
 
-template <typename index_at> py::object get_member(index_at const& index, key_t key, scalar_kind_t scalar_kind) {
+template <typename index_at> py::tuple get_many(index_at const& index, py::buffer keys, scalar_kind_t scalar_kind) {
     if (scalar_kind == scalar_kind_t::f32_k)
-        return get_typed_member<f32_t>(index, key);
+        return get_typed_vectors_for_keys<f32_t>(index, keys);
     else if (scalar_kind == scalar_kind_t::f64_k)
-        return get_typed_member<f64_t>(index, key);
+        return get_typed_vectors_for_keys<f64_t>(index, keys);
     else if (scalar_kind == scalar_kind_t::f16_k)
-        return get_typed_member<f16_t, std::uint16_t>(index, key);
+        return get_typed_vectors_for_keys<f16_t, std::uint16_t>(index, keys);
     else if (scalar_kind == scalar_kind_t::i8_k)
-        return get_typed_member<i8_bits_t, std::int8_t>(index, key);
+        return get_typed_vectors_for_keys<i8_bits_t, std::int8_t>(index, keys);
     else if (scalar_kind == scalar_kind_t::b1x8_k)
-        return get_typed_member<b1x8_t, std::uint8_t>(index, key);
+        return get_typed_vectors_for_keys<b1x8_t, std::uint8_t>(index, keys);
     else
         throw std::invalid_argument("Incompatible scalars in the query matrix!");
 }
 
-template <typename index_at> py::array_t<key_t> get_keys(index_at const& index, std::size_t offset, std::size_t limit) {
-    limit = std::min(index.size(), limit);
-    py::array_t<key_t> result_py(static_cast<Py_ssize_t>(limit));
-    auto result_py1d = result_py.template mutable_unchecked<1>();
-    index.export_keys(&result_py1d(0), offset, limit);
-    return result_py;
-}
-
-template <typename index_at> py::array_t<key_t> get_all_keys(index_at const& index) {
-    return get_keys(index, 0, index.size());
-}
-
-template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
-    if (begin == end)
-        return false;
-    element_at const* last = begin;
-    begin++;
-    for (; begin != end; ++begin, ++last) {
-        if (*begin == *last)
-            return true;
-    }
-    return false;
-}
-
 PYBIND11_MODULE(compiled, m) {
     m.doc() = "Smaller & Faster Single-File Vector Search Engine from Unum";
 
@@ -625,7 +833,7 @@ PYBIND11_MODULE(compiled, m) {
     py::enum_<metric_kind_t>(m, "MetricKind")
         .value("Unknown", metric_kind_t::unknown_k)
 
-        .value("IP", metric_kind_t::cos_k)
+        .value("IP", metric_kind_t::ip_k)
         .value("Cos", metric_kind_t::cos_k)
         .value("L2sq", metric_kind_t::l2sq_k)
 
@@ -637,7 +845,7 @@ PYBIND11_MODULE(compiled, m) {
         .value("Sorensen", metric_kind_t::sorensen_k)
 
         .value("Cosine", metric_kind_t::cos_k)
-        .value("InnerProduct", metric_kind_t::cos_k);
+        .value("InnerProduct", metric_kind_t::ip_k);
 
     py::enum_<scalar_kind_t>(m, "ScalarKind")
         .value("Unknown", scalar_kind_t::unknown_k)
@@ -659,7 +867,8 @@ PYBIND11_MODULE(compiled, m) {
 
     m.def("index_dense_metadata", [](std::string const& path) -> py::dict {
         index_dense_metadata_result_t meta = index_dense_metadata(path.c_str());
-        meta.error.raise();
+        forward_error(meta);
+
         index_dense_head_t const& head = meta.head;
 
         py::dict result;
@@ -704,40 +913,87 @@ PYBIND11_MODULE(compiled, m) {
           py::arg("expansion_search") = default_expansion_search(),        //
           py::arg("metric_kind") = metric_kind_t::cos_k,                   //
           py::arg("metric_signature") = metric_signature_t::array_array_k, //
-          py::arg("metric_pointer") = 0                                    //
+          py::arg("metric_pointer") = 0,                                   //
+          py::arg("multi") = false                                         //
     );
 
-    i.def(                                           //
-        "add", &add_many_to_index<dense_index_py_t>, //
-        py::arg("keys"),                             //
-        py::arg("vectors"),                          //
-        py::kw_only(),                               //
-        py::arg("copy") = true,                      //
-        py::arg("threads") = 0                       //
+    i.def(                                                //
+        "add_many", &add_many_to_index<dense_index_py_t>, //
+        py::arg("keys"),                                  //
+        py::arg("vectors"),                               //
+        py::kw_only(),                                    //
+        py::arg("copy") = true,                           //
+        py::arg("threads") = 0                            //
     );
 
-    i.def(                                                 //
-        "search", &search_many_in_index<dense_index_py_t>, //
-        py::arg("query"),                                  //
-        py::arg("count") = 10,                             //
-        py::arg("exact") = false,                          //
-        py::arg("threads") = 0                             //
+    i.def(                                                      //
+        "search_many", &search_many_in_index<dense_index_py_t>, //
+        py::arg("queries"),                                     //
+        py::arg("count") = 10,                                  //
+        py::arg("exact") = false,                               //
+        py::arg("threads") = 0                                  //
+    );
+
+    i.def(                                                     //
+        "cluster_vectors", &cluster_vectors<dense_index_py_t>, //
+        py::arg("queries"),                                    //
+        py::arg("min_count") = 0,                              //
+        py::arg("max_count") = 0,                              //
+        py::arg("threads") = 0                                 //
+    );
+
+    i.def(                                               //
+        "cluster_keys", &cluster_keys<dense_index_py_t>, //
+        py::arg("queries"),                              //
+        py::arg("min_count") = 0,                        //
+        py::arg("max_count") = 0,                        //
+        py::arg("threads") = 0                           //
     );
 
     i.def(
-        "rename",
-        [](dense_index_py_t& index, key_t from, key_t to) -> bool {
+        "rename_one_to_one",
+        [](dense_index_py_t& index, dense_key_t from, dense_key_t to) -> bool {
             dense_labeling_result_t result = index.rename(from, to);
-            result.error.raise();
+            forward_error(result);
             return result.completed;
         },
-        py::arg("from"), py::arg("to"));
+        py::arg("from_"), py::arg("to"));
 
     i.def(
-        "remove",
-        [](dense_index_py_t& index, key_t key, bool compact, std::size_t threads) -> bool {
+        "rename_many_to_many",
+        [](dense_index_py_t& index, std::vector<dense_key_t> const& from,
+           std::vector<dense_key_t> const& to) -> std::vector<bool> {
+            if (from.size() != to.size())
+                throw std::invalid_argument("Sizes of `from` and `to` arrays don't match!");
+
+            std::vector<bool> results(from.size(), false);
+            for (std::size_t i = 0; i != from.size(); ++i) {
+                dense_labeling_result_t result = index.rename(from[i], to[i]);
+                results[i] = result.completed;
+                forward_error(result);
+            }
+            return results;
+        },
+        py::arg("from_"), py::arg("to"));
+
+    i.def(
+        "rename_many_to_one",
+        [](dense_index_py_t& index, std::vector<dense_key_t> const& from, dense_key_t to) -> std::vector<bool> {
+            std::vector<bool> results(from.size(), false);
+            for (std::size_t i = 0; i != from.size(); ++i) {
+                dense_labeling_result_t result = index.rename(from[i], to);
+                results[i] = result.completed;
+                forward_error(result);
+            }
+            return results;
+        },
+        py::arg("from_"), py::arg("to"));
+
+    i.def(
+        "remove_one",
+        [](dense_index_py_t& index, dense_key_t key, bool compact, std::size_t threads) -> bool {
             dense_labeling_result_t result = index.remove(key);
-            result.error.raise();
+            forward_error(result);
             if (!compact)
                 return result.completed;
 
@@ -752,10 +1008,11 @@ PYBIND11_MODULE(compiled, m) {
         py::arg("key"), py::arg("compact"), py::arg("threads"));
 
     i.def(
-        "remove",
-        [](dense_index_py_t& index, std::vector<key_t> const& keys, bool compact, std::size_t threads) -> std::size_t {
+        "remove_many",
+        [](dense_index_py_t& index, std::vector<dense_key_t> const& keys, bool compact,
+           std::size_t threads) -> std::size_t {
             dense_labeling_result_t result = index.remove(keys.begin(), keys.end());
-            result.error.raise();
+            forward_error(result);
             if (!compact)
                 return result.completed;
 
@@ -778,8 +1035,7 @@ PYBIND11_MODULE(compiled, m) {
     i.def_property_readonly( //
         "dtype", [](dense_index_py_t const& index) -> scalar_kind_t { return index.scalar_kind(); });
     i.def_property_readonly( //
-        "memory_usage", [](dense_index_py_t const& index) -> std::size_t { return index.memory_usage(); },
-        py::call_guard<py::gil_scoped_release>());
+        "memory_usage", [](dense_index_py_t const& index) -> std::size_t { return index.memory_usage(); });
 
     i.def_property("expansion_add", &dense_index_py_t::expansion_add, &dense_index_py_t::change_expansion_add);
     i.def_property("expansion_search", &dense_index_py_t::expansion_search, &dense_index_py_t::change_expansion_search);
@@ -792,7 +1048,7 @@ PYBIND11_MODULE(compiled, m) {
             std::size_t dimensions = index.dimensions();
             metric_t metric =  //
                 metric_uintptr //
-                    ? udf(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
+                    ? wrap_user_defined_metric(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
                     : metric_t(dimensions, metric_kind, scalar_kind);
             index.change_metric(std::move(metric));
         },
@@ -801,21 +1057,95 @@ PYBIND11_MODULE(compiled, m) {
         py::arg("metric_pointer") = 0                                    //
     );
 
-    i.def_property_readonly("keys", &get_all_keys<dense_index_py_t>);
-    i.def("get_keys", &get_keys<dense_index_py_t>, py::arg("offset") = 0,
-          py::arg("limit") = std::numeric_limits<std::size_t>::max());
-    i.def("__contains__", &dense_index_py_t::contains);
-    i.def("__getitem__", &get_member<dense_index_py_t>, py::arg("key"), py::arg("dtype") = scalar_kind_t::f32_k);
-
-    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("reset", &reset_index<dense_index_py_t>, 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("compact", &compact_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>());
+    i.def_property_readonly("hardware_acceleration", [](dense_index_py_t const& index) -> py::str {
+        return isa_name(index.metric().isa_kind());
+    });
+
+    i.def("contains_one", &dense_index_py_t::contains);
+    i.def("count_one", &dense_index_py_t::count);
+
+    i.def( //
+        "contains_many",
+        [](dense_index_py_t const& index, py::array_t<dense_key_t> const& keys_py) -> py::array_t<bool> {
+            py::array_t<bool> results_py(keys_py.size());
+            auto results_py1d = results_py.template mutable_unchecked<1>();
+            auto keys_py1d = keys_py.template unchecked<1>();
+            for (Py_ssize_t task_idx = 0; task_idx != keys_py.size(); ++task_idx)
+                results_py1d(task_idx) = index.contains(keys_py1d(task_idx));
+            return results_py;
+        });
+
+    i.def( //
+        "count_many",
+        [](dense_index_py_t const& index, py::array_t<dense_key_t> const& keys_py) -> py::array_t<std::size_t> {
+            py::array_t<std::size_t> results_py(keys_py.size());
+            auto results_py1d = results_py.template mutable_unchecked<1>();
+            auto keys_py1d = keys_py.template unchecked<1>();
+            for (Py_ssize_t task_idx = 0; task_idx != keys_py.size(); ++task_idx)
+                results_py1d(task_idx) = index.count(keys_py1d(task_idx));
+            return results_py;
+        });
+
+    i.def( //
+        "pairwise_distances",
+        [](dense_index_py_t const& index, py::array_t<dense_key_t> const& left_py,
+           py::array_t<dense_key_t> const& right_py) -> py::array_t<distance_t> {
+            py::array_t<distance_t> results_py(left_py.size());
+            auto results_py1d = results_py.template mutable_unchecked<1>();
+            auto left_py1d = left_py.template unchecked<1>();
+            auto right_py1d = right_py.template unchecked<1>();
+            for (Py_ssize_t task_idx = 0; task_idx != left_py.size(); ++task_idx)
+                results_py1d(task_idx) = index.distance_between(left_py1d(task_idx), right_py1d(task_idx)).min;
+            return results_py;
+        });
+
+    i.def( //
+        "pairwise_distance", [](dense_index_py_t const& index, dense_key_t left, dense_key_t right) -> distance_t {
+            return index.distance_between(left, right).min;
+        });
+
+    i.def("get_many", &get_many<dense_index_py_t>, py::arg("keys"), py::arg("dtype") = scalar_kind_t::f32_k);
+
+    i.def(
+        "get_keys_in_slice",
+        [](dense_index_py_t const& index, std::size_t offset, std::size_t limit) -> py::array_t<dense_key_t> {
+            limit = std::min(index.size(), limit);
+            py::array_t<dense_key_t> result_py(static_cast<Py_ssize_t>(limit));
+            auto result_py1d = result_py.template mutable_unchecked<1>();
+            index.export_keys(&result_py1d(0), offset, limit);
+            return result_py;
+        },
+        py::arg("offset") = 0, py::arg("limit") = std::numeric_limits<std::size_t>::max());
+
+    i.def(
+        "get_keys_at_offsets",
+        [](dense_index_py_t const& index, py::array_t<Py_ssize_t> const& offsets_py) -> py::array_t<dense_key_t> {
+            py::array_t<dense_key_t> result_py(offsets_py.size());
+            auto result_py1d = result_py.template mutable_unchecked<1>();
+            auto offsets_py1d = offsets_py.template unchecked<1>();
+            for (Py_ssize_t task_idx = 0; task_idx != offsets_py.size(); ++task_idx)
+                index.export_keys(&result_py1d(task_idx), offsets_py1d(task_idx), 1);
+            return result_py;
+        },
+        py::arg("offsets"));
+
+    i.def(
+        "get_key_at_offset",
+        [](dense_index_py_t const& index, std::size_t offset) -> dense_key_t {
+            dense_key_t result;
+            index.export_keys(&result, offset, 1);
+            return result;
+        },
+        py::arg("offset"));
+
+    i.def("save", &save_index<dense_index_py_t>, py::arg("path"));
+    i.def("load", &load_index<dense_index_py_t>, py::arg("path"));
+    i.def("view", &view_index<dense_index_py_t>, py::arg("path"));
+    i.def("reset", &reset_index<dense_index_py_t>);
+    i.def("clear", &clear_index<dense_index_py_t>);
+    i.def("copy", &copy_index, py::kw_only(), py::arg("copy") = true);
+    i.def("compact", &compact_index);
+    i.def("join", &join_index, py::arg("other"), py::arg("max_proposals") = 0, py::arg("exact") = false);
 
     using punned_index_stats_t = typename dense_index_py_t::stats_t;
     auto i_stats = py::class_<punned_index_stats_t>(m, "IndexStats");
@@ -831,14 +1161,14 @@ PYBIND11_MODULE(compiled, m) {
     auto is = py::class_<dense_indexes_py_t>(m, "Indexes");
     is.def(py::init());
     is.def("__len__", &dense_indexes_py_t::size);
-    is.def("add", &dense_indexes_py_t::add);
-    is.def("add_paths", &dense_indexes_py_t::add_paths, py::arg("paths"), py::arg("view") = true,
+    is.def("merge", &dense_indexes_py_t::merge);
+    is.def("merge_paths", &dense_indexes_py_t::merge_paths, py::arg("paths"), py::arg("view") = true,
            py::arg("threads") = 0);
-    is.def(                                                  //
-        "search", &search_many_in_index<dense_indexes_py_t>, //
-        py::arg("query"),                                    //
-        py::arg("count") = 10,                               //
-        py::arg("exact") = false,                            //
-        py::arg("threads") = 0                               //
+    is.def(                                                       //
+        "search_many", &search_many_in_index<dense_indexes_py_t>, //
+        py::arg("query"),                                         //
+        py::arg("count") = 10,                                    //
+        py::arg("exact") = false,                                 //
+        py::arg("threads") = 0                                    //
     );
 }