Add HNSW removal with a deleted mark

.github/workflows/test.yaml

@@ -206,7 +206,7 @@ jobs:
           simd_filter+=':RaBitQKMeansTest.FinalModesKeepFlatLloydAssignmentsAcrossClusterCounts'
           simd_filter+=':IvfPersistenceTest.ReloadsRawAndQuantizedStorage'
           simd_filter+=':HnswSaveTest.SuccessfulSavePreservesSearchAfterLoading'
-          simd_filter+=':IvfAddTest.*:IvfRemoveTest.*'
+          simd_filter+=':IvfAddTest.*:IvfRemoveTest.*:HnswRemoveTest.*'
           for cpu_model in hsw icl; do
             RABITQ_TEST_CPU="$cpu_model" "$RUNNER_TEMP/sde/sde64" "-$cpu_model" -- \
               build-install/tests/rabitq_tests --gtest_filter="$simd_filter"

DEVELOPMENT.md

@@ -366,6 +366,21 @@ never appear in results. Keep this convention when touching IVF code:
 - `IvfRemoveTest.ExcludesRemovedPointsOnEveryScanPath`, `IvfRemoveTest.SurvivesSaveLoadAndLaterAdds`,
   and `FastScanPackingTest.UnpackInvertsPackIncludingTailsAndDoesNotOverrun` guard these rules.
 
+HNSW `remove` sets the high bit (`kRemovedMask`) of the stored cluster ID. `construct` and `load`
+cap the number of clusters at 2^31, so no index written before `remove` existed has that bit set and
+nothing in an old file is reinterpreted. A file that has removals is rejected by releases 0.3.7
+through 0.5.1, whose `load` range-checks every cluster ID; that rejection is the compatibility path.
+This mark is unrelated to the `SearchBuffer` ID marker. Keep these rules when touching HNSW code:
+
+- Read cluster IDs only through `get_clusterid_by_internalid`, which masks the bit. Estimators,
+  routing, and the `add` reconstruction index centroids with it.
+- A removed point stays in the graph: base-layer search still inserts it into `candidate_set`, and
+  `add` may link to it. Only the result buffer skips it, including when it is the entry point.
+- `load` masks the bit before range-checking the cluster ID.
+- `HnswRemoveTest.SearchReachesLivePointsThroughRemovedOnes`,
+  `HnswRemoveTest.SurvivesSaveLoadResizeAndLaterAdds`, and
+  `HnswRemoveTest.LoadRejectsAnOutOfRangeClusterEvenWithTheMark` guard these rules.
+
 ### Change Python bindings
 
 Shared NumPy and string conversion helpers live in `python_bindings/bindings_common.hpp`. Register

docs/docs/index/hnsw.md

@@ -47,7 +47,7 @@ Each indexed element is stored in the following layout:
 ```
 [number of edges]
 [edges]
-[cluster ID]
+[cluster ID]       (its high bit marks a removed point, see Removing points)
 [external label]
 [BinData (1-bit * dim + factors)]
 [ExData (ex-bits * dim + factors)]
@@ -99,6 +99,12 @@ Repeat until `candidate_set` is empty:
 
 The search terminates when `candidate_set` is empty.
 
+A point hidden by `remove` never enters `boundedKNN`, so search never returns it,
+but it still enters `candidate_set` and the search keeps walking through it. A
+removed entry point keeps the estimate already computed for it (the full-bits one
+when nbits > 1). A removed neighbor is not refined and enters `candidate_set` with
+its 1-bit estimate.
+
 ## Updating an Index
 
 A constructed or loaded index can gain points without the original data. The
@@ -176,3 +182,54 @@ never retrains them. If the added vectors come from a different distribution, or
 the index grows many times beyond its original size, recall can drop. Rebuild
 from the original data with new centroids when recall matters more than the cost
 of a rebuild.
+
+### Removing points
+
+```c++
+size_t HierarchicalNSW::remove(const PID* labels, size_t n);
+```
+
+- **labels**, **n**: Labels of the points to remove, as returned by `add` and
+  `search`. Every label must be in the index; nothing is removed if one is not.
+  `remove` returns how many points were newly removed, so repeating a label or a
+  call is safe.
+
+In Python:
+
+```python
+removed = index.remove(ids)   # ids in [0, num_points)
+```
+
+A removed point keeps its codes, its links, and its place in the graph. Search
+still walks through it and `add` may link new points to it, which keeps the graph
+connected, as in hnswlib. Search never returns it, so a query can get fewer than
+`k` results once points are removed: in Python the missing slots hold `kPidMax`
+(`2**32 - 1`) with an infinite distance. Removed points still count in
+`num_points()` and cannot be restored. It is not safe to call `remove` while
+another thread searches or adds to the same index.
+
+#### Choosing `ef` after removals
+
+Removed points still take slots in the `efSearch` candidate set, so a search with
+`efSearch` close to `k` finds fewer live points as removals accumulate. `search`
+uses the `efSearch` it is given. To keep about the same number of live candidates
+as before the removals, scale it by the share of live points:
+
+```text
+efSearch * num_points() / (num_points() - removed points)
+```
+
+Callers that track how many points they removed can apply this directly.
+
+#### How removal is stored
+
+`remove` sets the high bit of the point's stored cluster ID. `construct` and `load`
+limit the number of clusters to 2^31, so no real cluster ID has that bit, and
+every reader of the cluster ID masks it off. The mark lives in the base layer, so
+it survives `save`, `load`, and `resize`, and a file without removals is
+unchanged.
+
+A file that has removals cannot be read by earlier releases. Releases 0.3.7
+through 0.5.1 check every cluster ID on `load` and reject the file as invalid
+instead of returning removed points. Releases before 0.3.7 do not check cluster
+IDs, so do not open such a file with them.

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -105,6 +105,23 @@ class HierarchicalNSW {
      * @param new_max_elements The new capacity, at least `num_points()`
      */
     void resize(size_t);
+
+    /**
+     * @brief Exclude points from all later search results.
+     *
+     * A removed point keeps its codes and its place in the graph: search and `add`
+     * still traverse it and may link new points to it, but search never returns it,
+     * so a query may get fewer than k results. A removed point still counts in
+     * `num_points()`. Removal is stored in the index, so it survives save and load,
+     * and is idempotent. Removed points cannot be restored. Every label is checked
+     * before anything changes. Not safe to call while another thread searches or adds.
+     *
+     * @param labels Labels of the points to remove, as returned by `add` or `search`
+     * @param n Number of labels
+     * @return Number of points newly removed
+     */
+    size_t remove(const PID*, size_t);
+
     std::vector<std::vector<std::pair<float, PID>>> search(
         const float*, size_t, size_t, size_t, size_t
     );
@@ -181,6 +198,13 @@ class HierarchicalNSW {
     );
 
     static constexpr PID kMaxLabelOperationLock = 65536;
+    // `remove` sets the high bit of a point's stored cluster ID. construct and load
+    // keep num_cluster_ at or below 2^31, so a real cluster ID never has it.
+    static constexpr PID kRemovedMask = 0x80000000U;
+    static_assert(
+        buffer::kSearchBufferMaxPointCount <= kRemovedMask,
+        "cluster IDs must leave the removal bit free"
+    );
     size_t max_elements_{0};
     mutable std::atomic<size_t> cur_element_count_{0};  // current number of elements
     size_t size_data_per_element_{0};
@@ -313,13 +337,22 @@ class HierarchicalNSW {
         );
     }
 
-    PID get_clusterid_by_internalid(PID internal_id) const {
+    // The stored cluster ID with the removal mark, as save and load see it.
+    PID get_raw_clusterid(PID internal_id) const {
         return *(reinterpret_cast<PID*>(
             data_level0_memory_ + (internal_id * size_data_per_element_) +
             size_links_level0_
         ));
     }
 
+    PID get_clusterid_by_internalid(PID internal_id) const {
+        return get_raw_clusterid(internal_id) & ~kRemovedMask;
+    }
+
+    bool is_removed(PID internal_id) const {
+        return (get_raw_clusterid(internal_id) & kRemovedMask) != 0;
+    }
+
     char* get_clusterid_pt(PID internal_id) const {
         return reinterpret_cast<char*>(
             data_level0_memory_ + (internal_id * size_data_per_element_) +
@@ -716,7 +749,8 @@ inline void HierarchicalNSW::load(const char* filename) {
         std::memcpy(&level0_degree, row, sizeof(level0_degree));
         std::memcpy(&cluster_id, row + loaded.size_links_level0_, sizeof(cluster_id));
         std::memcpy(&label, row + loaded.label_offset_, sizeof(label));
-        if (cluster_id >= loaded.num_cluster_ || level0_degree > loaded.maxM0_) {
+        if ((cluster_id & ~kRemovedMask) >= loaded.num_cluster_ ||
+            level0_degree > loaded.maxM0_) {
             invalid_file();
         }
         for (size_t edge = 0; edge < level0_degree; ++edge) {
@@ -1364,6 +1398,43 @@ inline std::vector<std::vector<std::pair<float, PID>>> HierarchicalNSW::search(
     return results;
 }
 
+inline size_t HierarchicalNSW::remove(const PID* labels, size_t n) {
+    if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr ||
+        rotator_ == nullptr || num_cluster_ == 0) {
+        throw std::logic_error("HNSW index must be constructed or loaded before remove");
+    }
+    if (n == 0) {
+        return 0;
+    }
+    if (labels == nullptr) {
+        throw std::invalid_argument("HNSW remove labels must not be null");
+    }
+
+    // Resolve every label before marking any point, so an unknown label removes
+    // nothing.
+    std::vector<PID> internal_ids(n);
+    {
+        std::unique_lock<std::mutex> lock_table(label_lookup_lock_);
+        for (size_t i = 0; i < n; ++i) {
+            const auto found = label_lookup_.find(labels[i]);
+            if (found == label_lookup_.end()) {
+                throw std::invalid_argument("HNSW remove label is not in the index");
+            }
+            internal_ids[i] = found->second;
+        }
+    }
+
+    size_t removed = 0;
+    for (const PID id : internal_ids) {
+        if (!is_removed(id)) {
+            const PID marked = get_raw_clusterid(id) | kRemovedMask;
+            std::memcpy(get_clusterid_pt(id), &marked, sizeof(PID));
+            ++removed;
+        }
+    }
+    return removed;
+}
+
 inline maxheap<std::pair<float, PID>> HierarchicalNSW::search_knn(
     const float* rotated_query, size_t TOPK
 ) {
@@ -1495,12 +1566,13 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect(
     float est_dist = start_estimate_record.est_dist;
     float low_dist = start_estimate_record.low_dist;
 
-    // Insert initial candidate.
-    boundedKNN.insert({ResultRecord(est_dist, low_dist), ep_id});
+    // Insert initial candidate. A removed point is still explored, never returned.
+    if (!is_removed(ep_id)) {
+        boundedKNN.insert({ResultRecord(est_dist, low_dist), ep_id});
+        distk = est_dist;
+    }
     candidate_set.insert(ep_id, est_dist);
 
-    distk = est_dist;
-
     vl->set(ep_id);
 
     const size_t prefetch_size = (((padded_dim_ / 8) + 63) / 64) + 1;
@@ -1538,7 +1610,9 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect(
                 q_to_centroids, query_wrapper, candidate_id, candest
             );
 
-            bool flag_update_KNNs = boundedKNN.size() < TOPK || candest.low_dist < distk;
+            // A removed neighbor keeps its binary estimate and stays in candidate_set.
+            bool flag_update_KNNs = !is_removed(static_cast<PID>(candidate_id)) &&
+                                    (boundedKNN.size() < TOPK || candest.low_dist < distk);
 
             if (flag_update_KNNs) {
                 // Compute the full estimate if promising.

python_bindings/hnsw_bindings.cpp

@@ -164,6 +164,25 @@ class HnswIndex {
         max_elements_ = index_->max_elements();
     }
 
+    size_t remove(py::handle ids) {
+        auto ids_array = ensure_1d_integer_array(ids, "ids");
+        if (!built_) {
+            throw std::runtime_error("HnswIndex must be built or loaded before remove");
+        }
+        // Labels are dense: construct gives 0..n-1 and add continues from num_points.
+        const size_t count = index_->num_points();
+        std::vector<rabitqlib::PID> to_remove;
+        to_remove.reserve(static_cast<size_t>(ids_array.shape(0)));
+        for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) {
+            const int64_t id = ids_array.data()[i];
+            if (id < 0 || static_cast<uint64_t>(id) >= count) {
+                throw std::invalid_argument("ids must be in [0, num_points)");
+            }
+            to_remove.push_back(static_cast<rabitqlib::PID>(id));
+        }
+        return index_->remove(to_remove.data(), to_remove.size());
+    }
+
     [[nodiscard]] size_t num_points() const { return built_ ? index_->num_points() : 0; }
 
     py::tuple search(py::handle queries, size_t k, size_t ef = 0, size_t num_threads = 1) {
@@ -303,6 +322,14 @@ void register_hnsw(py::module_& m) {
             py::arg("fast_quantization") = false
         )
         .def("resize", &HnswIndex::resize, py::arg("max_elements"))
+        .def(
+            "remove",
+            &HnswIndex::remove,
+            py::arg("ids"),
+            "Exclude ids from later search results and return how many were newly "
+            "removed. Removed points stay in the graph and in num_points, and cannot "
+            "be restored."
+        )
         .def("save", &HnswIndex::save, py::arg("path"))
         .def_static("load", &HnswIndex::load, py::arg("path"))
         .def_property_readonly("dim", &HnswIndex::dim)

tests/failures/hnsw_add_allocation_test.cpp

@@ -224,5 +224,30 @@ TEST(HnswAddAllocationTest, EveryAllocationFailureLeavesUsableIndex) {
     }
     std::filesystem::remove(path);
 }
+
+TEST(HnswRemoveAllocationTest, FailureRemovesNothing) {
+    Fixture fixture;
+    HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1);
+    auto ids = fixture.clusters;
+    index.construct(
+        1, fixture.centroid.data(), kCapacity, fixture.data.data(), ids.data(), 1, false
+    );
+    const auto before =
+        index.search(fixture.data.data(), kCapacity, kCapacity, kCapacity, 1);
+    const std::array<PID, 3> labels{0, 4, 7};
+    bool failed = false;
+    {
+        // The only allocation in remove holds the resolved internal ids.
+        AllocationFailure failure(0, labels.size() * sizeof(PID));
+        try {
+            index.remove(labels.data(), labels.size());
+        } catch (const std::bad_alloc&) { failed = true; }
+    }
+    ASSERT_TRUE(failed);
+    EXPECT_EQ(
+        index.search(fixture.data.data(), kCapacity, kCapacity, kCapacity, 1), before
+    );
+    EXPECT_EQ(index.remove(labels.data(), labels.size()), labels.size());
+}
 }  // namespace
 }  // namespace rabitqlib::hnsw