From 0758b54079a49d96046f9b5676d698da093daeeb Mon Sep 17 00:00:00 2001 From: pma1999 Date: Thu, 1 Oct 2026 12:03:24 +0200 Subject: [PATCH 1/4] Remove points from an HNSW index with a deleted mark remove sets the high bit of each point's stored cluster ID. construct and load keep the cluster count at or below 2^31, so a real cluster ID never has that bit, and the getter masks it for every existing reader. Search keeps marked points out of boundedKNN but still adds them to candidate_set, so it walks through them as before. add still links new points to them. The arguments are external labels; every one is checked before any point is marked, as in IVF::remove. load masks the bit before its range check. --- .github/workflows/test.yaml | 2 +- include/rabitqlib/index/hnsw/hnsw.hpp | 88 ++++- tests/failures/hnsw_add_allocation_test.cpp | 25 ++ tests/unit/rabitqlib/index/hnsw_test.cpp | 371 ++++++++++++++++++++ 4 files changed, 478 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 19d260b6..ebfce9da 100644 --- a/.github/workflows/test.yaml +++ b/.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" diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 4cb1d218..c1b30433 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/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>> 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 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( 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( 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>> 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 internal_ids(n); + { + std::unique_lock 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> 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(candidate_id)) && + (boundedKNN.size() < TOPK || candest.low_dist < distk); if (flag_update_KNNs) { // Compute the full estimate if promising. diff --git a/tests/failures/hnsw_add_allocation_test.cpp b/tests/failures/hnsw_add_allocation_test.cpp index 404e8452..fd4e8ab8 100644 --- a/tests/failures/hnsw_add_allocation_test.cpp +++ b/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 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 diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 732f3aed..d55c6b4a 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -5,8 +5,11 @@ #include #include #include +#include #include #include +#include +#include #include #include #include @@ -39,6 +42,32 @@ struct HnswPruningTestAccess { index.mutually_connect_quant(incoming, candidates, 0); return {neighbors + 1, neighbors + 1 + index.get_list_count(neighbors)}; } + + static PID entry_point_label(const HierarchicalNSW& index) { + return index.get_external_label(index.enterpoint_node_); + } + + static PID internal_id(const HierarchicalNSW& index, PID label) { + return index.label_lookup_.at(label); + } + + static void set_raw_cluster_id(HierarchicalNSW& index, PID label, PID value) { + std::memcpy( + index.get_clusterid_pt(index.label_lookup_.at(label)), &value, sizeof(PID) + ); + } + + // Give internal slot i the label count - 1 - i, as parallel construction can + // when threads take slots out of order. + static void reverse_labels(HierarchicalNSW& index) { + const size_t count = index.num_points(); + index.label_lookup_.clear(); + for (size_t i = 0; i < count; ++i) { + const auto label = static_cast(count - 1 - i); + index.set_external_label(static_cast(i), label); + index.label_lookup_.emplace(label, static_cast(i)); + } + } }; namespace { @@ -737,5 +766,347 @@ TEST_F(HnswSaveTest, RejectsTruncatedRotatorWithoutLosingExistingIndex) { EXPECT_EQ(loaded.search(data_.data(), 1, 2, kCount, 1), expected); } +using Results = std::vector>>; + +struct RemoveCase { + MetricType metric; + size_t bits; +}; + +std::vector remove_cases() { + std::vector cases; + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + for (const size_t bits : {1U, 2U, 4U, 8U, 9U}) { + cases.push_back({metric, bits}); + } + } + return cases; +} + +// Constructs the first `count` points of `fixture` into an index that has room +// for `capacity`. +std::unique_ptr build_for_remove( + AddFixture& fixture, + size_t count, + size_t capacity, + const RemoveCase& config, + size_t num_threads = 1, + size_t M = 8 +) { + auto index = std::make_unique( + capacity, AddFixture::kDim, config.bits, M, 50, 100, config.metric + ); + index->construct( + 1, + fixture.centroid.data(), + count, + fixture.data.data(), + fixture.cluster_ids.data(), + num_threads, + false + ); + return index; +} + +// With k and ef at the point count, the search visits every reachable point. +Results search_everything( + HierarchicalNSW& index, const AddFixture& fixture, size_t queries +) { + const size_t count = index.num_points(); + return index.search(fixture.data.data(), queries, count, count, 1); +} + +// `before` without the removed labels, in the same order. +Results without_removed(const Results& before, const std::vector& removed) { + Results kept(before.size()); + for (size_t q = 0; q < before.size(); ++q) { + std::copy_if( + before[q].begin(), + before[q].end(), + std::back_inserter(kept[q]), + [&](const std::pair& hit) { return !removed[hit.second]; } + ); + } + return kept; +} + +std::vector saved_bytes(const HierarchicalNSW& index, const std::string& path) { + index.save(path.c_str()); + std::ifstream file(path, std::ios::binary); + return {std::istreambuf_iterator(file), std::istreambuf_iterator()}; +} + +TEST(HnswRemoveTest, ExcludesRemovedPointsAndKeepsSurvivors) { + constexpr size_t kCount = 120; + constexpr size_t kQueries = 8; + for (const auto& config : remove_cases()) { + SCOPED_TRACE( + ::testing::Message() + << "metric " << static_cast(config.metric) << " bits " << config.bits + ); + AddFixture fixture(kCount, 59); + auto index = build_for_remove(fixture, kCount, kCount, config); + const Results before = search_everything(*index, fixture, kQueries); + + // Every third point, the last one, and a repeated label. + std::vector labels; + std::vector removed(kCount, false); + for (PID label = 0; label < kCount; label += 3) { + labels.push_back(label); + removed[label] = true; + } + labels.push_back(kCount - 1); + removed[kCount - 1] = true; + labels.push_back(0); + const auto unique = + static_cast(std::count(removed.begin(), removed.end(), true)); + + EXPECT_EQ(index->remove(labels.data(), labels.size()), unique); + EXPECT_EQ(index->remove(labels.data(), labels.size()), 0U); + EXPECT_EQ(index->num_points(), kCount); + + // Survivors keep their order and their distances. + EXPECT_EQ( + search_everything(*index, fixture, kQueries), without_removed(before, removed) + ); + + // A small ef still fills k with live points only. + const auto top = index->search(fixture.data.data(), kQueries, 10, 10, 1); + for (const auto& hits : top) { + EXPECT_EQ(hits.size(), 10U); + for (const auto& hit : hits) { + EXPECT_FALSE(removed[hit.second]) << hit.second; + EXPECT_TRUE(std::isfinite(hit.first)); + } + } + } +} + +TEST(HnswRemoveTest, RemovingTheEntryPointKeepsSearchWorking) { + constexpr size_t kCount = 160; + for (const auto& config : remove_cases()) { + SCOPED_TRACE( + ::testing::Message() + << "metric " << static_cast(config.metric) << " bits " << config.bits + ); + AddFixture fixture(kCount, 61); + auto index = build_for_remove(fixture, kCount, kCount, config); + const PID entry = HnswPruningTestAccess::entry_point_label(*index); + ASSERT_EQ(index->remove(&entry, 1), 1U); + + // Every search starts at the removed entry point and must still return k + // live points without it. + const auto results = index->search(fixture.data.data(), kCount, 10, 10, 1); + for (const auto& hits : results) { + EXPECT_EQ(hits.size(), 10U); + for (const auto& hit : hits) { + EXPECT_NE(hit.second, entry); + } + } + const Results everything = search_everything(*index, fixture, 1); + EXPECT_EQ(everything[0].size(), kCount - 1); + } +} + +TEST(HnswRemoveTest, SearchReachesLivePointsThroughRemovedOnes) { + constexpr size_t kCount = 200; + constexpr size_t kStride = 10; + for (const auto& config : remove_cases()) { + SCOPED_TRACE( + ::testing::Message() + << "metric " << static_cast(config.metric) << " bits " << config.bits + ); + AddFixture fixture(kCount, 67); + auto index = build_for_remove(fixture, kCount, kCount, config, 1, 4); + + // Keep one point in ten. Most paths between the survivors now pass through + // removed points, which search must still walk through. + std::vector labels; + std::vector live; + for (PID label = 0; label < kCount; ++label) { + (label % kStride == 0 ? live : labels).push_back(label); + } + ASSERT_EQ(index->remove(labels.data(), labels.size()), labels.size()); + + const Results results = search_everything(*index, fixture, kCount); + for (size_t q = 0; q < kCount; q += kStride) { + std::vector found; + for (const auto& hit : results[q]) { + found.push_back(hit.second); + } + std::sort(found.begin(), found.end()); + EXPECT_EQ(found, live) << "query " << q; + } + } +} + +TEST(HnswRemoveTest, UsesLabelsNotInternalIds) { + constexpr size_t kCount = 256; + const RemoveCase config{METRIC_L2, 4}; + AddFixture fixture(kCount, 71); + auto index = build_for_remove(fixture, kCount, kCount, config); + HnswPruningTestAccess::reverse_labels(*index); + + // Every fifth label below kCount / 2: no removed label shares a slot with + // another removed label, so treating labels as slots would hide other points. + std::vector labels; + std::vector removed(kCount, false); + for (PID label = 0; label < kCount / 2; label += 5) { + labels.push_back(label); + removed[label] = true; + } + ASSERT_EQ(index->remove(labels.data(), labels.size()), labels.size()); + + const Results results = search_everything(*index, fixture, 1); + std::vector found(kCount, false); + for (const auto& hit : results[0]) { + found[hit.second] = true; + } + for (PID label = 0; label < kCount; ++label) { + EXPECT_EQ(found[label], !removed[label]) << "label " << label; + } +} + +TEST(HnswRemoveTest, SurvivesSaveLoadResizeAndLaterAdds) { + constexpr size_t kBuilt = 96; + constexpr size_t kAdded = 32; + constexpr size_t kTotal = kBuilt + kAdded; + const std::string path = ::testing::TempDir() + "rabitq_hnsw_remove_roundtrip.index"; + for (const auto& config : remove_cases()) { + SCOPED_TRACE( + ::testing::Message() + << "metric " << static_cast(config.metric) << " bits " << config.bits + ); + AddFixture fixture(kTotal, 73); + auto index = build_for_remove(fixture, kBuilt, kBuilt, config); + std::vector labels; + std::vector removed(kTotal, false); + for (PID label = 1; label < kBuilt; label += 4) { + labels.push_back(label); + removed[label] = true; + } + ASSERT_EQ(index->remove(labels.data(), labels.size()), labels.size()); + const Results expected = search_everything(*index, fixture, 4); + + index->save(path.c_str()); + HierarchicalNSW loaded; + loaded.load(path.c_str()); + EXPECT_EQ(search_everything(loaded, fixture, 4), expected); + EXPECT_EQ(loaded.remove(labels.data(), labels.size()), 0U); + + // Resizing copies the marks with the rest of the base layer. + index->resize(kTotal); + loaded.resize(kTotal); + EXPECT_EQ(search_everything(loaded, fixture, 4), expected); + + // New points can be linked through removed ones, and both copies grow the + // same way. + const float* extra = fixture.data.data() + (kBuilt * AddFixture::kDim); + index->add(extra, kAdded, fixture.cluster_ids.data()); + loaded.add(extra, kAdded, fixture.cluster_ids.data()); + const Results grown = search_everything(*index, fixture, kTotal); + EXPECT_EQ(search_everything(loaded, fixture, kTotal), grown); + for (size_t q = 0; q < kTotal; ++q) { + ASSERT_EQ(grown[q].size(), kTotal - labels.size()); + for (const auto& hit : grown[q]) { + EXPECT_FALSE(removed[hit.second]) << hit.second; + } + } + const auto top = index->search(extra, kAdded, 1, kTotal, 1); + for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(top[i][0].second, kBuilt + i); + } + } + std::remove(path.c_str()); +} + +TEST(HnswRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { + constexpr size_t kCount = 64; + const PID one = 3; + HierarchicalNSW unbuilt; + try { + unbuilt.remove(&one, 1); + FAIL() << "Removing from an index that was never built must fail"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "HNSW index must be constructed or loaded before remove" + ); + } + HierarchicalNSW unconstructed(kCount, AddFixture::kDim, 4, 8, 50); + EXPECT_THROW(unconstructed.remove(&one, 1), std::logic_error); + + AddFixture fixture(kCount, 79); + auto index = build_for_remove(fixture, kCount, kCount, {METRIC_L2, 4}); + const std::string path = ::testing::TempDir() + "rabitq_hnsw_remove_rejects.index"; + const auto expected = saved_bytes(*index, path); + + const std::vector unknown{3, kCount}; + try { + index->remove(unknown.data(), unknown.size()); + FAIL() << "An unknown label must be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW remove label is not in the index"); + } + EXPECT_THROW(index->remove(nullptr, 1), std::invalid_argument); + EXPECT_EQ(index->remove(nullptr, 0), 0U); + EXPECT_EQ(saved_bytes(*index, path), expected); + + // Label 3 was valid in the rejected call, so it is still there to remove. + EXPECT_EQ(index->remove(&one, 1), 1U); + std::remove(path.c_str()); +} + +TEST(HnswRemoveTest, RemovingEveryPointReturnsNoResults) { + constexpr size_t kBuilt = 64; + constexpr size_t kAdded = 16; + AddFixture fixture(kBuilt + kAdded, 83); + auto index = build_for_remove(fixture, kBuilt, kBuilt, {METRIC_IP, 1}); + std::vector labels(kBuilt); + for (PID label = 0; label < kBuilt; ++label) { + labels[label] = label; + } + ASSERT_EQ(index->remove(labels.data(), labels.size()), kBuilt); + for (const auto& hits : index->search(fixture.data.data(), 4, 10, 10, 1)) { + EXPECT_TRUE(hits.empty()); + } + for (const auto& hits : search_everything(*index, fixture, 4)) { + EXPECT_TRUE(hits.empty()); + } + + // The graph of removed points still carries new ones. + index->resize(kBuilt + kAdded); + const float* extra = fixture.data.data() + (kBuilt * AddFixture::kDim); + index->add(extra, kAdded, fixture.cluster_ids.data()); + const auto results = index->search(extra, kAdded, 1, kBuilt + kAdded, 1); + for (size_t i = 0; i < kAdded; ++i) { + ASSERT_EQ(results[i].size(), 1U); + EXPECT_GE(results[i][0].second, kBuilt); + } +} + +TEST(HnswRemoveTest, LoadRejectsAnOutOfRangeClusterEvenWithTheMark) { + constexpr size_t kCount = 32; + AddFixture fixture(kCount, 89); + auto index = build_for_remove(fixture, kCount, kCount, {METRIC_L2, 4}); + const PID one = 5; + ASSERT_EQ(index->remove(&one, 1), 1U); + + const std::string path = ::testing::TempDir() + "rabitq_hnsw_remove_bad_cluster.index"; + index->save(path.c_str()); + HierarchicalNSW loaded; + ASSERT_NO_THROW(loaded.load(path.c_str())); + + // Cluster 1 does not exist in a one-cluster index, marked or not. + HnswPruningTestAccess::set_raw_cluster_id(*index, one, 0x80000001U); + index->save(path.c_str()); + try { + loaded.load(path.c_str()); + FAIL() << "A marked out-of-range cluster ID must be rejected"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "HNSW: invalid or truncated index file"); + } + std::remove(path.c_str()); +} + } // namespace } // namespace rabitqlib::hnsw From 55fa9fe65cd8a200137cb34a6dc069e9234aaa34 Mon Sep 17 00:00:00 2001 From: pma1999 Date: Thu, 1 Oct 2026 12:03:24 +0200 Subject: [PATCH 2/4] Expose HNSW remove to Python --- python_bindings/hnsw_bindings.cpp | 27 +++++++ tests/python/test_hnsw.py | 124 ++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+) diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 5a3e1cf2..b7f304dd 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/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 to_remove; + to_remove.reserve(static_cast(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(id) >= count) { + throw std::invalid_argument("ids must be in [0, num_points)"); + } + to_remove.push_back(static_cast(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) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index 82ee4194..d479f25d 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -312,3 +312,127 @@ def test_added_points_survive_save_and_load( after_ids, after_dists = loaded.search(base_data, k=_TOPK, ef=_EF) np.testing.assert_array_equal(before_ids, after_ids) np.testing.assert_allclose(before_dists, after_dists, rtol=1e-5) + + +# ── remove ──────────────────────────────────────────────────────────────────── + +_NO_ID = np.iinfo(np.uint32).max + + +def _removable(nbits=4, metric="l2", count=120, dim=64, seed=61): + """A small built index with room for as many points again, plus those points.""" + rng = np.random.default_rng(seed) + data = rng.standard_normal((count * 2, dim)).astype(np.float32) + centroids = data[:count].mean(axis=0, keepdims=True) + idx = HnswIndex(dim, count, M=8, ef_construction=50, nbits=nbits, metric=metric) + idx.build(data[:count], centroids, np.zeros(count, dtype=np.uint32)) + return idx, data[:count], data[count:] + + +def _everything(index, queries): + """With k and ef at the point count, search visits every reachable point.""" + return index.search(queries, k=index.num_points, ef=index.num_points) + + +@pytest.mark.parametrize("nbits", [1, 2, 4, 8, 9]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_remove_excludes_points_and_survives_reload(tmp_path, nbits, metric): + idx, old, new = _removable(nbits=nbits, metric=metric) + count = len(old) + queries = new[:5] + before_ids, before_dists = _everything(idx, queries) + + removed = np.array([0, 5, 6, count - 1, 5]) + assert idx.remove(removed) == 4 + assert idx.remove(removed) == 0 + assert idx.remove([]) == 0 + assert idx.num_points == count + + def check(index): + ids, dists = _everything(index, queries) + alive = count - 4 + for q in range(len(queries)): + keep = ~np.isin(before_ids[q], removed) + np.testing.assert_array_equal(ids[q, :alive], before_ids[q][keep]) + np.testing.assert_array_equal(dists[q, :alive], before_dists[q][keep]) + assert np.all(ids[q, alive:] == _NO_ID) + assert np.all(np.isinf(dists[q, alive:])) + top, top_dists = index.search(queries, k=10, ef=10) + assert not np.isin(top, removed).any() + assert np.isfinite(top_dists).all() + + check(idx) + path = tmp_path / "removed.index" + idx.save(str(path)) + loaded = HnswIndex.load(str(path)) + check(loaded) + assert loaded.remove(removed) == 0 + + # New points link through removed ones, which stay hidden. + for index in (idx, loaded): + index.resize(count + len(new)) + index.add(new) + ids, dists = _everything(loaded, queries) + assert not np.isin(ids, removed).any() + assert np.count_nonzero(ids[0] != _NO_ID) == count + len(new) - 4 + expected_ids, expected_dists = _everything(idx, queries) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(dists, expected_dists) + + +def test_remove_rejects_invalid_ids_without_removing_anything(): + idx, old, new = _removable() + count = len(old) + before = _everything(idx, new[:3]) + # 2**32 + 1 would wrap to point 1 if it were cast to uint32 first. + for bad in (-1, count, 2**32 + 1): + with pytest.raises(ValueError, match="ids"): + idx.remove([0, 1, bad]) + with pytest.raises(ValueError, match="1D"): + idx.remove(np.zeros((2, 2), dtype=np.int64)) + after = _everything(idx, new[:3]) + np.testing.assert_array_equal(before[0], after[0]) + np.testing.assert_array_equal(before[1], after[1]) + assert idx.remove([0, 1]) == 2 + + unbuilt = HnswIndex(64, 4, M=8, ef_construction=50, nbits=4) + with pytest.raises(RuntimeError, match="built or loaded"): + unbuilt.remove([0]) + + +def test_removing_every_point_returns_only_sentinels(): + idx, old, new = _removable() + assert idx.remove(np.arange(len(old))) == len(old) + ids, dists = idx.search(new[:2], k=5, ef=10) + assert np.all(ids == _NO_ID) + assert np.all(np.isinf(dists)) + + +@pytest.mark.parametrize( + "bad", + [ + [2.7], + [-0.5], + # Integral floats are refused too: the dtype is the contract. + np.array([1.0, 2.0]), + np.array([True, False]), # a boolean mask is not a list of ids + ["3"], + np.array([1, 2], dtype=object), + ], +) +def test_remove_refuses_non_integer_ids(bad): + idx, old, new = _removable() + before = _everything(idx, new[:3]) + with pytest.raises(ValueError, match="integers"): + idx.remove(bad) + after = _everything(idx, new[:3]) + np.testing.assert_array_equal(before[0], after[0]) + + +def test_remove_accepts_any_integer_dtype_and_empty_input(): + idx, _, _ = _removable() + assert idx.remove([]) == 0 + assert idx.remove(np.array([], dtype=np.float64)) == 0 # no dtype to check + assert idx.remove([1, 2]) == 2 + assert idx.remove(np.array([3, 4], dtype=np.uint8)) == 2 + assert idx.remove(np.array([5], dtype=np.int16)) == 1 From fc475cf4bac040ebaa1cc103930b7b6b357f96c3 Mon Sep 17 00:00:00 2001 From: pma1999 Date: Thu, 1 Oct 2026 12:03:24 +0200 Subject: [PATCH 3/4] Document removing points from an HNSW index Cover the API, the deleted mark in the file layout, scaling ef after removals as a caller choice, and which older releases refuse files with removed points. --- DEVELOPMENT.md | 15 +++++++++++ docs/docs/index/hnsw.md | 59 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 2543c4b0..f5ff34fc 100644 --- a/DEVELOPMENT.md +++ b/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 diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 69ba7050..17a4c0f5 100644 --- a/docs/docs/index/hnsw.md +++ b/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. From c8e7507a0ac9e2a256bb02af4cd0a39237689bff Mon Sep 17 00:00:00 2001 From: pma1999 Date: Thu, 1 Oct 2026 12:47:46 +0200 Subject: [PATCH 4/4] Compare HNSW remove tests with the index before removal add can leave a point unreachable from the entry point, with or without removals, so a test that expects every point to be found fails at random across rotators. Compare each search with the same index before the removal, or with a copy that removes nothing, which is the property remove has to keep. --- tests/python/test_hnsw.py | 19 ++++-- tests/unit/rabitqlib/index/hnsw_test.cpp | 79 +++++++++++++----------- 2 files changed, 57 insertions(+), 41 deletions(-) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index d479f25d..fd8e7b2e 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -341,6 +341,8 @@ def test_remove_excludes_points_and_survives_reload(tmp_path, nbits, metric): count = len(old) queries = new[:5] before_ids, before_dists = _everything(idx, queries) + plain_path = tmp_path / "plain.index" + idx.save(str(plain_path)) removed = np.array([0, 5, 6, count - 1, 5]) assert idx.remove(removed) == 4 @@ -368,16 +370,25 @@ def check(index): check(loaded) assert loaded.remove(removed) == 0 - # New points link through removed ones, which stay hidden. - for index in (idx, loaded): + # New points link through removed ones, which stay hidden. add can leave a + # point unreachable with or without removals, so compare with a copy saved + # before the removal instead of expecting every point. + plain = HnswIndex.load(str(plain_path)) + for index in (idx, loaded, plain): index.resize(count + len(new)) index.add(new) ids, dists = _everything(loaded, queries) - assert not np.isin(ids, removed).any() - assert np.count_nonzero(ids[0] != _NO_ID) == count + len(new) - 4 expected_ids, expected_dists = _everything(idx, queries) np.testing.assert_array_equal(ids, expected_ids) np.testing.assert_array_equal(dists, expected_dists) + plain_ids, plain_dists = _everything(plain, queries) + for q in range(len(queries)): + keep = ~np.isin(plain_ids[q], removed) & (plain_ids[q] != _NO_ID) + found = np.count_nonzero(keep) + assert found >= 0.95 * (count + len(new) - 4) + np.testing.assert_array_equal(ids[q, :found], plain_ids[q][keep]) + np.testing.assert_array_equal(dists[q, :found], plain_dists[q][keep]) + assert np.all(ids[q, found:] == _NO_ID) def test_remove_rejects_invalid_ids_without_removing_anything(): diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index d55c6b4a..badd21e4 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -808,7 +808,9 @@ std::unique_ptr build_for_remove( return index; } -// With k and ef at the point count, the search visits every reachable point. +// With k and ef at the point count, the search returns every point reachable +// from the entry point. HNSW does not promise that every point is reachable, so +// the tests compare against the same index without removals. Results search_everything( HierarchicalNSW& index, const AddFixture& fixture, size_t queries ) { @@ -830,6 +832,18 @@ Results without_removed(const Results& before, const std::vector& removed) return kept; } +// A copy through save and load: the same rotator, codes and graph. +std::unique_ptr copy_of( + const HierarchicalNSW& index, const std::string& name +) { + const std::string path = ::testing::TempDir() + name; + index.save(path.c_str()); + auto copy = std::make_unique(); + copy->load(path.c_str()); + std::remove(path.c_str()); + return copy; +} + std::vector saved_bytes(const HierarchicalNSW& index, const std::string& path) { index.save(path.c_str()); std::ifstream file(path, std::ios::binary); @@ -891,7 +905,10 @@ TEST(HnswRemoveTest, RemovingTheEntryPointKeepsSearchWorking) { ); AddFixture fixture(kCount, 61); auto index = build_for_remove(fixture, kCount, kCount, config); + const Results before = search_everything(*index, fixture, 1); const PID entry = HnswPruningTestAccess::entry_point_label(*index); + std::vector removed(kCount, false); + removed[entry] = true; ASSERT_EQ(index->remove(&entry, 1), 1U); // Every search starts at the removed entry point and must still return k @@ -903,8 +920,7 @@ TEST(HnswRemoveTest, RemovingTheEntryPointKeepsSearchWorking) { EXPECT_NE(hit.second, entry); } } - const Results everything = search_everything(*index, fixture, 1); - EXPECT_EQ(everything[0].size(), kCount - 1); + EXPECT_EQ(search_everything(*index, fixture, 1), without_removed(before, removed)); } } @@ -918,25 +934,24 @@ TEST(HnswRemoveTest, SearchReachesLivePointsThroughRemovedOnes) { ); AddFixture fixture(kCount, 67); auto index = build_for_remove(fixture, kCount, kCount, config, 1, 4); + const Results before = search_everything(*index, fixture, kCount); // Keep one point in ten. Most paths between the survivors now pass through - // removed points, which search must still walk through. + // removed points, which search must still walk through to reach every + // survivor it reached before. std::vector labels; - std::vector live; + std::vector removed(kCount, false); for (PID label = 0; label < kCount; ++label) { - (label % kStride == 0 ? live : labels).push_back(label); + if (label % kStride != 0) { + labels.push_back(label); + removed[label] = true; + } } ASSERT_EQ(index->remove(labels.data(), labels.size()), labels.size()); - const Results results = search_everything(*index, fixture, kCount); - for (size_t q = 0; q < kCount; q += kStride) { - std::vector found; - for (const auto& hit : results[q]) { - found.push_back(hit.second); - } - std::sort(found.begin(), found.end()); - EXPECT_EQ(found, live) << "query " << q; - } + EXPECT_EQ( + search_everything(*index, fixture, kCount), without_removed(before, removed) + ); } } @@ -946,6 +961,7 @@ TEST(HnswRemoveTest, UsesLabelsNotInternalIds) { AddFixture fixture(kCount, 71); auto index = build_for_remove(fixture, kCount, kCount, config); HnswPruningTestAccess::reverse_labels(*index); + const Results before = search_everything(*index, fixture, 1); // Every fifth label below kCount / 2: no removed label shares a slot with // another removed label, so treating labels as slots would hide other points. @@ -957,14 +973,7 @@ TEST(HnswRemoveTest, UsesLabelsNotInternalIds) { } ASSERT_EQ(index->remove(labels.data(), labels.size()), labels.size()); - const Results results = search_everything(*index, fixture, 1); - std::vector found(kCount, false); - for (const auto& hit : results[0]) { - found[hit.second] = true; - } - for (PID label = 0; label < kCount; ++label) { - EXPECT_EQ(found[label], !removed[label]) << "label " << label; - } + EXPECT_EQ(search_everything(*index, fixture, 1), without_removed(before, removed)); } TEST(HnswRemoveTest, SurvivesSaveLoadResizeAndLaterAdds) { @@ -979,6 +988,7 @@ TEST(HnswRemoveTest, SurvivesSaveLoadResizeAndLaterAdds) { ); AddFixture fixture(kTotal, 73); auto index = build_for_remove(fixture, kBuilt, kBuilt, config); + const auto plain = copy_of(*index, "rabitq_hnsw_remove_plain.index"); std::vector labels; std::vector removed(kTotal, false); for (PID label = 1; label < kBuilt; label += 4) { @@ -999,23 +1009,18 @@ TEST(HnswRemoveTest, SurvivesSaveLoadResizeAndLaterAdds) { loaded.resize(kTotal); EXPECT_EQ(search_everything(loaded, fixture, 4), expected); - // New points can be linked through removed ones, and both copies grow the - // same way. + // New points can be linked through removed ones. All three copies grow the + // same graph, and the removed points stay hidden in the two that have them. const float* extra = fixture.data.data() + (kBuilt * AddFixture::kDim); - index->add(extra, kAdded, fixture.cluster_ids.data()); - loaded.add(extra, kAdded, fixture.cluster_ids.data()); + plain->resize(kTotal); + for (HierarchicalNSW* copy : {index.get(), &loaded, plain.get()}) { + copy->add(extra, kAdded, fixture.cluster_ids.data()); + } const Results grown = search_everything(*index, fixture, kTotal); EXPECT_EQ(search_everything(loaded, fixture, kTotal), grown); - for (size_t q = 0; q < kTotal; ++q) { - ASSERT_EQ(grown[q].size(), kTotal - labels.size()); - for (const auto& hit : grown[q]) { - EXPECT_FALSE(removed[hit.second]) << hit.second; - } - } - const auto top = index->search(extra, kAdded, 1, kTotal, 1); - for (size_t i = 0; i < kAdded; ++i) { - EXPECT_EQ(top[i][0].second, kBuilt + i); - } + EXPECT_EQ( + grown, without_removed(search_everything(*plain, fixture, kTotal), removed) + ); } std::remove(path.c_str()); }