diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index e98b5093..8e8ed6f3 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -458,6 +458,9 @@ inline HierarchicalNSW::HierarchicalNSW( , element_levels_(max_elements) , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) { validate_metric_type(metric_type); + if (M < 2) { + throw std::invalid_argument("HNSW M must be at least 2"); + } max_elements_ = max_elements; dim_ = dim; rotator_.reset(choose_rotator( @@ -652,7 +655,7 @@ inline void HierarchicalNSW::load(const char* filename) { loaded.padded_dim_ != round_up_to_multiple(loaded.dim_, 64) || loaded.ex_bits_ > 8 || (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || - loaded.M_ == 0 || loaded.M_ > 10000 || loaded.maxM_ != loaded.M_ || + loaded.M_ < 2 || loaded.M_ > 10000 || loaded.maxM_ != loaded.M_ || loaded.maxM0_ != 2 * loaded.M_ || loaded.ef_construction_ < loaded.M_ || (element_count == 0 && (loaded.maxlevel_ != -1 || loaded.enterpoint_node_ != kPidMax)) || @@ -676,8 +679,7 @@ inline void HierarchicalNSW::load(const char* filename) { loaded.offsetExData_ != expected_ex_offset || loaded.size_data_per_element_ != expected_stride || loaded.size_links_per_element_ != (loaded.maxM_ + 1) * sizeof(PID) || - (loaded.M_ > 1 && (!std::isfinite(loaded.mult_) || loaded.mult_ <= 0.0)) || - (loaded.M_ == 1 && !std::isinf(loaded.mult_))) { + !std::isfinite(loaded.mult_) || loaded.mult_ <= 0.0) { invalid_file(); } @@ -898,6 +900,9 @@ inline void HierarchicalNSW::construct( size_t num_threads = 0, bool faster = false ) { + if (num_cluster_ != 0 || cur_element_count_ != 0) { + throw std::logic_error("HNSW index is already constructed or loaded"); + } if (cluster_num == 0 || cluster_num > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("HNSW cluster count is out of range"); } @@ -917,32 +922,55 @@ inline void HierarchicalNSW::construct( } } - num_cluster_ = cluster_num; - const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float); - centroids_memory_ = memory::huge_page_allocate(centroids_bytes); - if (centroids_memory_ == nullptr) { - throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids"); - } + try { + num_cluster_ = cluster_num; + const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float); + centroids_memory_ = memory::huge_page_allocate(centroids_bytes); + if (centroids_memory_ == nullptr) { + throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids" + ); + } - for (size_t i = 0; i < cluster_num; ++i) { - this->rotator_->rotate( - centroids + (i * dim_), - reinterpret_cast(centroids_memory_) + (i * padded_dim_) - ); - } + for (size_t i = 0; i < cluster_num; ++i) { + this->rotator_->rotate( + centroids + (i * dim_), + reinterpret_cast(centroids_memory_) + (i * padded_dim_) + ); + } - quant::RabitqConfig config; - if (faster) { - config = quant::faster_config(padded_dim_, ex_bits_ + 1); - } + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } - rawDataPtr_ = data; - rabitqlib::ivf::parallel_for( - 0, - data_num, - num_threads, - [&](size_t idx, size_t /*threadId*/) { add_point(idx, cluster_ids[idx], config); } - ); + rawDataPtr_ = data; + rabitqlib::ivf::parallel_for( + 0, + data_num, + num_threads, + [&](size_t idx, size_t /*threadId*/) { + add_point(idx, cluster_ids[idx], config); + } + ); + } catch (...) { + // parallel_for joins every worker before propagating an exception. + // Discard the partial graph, retaining capacity and rotation for a retry. + for (size_t i = 0; i < cur_element_count_; ++i) { + std::free(linkLists_[i]); + linkLists_[i] = nullptr; + element_levels_[i] = 0; + } + cur_element_count_ = 0; + label_lookup_.clear(); + enterpoint_node_ = kPidMax; + maxlevel_ = -1; + memory::aligned_deallocate(centroids_memory_); + centroids_memory_ = nullptr; + num_cluster_ = 0; + rawDataPtr_ = nullptr; + throw; + } + rawDataPtr_ = nullptr; } inline void HierarchicalNSW::add_point( @@ -967,6 +995,7 @@ inline void HierarchicalNSW::add_point( } cur_c = cur_element_count_; + linkLists_[cur_c] = nullptr; cur_element_count_++; label_lookup_[label] = cur_c; curlevel = get_random_level(mult_); @@ -980,10 +1009,10 @@ inline void HierarchicalNSW::add_point( element_levels_[cur_c] = curlevel; std::unique_lock templock(global_); int maxlevelcopy = maxlevel_; + PID curr_obj = enterpoint_node_; if (curlevel <= maxlevelcopy) { templock.unlock(); } - PID curr_obj = enterpoint_node_; // initialize the current memory. memset( @@ -1070,7 +1099,7 @@ inline void HierarchicalNSW::add_point( inline maxheap> HierarchicalNSW::search_base_layer( PID ep_id, PID cur_c, int layer ) { - VisitedSet* vl = visited_list_pool_->get_free_vislist(); + std::unique_ptr vl(visited_list_pool_->get_free_vislist()); maxheap> top_candidates; minheap> candidate_set; @@ -1143,7 +1172,8 @@ inline maxheap> HierarchicalNSW::search_base_layer( } } } - visited_list_pool_->release_vis_list(vl); + visited_list_pool_->release_vis_list(vl.get()); + vl.release(); return top_candidates; } diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index 025c0297..4558ff29 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -1,5 +1,5 @@ // Insertion into a built index: the construct path with get_quant_dist in place -// of get_data_dist, because rawDataPtr_ dangles once construct returns. +// of get_data_dist, because raw input is borrowed only during construct. // // add routes, draws every level, allocates every link list and quantizes every // point before reserving one block of slots and linking in parallel. Levels come @@ -356,10 +356,10 @@ inline void HierarchicalNSW::add_point_quant( std::unique_lock templock(global_); int maxlevelcopy = maxlevel_; + PID curr_obj = enterpoint_node_; if (curlevel <= maxlevelcopy) { templock.unlock(); } - PID curr_obj = enterpoint_node_; const PID previous_entry_point = curr_obj; linkLists_[cur_c] = link_list.release(); diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index b7f304dd..703f17a9 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -84,8 +84,6 @@ class HnswIndex { throw std::invalid_argument("cluster_ids contains an out-of-range value"); } } - num_clusters_ = num_clusters; - // Ensure cluster_ids are writable for the C++ API by making a copy std::vector cluster_ids_vec( static_cast(cluster_ids_array.shape(0)) @@ -105,6 +103,7 @@ class HnswIndex { num_threads, fast_quantization ); + num_clusters_ = num_clusters; built_ = true; } diff --git a/tests/failures/hnsw_add_allocation_test.cpp b/tests/failures/hnsw_add_allocation_test.cpp index fd4e8ab8..dbe9af55 100644 --- a/tests/failures/hnsw_add_allocation_test.cpp +++ b/tests/failures/hnsw_add_allocation_test.cpp @@ -86,6 +86,63 @@ struct Fixture { } }; +TEST(HnswConstructionAllocationTest, EveryAllocationFailureLeavesEmptyReusableIndex) { + Fixture fixture; + for (const size_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + size_t failures = 0; + bool succeeded = false; + for (size_t allocation = 0; allocation < 512; ++allocation) { + SCOPED_TRACE(allocation); + HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); + auto clusters = fixture.clusters; + bool failed = false; + { + AllocationFailure failure(allocation); + try { + index.construct( + 1, + fixture.centroid.data(), + kAdded, + fixture.data.data(), + clusters.data(), + threads, + false + ); + } catch (const std::bad_alloc&) { failed = true; } + } + if (!failed) { + succeeded = true; + break; + } + ++failures; + EXPECT_EQ(index.num_points(), 0U); + EXPECT_EQ(index.num_clusters(), 0U); + const auto empty = index.search(fixture.data.data(), 1, 1, kCapacity, 1); + ASSERT_EQ(empty.size(), 1U); + EXPECT_TRUE(empty[0].empty()); + ASSERT_NO_THROW(index.construct( + 1, + fixture.centroid.data(), + kAdded, + fixture.data.data(), + clusters.data(), + threads, + false + )); + EXPECT_EQ(index.num_points(), kAdded); + EXPECT_EQ(index.num_clusters(), 1U); + const auto results = index.search(fixture.data.data(), kAdded, 1, kCapacity, 1); + for (size_t i = 0; i < kAdded; ++i) { + ASSERT_EQ(results[i].size(), 1U); + EXPECT_EQ(results[i][0].second, i); + } + } + EXPECT_GT(failures, 0U); + EXPECT_TRUE(succeeded); + } +} + TEST(HnswAddAllocationTest, AllocatesReturnedIdsBeforeMutation) { Fixture fixture; HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index 3bb12b0f..fbff67bc 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -29,6 +29,33 @@ def test_is_built(built_hnsw): assert built_hnsw.is_built +@pytest.mark.parametrize("m", [0, 1]) +def test_rejects_too_small_m(m): + with pytest.raises(ValueError, match="^HNSW M must be at least 2$"): + HnswIndex(DIM, N_VECTORS, M=m) + + +@pytest.mark.parametrize("reload", [False, True]) +def test_rejected_rebuild_preserves_index(base_data, clusters, tmp_path, reload): + idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + centroids, cluster_ids = clusters + idx.build(base_data, centroids, cluster_ids) + if reload: + path = str(tmp_path / "hnsw.index") + idx.save(path) + idx = HnswIndex.load(path) + before = idx.search(base_data[:10], k=3, ef=_EF) + with pytest.raises( + RuntimeError, match="^HNSW index is already constructed or loaded$" + ): + idx.build(base_data, centroids[:1], np.zeros(N_VECTORS, dtype=np.uint32)) + assert idx.is_built + assert idx.num_clusters == N_CLUSTERS + after = idx.search(base_data[:10], k=3, ef=_EF) + for actual, expected in zip(after, before, strict=True): + np.testing.assert_array_equal(actual, expected) + + def test_properties(built_hnsw): assert built_hnsw.dim == DIM assert built_hnsw.nbits == 4 diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index ec862ada..11efd703 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -645,6 +645,19 @@ TEST(HnswAddTest, AddIntoAnEmptyIndexBuildsFromScratch) { EXPECT_EQ(exact, kCount); } +TEST(HnswConstructionTest, RejectsTooSmallM) { + for (const size_t M : {0U, 1U}) { + SCOPED_TRACE(M); + try { + HierarchicalNSW index(8, 64, 4, M, 10); + FAIL() << "M below 2 must be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW M must be at least 2"); + } + } + EXPECT_NO_THROW(HierarchicalNSW(8, 64, 4, 2, 10)); +} + TEST(HnswConstructionTest, RejectsInvalidInputsBeforeChangingIndex) { constexpr size_t kDim = 64; std::vector data(kDim, 1.0F); @@ -705,6 +718,40 @@ TEST_F(HnswSaveTest, RejectsUnopenableDestination) { } } +TEST_F(HnswSaveTest, RejectedRebuildPreservesConstructedAndLoadedIndexes) { + // Multiple original centroids expose a rejected rebuild replacing the + // centroid buffer with a smaller one while retaining the old cluster IDs. + std::vector centroids(2 * kDim, 0.0F); + centroids[kDim] = 1.0F; + for (size_t i = 0; i < kCount; ++i) { + cluster_ids_[i] = i % 2; + } + HierarchicalNSW original(kCount, kDim, 4, 4, 10); + original.construct(2, centroids.data(), kCount, data_.data(), cluster_ids_.data(), 1); + original.save(path_.c_str()); + HierarchicalNSW loaded; + loaded.load(path_.c_str()); + std::vector new_clusters(kCount, 0); + for (auto* index : {&original, &loaded}) { + const auto before = index->search(data_.data(), kCount, 2, kCount, 1); + try { + index->construct( + 1, centroid_.data(), kCount, data_.data(), new_clusters.data(), 1 + ); + FAIL() << "A second construction must be rejected"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW index is already constructed or loaded"); + } + EXPECT_EQ(index->num_points(), kCount); + EXPECT_EQ(index->num_clusters(), 2U); + EXPECT_EQ(index->search(data_.data(), kCount, 2, kCount, 1), before); + index->save(path_.c_str()); + HierarchicalNSW roundtrip; + roundtrip.load(path_.c_str()); + EXPECT_EQ(roundtrip.search(data_.data(), kCount, 2, kCount, 1), before); + } +} + TEST_F(HnswSaveTest, ReportsWriteOrCloseFailure) { if (!std::filesystem::exists("/dev/full")) { GTEST_SKIP() << "/dev/full is unavailable";