From af840cd29a2f7567c5faa517638f0fbb0a4ff358 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:14:13 +0800 Subject: [PATCH 01/12] Add an inverse for the packed extra-bit codes --- .../rabitqlib/quantization/pack_excode.hpp | 10 ++++ .../rabitqlib/simd/pack_excode_dispatch.hpp | 5 ++ src/simd/portable_generic.cpp | 46 +++++++++++++++++++ .../rabitqlib/utils/portable_kernel_test.cpp | 8 ++++ 4 files changed, 69 insertions(+) diff --git a/include/rabitqlib/quantization/pack_excode.hpp b/include/rabitqlib/quantization/pack_excode.hpp index 044a43c5..75b39719 100644 --- a/include/rabitqlib/quantization/pack_excode.hpp +++ b/include/rabitqlib/quantization/pack_excode.hpp @@ -85,4 +85,14 @@ inline void packing_rabitqplus_code( throw std::invalid_argument("ex_bits must be in [1, 8]"); } } + +/** + * @brief Inverse of packing_rabitqplus_code. Scalar for every width; the packed + * layout does not depend on which backend wrote it. + */ +inline void unpacking_rabitqplus_code( + const uint8_t* o_compact, uint8_t* o_raw, size_t dim, size_t ex_bits +) { + ::rabitqlib::simd::unpacking_excode_generic(o_compact, o_raw, dim, ex_bits); +} } // namespace rabitqlib::quant::rabitq_impl::ex_bits diff --git a/include/rabitqlib/simd/pack_excode_dispatch.hpp b/include/rabitqlib/simd/pack_excode_dispatch.hpp index a1765fb3..a73117c7 100644 --- a/include/rabitqlib/simd/pack_excode_dispatch.hpp +++ b/include/rabitqlib/simd/pack_excode_dispatch.hpp @@ -17,6 +17,11 @@ void packing_6bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_7bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +// Inverse of the packers above, for bits in [1, 8]. +void unpacking_excode_generic( + const uint8_t* o_compact, uint8_t* o_raw, size_t dim, size_t bits +); + void packing_2bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_2bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp index 2f78e400..d4804165 100644 --- a/src/simd/portable_generic.cpp +++ b/src/simd/portable_generic.cpp @@ -53,6 +53,19 @@ void pack(const uint8_t* raw, uint8_t* compact, size_t dim) { } } +// Inverse of pack, over the same mapping. Not a search path, so scalar only. +template +void unpack(const uint8_t* compact, uint8_t* raw, size_t dim) { + for (size_t i = 0; i < dim; ++i) { + unsigned code = 0; + for (size_t b = 0; b < Bits; ++b) { + const size_t pos = packed_bit(i, b); + code |= ((compact[pos / 8] >> (pos % 8)) & 1U) << b; + } + raw[i] = static_cast(code); + } +} + template float excode_ip(const float* query, const uint8_t* compact, size_t dim) { double sum = 0; @@ -102,6 +115,39 @@ RABITQ_PACK(6) RABITQ_PACK(7) #undef RABITQ_PACK +void unpacking_excode_generic( + const uint8_t* compact, uint8_t* raw, size_t dim, size_t bits +) { + switch (bits) { + case 1: + unpack<1>(compact, raw, dim); + break; + case 2: + unpack<2>(compact, raw, dim); + break; + case 3: + unpack<3>(compact, raw, dim); + break; + case 4: + unpack<4>(compact, raw, dim); + break; + case 5: + unpack<5>(compact, raw, dim); + break; + case 6: + unpack<6>(compact, raw, dim); + break; + case 7: + unpack<7>(compact, raw, dim); + break; + case 8: + unpack<8>(compact, raw, dim); + break; + default: + throw std::invalid_argument("ex_bits must be in [1, 8]"); + } +} + namespace excode_ipimpl { #define RABITQ_IP(Block, Bits) \ float ip##Block##_fxu##Bits##_generic( \ diff --git a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp index 7d781ce3..30deaac4 100644 --- a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp +++ b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp @@ -53,6 +53,14 @@ TEST(PortableKernels, PackingMatchesDispatchedBytesAndDotProducts) { pack[bits - 2](raw.data(), scalar.data(), dim); EXPECT_EQ(scalar, dispatched); } + + // The inverse has to return exactly what was packed, whichever + // backend did the packing. + std::vector roundtrip(dim, 0xFF); + quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( + dispatched.data(), roundtrip.data(), dim, bits + ); + EXPECT_EQ(roundtrip, raw); EXPECT_EQ(ip[bits - 1](query.data(), dispatched.data(), dim), expected); EXPECT_EQ( select_excode_ipfunc(bits)(query.data(), dispatched.data(), dim), expected From 7bc418f61457a99aae563d54c12d811757c89a08 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:14:19 +0800 Subject: [PATCH 02/12] Add HNSW insertion after construction using the stored codes --- include/rabitqlib/index/hnsw/hnsw.hpp | 16 + include/rabitqlib/index/hnsw/hnsw_quant.hpp | 491 ++++++++++++++++++++ tests/unit/rabitqlib/index/hnsw_test.cpp | 262 +++++++++++ 3 files changed, 769 insertions(+) create mode 100644 include/rabitqlib/index/hnsw/hnsw_quant.hpp diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 5c778f46..f4f49105 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -61,11 +61,16 @@ class HierarchicalNSW { [[nodiscard]] size_t ef_construction() const { return ef_construction_; } [[nodiscard]] MetricType metric_type() const { return metric_type_; } [[nodiscard]] size_t max_elements() const { return max_elements_; } + [[nodiscard]] size_t num_points() const { return cur_element_count_; } void save(const char*) const; void load(const char*); void construct(size_t, const float*, size_t, const float*, PID*, size_t, bool); + + // Inserts points using the stored codes, labelling them num_points() + i. + // Single-threaded, and throws rather than growing past max_elements(). + std::vector add(const float*, size_t, const PID*, bool = false); std::vector>> search( const float*, size_t, size_t, size_t, size_t ); @@ -355,6 +360,14 @@ class HierarchicalNSW { void add_point(PID, PID, const quant::RabitqConfig&); + // Post-construction insertion, all defined in hnsw_quant.hpp. + void reconstruct_rotated(PID, float*) const; + float get_quant_dist(PID target, PID query) const; + maxheap> search_base_layer_quant(PID, PID, int); + void get_neighbors_by_heuristic2_quant(maxheap>&, size_t); + PID mutually_connect_quant(PID, maxheap>&, int); + PID add_point_quant(const float*, PID, const quant::RabitqConfig&); + maxheap> search_base_layer(PID, PID, int); PID mutually_connect_new_element(PID, maxheap>&, int); @@ -1520,3 +1533,6 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( } } // namespace rabitqlib::hnsw + +// Included last: these definitions need the complete class. +#include "rabitqlib/index/hnsw/hnsw_quant.hpp" // IWYU pragma: export diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp new file mode 100644 index 00000000..d4991e3a --- /dev/null +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -0,0 +1,491 @@ +// The construct path with get_quant_dist in place of get_data_dist: rawDataPtr_ +// dangles once construct returns, so these link using the stored codes. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/pack_excode.hpp" +#include "rabitqlib/quantization/rabitq.hpp" + +namespace rabitqlib::hnsw { + +namespace detail { +// The reconstructed query of get_quant_dist, reused while the query does not +// move. Cleared at the start of every insert so it cannot outlive its index. +struct QuantQuery { + PID id = kPidMax; + std::vector rotated; + std::optional> prepared; + + void clear() { + id = kPidMax; + prepared.reset(); + } +}; + +inline QuantQuery& quant_query() { + thread_local QuantQuery query; + return query; +} +} // namespace detail + +// Code per dimension is (sign << ex_bits) + ex, up to nine bits, scaled by +// f_rescale_ex. +inline void HierarchicalNSW::reconstruct_rotated(PID id, float* out) const { + thread_local std::vector combined; + thread_local std::vector ex_codes; + combined.resize(padded_dim_); + + const ConstBinDataMap bin(get_bindata_by_internalid(id), padded_dim_); + const uint8_t* sign_code = bin.bin_code(); + for (size_t word = 0; word * 64 < padded_dim_; ++word) { + uint64_t bits = 0; + std::memcpy(&bits, sign_code + (word * sizeof(uint64_t)), sizeof(bits)); + const size_t base = word * 64; + const size_t span = std::min(64, padded_dim_ - base); + for (size_t j = 0; j < span; ++j) { + combined[base + j] = static_cast((bits >> (63 - j)) & 1ULL); + } + } + + const auto* centroid = reinterpret_cast(centroids_memory_) + + (get_clusterid_by_internalid(id) * padded_dim_); + + if (ex_bits_ == 0) { + quant::reconstruct_full_vec( + combined.data(), centroid, padded_dim_, 1, bin.f_rescale(), out, metric_type_ + ); + return; + } + + ex_codes.resize(padded_dim_); + const ConstExDataMap ex(get_exdata_by_internalid(id), padded_dim_, ex_bits_); + quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( + ex.ex_code(), ex_codes.data(), padded_dim_, ex_bits_ + ); + for (size_t d = 0; d < padded_dim_; ++d) { + combined[d] = static_cast((combined[d] << ex_bits_) + ex_codes[d]); + } + + quant::reconstruct_full_vec( + combined.data(), + centroid, + padded_dim_, + ex_bits_ + 1, + ex.f_rescale_ex(), + out, + metric_type_ + ); +} + +// Only `query` is reconstructed; `target` is scored from its codes, as QGBuilder +// does. Callers hold `query` fixed across a loop, so it is prepared once. +inline float HierarchicalNSW::get_quant_dist(PID target, PID query) const { + detail::QuantQuery& cached = detail::quant_query(); + if (cached.id != query || !cached.prepared.has_value()) { + cached.rotated.resize(padded_dim_); + reconstruct_rotated(query, cached.rotated.data()); + cached.prepared.emplace( + cached.rotated.data(), padded_dim_, ex_bits_, query_config_, metric_type_ + ); + cached.id = query; + } + + const auto* centroid = reinterpret_cast(centroids_memory_) + + (get_clusterid_by_internalid(target) * padded_dim_); + const float dist_sqr = euclidean_sqr(cached.rotated.data(), centroid, padded_dim_); + const float g_error = std::sqrt(dist_sqr); + const float g_add = (metric_type_ == METRIC_IP) + ? -dot_product(cached.rotated.data(), centroid, padded_dim_) + : dist_sqr; + + float est_dist = 0; + float low_dist = 0; + float ip_x0_qr = 0; + if (ex_bits_ == 0) { + split_single_estdist( + get_bindata_by_internalid(target), + *cached.prepared, + padded_dim_, + ip_x0_qr, + est_dist, + low_dist, + g_add, + g_error + ); + } else { + split_single_fulldist( + get_bindata_by_internalid(target), + get_exdata_by_internalid(target), + ip_func_, + *cached.prepared, + padded_dim_, + ex_bits_, + est_dist, + low_dist, + ip_x0_qr, + g_add, + g_error + ); + } + return est_dist; +} + +inline maxheap> HierarchicalNSW::search_base_layer_quant( + PID ep_id, PID cur_c, int layer +) { + VisitedSet* vl = visited_list_pool_->get_free_vislist(); + + maxheap> top_candidates; + minheap> candidate_set; + + float lower_bound = get_quant_dist(ep_id, cur_c); + top_candidates.emplace(lower_bound, ep_id); + candidate_set.emplace(lower_bound, ep_id); + vl->set(ep_id); + + while (!candidate_set.empty()) { + std::pair curr_el_pair = candidate_set.top(); + if (curr_el_pair.first > lower_bound && top_candidates.size() == ef_construction_) { + break; + } + candidate_set.pop(); + + PID cur_node_num = curr_el_pair.second; + std::unique_lock lock(link_list_locks_[cur_node_num]); + + PID* data = + (layer == 0) ? get_linklist0(cur_node_num) : get_linklist(cur_node_num, layer); + size_t size = get_list_count(data); + auto* datal = data + 1; + + for (size_t j = 0; j < size; j++) { + PID candidate_id = datal[j]; + if (candidate_id >= cur_element_count_) { + throw std::runtime_error("cand error"); + } + if (vl->get(candidate_id)) { + continue; + } + vl->set(candidate_id); + + float dist1 = get_quant_dist(candidate_id, cur_c); + if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { + candidate_set.emplace(dist1, candidate_id); + top_candidates.emplace(dist1, candidate_id); + if (top_candidates.size() > ef_construction_) { + top_candidates.pop(); + } + if (!top_candidates.empty()) { + lower_bound = top_candidates.top().first; + } + } + } + } + visited_list_pool_->release_vis_list(vl); + return top_candidates; +} + +inline void HierarchicalNSW::get_neighbors_by_heuristic2_quant( + maxheap>& top_candidates, size_t M +) { + if (top_candidates.size() < M) { + return; + } + + minheap> queue_closest; + std::vector> return_list; + while (top_candidates.size() > 0) { + queue_closest.emplace(top_candidates.top()); + top_candidates.pop(); + } + + while (queue_closest.size() > 0) { + if (return_list.size() >= M) { + break; + } + std::pair current_pair = queue_closest.top(); + float dist_to_query = current_pair.first; + queue_closest.pop(); + bool good = true; + + for (std::pair second_pair : return_list) { + float curdist = get_quant_dist(second_pair.second, current_pair.second); + if (curdist < dist_to_query) { + good = false; + break; + } + } + if (good) { + return_list.push_back(current_pair); + } + } + + for (std::pair current_pair : return_list) { + top_candidates.emplace(current_pair); + } +} + +inline PID HierarchicalNSW::mutually_connect_quant( + PID cur_c, maxheap>& top_candidates, int level +) { + size_t max_m = level > 0 ? maxM_ : maxM0_; + get_neighbors_by_heuristic2_quant(top_candidates, M_); + if (top_candidates.size() > M_) { + throw std::runtime_error( + "Should be not be more than M_ candidates returned by the heuristic" + ); + } + + std::vector selected_neighbors; + selected_neighbors.reserve(M_); + while (top_candidates.size() > 0) { + selected_neighbors.push_back(top_candidates.top().second); + top_candidates.pop(); + } + if (selected_neighbors.empty()) { + throw std::runtime_error("No neighbors selected for the new element"); + } + PID next_closest_entry_point = selected_neighbors.back(); + + { + PID* ll_cur = (level == 0) ? get_linklist0(cur_c) : get_linklist(cur_c, level); + if (*ll_cur > 0) { + throw std::runtime_error( + "The newly inserted element should have blank link list" + ); + } + + set_list_count(ll_cur, selected_neighbors.size()); + auto* data = ll_cur + 1; + for (size_t idx = 0; idx < selected_neighbors.size(); idx++) { + if (data[idx] != 0) { + throw std::runtime_error("Possible memory corruption"); + } + if (level > element_levels_[selected_neighbors[idx]]) { + throw std::runtime_error("Trying to make a link on a non-existent level"); + } + data[idx] = selected_neighbors[idx]; + } + } + + for (auto selected_neighbor : selected_neighbors) { + std::unique_lock lock(link_list_locks_[selected_neighbor]); + + PID* ll_other = (level == 0) ? get_linklist0(selected_neighbor) + : get_linklist(selected_neighbor, level); + size_t sz_link_list_other = get_list_count(ll_other); + + if (sz_link_list_other > max_m) { + throw std::runtime_error("Bad value of sz_link_list_other"); + } + if (selected_neighbor == cur_c) { + throw std::runtime_error("Trying to connect an element to itself"); + } + if (level > element_levels_[selected_neighbor]) { + throw std::runtime_error("Trying to make a link on a non-existent level"); + } + + auto* data = ll_other + 1; + + bool is_cur_c_present = false; + for (size_t j = 0; j < sz_link_list_other; j++) { + if (data[j] == cur_c) { + is_cur_c_present = true; + break; + } + } + + if (!is_cur_c_present) { + if (sz_link_list_other < max_m) { + data[sz_link_list_other] = cur_c; + set_list_count(ll_other, sz_link_list_other + 1); + } else { + float d_max = get_quant_dist(selected_neighbor, cur_c); + maxheap> candidates; + candidates.emplace(d_max, cur_c); + for (size_t j = 0; j < sz_link_list_other; j++) { + candidates.emplace(get_quant_dist(data[j], selected_neighbor), data[j]); + } + + get_neighbors_by_heuristic2_quant(candidates, max_m); + + int indx = 0; + while (candidates.size() > 0) { + data[indx] = candidates.top().second; + candidates.pop(); + indx++; + } + set_list_count(ll_other, indx); + } + } + } + + return next_closest_entry_point; +} + +inline PID HierarchicalNSW::add_point_quant( + const float* vec, PID cluster_id, const quant::RabitqConfig& config +) { + detail::quant_query().clear(); + int curlevel = get_random_level(mult_); + + // Allocated before the slot is claimed so a failure leaves the index untouched. + std::unique_ptr link_list(nullptr, std::free); + if (curlevel > 0) { + auto* raw = + static_cast(std::calloc(1, (size_links_per_element_ * curlevel) + 1)); + if (raw == nullptr) { + throw std::runtime_error( + "Not enough memory: add_point_quant failed to allocate linklist" + ); + } + link_list.reset(raw); + } + + PID cur_c = 0; + PID label = 0; + { + std::unique_lock lock_table(label_lookup_lock_); + if (cur_element_count_ >= max_elements_) { + throw std::runtime_error("The number of elements exceeds the specified limit"); + } + cur_c = static_cast(cur_element_count_.load()); + label = cur_c; + if (label_lookup_.find(label) != label_lookup_.end()) { + throw std::runtime_error("Label already present"); + } + cur_element_count_++; + label_lookup_[label] = cur_c; + } + + std::unique_lock lock_el(link_list_locks_[cur_c]); + element_levels_[cur_c] = curlevel; + + std::unique_lock templock(global_); + int maxlevelcopy = maxlevel_; + if (curlevel <= maxlevelcopy) { + templock.unlock(); + } + PID curr_obj = enterpoint_node_; + + std::memset( + data_level0_memory_ + (cur_c * size_data_per_element_), 0, size_data_per_element_ + ); + std::memcpy(get_external_label_pt(cur_c), &label, sizeof(PID)); + std::memcpy(get_clusterid_pt(cur_c), &cluster_id, sizeof(PID)); + + std::vector rotated_data(padded_dim_); + rotator_->rotate(vec, rotated_data.data()); + quant::quantize_split_single( + rotated_data.data(), + reinterpret_cast(centroids_memory_) + (cluster_id * padded_dim_), + padded_dim_, + ex_bits_, + get_bindata_by_internalid(cur_c), + get_exdata_by_internalid(cur_c), + metric_type_, + config + ); + + if (curlevel > 0) { + linkLists_[cur_c] = link_list.release(); + } + + if (static_cast(curr_obj) != -1) { + if (curlevel < maxlevelcopy) { + float curdist = get_quant_dist(curr_obj, cur_c); + for (int level = maxlevelcopy; level > curlevel; level--) { + bool changed = true; + while (changed) { + changed = false; + std::unique_lock lock(link_list_locks_[curr_obj]); + PID* data = get_linklist(curr_obj, level); + int size = get_list_count(data); + auto* datal = data + 1; + for (int i = 0; i < size; i++) { + PID cand = datal[i]; + if (cand >= cur_element_count_) { + throw std::runtime_error("cand error"); + } + float d = get_quant_dist(cand, cur_c); + if (d < curdist) { + curdist = d; + curr_obj = cand; + changed = true; + } + } + } + } + } + + for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { + maxheap> top_candidates = + search_base_layer_quant(curr_obj, cur_c, level); + curr_obj = mutually_connect_quant(cur_c, top_candidates, level); + } + } else { + enterpoint_node_ = cur_c; + maxlevel_ = curlevel; + } + + if (curlevel > maxlevelcopy) { + enterpoint_node_ = cur_c; + maxlevel_ = curlevel; + } + return label; +} + +inline std::vector HierarchicalNSW::add( + const float* data, size_t n, const PID* cluster_ids, bool faster +) { + 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 add"); + } + if (n == 0) { + return {}; + } + if (data == nullptr || cluster_ids == nullptr) { + throw std::invalid_argument("HNSW add inputs must not be null"); + } + if (n > max_elements_ - cur_element_count_) { + throw std::invalid_argument("HNSW add: not enough capacity"); + } + if (n > buffer::kSearchBufferMaxPointCount - cur_element_count_) { + throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); + } + for (size_t i = 0; i < n; ++i) { + if (cluster_ids[i] >= num_cluster_) { + throw std::invalid_argument("HNSW cluster ID is out of range"); + } + } + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + std::vector labels; + labels.reserve(n); + for (size_t i = 0; i < n; ++i) { + labels.push_back(add_point_quant(data + (i * dim_), cluster_ids[i], config)); + } + return labels; +} + +} // namespace rabitqlib::hnsw diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 459bfad4..698efd24 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -77,6 +77,268 @@ TEST(HnswConstructionTest, ParallelConstructionProducesSearchableIndex) { } } +// Builds an index over the first `built` of `total` random vectors, leaving the +// rest as capacity for add. +struct AddFixture { + static constexpr size_t kDim = 64; + std::vector data; + std::vector centroid; + std::vector cluster_ids; + + AddFixture(size_t total, unsigned seed) + : data(total * kDim), centroid(kDim, 0.0F), cluster_ids(total, 0) { + std::mt19937 generator(seed); + std::uniform_real_distribution distribution(-1, 1); + for (auto& value : data) { + value = distribution(generator); + } + } +}; + +TEST(HnswAddTest, AddedPointsAreSearchable) { + constexpr size_t kBuilt = 256; + constexpr size_t kAdded = 64; + constexpr size_t kTotal = kBuilt + kAdded; + AddFixture fixture(kTotal, 7); + + HierarchicalNSW index(kTotal, AddFixture::kDim, 4, 8, 50); + index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + ); + ASSERT_EQ(index.num_points(), kBuilt); + + const auto labels = index.add( + fixture.data.data() + (kBuilt * AddFixture::kDim), + kAdded, + fixture.cluster_ids.data(), + false + ); + + ASSERT_EQ(labels.size(), kAdded); + for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(labels[i], kBuilt + i); + } + EXPECT_EQ(index.num_points(), kTotal); + + // Every added vector must find itself when used as its own query. + const auto results = index.search( + fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, 1, kTotal, 1 + ); + ASSERT_EQ(results.size(), kAdded); + size_t exact = 0; + for (size_t i = 0; i < kAdded; ++i) { + ASSERT_EQ(results[i].size(), 1U); + EXPECT_TRUE(std::isfinite(results[i][0].first)); + exact += (results[i][0].second == kBuilt + i) ? 1 : 0; + } + EXPECT_EQ(exact, kAdded); + + // The points placed by construct must still be reachable afterwards. + const auto old_results = index.search(fixture.data.data(), 16, 1, kTotal, 1); + for (size_t i = 0; i < 16; ++i) { + EXPECT_EQ(old_results[i][0].second, i); + } +} + +TEST(HnswAddTest, SurvivesSaveAndLoad) { + constexpr size_t kBuilt = 128; + constexpr size_t kAdded = 32; + constexpr size_t kTotal = kBuilt + kAdded; + AddFixture fixture(kTotal, 11); + + const std::filesystem::path path = + std::filesystem::temp_directory_path() / "hnsw_add_roundtrip.index"; + + { + HierarchicalNSW index(kTotal, AddFixture::kDim, 4, 8, 50); + index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + ); + index.add( + fixture.data.data() + (kBuilt * AddFixture::kDim), + kAdded, + fixture.cluster_ids.data(), + false + ); + index.save(path.string().c_str()); + } + + HierarchicalNSW loaded; + loaded.load(path.string().c_str()); + EXPECT_EQ(loaded.num_points(), kTotal); + EXPECT_EQ(loaded.max_elements(), kTotal); + + const auto results = loaded.search( + fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, 1, kTotal, 1 + ); + for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(results[i][0].second, kBuilt + i); + } + std::filesystem::remove(path); +} + +TEST(HnswAddTest, RejectsBadInputWithoutDisturbingTheIndex) { + constexpr size_t kBuilt = 64; + constexpr size_t kTotal = kBuilt + 8; + AddFixture fixture(kTotal, 13); + + HierarchicalNSW index(kTotal, AddFixture::kDim, 4, 8, 50); + index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + ); + + const float* extra = fixture.data.data() + (kBuilt * AddFixture::kDim); + std::vector bad_cluster(8, 5); + + EXPECT_THROW(index.add(nullptr, 8, fixture.cluster_ids.data()), std::invalid_argument); + EXPECT_THROW(index.add(extra, 8, nullptr), std::invalid_argument); + EXPECT_THROW(index.add(extra, 8, bad_cluster.data()), std::invalid_argument); + // More points than the remaining capacity. + EXPECT_THROW(index.add(extra, 9, fixture.cluster_ids.data()), std::invalid_argument); + + EXPECT_EQ(index.num_points(), kBuilt); + EXPECT_EQ(index.add(extra, 0, fixture.cluster_ids.data()).size(), 0U); + + // The index is still searchable and still holds only what construct placed. + const auto results = index.search(fixture.data.data(), 8, 1, kBuilt, 1); + for (size_t i = 0; i < 8; ++i) { + EXPECT_EQ(results[i][0].second, i); + } +} + +// Several clusters, so the per-point centroid actually varies and the correction +// terms have to pick the right one. Both metrics, since their g_add differ. +TEST(HnswAddTest, AddedPointsAreSearchableWithManyClusters) { + constexpr size_t kBuilt = 256; + constexpr size_t kAdded = 64; + constexpr size_t kTotal = kBuilt + kAdded; + constexpr size_t kClusters = 8; + + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + AddFixture fixture(kTotal, 23); + std::vector centroids(kClusters * AddFixture::kDim); + std::mt19937 generator(5); + std::uniform_real_distribution distribution(-1, 1); + for (auto& value : centroids) { + value = distribution(generator); + } + std::vector cluster_ids(kTotal); + for (size_t i = 0; i < kTotal; ++i) { + cluster_ids[i] = static_cast(i % kClusters); + } + + HierarchicalNSW index(kTotal, AddFixture::kDim, 4, 8, 50, 100, metric); + index.construct( + kClusters, + centroids.data(), + kBuilt, + fixture.data.data(), + cluster_ids.data(), + 1, + false + ); + const auto labels = index.add( + fixture.data.data() + (kBuilt * AddFixture::kDim), + kAdded, + cluster_ids.data() + kBuilt, + false + ); + ASSERT_EQ(labels.size(), kAdded); + EXPECT_EQ(index.num_points(), kTotal); + + const auto results = index.search( + fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, 1, kTotal, 1 + ); + size_t exact = 0; + for (size_t i = 0; i < kAdded; ++i) { + exact += (results[i][0].second == kBuilt + i) ? 1 : 0; + } + EXPECT_EQ(exact, kAdded); + } +} + +// Two indexes with different shapes, used one after the other: the per-thread +// query scratch must not carry over between them. +TEST(HnswAddTest, DoesNotReuseScratchAcrossIndexes) { + constexpr size_t kCount = 64; + for (const size_t dim : {64U, 128U}) { + std::vector data(kCount * dim); + std::mt19937 generator(31); + std::uniform_real_distribution distribution(-1, 1); + for (auto& value : data) { + value = distribution(generator); + } + std::vector centroid(dim, 0.0F); + std::vector cluster_ids(kCount, 0); + + HierarchicalNSW index(kCount, dim, 4, 8, 50); + index.construct( + 1, centroid.data(), kCount / 2, data.data(), cluster_ids.data(), 1, false + ); + index.add( + data.data() + ((kCount / 2) * dim), kCount / 2, cluster_ids.data(), false + ); + + const auto results = index.search(data.data(), kCount, 1, kCount, 1); + for (size_t i = 0; i < kCount; ++i) { + EXPECT_EQ(results[i][0].second, i); + } + } +} + +TEST(HnswAddTest, AddIntoAnEmptyIndexBuildsFromScratch) { + constexpr size_t kCount = 96; + AddFixture fixture(kCount, 17); + + HierarchicalNSW index(kCount, AddFixture::kDim, 4, 8, 50); + // construct one point so the centroids and rotator exist, then add the rest. + index.construct( + 1, + fixture.centroid.data(), + 1, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + ); + const auto labels = index.add( + fixture.data.data() + AddFixture::kDim, + kCount - 1, + fixture.cluster_ids.data(), + false + ); + ASSERT_EQ(labels.size(), kCount - 1); + EXPECT_EQ(index.num_points(), kCount); + + const auto results = index.search(fixture.data.data(), kCount, 1, kCount, 1); + size_t exact = 0; + for (size_t i = 0; i < kCount; ++i) { + exact += (results[i][0].second == i) ? 1 : 0; + } + // A graph built entirely from one-bit reconstructions still has to route a + // vector to itself; the codes distinguish points far better than they order + // near neighbors. + EXPECT_EQ(exact, kCount); +} + TEST(HnswConstructionTest, RejectsInvalidInputsBeforeChangingIndex) { constexpr size_t kDim = 64; std::vector data(kDim, 1.0F); From 06b691c206e1e5cc27c2708e43acd655a13d9d3b Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:23:22 +0800 Subject: [PATCH 03/12] Route new HNSW points to their nearest centroid when no clusters are given --- include/rabitqlib/index/hnsw/hnsw.hpp | 5 ++- include/rabitqlib/index/hnsw/hnsw_quant.hpp | 33 +++++++++++--- tests/unit/rabitqlib/index/hnsw_test.cpp | 48 ++++++++++++++++++++- 3 files changed, 79 insertions(+), 7 deletions(-) diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index f4f49105..f036f3ff 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -62,6 +62,9 @@ class HierarchicalNSW { [[nodiscard]] MetricType metric_type() const { return metric_type_; } [[nodiscard]] size_t max_elements() const { return max_elements_; } [[nodiscard]] size_t num_points() const { return cur_element_count_; } + [[nodiscard]] PID cluster_id_of(PID internal_id) const { + return get_clusterid_by_internalid(internal_id); + } void save(const char*) const; void load(const char*); @@ -70,7 +73,7 @@ class HierarchicalNSW { // Inserts points using the stored codes, labelling them num_points() + i. // Single-threaded, and throws rather than growing past max_elements(). - std::vector add(const float*, size_t, const PID*, bool = false); + std::vector add(const float*, size_t, const PID* = nullptr, bool = false); std::vector>> search( const float*, size_t, size_t, size_t, size_t ); diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index d4991e3a..4153e588 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -460,7 +461,7 @@ inline std::vector HierarchicalNSW::add( if (n == 0) { return {}; } - if (data == nullptr || cluster_ids == nullptr) { + if (data == nullptr) { throw std::invalid_argument("HNSW add inputs must not be null"); } if (n > max_elements_ - cur_element_count_) { @@ -469,9 +470,31 @@ inline std::vector HierarchicalNSW::add( if (n > buffer::kSearchBufferMaxPointCount - cur_element_count_) { throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); } - for (size_t i = 0; i < n; ++i) { - if (cluster_ids[i] >= num_cluster_) { - throw std::invalid_argument("HNSW cluster ID is out of range"); + // Either the caller names the cluster of every point, or each one goes to its + // nearest centroid, ordered the way a query is routed. + std::vector assigned(n, 0); + if (cluster_ids != nullptr) { + for (size_t i = 0; i < n; ++i) { + if (cluster_ids[i] >= num_cluster_) { + throw std::invalid_argument("HNSW cluster ID is out of range"); + } + assigned[i] = cluster_ids[i]; + } + } else { + const auto* centroids = reinterpret_cast(centroids_memory_); + std::vector rotated(padded_dim_); + for (size_t i = 0; i < n; ++i) { + rotator_->rotate(data + (i * dim_), rotated.data()); + float best = std::numeric_limits::max(); + for (size_t c = 0; c < num_cluster_; ++c) { + const float dist = raw_dist_func_( + rotated.data(), centroids + (c * padded_dim_), padded_dim_ + ); + if (dist < best) { + best = dist; + assigned[i] = static_cast(c); + } + } } } @@ -483,7 +506,7 @@ inline std::vector HierarchicalNSW::add( std::vector labels; labels.reserve(n); for (size_t i = 0; i < n; ++i) { - labels.push_back(add_point_quant(data + (i * dim_), cluster_ids[i], config)); + labels.push_back(add_point_quant(data + (i * dim_), assigned[i], config)); } return labels; } diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 698efd24..e91f2e98 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -209,7 +209,6 @@ TEST(HnswAddTest, RejectsBadInputWithoutDisturbingTheIndex) { std::vector bad_cluster(8, 5); EXPECT_THROW(index.add(nullptr, 8, fixture.cluster_ids.data()), std::invalid_argument); - EXPECT_THROW(index.add(extra, 8, nullptr), std::invalid_argument); EXPECT_THROW(index.add(extra, 8, bad_cluster.data()), std::invalid_argument); // More points than the remaining capacity. EXPECT_THROW(index.add(extra, 9, fixture.cluster_ids.data()), std::invalid_argument); @@ -304,6 +303,53 @@ TEST(HnswAddTest, DoesNotReuseScratchAcrossIndexes) { } } +TEST(HnswAddTest, RoutesToTheNearestCentroidWhenNoClustersGiven) { + constexpr size_t kBuilt = 128; + constexpr size_t kAdded = 32; + constexpr size_t kTotal = kBuilt + kAdded; + constexpr size_t kClusters = 4; + AddFixture fixture(kTotal, 29); + + // Well-separated centroids, so the nearest one is unambiguous. + std::vector centroids(kClusters * AddFixture::kDim, 0.0F); + for (size_t c = 0; c < kClusters; ++c) { + centroids[(c * AddFixture::kDim) + c] = 50.0F; + } + std::vector cluster_ids(kTotal); + for (size_t i = 0; i < kTotal; ++i) { + cluster_ids[i] = static_cast(i % kClusters); + // Push each point hard towards its centroid. + for (size_t d = 0; d < AddFixture::kDim; ++d) { + fixture.data[(i * AddFixture::kDim) + d] += + centroids[(cluster_ids[i] * AddFixture::kDim) + d]; + } + } + + HierarchicalNSW index(kTotal, AddFixture::kDim, 4, 8, 50); + index.construct( + kClusters, + centroids.data(), + kBuilt, + fixture.data.data(), + cluster_ids.data(), + 1, + false + ); + const auto labels = + index.add(fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, nullptr); + ASSERT_EQ(labels.size(), kAdded); + + for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(index.cluster_id_of(labels[i]), cluster_ids[kBuilt + i]); + } + const auto results = index.search( + fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, 1, kTotal, 1 + ); + for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(results[i][0].second, kBuilt + i); + } +} + TEST(HnswAddTest, AddIntoAnEmptyIndexBuildsFromScratch) { constexpr size_t kCount = 96; AddFixture fixture(kCount, 17); From a844684f78e7c1b6e97ae816b6cfddb6d9b5299c Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:24:03 +0800 Subject: [PATCH 04/12] Grow HNSW capacity with resize --- include/rabitqlib/index/hnsw/hnsw.hpp | 3 + include/rabitqlib/index/hnsw/hnsw_quant.hpp | 56 ++++++++++++++ tests/unit/rabitqlib/index/hnsw_test.cpp | 81 +++++++++++++++++++++ 3 files changed, 140 insertions(+) diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index f036f3ff..2939267e 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -74,6 +74,9 @@ class HierarchicalNSW { // Inserts points using the stored codes, labelling them num_points() + i. // Single-threaded, and throws rather than growing past max_elements(). std::vector add(const float*, size_t, const PID* = nullptr, bool = false); + + // Grows max_elements(). Invalidates every pointer into the index. + void resize(size_t); std::vector>> search( const float*, size_t, size_t, size_t, size_t ); diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index 4153e588..9bd93377 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -451,6 +451,62 @@ inline PID HierarchicalNSW::add_point_quant( return label; } +// Grows the capacity. Every pointer into the level-0 block is invalidated, so +// this must not run while another thread searches. +inline void HierarchicalNSW::resize(size_t new_max_elements) { + if (new_max_elements < cur_element_count_) { + throw std::invalid_argument("HNSW resize: capacity below the element count"); + } + if (new_max_elements > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("HNSW resize: capacity exceeds the supported ID range"); + } + if (new_max_elements == max_elements_) { + return; + } + if (new_max_elements > std::numeric_limits::max() / size_data_per_element_) { + throw std::invalid_argument("HNSW resize: capacity exceeds addressable storage"); + } + + // Everything new is built before anything old is released, so a failure + // leaves the index as it was. realloc is not an option: the level-0 block + // comes from aligned_alloc, which realloc does not keep aligned. + const size_t level0_bytes = new_max_elements * size_data_per_element_; + std::unique_ptr level0( + memory::huge_page_allocate(level0_bytes), + [](char* ptr) { memory::aligned_deallocate(ptr); } + ); + if (level0 == nullptr) { + throw std::bad_alloc(); + } + std::memset(level0.get(), 0, level0_bytes); + std::memcpy( + level0.get(), data_level0_memory_, cur_element_count_ * size_data_per_element_ + ); + + std::unique_ptr links( + static_cast(std::calloc(new_max_elements, sizeof(char*))), + [](char** ptr) { std::free(ptr); } + ); + if (links == nullptr) { + throw std::bad_alloc(); + } + std::memcpy(links.get(), linkLists_, cur_element_count_ * sizeof(char*)); + + std::vector levels(element_levels_); + levels.resize(new_max_elements); + std::vector locks(new_max_elements); + auto pool = std::make_unique(1, new_max_elements); + + memory::aligned_deallocate(data_level0_memory_); + data_level0_memory_ = level0.release(); + std::free(reinterpret_cast(linkLists_)); + linkLists_ = links.release(); + element_levels_.swap(levels); + link_list_locks_.swap(locks); + visited_list_pool_ = std::move(pool); + max_elements_ = new_max_elements; +} + inline std::vector HierarchicalNSW::add( const float* data, size_t n, const PID* cluster_ids, bool faster ) { diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index e91f2e98..8b7b294e 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -350,6 +350,87 @@ TEST(HnswAddTest, RoutesToTheNearestCentroidWhenNoClustersGiven) { } } +TEST(HnswResizeTest, GrowsCapacityAndKeepsTheGraph) { + constexpr size_t kBuilt = 128; + constexpr size_t kAdded = 64; + constexpr size_t kTotal = kBuilt + kAdded; + AddFixture fixture(kTotal, 37); + + HierarchicalNSW index(kBuilt, AddFixture::kDim, 4, 8, 50); + index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + ); + ASSERT_EQ(index.max_elements(), kBuilt); + + // Full: adding anything must fail, and leave the index usable. + EXPECT_THROW( + index.add(fixture.data.data() + (kBuilt * AddFixture::kDim), 1, nullptr), + std::invalid_argument + ); + EXPECT_EQ(index.num_points(), kBuilt); + + EXPECT_THROW(index.resize(kBuilt - 1), std::invalid_argument); + index.resize(kTotal); + EXPECT_EQ(index.max_elements(), kTotal); + EXPECT_EQ(index.num_points(), kBuilt); + + // Points placed before the resize survive it. + const auto before = index.search(fixture.data.data(), 16, 1, kBuilt, 1); + for (size_t i = 0; i < 16; ++i) { + EXPECT_EQ(before[i][0].second, i); + } + + index.add(fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, nullptr); + EXPECT_EQ(index.num_points(), kTotal); + + const auto results = index.search(fixture.data.data(), kTotal, 1, kTotal, 1); + for (size_t i = 0; i < kTotal; ++i) { + EXPECT_EQ(results[i][0].second, i); + } +} + +TEST(HnswResizeTest, SurvivesSaveAndLoad) { + constexpr size_t kBuilt = 64; + constexpr size_t kTotal = 96; + AddFixture fixture(kTotal, 41); + + const std::filesystem::path path = + std::filesystem::temp_directory_path() / "hnsw_resize_roundtrip.index"; + { + HierarchicalNSW index(kBuilt, AddFixture::kDim, 4, 8, 50); + index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + ); + index.resize(kTotal); + index.add( + fixture.data.data() + (kBuilt * AddFixture::kDim), kTotal - kBuilt, nullptr + ); + index.save(path.string().c_str()); + } + + HierarchicalNSW loaded; + loaded.load(path.string().c_str()); + EXPECT_EQ(loaded.max_elements(), kTotal); + EXPECT_EQ(loaded.num_points(), kTotal); + const auto results = loaded.search(fixture.data.data(), kTotal, 1, kTotal, 1); + for (size_t i = 0; i < kTotal; ++i) { + EXPECT_EQ(results[i][0].second, i); + } + std::filesystem::remove(path); +} + TEST(HnswAddTest, AddIntoAnEmptyIndexBuildsFromScratch) { constexpr size_t kCount = 96; AddFixture fixture(kCount, 17); From 0ba0a894c5731b7ca6342edcf3ab7abf16891a4e Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:34:40 +0800 Subject: [PATCH 05/12] Expose HNSW add and resize to Python --- python_bindings/hnsw_bindings.cpp | 64 +++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 559237dd..81c55e1e 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -107,6 +108,60 @@ class HnswIndex { built_ = true; } + py::array_t add( + py::handle data, const py::object& cluster_ids, bool fast_quantization = false + ) { + auto data_array = ensure_2d_array(data, "data"); + if (!built_) { + throw std::runtime_error("HnswIndex must be built or loaded before add"); + } + if (static_cast(data_array.shape(1)) != dim_) { + throw std::invalid_argument("data dimension does not match index dim"); + } + const auto rows = static_cast(data_array.shape(0)); + + // Validate as int64 first: casting straight to uint32 would wrap negative + // or oversized values onto valid clusters. + std::vector assigned; + if (!cluster_ids.is_none()) { + auto cluster_ids_array = ensure_1d_integer_array(cluster_ids, "cluster_ids"); + if (static_cast(cluster_ids_array.shape(0)) != rows) { + throw std::invalid_argument( + "cluster_ids length must match number of rows in data" + ); + } + assigned.reserve(rows); + for (py::ssize_t i = 0; i < cluster_ids_array.shape(0); ++i) { + const int64_t id = cluster_ids_array.data()[i]; + if (id < 0 || static_cast(id) >= num_clusters_) { + throw std::invalid_argument("cluster_ids must be in [0, num_clusters)"); + } + assigned.push_back(static_cast(id)); + } + } + + // Allocated before the index is touched. + auto ids = py::array_t(static_cast(rows)); + const auto labels = index_->add( + data_array.data(), + rows, + assigned.empty() ? nullptr : assigned.data(), + fast_quantization + ); + std::copy(labels.begin(), labels.end(), ids.mutable_data()); + return ids; + } + + void resize(size_t new_max_elements) { + if (!built_) { + throw std::runtime_error("HnswIndex must be built or loaded before resize"); + } + index_->resize(new_max_elements); + max_elements_ = index_->max_elements(); + } + + [[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) { auto query_array = ensure_2d_array(queries, "queries"); if (!built_) { @@ -235,10 +290,19 @@ void register_hnsw(py::module_& m) { py::arg("ef") = 0, py::arg("num_threads") = 1 ) + .def( + "add", + &HnswIndex::add, + py::arg("data"), + py::arg("cluster_ids") = py::none(), + py::arg("fast_quantization") = false + ) + .def("resize", &HnswIndex::resize, py::arg("max_elements")) .def("save", &HnswIndex::save, py::arg("path")) .def_static("load", &HnswIndex::load, py::arg("path")) .def_property_readonly("dim", &HnswIndex::dim) .def_property_readonly("max_elements", &HnswIndex::max_elements) + .def_property_readonly("num_points", &HnswIndex::num_points) .def_property_readonly("nbits", &HnswIndex::nbits) .def_property_readonly("num_clusters", &HnswIndex::num_clusters) .def_property_readonly("is_built", &HnswIndex::is_built) From 67c6ca89ed119125ccbc2c75faeba948dde705af Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:34:40 +0800 Subject: [PATCH 06/12] Test HNSW add and resize through the Python bindings --- tests/python/test_hnsw.py | 129 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index f6959699..ef1f67fe 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -179,3 +179,132 @@ def test_load_rejects_count_larger_than_capacity(built_hnsw, tmp_path): index_file.write((1).to_bytes(np.dtype(np.uintp).itemsize, byteorder="little")) with pytest.raises(RuntimeError, match="HNSW"): HnswIndex.load(str(path)) + + +# ── add and resize ──────────────────────────────────────────────────────────── + + +@pytest.fixture +def partial_hnsw(base_data, clusters): + """Built from the first 400 vectors, with room for the remaining 100.""" + idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + centroids, cluster_ids = clusters + idx.build(base_data[:400], centroids, cluster_ids[:400]) + return idx + + +def test_add_returns_dense_labels(partial_hnsw, base_data, clusters): + _, cluster_ids = clusters + assert partial_hnsw.num_points == 400 + ids = partial_hnsw.add(base_data[400:], cluster_ids=cluster_ids[400:]) + np.testing.assert_array_equal(ids, np.arange(400, N_VECTORS)) + assert partial_hnsw.num_points == N_VECTORS + + +def test_add_makes_points_searchable(partial_hnsw, base_data, clusters): + # A sparse graph can isolate the odd point, as it can after build, so this + # asks for near-perfect self-retrieval rather than exact. + _, cluster_ids = clusters + partial_hnsw.add(base_data[400:], cluster_ids=cluster_ids[400:]) + ids, distances = partial_hnsw.search(base_data[400:], k=1, ef=N_VECTORS) + found = np.count_nonzero(ids[:, 0] == np.arange(400, N_VECTORS)) + assert found >= 0.95 * (N_VECTORS - 400) + assert np.isfinite(distances).all() + + +def test_add_into_a_dense_graph_retrieves_every_point(base_data, clusters): + # Keeping every neighbor removes the pruning that isolates points, so the + # added points must all be found. + centroids, cluster_ids = clusters + idx = HnswIndex(DIM, N_VECTORS, M=N_VECTORS, ef_construction=N_VECTORS, nbits=4) + idx.build(base_data[:400], centroids, cluster_ids[:400]) + idx.add(base_data[400:], cluster_ids=cluster_ids[400:]) + ids, _ = idx.search(base_data[400:], k=1, ef=N_VECTORS) + np.testing.assert_array_equal(ids[:, 0], np.arange(400, N_VECTORS)) + + +def test_add_without_cluster_ids_routes(partial_hnsw, base_data, clusters): + _, cluster_ids = clusters + ids = partial_hnsw.add(base_data[400:]) + np.testing.assert_array_equal(ids, np.arange(400, N_VECTORS)) + # Round-robin clusters, so routing should recover the original assignment for + # most points; the centroids are close together and a few land elsewhere. + found, _ = partial_hnsw.search(base_data[400:], k=1, ef=N_VECTORS) + assert np.count_nonzero(found[:, 0] == np.arange(400, N_VECTORS)) >= 0.95 * 100 + + +def test_add_preserves_earlier_points(partial_hnsw, base_data, clusters): + _, cluster_ids = clusters + before, _ = partial_hnsw.search(base_data[:10], k=1, ef=N_VECTORS) + partial_hnsw.add(base_data[400:], cluster_ids=cluster_ids[400:]) + after, _ = partial_hnsw.search(base_data[:10], k=1, ef=N_VECTORS) + np.testing.assert_array_equal(before[:, 0], after[:, 0]) + + +def test_add_beyond_capacity_raises(partial_hnsw, base_data): + with pytest.raises(ValueError, match="capacity"): + partial_hnsw.add(base_data) + assert partial_hnsw.num_points == 400 + + +def test_add_rejects_bad_cluster_ids(partial_hnsw, base_data): + rows = N_VECTORS - 400 + with pytest.raises(ValueError, match="cluster_ids"): + partial_hnsw.add(base_data[400:], cluster_ids=np.full(rows, N_CLUSTERS)) + with pytest.raises(ValueError, match="cluster_ids"): + partial_hnsw.add(base_data[400:], cluster_ids=-np.ones(rows, dtype=np.int64)) + with pytest.raises(ValueError, match="cluster_ids"): + partial_hnsw.add(base_data[400:], cluster_ids=np.zeros(rows - 1, dtype=np.int64)) + assert partial_hnsw.num_points == 400 + + +def test_add_rejects_wrong_dimension(partial_hnsw): + with pytest.raises(ValueError, match="dimension"): + partial_hnsw.add(np.zeros((4, DIM + 1), dtype=np.float32)) + + +def test_add_before_build_raises(base_data): + idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + with pytest.raises(RuntimeError, match="built or loaded"): + idx.add(base_data[:4]) + + +def test_resize_grows_capacity_and_allows_more(base_data, clusters): + centroids, cluster_ids = clusters + idx = HnswIndex(DIM, 400, M=8, ef_construction=50, nbits=4) + idx.build(base_data[:400], centroids, cluster_ids[:400]) + assert idx.max_elements == 400 + + with pytest.raises(ValueError, match="capacity"): + idx.add(base_data[400:]) + + idx.resize(N_VECTORS) + assert idx.max_elements == N_VECTORS + assert idx.num_points == 400 + + idx.add(base_data[400:], cluster_ids=cluster_ids[400:]) + assert idx.num_points == N_VECTORS + ids, _ = idx.search(base_data[:10], k=1, ef=N_VECTORS) + np.testing.assert_array_equal(ids[:, 0], np.arange(10)) + + +def test_resize_below_element_count_raises(partial_hnsw): + with pytest.raises(ValueError, match="element count"): + partial_hnsw.resize(10) + assert partial_hnsw.max_elements == N_VECTORS + + +def test_added_points_survive_save_and_load(partial_hnsw, base_data, clusters, tmp_path): + _, cluster_ids = clusters + partial_hnsw.add(base_data[400:], cluster_ids=cluster_ids[400:]) + before_ids, before_dists = partial_hnsw.search(base_data, k=_TOPK, ef=_EF) + + path = str(tmp_path / "hnsw-added.index") + partial_hnsw.save(path) + loaded = HnswIndex.load(path) + + assert loaded.num_points == N_VECTORS + assert loaded.max_elements == N_VECTORS + 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) From 8d26aac095a5a82cbdeb9a0b9a9ba000ff5fb738 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:34:40 +0800 Subject: [PATCH 07/12] Document updating an HNSW index --- docs/docs/index/hnsw.md | 75 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index f3e363ab..67ea7c29 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -98,3 +98,78 @@ Repeat until `candidate_set` is empty: - Insert the neighbor into `candidate_set` with its (possibly refined) estimated distance. The search terminates when `candidate_set` is empty. + +## Updating an Index + +A constructed or loaded index can gain points without the original data. The +index file format does not change. + +```c++ +std::vector HierarchicalNSW::add( + const float* data, + size_t n, + const PID* cluster_ids = nullptr, + bool faster = false +); + +void HierarchicalNSW::resize(size_t new_max_elements); +``` + +- **data**, **n**: `n` new vectors in the original coordinates. Vector `i` receives + the label `num_points() + i`, so labels stay dense, and `add` returns them. +- **cluster_ids**: The cluster of each new vector, in `[0, num_clusters)`. When it + is `nullptr`, each vector goes to its nearest centroid, chosen the same way a + query is routed. +- **faster**: Same as in `construct`. +- **new_max_elements**: The new capacity. It must be at least `num_points()`. + +In Python: + +```python +new_ids = index.add(vectors) # route to the nearest centroids +new_ids = index.add(vectors, cluster_ids=labels) # or choose the clusters +index.resize(index.max_elements + 100_000) +``` + +It is not safe to call `add` or `resize` while another thread searches the same +index, and `add` inserts one point at a time rather than in parallel. + +### Capacity + +`add` never grows the index: it throws once `num_points()` reaches +`max_elements()`, and `resize` is the only way to raise that. An index built with +spare capacity absorbs points with no reallocation at all, so sizing +`max_elements` ahead of time is the cheap path. + +`resize` rebuilds the base-layer storage and copies it, so it needs memory for +both copies while it runs and costs time proportional to the whole index. It +invalidates every pointer into the index. The per-node upper-layer link lists are +not copied, only their pointer table. + +### How added points are linked + +`construct` borrows the caller's vectors and links points using exact distances +between them. That pointer is not retained, so `add` cannot use it, and instead +scores candidates with the estimator the search already uses: the vector being +compared against is rebuilt from its stored RaBitQ code, sign bit and extra bits +together, and the other side is scored straight from its codes. This is the same +approach quantized SymphonyQG construction takes. + +Graph quality therefore depends on how a point arrived. On SIFT-1M restricted to +100,000 vectors (M = 32, 9 bits, 500 queries, recall@10 against exact +neighbors), an index built entirely by `construct` and one whose last fifth +arrived through `add` are indistinguishable: + +| built by | ef 10 | ef 50 | ef 100 | ef 200 | +| --- | --- | --- | --- | --- | +| all `construct` | 0.829 | 0.982 | 0.992 | 0.993 | +| 80% `construct`, 20% `add` | 0.826 | 0.981 | 0.991 | 0.993 | +| 50% `construct`, 50% `add` | 0.828 | 0.983 | 0.993 | 0.995 | + +### Recall after `add` + +The centroids and the rotation are fixed when the index is constructed, so `add` +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. From 0b3b2b16b34a4b78eea4b94ef5c3a31e5bb39c37 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 30 Sep 2026 21:41:44 +0800 Subject: [PATCH 08/12] Format the HNSW Python tests --- tests/python/test_hnsw.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index ef1f67fe..82ee4194 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -254,7 +254,9 @@ def test_add_rejects_bad_cluster_ids(partial_hnsw, base_data): with pytest.raises(ValueError, match="cluster_ids"): partial_hnsw.add(base_data[400:], cluster_ids=-np.ones(rows, dtype=np.int64)) with pytest.raises(ValueError, match="cluster_ids"): - partial_hnsw.add(base_data[400:], cluster_ids=np.zeros(rows - 1, dtype=np.int64)) + partial_hnsw.add( + base_data[400:], cluster_ids=np.zeros(rows - 1, dtype=np.int64) + ) assert partial_hnsw.num_points == 400 @@ -294,7 +296,9 @@ def test_resize_below_element_count_raises(partial_hnsw): assert partial_hnsw.max_elements == N_VECTORS -def test_added_points_survive_save_and_load(partial_hnsw, base_data, clusters, tmp_path): +def test_added_points_survive_save_and_load( + partial_hnsw, base_data, clusters, tmp_path +): _, cluster_ids = clusters partial_hnsw.add(base_data[400:], cluster_ids=cluster_ids[400:]) before_ids, before_dists = partial_hnsw.search(base_data, k=_TOPK, ef=_EF) From e90665f8a4084ee0394a174bedb1c887ec71f2ea Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Thu, 1 Oct 2026 01:01:51 +0800 Subject: [PATCH 09/12] Insert HNSW points in parallel --- include/rabitqlib/index/hnsw/hnsw.hpp | 48 ++++++-- include/rabitqlib/index/hnsw/hnsw_quant.hpp | 119 +++++++++++++------- python_bindings/hnsw_bindings.cpp | 9 +- tests/unit/rabitqlib/index/hnsw_test.cpp | 38 +++++++ 4 files changed, 157 insertions(+), 57 deletions(-) diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 2939267e..865739e1 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -71,11 +71,35 @@ class HierarchicalNSW { void construct(size_t, const float*, size_t, const float*, PID*, size_t, bool); - // Inserts points using the stored codes, labelling them num_points() + i. - // Single-threaded, and throws rather than growing past max_elements(). - std::vector add(const float*, size_t, const PID* = nullptr, bool = false); + /** + * @brief Insert points into a constructed or loaded index. + * + * Linked with the stored quantized codes, not the exact vectors construct uses, + * which are not kept. Centroids and rotation are never retrained. Point `i` gets + * the label `num_points() + i` whatever `num_threads` is. Throws rather than + * growing past `max_elements()`. Not safe to call while another thread searches. + * + * @param data New points, `n * dim` floats in the original coordinates + * @param n Number of new points + * @param cluster_ids Cluster of each point in `[0, num_clusters)`, or nullptr to + * route each one to its nearest centroid + * @param faster Same meaning as in `construct` + * @param num_threads Threads used to route and insert + * @return The label given to each new point + */ + std::vector add( + const float*, size_t, const PID* = nullptr, bool = false, size_t = 1 + ); - // Grows max_elements(). Invalidates every pointer into the index. + /** + * @brief Grow the number of points the index can hold. + * + * Reallocates and copies the base layer, so it needs room for both copies and + * costs time proportional to the whole index. Unchanged if it throws, and + * invalidates every pointer into the index. + * + * @param new_max_elements The new capacity, at least `num_points()` + */ void resize(size_t); std::vector>> search( const float*, size_t, size_t, size_t, size_t @@ -366,19 +390,19 @@ class HierarchicalNSW { void add_point(PID, PID, const quant::RabitqConfig&); - // Post-construction insertion, all defined in hnsw_quant.hpp. - void reconstruct_rotated(PID, float*) const; - float get_quant_dist(PID target, PID query) const; - maxheap> search_base_layer_quant(PID, PID, int); - void get_neighbors_by_heuristic2_quant(maxheap>&, size_t); - PID mutually_connect_quant(PID, maxheap>&, int); - PID add_point_quant(const float*, PID, const quant::RabitqConfig&); - maxheap> search_base_layer(PID, PID, int); PID mutually_connect_new_element(PID, maxheap>&, int); void get_neighbors_by_heuristic2(maxheap>&, size_t); + + // The same steps once the original data is gone, defined in hnsw_quant.hpp. + void reconstruct_rotated(PID, float*) const; + float get_quant_dist(PID target, PID query) const; + void add_point_quant(PID, int, char*, const float*, PID, const quant::RabitqConfig&); + maxheap> search_base_layer_quant(PID, PID, int); + PID mutually_connect_quant(PID, maxheap>&, int); + void get_neighbors_by_heuristic2_quant(maxheap>&, size_t); }; inline HierarchicalNSW::HierarchicalNSW( diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index 9bd93377..bc529227 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -1,5 +1,10 @@ -// The construct path with get_quant_dist in place of get_data_dist: rawDataPtr_ -// dangles once construct returns, so these link using the stored codes. +// Insertion into a built index: the construct path with get_quant_dist in place +// of get_data_dist, because rawDataPtr_ dangles once construct returns. +// +// add routes, draws every level, allocates every link list, then reserves one +// block of slots before inserting in parallel. Levels come first because the +// generator is shared; the slots are one block so labels do not depend on which +// thread finishes first. #pragma once #include @@ -339,40 +344,18 @@ inline PID HierarchicalNSW::mutually_connect_quant( return next_closest_entry_point; } -inline PID HierarchicalNSW::add_point_quant( - const float* vec, PID cluster_id, const quant::RabitqConfig& config +// The slot, the level and the link list are all decided by the caller, so this +// can run on many threads at once and the labels do not depend on their order. +inline void HierarchicalNSW::add_point_quant( + PID cur_c, + int curlevel, + char* link_list, + const float* vec, + PID cluster_id, + const quant::RabitqConfig& config ) { detail::quant_query().clear(); - int curlevel = get_random_level(mult_); - - // Allocated before the slot is claimed so a failure leaves the index untouched. - std::unique_ptr link_list(nullptr, std::free); - if (curlevel > 0) { - auto* raw = - static_cast(std::calloc(1, (size_links_per_element_ * curlevel) + 1)); - if (raw == nullptr) { - throw std::runtime_error( - "Not enough memory: add_point_quant failed to allocate linklist" - ); - } - link_list.reset(raw); - } - - PID cur_c = 0; - PID label = 0; - { - std::unique_lock lock_table(label_lookup_lock_); - if (cur_element_count_ >= max_elements_) { - throw std::runtime_error("The number of elements exceeds the specified limit"); - } - cur_c = static_cast(cur_element_count_.load()); - label = cur_c; - if (label_lookup_.find(label) != label_lookup_.end()) { - throw std::runtime_error("Label already present"); - } - cur_element_count_++; - label_lookup_[label] = cur_c; - } + const PID label = cur_c; std::unique_lock lock_el(link_list_locks_[cur_c]); element_levels_[cur_c] = curlevel; @@ -404,7 +387,7 @@ inline PID HierarchicalNSW::add_point_quant( ); if (curlevel > 0) { - linkLists_[cur_c] = link_list.release(); + linkLists_[cur_c] = link_list; } if (static_cast(curr_obj) != -1) { @@ -448,7 +431,6 @@ inline PID HierarchicalNSW::add_point_quant( enterpoint_node_ = cur_c; maxlevel_ = curlevel; } - return label; } // Grows the capacity. Every pointer into the level-0 block is invalidated, so @@ -508,7 +490,7 @@ inline void HierarchicalNSW::resize(size_t new_max_elements) { } inline std::vector HierarchicalNSW::add( - const float* data, size_t n, const PID* cluster_ids, bool faster + const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads ) { if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || rotator_ == nullptr || num_cluster_ == 0) { @@ -538,8 +520,8 @@ inline std::vector HierarchicalNSW::add( } } else { const auto* centroids = reinterpret_cast(centroids_memory_); - std::vector rotated(padded_dim_); - for (size_t i = 0; i < n; ++i) { + rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { + std::vector rotated(padded_dim_); rotator_->rotate(data + (i * dim_), rotated.data()); float best = std::numeric_limits::max(); for (size_t c = 0; c < num_cluster_; ++c) { @@ -551,7 +533,7 @@ inline std::vector HierarchicalNSW::add( assigned[i] = static_cast(c); } } - } + }); } quant::RabitqConfig config; @@ -559,10 +541,61 @@ inline std::vector HierarchicalNSW::add( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } - std::vector labels; - labels.reserve(n); + // Drawn here because the generator is shared: the inserts below run in + // parallel and must not touch it. + std::vector levels(n); + for (size_t i = 0; i < n; ++i) { + levels[i] = get_random_level(mult_); + } + + // Every link list is allocated before any of them is installed, so running + // out of memory leaves the index untouched. + std::vector> link_lists; + link_lists.reserve(n); + for (size_t i = 0; i < n; ++i) { + char* raw = nullptr; + if (levels[i] > 0) { + raw = + static_cast(std::calloc(1, (size_links_per_element_ * levels[i]) + 1) + ); + if (raw == nullptr) { + throw std::runtime_error( + "Not enough memory: HNSW add failed to allocate a link list" + ); + } + } + link_lists.emplace_back(raw, std::free); + } + + // One contiguous block of slots, so a point keeps its label whatever order + // the threads finish in. + PID first = 0; + { + std::unique_lock lock_table(label_lookup_lock_); + if (cur_element_count_ + n > max_elements_) { + throw std::runtime_error("The number of elements exceeds the specified limit"); + } + first = static_cast(cur_element_count_.load()); + for (size_t i = 0; i < n; ++i) { + label_lookup_[static_cast(first + i)] = static_cast(first + i); + } + cur_element_count_ += n; + } + + rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { + add_point_quant( + static_cast(first + i), + levels[i], + link_lists[i].release(), + data + (i * dim_), + assigned[i], + config + ); + }); + + std::vector labels(n); for (size_t i = 0; i < n; ++i) { - labels.push_back(add_point_quant(data + (i * dim_), assigned[i], config)); + labels[i] = static_cast(first + i); } return labels; } diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 81c55e1e..5a3e1cf2 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -109,7 +109,10 @@ class HnswIndex { } py::array_t add( - py::handle data, const py::object& cluster_ids, bool fast_quantization = false + py::handle data, + const py::object& cluster_ids, + size_t num_threads = 1, + bool fast_quantization = false ) { auto data_array = ensure_2d_array(data, "data"); if (!built_) { @@ -146,7 +149,8 @@ class HnswIndex { data_array.data(), rows, assigned.empty() ? nullptr : assigned.data(), - fast_quantization + fast_quantization, + num_threads ); std::copy(labels.begin(), labels.end(), ids.mutable_data()); return ids; @@ -295,6 +299,7 @@ void register_hnsw(py::module_& m) { &HnswIndex::add, py::arg("data"), py::arg("cluster_ids") = py::none(), + py::arg("num_threads") = 1, py::arg("fast_quantization") = false ) .def("resize", &HnswIndex::resize, py::arg("max_elements")) diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 8b7b294e..0f774481 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -431,6 +431,44 @@ TEST(HnswResizeTest, SurvivesSaveAndLoad) { std::filesystem::remove(path); } +// Labels must not depend on the order threads happen to finish in, and every +// added point must still be reachable. +TEST(HnswAddTest, ParallelAddKeepsLabelsAndTheGraph) { + constexpr size_t kBuilt = 256; + constexpr size_t kAdded = 256; + constexpr size_t kTotal = kBuilt + kAdded; + AddFixture fixture(kTotal, 53); + + HierarchicalNSW index(kTotal, AddFixture::kDim, 4, kTotal, kTotal); + index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 8, + false + ); + const auto labels = index.add( + fixture.data.data() + (kBuilt * AddFixture::kDim), + kAdded, + fixture.cluster_ids.data(), + false, + 8 + ); + + ASSERT_EQ(labels.size(), kAdded); + for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(labels[i], kBuilt + i); + } + EXPECT_EQ(index.num_points(), kTotal); + + const auto results = index.search(fixture.data.data(), kTotal, 1, kTotal, 1); + for (size_t i = 0; i < kTotal; ++i) { + EXPECT_EQ(results[i][0].second, i); + } +} + TEST(HnswAddTest, AddIntoAnEmptyIndexBuildsFromScratch) { constexpr size_t kCount = 96; AddFixture fixture(kCount, 17); From 0d5701cfd94b04a2e92f3920082d3b2ea9a85d82 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Thu, 1 Oct 2026 01:11:37 +0800 Subject: [PATCH 10/12] Drop the unreproducible recall table from the HNSW docs --- docs/docs/index/hnsw.md | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 67ea7c29..014722fd 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -109,7 +109,8 @@ std::vector HierarchicalNSW::add( const float* data, size_t n, const PID* cluster_ids = nullptr, - bool faster = false + bool faster = false, + size_t num_threads = 1 ); void HierarchicalNSW::resize(size_t new_max_elements); @@ -120,7 +121,9 @@ void HierarchicalNSW::resize(size_t new_max_elements); - **cluster_ids**: The cluster of each new vector, in `[0, num_clusters)`. When it is `nullptr`, each vector goes to its nearest centroid, chosen the same way a query is routed. -- **faster**: Same as in `construct`. +- **faster**, **num_threads**: Same as in `construct`. A point keeps the label + `num_points() + i` whatever order the threads finish in, because the slots are + reserved as one block before the inserts start. - **new_max_elements**: The new capacity. It must be at least `num_points()`. In Python: @@ -132,7 +135,7 @@ index.resize(index.max_elements + 100_000) ``` It is not safe to call `add` or `resize` while another thread searches the same -index, and `add` inserts one point at a time rather than in parallel. +index. ### Capacity @@ -155,16 +158,7 @@ compared against is rebuilt from its stored RaBitQ code, sign bit and extra bits together, and the other side is scored straight from its codes. This is the same approach quantized SymphonyQG construction takes. -Graph quality therefore depends on how a point arrived. On SIFT-1M restricted to -100,000 vectors (M = 32, 9 bits, 500 queries, recall@10 against exact -neighbors), an index built entirely by `construct` and one whose last fifth -arrived through `add` are indistinguishable: - -| built by | ef 10 | ef 50 | ef 100 | ef 200 | -| --- | --- | --- | --- | --- | -| all `construct` | 0.829 | 0.982 | 0.992 | 0.993 | -| 80% `construct`, 20% `add` | 0.826 | 0.981 | 0.991 | 0.993 | -| 50% `construct`, 50% `add` | 0.828 | 0.983 | 0.993 | 0.995 | +Graph quality therefore depends on how a point arrived. ### Recall after `add` From 0a99549ecef5df8ecf6d7aa22b94350c4502d47a Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 1 Oct 2026 14:07:54 +0800 Subject: [PATCH 11/12] Fix HNSW insertion exception safety and neighbor pruning Prepare returned IDs, quantized rows, and owned link lists before publishing new points. Keep the index valid if graph linking fails, and document the partial-insertion behavior. Reject resizing an uninitialized index and score all pruning candidates with the same query direction. Add allocation-failure, resize, and pruning regression coverage. Validation: 272 C++ tests; 633 Python tests passed, 2 skipped; full core and binding clang-tidy; focused test clang-tidy; formatting; strict MkDocs build. Focused pruning regression also passed ASan/UBSan with leak detection disabled. Benchmarks omitted because these changes fix correctness; no performance improvement is claimed. --- docs/docs/index/hnsw.md | 9 + include/rabitqlib/index/hnsw/hnsw.hpp | 7 +- include/rabitqlib/index/hnsw/hnsw_quant.hpp | 198 +++++++++-------- tests/CMakeLists.txt | 6 + tests/failures/hnsw_add_allocation_test.cpp | 228 ++++++++++++++++++++ tests/unit/rabitqlib/index/hnsw_test.cpp | 104 +++++++++ 6 files changed, 464 insertions(+), 88 deletions(-) create mode 100644 tests/failures/hnsw_add_allocation_test.cpp diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 014722fd..69ba7050 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -125,6 +125,8 @@ void HierarchicalNSW::resize(size_t new_max_elements); `num_points() + i` whatever order the threads finish in, because the slots are reserved as one block before the inserts start. - **new_max_elements**: The new capacity. It must be at least `num_points()`. + In C++, initialize the index through its parameterized constructor or `load` + before calling `resize`. In Python: @@ -137,6 +139,13 @@ index.resize(index.max_elements + 100_000) It is not safe to call `add` or `resize` while another thread searches the same index. +`add` prepares the result IDs, quantized vectors, and link-list allocations before +changing the graph or point count. A failure during preparation leaves those +unchanged. A failure during graph linking can leave the batch counted in +`num_points()` with only some points reachable through the graph. The index +remains safe to search, save, load, and destroy; inspect `num_points()` before +retrying, or rebuild to restore full reachability. + ### Capacity `add` never grows the index: it throws once `num_points()` reaches diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 865739e1..4cb1d218 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -78,6 +78,9 @@ class HierarchicalNSW { * which are not kept. Centroids and rotation are never retrained. Point `i` gets * the label `num_points() + i` whatever `num_threads` is. Throws rather than * growing past `max_elements()`. Not safe to call while another thread searches. + * Preparation failures leave the graph and point count unchanged. A failure + * during linking may leave a partially linked batch counted in the index, + * which remains safe to search, save, load, and destroy. * * @param data New points, `n * dim` floats in the original coordinates * @param n Number of new points @@ -97,6 +100,7 @@ class HierarchicalNSW { * Reallocates and copies the base layer, so it needs room for both copies and * costs time proportional to the whole index. Unchanged if it throws, and * invalidates every pointer into the index. + * Requires initialization through the parameterized constructor or `load`. * * @param new_max_elements The new capacity, at least `num_points()` */ @@ -159,6 +163,7 @@ class HierarchicalNSW { }; private: + friend struct HnswPruningTestAccess; friend maxheap> detail::search_knn_neon( HierarchicalNSW&, const float*, size_t ); @@ -399,7 +404,7 @@ class HierarchicalNSW { // The same steps once the original data is gone, defined in hnsw_quant.hpp. void reconstruct_rotated(PID, float*) const; float get_quant_dist(PID target, PID query) const; - void add_point_quant(PID, int, char*, const float*, PID, const quant::RabitqConfig&); + void add_point_quant(PID, int, std::unique_ptr&); maxheap> search_base_layer_quant(PID, PID, int); PID mutually_connect_quant(PID, maxheap>&, int); void get_neighbors_by_heuristic2_quant(maxheap>&, size_t); diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index bc529227..025c0297 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -1,10 +1,10 @@ // Insertion into a built index: the construct path with get_quant_dist in place // of get_data_dist, because rawDataPtr_ dangles once construct returns. // -// add routes, draws every level, allocates every link list, then reserves one -// block of slots before inserting in parallel. Levels come first because the -// generator is shared; the slots are one block so labels do not depend on which -// thread finishes first. +// 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 +// first because the generator is shared; the slots are one block so labels do +// not depend on which thread finishes first. #pragma once #include @@ -154,7 +154,7 @@ inline float HierarchicalNSW::get_quant_dist(PID target, PID query) const { inline maxheap> HierarchicalNSW::search_base_layer_quant( 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; @@ -202,7 +202,8 @@ inline maxheap> HierarchicalNSW::search_base_layer_quant( } } } - visited_list_pool_->release_vis_list(vl); + visited_list_pool_->release_vis_list(vl.get()); + vl.release(); return top_candidates; } @@ -321,7 +322,7 @@ inline PID HierarchicalNSW::mutually_connect_quant( data[sz_link_list_other] = cur_c; set_list_count(ll_other, sz_link_list_other + 1); } else { - float d_max = get_quant_dist(selected_neighbor, cur_c); + float d_max = get_quant_dist(cur_c, selected_neighbor); maxheap> candidates; candidates.emplace(d_max, cur_c); for (size_t j = 0; j < sz_link_list_other; j++) { @@ -347,18 +348,11 @@ inline PID HierarchicalNSW::mutually_connect_quant( // The slot, the level and the link list are all decided by the caller, so this // can run on many threads at once and the labels do not depend on their order. inline void HierarchicalNSW::add_point_quant( - PID cur_c, - int curlevel, - char* link_list, - const float* vec, - PID cluster_id, - const quant::RabitqConfig& config + PID cur_c, int curlevel, std::unique_ptr& link_list ) { detail::quant_query().clear(); - const PID label = cur_c; std::unique_lock lock_el(link_list_locks_[cur_c]); - element_levels_[cur_c] = curlevel; std::unique_lock templock(global_); int maxlevelcopy = maxlevel_; @@ -366,76 +360,80 @@ inline void HierarchicalNSW::add_point_quant( templock.unlock(); } PID curr_obj = enterpoint_node_; + const PID previous_entry_point = curr_obj; - std::memset( - data_level0_memory_ + (cur_c * size_data_per_element_), 0, size_data_per_element_ - ); - std::memcpy(get_external_label_pt(cur_c), &label, sizeof(PID)); - std::memcpy(get_clusterid_pt(cur_c), &cluster_id, sizeof(PID)); - - std::vector rotated_data(padded_dim_); - rotator_->rotate(vec, rotated_data.data()); - quant::quantize_split_single( - rotated_data.data(), - reinterpret_cast(centroids_memory_) + (cluster_id * padded_dim_), - padded_dim_, - ex_bits_, - get_bindata_by_internalid(cur_c), - get_exdata_by_internalid(cur_c), - metric_type_, - config - ); + linkLists_[cur_c] = link_list.release(); + element_levels_[cur_c] = curlevel; - if (curlevel > 0) { - linkLists_[cur_c] = link_list; - } - - if (static_cast(curr_obj) != -1) { - if (curlevel < maxlevelcopy) { - float curdist = get_quant_dist(curr_obj, cur_c); - for (int level = maxlevelcopy; level > curlevel; level--) { - bool changed = true; - while (changed) { - changed = false; - std::unique_lock lock(link_list_locks_[curr_obj]); - PID* data = get_linklist(curr_obj, level); - int size = get_list_count(data); - auto* datal = data + 1; - for (int i = 0; i < size; i++) { - PID cand = datal[i]; - if (cand >= cur_element_count_) { - throw std::runtime_error("cand error"); - } - float d = get_quant_dist(cand, cur_c); - if (d < curdist) { - curdist = d; - curr_obj = cand; - changed = true; + try { + if (static_cast(curr_obj) != -1) { + if (curlevel < maxlevelcopy) { + float curdist = get_quant_dist(curr_obj, cur_c); + for (int level = maxlevelcopy; level > curlevel; level--) { + bool changed = true; + while (changed) { + changed = false; + std::unique_lock lock(link_list_locks_[curr_obj]); + PID* data = get_linklist(curr_obj, level); + int size = get_list_count(data); + auto* datal = data + 1; + for (int i = 0; i < size; i++) { + PID cand = datal[i]; + if (cand >= cur_element_count_) { + throw std::runtime_error("cand error"); + } + float d = get_quant_dist(cand, cur_c); + if (d < curdist) { + curdist = d; + curr_obj = cand; + changed = true; + } } } } } - } - for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { - maxheap> top_candidates = - search_base_layer_quant(curr_obj, cur_c, level); - curr_obj = mutually_connect_quant(cur_c, top_candidates, level); + for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { + maxheap> top_candidates = + search_base_layer_quant(curr_obj, cur_c, level); + curr_obj = mutually_connect_quant(cur_c, top_candidates, level); + } + } else { + enterpoint_node_ = cur_c; + maxlevel_ = curlevel; } - } else { - enterpoint_node_ = cur_c; - maxlevel_ = curlevel; - } - if (curlevel > maxlevelcopy) { - enterpoint_node_ = cur_c; - maxlevel_ = curlevel; + if (curlevel > maxlevelcopy) { + enterpoint_node_ = cur_c; + maxlevel_ = curlevel; + } + } catch (...) { + // Upper layers may already expose this node. Keep a route to the old + // graph if base-layer linking was interrupted before choosing neighbors. + PID* base_links = get_linklist0(cur_c); + if (previous_entry_point != kPidMax && get_list_count(base_links) == 0) { + base_links[1] = previous_entry_point; + set_list_count(base_links, 1); + } + // An interrupted promotion has no incoming links above the previous + // maximum level. Discard those levels so save/load sees a valid graph. + if (curlevel > maxlevelcopy) { + element_levels_[cur_c] = std::max(0, maxlevelcopy); + if (element_levels_[cur_c] == 0) { + std::free(linkLists_[cur_c]); + linkLists_[cur_c] = nullptr; + } + } + throw; } } // Grows the capacity. Every pointer into the level-0 block is invalidated, so // this must not run while another thread searches. inline void HierarchicalNSW::resize(size_t new_max_elements) { + if (size_data_per_element_ == 0) { + throw std::logic_error("HNSW index must be initialized before resize"); + } if (new_max_elements < cur_element_count_) { throw std::invalid_argument("HNSW resize: capacity below the element count"); } @@ -508,6 +506,7 @@ inline std::vector HierarchicalNSW::add( if (n > buffer::kSearchBufferMaxPointCount - cur_element_count_) { throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); } + std::vector labels(n); // Either the caller names the cluster of every point, or each one goes to its // nearest centroid, ordered the way a query is routed. std::vector assigned(n, 0); @@ -567,36 +566,61 @@ inline std::vector HierarchicalNSW::add( link_lists.emplace_back(raw, std::free); } - // One contiguous block of slots, so a point keeps its label whatever order - // the threads finish in. - PID first = 0; + // Prepare every code before publishing the slots. If quantization or worker + // creation fails, the graph and point count are unchanged. { std::unique_lock lock_table(label_lookup_lock_); if (cur_element_count_ + n > max_elements_) { throw std::runtime_error("The number of elements exceeds the specified limit"); } - first = static_cast(cur_element_count_.load()); + const auto first = static_cast(cur_element_count_.load()); + rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { + const auto id = static_cast(first + i); + labels[i] = id; + std::memset( + data_level0_memory_ + (id * size_data_per_element_), + 0, + size_data_per_element_ + ); + std::memcpy(get_external_label_pt(id), &id, sizeof(PID)); + std::memcpy(get_clusterid_pt(id), &assigned[i], sizeof(PID)); + std::vector rotated_data(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated_data.data()); + quant::quantize_split_single( + rotated_data.data(), + reinterpret_cast(centroids_memory_) + (assigned[i] * padded_dim_), + padded_dim_, + ex_bits_, + get_bindata_by_internalid(id), + get_exdata_by_internalid(id), + metric_type_, + config + ); + }); + size_t registered = 0; + try { + for (; registered < n; ++registered) { + label_lookup_[labels[registered]] = labels[registered]; + } + } catch (...) { + for (size_t i = 0; i < registered; ++i) { + label_lookup_.erase(labels[i]); + } + throw; + } + // Unstarted inserts are valid level-0 nodes with owned codes. Their + // upper-layer lists remain under RAII until an insert takes ownership. for (size_t i = 0; i < n; ++i) { - label_lookup_[static_cast(first + i)] = static_cast(first + i); + element_levels_[labels[i]] = 0; + linkLists_[labels[i]] = nullptr; } cur_element_count_ += n; } rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { - add_point_quant( - static_cast(first + i), - levels[i], - link_lists[i].release(), - data + (i * dim_), - assigned[i], - config - ); + add_point_quant(labels[i], levels[i], link_lists[i]); }); - std::vector labels(n); - for (size_t i = 0; i < n; ++i) { - labels[i] = static_cast(first + i); - } return labels; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 9c76145d..144499b8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -47,6 +47,12 @@ target_link_libraries(fastscan_no_simd_test PRIVATE GTest::gtest_main rabitq_hea rabitq_enable_warnings(fastscan_no_simd_test) gtest_discover_tests(fastscan_no_simd_test DISCOVERY_MODE PRE_TEST) +# Keep the global allocation-failure override isolated from the main suite. +add_executable(hnsw_add_allocation_test failures/hnsw_add_allocation_test.cpp) +target_link_libraries(hnsw_add_allocation_test PRIVATE GTest::gtest_main rabitq_headers) +rabitq_enable_warnings(hnsw_add_allocation_test) +gtest_discover_tests(hnsw_add_allocation_test DISCOVERY_MODE PRE_TEST) + # Message to verify files are found message(STATUS "Discovered test files:") foreach(TEST_SRC ${UNIT_TESTS} ${INTEG_TESTS}) diff --git a/tests/failures/hnsw_add_allocation_test.cpp b/tests/failures/hnsw_add_allocation_test.cpp new file mode 100644 index 00000000..404e8452 --- /dev/null +++ b/tests/failures/hnsw_add_allocation_test.cpp @@ -0,0 +1,228 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/index/hnsw/hnsw.hpp" + +namespace { +std::atomic fail_enabled{false}; +std::atomic allocations_before_failure{0}; +size_t allocation_size = 0; +const rabitqlib::hnsw::HierarchicalNSW* watched_index = nullptr; +size_t watched_count = 0; + +class AllocationFailure { + public: + explicit AllocationFailure(size_t before_failure, size_t size = 0) { + allocations_before_failure = before_failure; + allocation_size = size; + fail_enabled = true; + } + ~AllocationFailure() { fail_enabled = false; } +}; +} // namespace + +// This executable is separate from the normal suite so fault injection cannot +// affect unrelated tests. Each failure is one-shot, allowing cleanup to allocate. +GTEST_NO_INLINE_ void* operator new(size_t size) { + if (fail_enabled && (allocation_size == 0 || allocation_size == size) && + (watched_index == nullptr || watched_index->num_points() == watched_count)) { + if (allocations_before_failure.fetch_sub(1) == 0) { + fail_enabled = false; + throw std::bad_alloc(); + } + } + if (void* ptr = std::malloc(size == 0 ? 1 : size)) { + return ptr; + } + throw std::bad_alloc(); +} + +GTEST_NO_INLINE_ void* operator new(size_t size, const std::nothrow_t&) noexcept { + try { + return ::operator new(size); + } catch (...) { return nullptr; } +} + +// The analyzer models delete as freeing before entering these replacements. +GTEST_NO_INLINE_ void operator delete(void* ptr) noexcept { + std::free(ptr); // NOLINT(clang-analyzer-cplusplus.NewDelete) +} +GTEST_NO_INLINE_ void operator delete(void* ptr, size_t) noexcept { + std::free(ptr); // NOLINT(clang-analyzer-cplusplus.NewDelete) +} + +namespace rabitqlib::hnsw { +namespace { +constexpr size_t kDim = 64; +constexpr size_t kAdded = 3; +constexpr size_t kCapacity = 8; + +struct Fixture { + std::array data{}; + std::array centroid{}; + std::array clusters{}; + + Fixture() { + std::mt19937 generator(42); + std::uniform_real_distribution distribution(-1, 1); + for (auto& value : data) { + value = distribution(generator); + } + } + + void build(HierarchicalNSW& index) const { + auto ids = clusters; + index.construct(1, centroid.data(), 1, data.data(), ids.data(), 1, false); + } +}; + +TEST(HnswAddAllocationTest, AllocatesReturnedIdsBeforeMutation) { + Fixture fixture; + HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); + fixture.build(index); + bool preparation_failed = false; + { + AllocationFailure failure(0, kAdded * sizeof(PID)); + try { + index.add(fixture.data.data() + kDim, kAdded, fixture.clusters.data()); + } catch (const std::bad_alloc&) { preparation_failed = true; } + } + ASSERT_TRUE(preparation_failed); + EXPECT_EQ(index.num_points(), 1U); + watched_index = &index; + watched_count = 1 + kAdded; + std::vector ids; + bool failed = false; + { + AllocationFailure failure(0, kAdded * sizeof(PID)); + try { + ids = index.add(fixture.data.data() + kDim, kAdded, fixture.clusters.data()); + } catch (const std::bad_alloc&) { failed = true; } + } + watched_index = nullptr; + EXPECT_FALSE(failed); + EXPECT_EQ(ids, (std::vector{1, 2, 3})); + EXPECT_EQ(index.num_points(), 1 + kAdded); +} + +TEST(HnswAddAllocationTest, PreparationFailurePreservesIndex) { + for (const size_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + Fixture fixture; + HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); + fixture.build(index); + const auto before = index.search(fixture.data.data(), 1, 1, kCapacity, 1); + bool failed = false; + { + AllocationFailure failure(0, kDim * sizeof(float)); + try { + index.add( + fixture.data.data() + kDim, + kAdded, + fixture.clusters.data(), + false, + threads + ); + } catch (const std::bad_alloc&) { failed = true; } + } + ASSERT_TRUE(failed); + EXPECT_EQ(index.num_points(), 1U); + EXPECT_EQ(index.search(fixture.data.data(), 1, 1, kCapacity, 1), before); + EXPECT_EQ( + index.add(fixture.data.data() + kDim, kAdded, fixture.clusters.data()), + (std::vector{1, 2, 3}) + ); + } +} + +TEST(HnswAddAllocationTest, EveryAllocationFailureLeavesUsableIndex) { + Fixture fixture; + const std::string path = + ::testing::TempDir() + "rabitq_hnsw_add_allocation_failure.index"; + // Exercise interrupted promotions from level 0 (freeing the upper list) and + // level 1 (retaining its lower part). Use this platform's default engine, + // just as HNSW does, so the selected seeds work across standard libraries. + std::array seeds{}; + for (size_t level = 0; level < seeds.size(); ++level) { + for (size_t seed = 1; seed < 10000; ++seed) { + std::default_random_engine generator(seed); + std::uniform_real_distribution distribution(0, 1); + const auto built_level = + static_cast(-std::log(distribution(generator)) / std::log(2.0)); + const auto added_level = + static_cast(-std::log(distribution(generator)) / std::log(2.0)); + if (built_level == level && added_level > level) { + seeds[level] = seed; + break; + } + } + ASSERT_NE(seeds[level], 0U); + } + for (const size_t seed : seeds) { + SCOPED_TRACE(seed); + for (const size_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + size_t failures = 0; + size_t linking_failures = 0; + bool succeeded = false; + for (size_t allocation = 0; allocation < 256; ++allocation) { + SCOPED_TRACE(allocation); + HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, seed); + fixture.build(index); + bool failed = false; + { + AllocationFailure failure(allocation); + try { + index.add( + fixture.data.data() + kDim, + kAdded, + fixture.clusters.data(), + false, + threads + ); + } catch (const std::bad_alloc&) { failed = true; } + } + failures += failed ? 1 : 0; + linking_failures += (failed && index.num_points() > 1) ? 1 : 0; + const auto results = index.search(fixture.data.data(), 1, 1, kCapacity, 1); + ASSERT_EQ(results.size(), 1U); + ASSERT_EQ(results[0].size(), 1U); + EXPECT_EQ(results[0][0].second, 0U); + ASSERT_NO_THROW(index.save(path.c_str())); + HierarchicalNSW loaded; + ASSERT_NO_THROW(loaded.load(path.c_str())); + EXPECT_EQ(loaded.num_points(), index.num_points()); + EXPECT_EQ(loaded.search(fixture.data.data(), 1, 1, kCapacity, 1), results); + const auto next_id = static_cast(loaded.num_points()); + EXPECT_EQ( + loaded.add( + fixture.data.data() + ((1 + kAdded) * kDim), + 1, + fixture.clusters.data() + ), + (std::vector{next_id}) + ); + if (!failed) { + succeeded = true; + break; + } + } + EXPECT_GT(failures, 0U); + EXPECT_GT(linking_failures, 0U); + EXPECT_TRUE(succeeded); + } + } + std::filesystem::remove(path); +} +} // namespace +} // namespace rabitqlib::hnsw diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 0f774481..8fa813d0 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -12,8 +13,68 @@ #include namespace rabitqlib::hnsw { + +struct HnswPruningTestAccess { + static void quantize(HierarchicalNSW& index, PID id, const float* rotated) { + const auto* centroid = reinterpret_cast(index.centroids_memory_); + quant::quantize_compact_one_bit( + rotated, centroid, index.padded_dim_, index.get_bindata_by_internalid(id) + ); + } + + static float distance(const HierarchicalNSW& index, PID target, PID query) { + return index.get_quant_dist(target, query); + } + + static std::vector connect(HierarchicalNSW& index, PID incoming) { + PID* neighbors = index.get_linklist0(0); + index.set_list_count(neighbors, index.maxM0_); + for (size_t i = 0; i < index.maxM0_; ++i) { + neighbors[i + 1] = static_cast(i + 1); + } + std::fill_n(index.get_linklist0(incoming), index.maxM0_ + 1, PID{0}); + detail::quant_query().clear(); + maxheap> candidates; + candidates.emplace(index.get_quant_dist(0, incoming), 0); + index.mutually_connect_quant(incoming, candidates, 0); + return {neighbors + 1, neighbors + 1 + index.get_list_count(neighbors)}; + } +}; + namespace { +TEST(HnswPruningTest, ScoresNewAndExistingNeighborsWithTheSameQuery) { + constexpr size_t kDim = 64; + constexpr size_t kCount = 6; + constexpr PID kIncoming = kCount - 1; + std::vector data(kCount * kDim, 2.0F); + std::fill_n(data.begin(), kDim, 1.0F); + // A sparse residual reconstructs with a much larger norm than its source. + // Scoring its codes and using its reconstruction as a query differ sharply. + std::fill_n(data.begin() + (kIncoming * kDim), kDim, 0.01F); + data[kIncoming * kDim] = 8.0F; + std::vector centroid(kDim, 0.0F); + std::vector clusters(kCount, 0); + HierarchicalNSW index(kCount, kDim, 1, 2, 10); + index.construct(1, centroid.data(), kCount, data.data(), clusters.data(), 1, false); + + // Set codes directly in the rotated domain to avoid random rotation changing + // the asymmetry, then fill node 0's base-layer list to force pruning. + for (PID id = 0; id < kCount; ++id) { + HnswPruningTestAccess::quantize(index, id, data.data() + (id * kDim)); + } + detail::quant_query().clear(); + const float existing_distance = HnswPruningTestAccess::distance(index, 1, 0); + EXPECT_LT(HnswPruningTestAccess::distance(index, kIncoming, 0), existing_distance); + EXPECT_GT(HnswPruningTestAccess::distance(index, 0, kIncoming), existing_distance); + + // The new point is closest when node 0 is the query. Its duplicate existing + // neighbors are pruned; the reversed estimate instead retains an old point. + EXPECT_EQ( + HnswPruningTestAccess::connect(index, kIncoming), std::vector{kIncoming} + ); +} + TEST(HnswConfigurationTest, RejectsUnsupportedMetric) { EXPECT_THROW( (HierarchicalNSW(8, 64, 1, 2, 10, 100, static_cast(255))), @@ -350,6 +411,47 @@ TEST(HnswAddTest, RoutesToTheNearestCentroidWhenNoClustersGiven) { } } +TEST(HnswResizeTest, RejectsUninitializedIndex) { + HierarchicalNSW index; + for (const size_t capacity : {10U, 0U}) { + SCOPED_TRACE(capacity); + try { + index.resize(capacity); + FAIL() << "Resizing an uninitialized index must fail"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW index must be initialized before resize"); + } + EXPECT_EQ(index.max_elements(), 0U); + EXPECT_EQ(index.num_points(), 0U); + } +} + +TEST(HnswResizeTest, AllowsResizeBeforeConstruction) { + constexpr size_t kBuilt = 64; + constexpr size_t kCapacity = 128; + AddFixture fixture(kBuilt, 43); + HierarchicalNSW index(kBuilt, AddFixture::kDim, 4, 8, 50); + + ASSERT_NO_THROW(index.resize(kCapacity)); + EXPECT_EQ(index.max_elements(), kCapacity); + EXPECT_EQ(index.num_points(), 0U); + ASSERT_NO_THROW(index.construct( + 1, + fixture.centroid.data(), + kBuilt, + fixture.data.data(), + fixture.cluster_ids.data(), + 1, + false + )); + const auto results = index.search(fixture.data.data(), 8, 1, kBuilt, 1); + ASSERT_EQ(results.size(), 8U); + for (size_t i = 0; i < results.size(); ++i) { + ASSERT_EQ(results[i].size(), 1U); + EXPECT_EQ(results[i][0].second, i); + } +} + TEST(HnswResizeTest, GrowsCapacityAndKeepsTheGraph) { constexpr size_t kBuilt = 128; constexpr size_t kAdded = 64; @@ -424,6 +526,8 @@ TEST(HnswResizeTest, SurvivesSaveAndLoad) { loaded.load(path.string().c_str()); EXPECT_EQ(loaded.max_elements(), kTotal); EXPECT_EQ(loaded.num_points(), kTotal); + ASSERT_NO_THROW(loaded.resize(kTotal + 8)); + EXPECT_EQ(loaded.max_elements(), kTotal + 8); const auto results = loaded.search(fixture.data.data(), kTotal, 1, kTotal, 1); for (size_t i = 0; i < kTotal; ++i) { EXPECT_EQ(results[i][0].second, i); From 153f391821eb269109b1f6f700831886f7d21907 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 1 Oct 2026 14:21:52 +0800 Subject: [PATCH 12/12] Fix flaky nearest-centroid routing test Keep exact checks for assigned cluster IDs, returned labels, and point count. Check finite, valid search results in the expected well-separated cluster instead of requiring exact point IDs from approximate search. The original test failed 13 of 500 local repetitions because random rotation can change ANN results. The updated test passed 1,000 repetitions. --- tests/unit/rabitqlib/index/hnsw_test.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 8fa813d0..732f3aed 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -399,15 +399,23 @@ TEST(HnswAddTest, RoutesToTheNearestCentroidWhenNoClustersGiven) { const auto labels = index.add(fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, nullptr); ASSERT_EQ(labels.size(), kAdded); + EXPECT_EQ(index.num_points(), kTotal); for (size_t i = 0; i < kAdded; ++i) { + EXPECT_EQ(labels[i], kBuilt + i); EXPECT_EQ(index.cluster_id_of(labels[i]), cluster_ids[kBuilt + i]); } + // ANN point IDs can vary with random rotation. Search should still return + // valid results in the query's well-separated cluster. const auto results = index.search( fixture.data.data() + (kBuilt * AddFixture::kDim), kAdded, 1, kTotal, 1 ); + ASSERT_EQ(results.size(), kAdded); for (size_t i = 0; i < kAdded; ++i) { - EXPECT_EQ(results[i][0].second, kBuilt + i); + ASSERT_EQ(results[i].size(), 1U); + EXPECT_TRUE(std::isfinite(results[i][0].first)); + ASSERT_LT(results[i][0].second, kTotal); + EXPECT_EQ(index.cluster_id_of(results[i][0].second), cluster_ids[kBuilt + i]); } }