usearch 0.10.0 → 0.11.0

package/src/CMakeLists.txt

@@ -69,6 +69,7 @@ include(clipp)
 set(USEARCH_PUNNED_INCLUDE_DIRS
     "${CMAKE_CURRENT_SOURCE_DIR}/../include"
     "${CMAKE_CURRENT_SOURCE_DIR}/../fp16/include"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../robin-map/include"
     "${CMAKE_CURRENT_SOURCE_DIR}/../simsimd/include"
     "${CMAKE_CURRENT_SOURCE_DIR}/"
 )

package/src/punned.hpp

@@ -1,13 +1,15 @@
 #pragma once
 #include <stdlib.h> // `aligned_alloc`
 
-#include <functional> // `std::function`
-#include <numeric>    // `std::iota`
-#include <stdexcept>  // `std::invalid_argument`
-#include <thread>     // `std::thread`
-#include <vector>     // `std::vector`
+#include <functional>   // `std::function`
+#include <numeric>      // `std::iota`
+#include <shared_mutex> // `std::shared_mutex`
+#include <stdexcept>    // `std::invalid_argument`
+#include <thread>       // `std::thread`
+#include <vector>       // `std::vector`
 
 #include <fp16/fp16.h>
+#include <tsl/robin_map.h>
 
 #include <usearch/usearch.hpp>
 
@@ -401,6 +403,8 @@ class punned_gt {
 
     using member_iterator_t = typename index_t::member_iterator_t;
     using member_citerator_t = typename index_t::member_citerator_t;
+    using member_ref_t = typename index_t::member_ref_t;
+    using member_cref_t = typename index_t::member_cref_t;
 
     std::size_t dimensions_ = 0;
     std::size_t casted_vector_bytes_ = 0;
@@ -414,6 +418,10 @@ class punned_gt {
         cast_t from_f16{};
         cast_t from_f32{};
         cast_t from_f64{};
+        cast_t to_f8{};
+        cast_t to_f16{};
+        cast_t to_f32{};
+        cast_t to_f64{};
     } casts_;
 
     punned_stateful_metric_t root_metric_;
@@ -421,6 +429,13 @@ class punned_gt {
     mutable std::vector<std::size_t> available_threads_;
     mutable std::mutex available_threads_mutex_;
 
+    using shared_mutex_t = std::mutex; // TODO: Find an OS-compatible solution
+    using shared_lock_t = std::unique_lock<shared_mutex_t>;
+    using unique_lock_t = std::unique_lock<shared_mutex_t>;
+
+    mutable shared_mutex_t lookup_table_mutex_;
+    tsl::robin_map<label_t, id_t> lookup_table_;
+
   public:
     using search_results_t = typename index_t::search_results_t;
     using add_result_t = typename index_t::add_result_t;
@@ -511,6 +526,11 @@ class punned_gt {
     search_results_t search_around(label_t hint, f32_t const* vector, std::size_t wanted, search_config_t config) const { return search_around_(hint, vector, wanted, config, casts_.from_f32); }
     search_results_t search_around(label_t hint, f64_t const* vector, std::size_t wanted, search_config_t config) const { return search_around_(hint, vector, wanted, config, casts_.from_f64); }
 
+    void reconstruct(label_t label, f8_bits_t* vector) const { return reconstruct_(label, vector, casts_.to_f8); }
+    void reconstruct(label_t label, f16_bits_t* vector) const { return reconstruct_(label, vector, casts_.to_f16); }
+    void reconstruct(label_t label, f32_t* vector) const { return reconstruct_(label, vector, casts_.to_f32); }
+    void reconstruct(label_t label, f64_t* vector) const { return reconstruct_(label, vector, casts_.to_f64); }
+
     static punned_gt ip(std::size_t dimensions, accuracy_t accuracy = accuracy_t::f16_k, config_t config = {}) { return make_(dimensions, accuracy, ip_metric_(dimensions, accuracy), make_casts_(accuracy), config); }
     static punned_gt l2sq(std::size_t dimensions, accuracy_t accuracy = accuracy_t::f16_k, config_t config = {}) { return make_(dimensions, accuracy, l2_metric_(dimensions, accuracy), make_casts_(accuracy), config); }
     static punned_gt cos(std::size_t dimensions, accuracy_t accuracy = accuracy_t::f32_k, config_t config = {}) { return make_(dimensions, accuracy, cos_metric_(dimensions, accuracy), make_casts_(accuracy), config); }
@@ -589,7 +609,12 @@ class punned_gt {
         if (casted)
             vector_data = casted_data, vector_bytes = casted_vector_bytes_, config.store_vector = true;
 
-        return typed_->add(label, {vector_data, vector_bytes}, config);
+        add_result_t result = typed_->add(label, {vector_data, vector_bytes}, config);
+        {
+            unique_lock_t lock(lookup_table_mutex_);
+            lookup_table_.emplace(label, result.id);
+        }
+        return result;
     }
 
     template <typename scalar_at>
@@ -624,6 +649,21 @@ class punned_gt {
         return typed_->search_around(static_cast<id_t>(hint), {vector_data, vector_bytes}, wanted, config);
     }
 
+    id_t lookup_id_(label_t label) const {
+        shared_lock_t lock(lookup_table_mutex_);
+        return lookup_table_.at(label);
+    }
+
+    template <typename scalar_at> void reconstruct_(label_t label, scalar_at* reconstructed, cast_t const& cast) const {
+        id_t id = lookup_id_(label);
+        member_citerator_t iterator = typed_->cbegin() + id;
+        member_cref_t member = *iterator;
+        byte_t const* casted_vector = reinterpret_cast<byte_t const*>(member.vector.data());
+        bool casted = cast(casted_vector, casted_vector_bytes_, (byte_t*)reconstructed);
+        if (!casted)
+            std::memcpy(reconstructed, casted_vector, casted_vector_bytes_);
+    }
+
     template <typename scalar_at> add_result_t add_(label_t label, scalar_at const* vector, cast_t const& cast) {
         thread_lock_t lock = thread_lock_();
         add_config_t add_config;
@@ -686,6 +726,10 @@ class punned_gt {
         result.from_f16 = cast_gt<f16_bits_t, to_scalar_at>{};
         result.from_f32 = cast_gt<f32_t, to_scalar_at>{};
         result.from_f64 = cast_gt<f64_t, to_scalar_at>{};
+        result.to_f8 = cast_gt<to_scalar_at, f8_bits_t>{};
+        result.to_f16 = cast_gt<to_scalar_at, f16_bits_t>{};
+        result.to_f32 = cast_gt<to_scalar_at, f32_t>{};
+        result.to_f64 = cast_gt<to_scalar_at, f64_t>{};
         return result;
     }
 

package/src/test.cpp

@@ -12,7 +12,7 @@ using namespace unum;
 using point_id_t = std::int64_t;
 
 template <typename scalar_at, typename index_at> void test3d(index_at&& index) {
-    using span_t = span_gt<scalar_at const>;
+    using view_t = span_gt<scalar_at const>;
     using distance_t = typename index_at::distance_t;
 
     scalar_at vec[3] = {10, 20, 15};
@@ -20,12 +20,12 @@ template <typename scalar_at, typename index_at> void test3d(index_at&& index) {
     scalar_at vec_b[3] = {16, 17, 18};
 
     index.reserve(10);
-    index.add(42, span_t{&vec[0], 3ul});
+    index.add(42, view_t{&vec[0], 3ul});
 
     // Default approximate search
     point_id_t matched_labels[10] = {0};
     distance_t matched_distances[10] = {0};
-    std::size_t matched_count = index.search(span_t{&vec[0], 3ul}, 5).dump_to(matched_labels, matched_distances);
+    std::size_t matched_count = index.search(view_t{&vec[0], 3ul}, 5).dump_to(matched_labels, matched_distances);
 
     assert(matched_count == 1);
     assert(matched_labels[0] == 42);
@@ -34,12 +34,12 @@ template <typename scalar_at, typename index_at> void test3d(index_at&& index) {
     // Add more entries
     search_config_t search_config;
     search_config.exact = true;
-    index.add(43, span_t{&vec_a[0], 3ul});
-    index.add(44, span_t{&vec_b[0], 3ul});
+    index.add(43, view_t{&vec_a[0], 3ul});
+    index.add(44, view_t{&vec_b[0], 3ul});
     assert(index.size() == 3);
 
     // Perform exact search
-    matched_count = index.search(span_t{&vec[0], 3ul}, 5, search_config).dump_to(matched_labels, matched_distances);
+    matched_count = index.search(view_t{&vec[0], 3ul}, 5, search_config).dump_to(matched_labels, matched_distances);
 
     // Validate scans
     std::size_t count = 0;
@@ -54,14 +54,14 @@ template <typename scalar_at, typename index_at> void test3d(index_at&& index) {
     // Search again over reconstructed index
     index.save("tmp.usearch");
     index.load("tmp.usearch");
-    matched_count = index.search(span_t{&vec[0], 3ul}, 5).dump_to(matched_labels, matched_distances);
+    matched_count = index.search(view_t{&vec[0], 3ul}, 5).dump_to(matched_labels, matched_distances);
     assert(matched_count == 3);
     assert(matched_labels[0] == 42);
     assert(std::abs(matched_distances[0]) < 0.01);
 
     // Search again over mapped index
     index.view("tmp.usearch");
-    matched_count = index.search(span_t{&vec[0], 3ul}, 5).dump_to(matched_labels, matched_distances);
+    matched_count = index.search(view_t{&vec[0], 3ul}, 5).dump_to(matched_labels, matched_distances);
     assert(matched_count == 3);
     assert(matched_labels[0] == 42);
     assert(std::abs(matched_distances[0]) < 0.01);
@@ -70,6 +70,24 @@ template <typename scalar_at, typename index_at> void test3d(index_at&& index) {
     assert(index.stats().max_edges > 0);
 }
 
+template <typename scalar_at, typename index_at> void test3d_punned(index_at&& index) {
+    using view_t = span_gt<scalar_at const>;
+    using span_t = span_gt<scalar_at>;
+    using distance_t = typename index_at::distance_t;
+
+    scalar_at vec[3] = {10, 20, 15};
+
+    index.reserve(10);
+    index.add(42, view_t{&vec[0], 3ul});
+
+    // Reconstruct
+    scalar_at vec_reconstructed[3] = {0, 0, 0};
+    index.reconstruct(42, span_t{&vec_reconstructed[0], 3ul});
+    assert(vec_reconstructed[0] == vec[0]);
+    assert(vec_reconstructed[1] == vec[1]);
+    assert(vec_reconstructed[2] == vec[2]);
+}
+
 int main(int, char**) {
 
     test3d<float>(index_gt<cos_gt<float>, point_id_t, std::uint32_t, float>{});
@@ -84,5 +102,8 @@ int main(int, char**) {
     test3d<double>(punned_small_t::cos(3));
     test3d<double>(punned_small_t::l2sq(3));
 
+    test3d_punned<float>(punned_small_t::cos(3));
+    test3d_punned<float>(punned_small_t::l2sq(3));
+
     return 0;
 }