Merge pull request #647 from unum-cloud/main-dev

Fix Isolating Removed Entries

Author: ashvardanian

.vscode/settings.json

@@ -135,7 +135,8 @@
     "xstring": "cpp",
     "xtr1common": "cpp",
     "xtree": "cpp",
-    "xutility": "cpp"
+    "xutility": "cpp",
+    "cfenv": "cpp"
   },
   "cSpell.words": [
     "allclose",

cpp/test.cpp

@@ -880,11 +880,11 @@ void test_exact_search(std::size_t dataset_count, std::size_t queries_count, std
     std::size_t dimensions = 32;
     metric_punned_t metric(dimensions, metric_kind_t::cos_k, scalar_kind<scalar_at>());
 
-    std::random_device rd;
-    std::mt19937 gen(rd());
-    std::uniform_real_distribution<> dis(0.0, 1.0);
+    std::random_device seed_source;
+    std::mt19937 generator(seed_source());
+    std::uniform_real_distribution<> distribution(0.0, 1.0); // ! We can't pass `scalar_at` to the distribution
     std::vector<scalar_at> dataset(dataset_count * dimensions);
-    std::generate(dataset.begin(), dataset.end(), [&] { return static_cast<scalar_at>(dis(gen)); });
+    std::generate(dataset.begin(), dataset.end(), [&] { return static_cast<scalar_at>(distribution(generator)); });
 
     exact_search_t search;
     auto results = search(                                                            //
@@ -1107,15 +1107,15 @@ void test_filtered_search() {
     constexpr std::size_t dimensions = 32;
     metric_punned_t metric(dimensions, metric_kind_t::cos_k);
 
-    std::random_device rd;
-    std::mt19937 gen(rd());
-    std::uniform_real_distribution<> dis(0.0, 1.0);
+    std::random_device seed_source;
+    std::mt19937 generator(seed_source());
+    std::uniform_real_distribution<float> distribution(0.0, 1.0);
     using vector_of_vectors_t = std::vector<std::vector<float>>;
 
     vector_of_vectors_t vector_of_vectors(dataset_count);
     for (auto& vector : vector_of_vectors) {
         vector.resize(dimensions);
-        std::generate(vector.begin(), vector.end(), [&] { return dis(gen); });
+        std::generate(vector.begin(), vector.end(), [&] { return distribution(generator); });
     }
 
     index_dense_t index = index_dense_t::make(metric);
@@ -1144,6 +1144,41 @@ void test_filtered_search() {
     }
 }
 
+void test_isolate() {
+    constexpr std::size_t dataset_count = 16;
+    constexpr std::size_t dimensions = 32;
+    metric_punned_t metric(dimensions, metric_kind_t::cos_k);
+
+    std::random_device seed_source;
+    std::mt19937 generator(seed_source());
+    std::uniform_real_distribution<float> distribution(0.0, 1.0);
+    using vector_of_vectors_t = std::vector<std::vector<float>>;
+
+    vector_of_vectors_t vector_of_vectors(dataset_count);
+    for (auto& vector : vector_of_vectors) {
+        vector.resize(dimensions);
+        std::generate(vector.begin(), vector.end(), [&] { return distribution(generator); });
+    }
+
+    index_dense_t index = index_dense_t::make(metric);
+    index.reserve(dataset_count);
+    for (std::size_t idx = 0; idx < dataset_count; ++idx) {
+        index.add(idx, vector_of_vectors[idx].data());
+    }
+    expect_eq(index.size(), dataset_count);
+
+    for (std::size_t idx = 0; idx < dataset_count; ++idx) {
+        if (idx % 2 == 0)
+            index.remove(idx);
+    }
+
+    auto result = index.isolate();
+    for (std::size_t idx = 0; idx < dataset_count; ++idx) {
+        auto result = index.search(vector_of_vectors[idx].data(), 16);
+        expect_eq(result.size(), dataset_count / 2);
+    }
+}
+
 int main(int, char**) {
     test_uint40();
     test_cosine<float, std::int64_t, uint40_t>(10, 10);
@@ -1221,5 +1256,6 @@ int main(int, char**) {
     test_strings<std::int64_t, slot32_t>();
 
     test_filtered_search();
+    test_isolate();
     return 0;
 }

include/usearch/index.hpp

@@ -2177,6 +2177,20 @@ class index_gt {
             misaligned_store<compressed_slot_t>(tape_ + shift(n), slot);
             misaligned_store<neighbors_count_t>(tape_, n + 1);
         }
+        template <typename allow_slot_at> std::size_t erase_if(allow_slot_at&& allow_slot) noexcept {
+            std::size_t old_count = misaligned_load<neighbors_count_t>(tape_);
+            std::size_t removed_count = 0;
+            for (std::size_t i = 0; i < old_count; ++i) {
+                compressed_slot_t slot = misaligned_load<compressed_slot_t>(tape_ + shift(i));
+                if (allow_slot(slot)) {
+                    removed_count++;
+                } else {
+                    misaligned_store<compressed_slot_t>(tape_ + shift(i - removed_count), slot);
+                }
+            }
+            misaligned_store<neighbors_count_t>(tape_, old_count - removed_count);
+            return removed_count;
+        }
     };
 
     /**
@@ -3654,14 +3668,10 @@ class index_gt {
             node_t node = node_at_(node_idx);
             for (level_t level = 0; level <= node.level(); ++level) {
                 neighbors_ref_t neighbors = neighbors_(node, level);
-                std::size_t old_size = neighbors.size();
-                neighbors.clear();
-                for (std::size_t i = 0; i != old_size; ++i) {
-                    compressed_slot_t neighbor_slot = neighbors[i];
+                neighbors.erase_if([&](compressed_slot_t neighbor_slot) {
                     node_t neighbor = node_at_(neighbor_slot);
-                    if (allow_member(member_cref_t{neighbor.ckey(), neighbor_slot}))
-                        neighbors.push_back(neighbor_slot);
-                }
+                    return !allow_member(member_cref_t{neighbor.ckey(), neighbor_slot});
+                });
             }
             ++processed;
             if (thread_idx == 0)

include/usearch/index_dense.hpp

@@ -1709,12 +1709,12 @@ class index_dense_gt {
     compaction_result_t isolate(executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
         compaction_result_t result;
         std::atomic<std::size_t> pruned_edges;
-        auto disallow = [&](member_cref_t const& member) noexcept {
+        auto allow = [&](member_cref_t const& member) noexcept {
             bool freed = member.key == free_key_;
             pruned_edges += freed;
-            return freed;
+            return !freed;
         };
-        typed_->isolate(disallow, std::forward<executor_at>(executor), std::forward<progress_at>(progress));
+        typed_->isolate(allow, std::forward<executor_at>(executor), std::forward<progress_at>(progress));
         result.pruned_edges = pruned_edges;
         return result;
     }