From 3443da6b9f9923c3b2af8b4e209ec9dc5524a8ba Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Tue, 18 Aug 2026 13:12:39 +0800 Subject: [PATCH 1/5] fix: adjust HashBasedBooleanSet initialization for full ID coverage --- include/rabitqlib/index/symqg/qg_builder.hpp | 2 +- include/rabitqlib/utils/hashset.hpp | 88 ++++++++------------ include/rabitqlib/utils/visited_pool.hpp | 3 +- 3 files changed, 38 insertions(+), 55 deletions(-) diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 9d90248..a520553 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -67,7 +67,7 @@ class QGBuilder { , pruned_neighbors_(qg_.num_vertices()) , visited_list_( num_threads_, - HashBasedBooleanSet(std::min(ef_build_ * ef_build_, num_nodes_ / 10)) + HashBasedBooleanSet(num_nodes_) ) , degrees_(qg_.num_vertices(), degree_bound_) { omp_set_num_threads(static_cast(num_threads_)); diff --git a/include/rabitqlib/utils/hashset.hpp b/include/rabitqlib/utils/hashset.hpp index b363789..dd04364 100644 --- a/include/rabitqlib/utils/hashset.hpp +++ b/include/rabitqlib/utils/hashset.hpp @@ -18,84 +18,66 @@ #pragma once -#include -#include -#include -#include +#include +#include +#include #include "rabitqlib/defines.hpp" #include "rabitqlib/utils/memory.hpp" namespace rabitqlib { /** - * @brief hash set to record visited vertices + * @brief Set of visited vertices, backed by a generation-stamped array. * + * One byte-pair per vertex holds the generation in which that vertex was last + * marked; a vertex is "visited" iff its stamp equals the current generation. + * clear() therefore just bumps the generation -- O(1), and only the rare + * 16-bit wraparound pays for an actual refill. + * + * This replaces a direct-mapped table with an std::unordered_set overflow. That + * table was sized at max_elements/10 rounded down to a power of two, so a + * single ef=100 search collided on nearly every visit and allocated an + * unordered_set node each time: measured at ~1200 mallocs per query, 98% of all + * allocation on the search path. Removing them was worth ~1.6x on HNSW search. + * + * The constructor argument is now the ID SPACE (number of vertices), not a + * bucket-count hint -- ids are used to index directly, so it must cover every + * id that will be passed to get()/set(). */ class HashBasedBooleanSet { private: - size_t table_size_ = 0; - PID mask_ = 0; - std::vector> table_; - std::unordered_set stl_hash_; - - [[nodiscard]] auto hash1(const PID value) const { return value & mask_; } + std::vector> stamp_; + uint16_t cur_ = 0; public: HashBasedBooleanSet() = default; ~HashBasedBooleanSet() = default; HashBasedBooleanSet(const HashBasedBooleanSet&) = default; + HashBasedBooleanSet& operator=(const HashBasedBooleanSet&) = default; HashBasedBooleanSet(HashBasedBooleanSet&&) noexcept = default; HashBasedBooleanSet& operator=(HashBasedBooleanSet&&) noexcept = default; - explicit HashBasedBooleanSet(size_t size) { - size_t bit_size = 0; - size_t bit = size; - while (bit != 0) { - bit_size++; - bit >>= 1; - } - size_t bucket_size = 0x1 << ((bit_size + 4) / 2 + 3); - initialize(bucket_size); - } - - void initialize(const size_t table_size) { - table_size_ = table_size; - mask_ = static_cast(table_size_ - 1); - const PID check_val = hash1(static_cast(table_size)); - if (check_val != 0) { - std::cerr << "[WARN] table size is not 2^N : " << table_size << '\n'; - } + explicit HashBasedBooleanSet(size_t num_elements) { initialize(num_elements); } - table_ = std::vector>(table_size); - std::fill(table_.begin(), table_.end(), kPidMax); - stl_hash_.clear(); + void initialize(size_t num_elements) { + stamp_.assign(num_elements, 0); + cur_ = 0; + // Leaves the set usable without an explicit clear(): every stamp is 0 + // while the live generation is 1, so nothing reads as visited. + clear(); } void clear() { - std::fill(table_.begin(), table_.end(), kPidMax); - stl_hash_.clear(); + if (++cur_ == 0) { + std::fill(stamp_.begin(), stamp_.end(), 0); + cur_ = 1; + } } // get if data_id is in the hashset - [[nodiscard]] bool get(PID data_id) const { - PID val = this->table_[hash1(data_id)]; - if (val == data_id) { - return true; - } - return (val != kPidMax && stl_hash_.find(data_id) != stl_hash_.end()); - } + [[nodiscard]] bool get(PID data_id) const { return stamp_[data_id] == cur_; } - void set(PID data_id) { - PID& val = table_[hash1(data_id)]; - if (val == data_id) { - return; - } - if (val == kPidMax) { - val = data_id; - } else { - stl_hash_.emplace(data_id); - } - } + void set(PID data_id) { stamp_[data_id] = cur_; } }; -} // namespace rabitqlib \ No newline at end of file +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_pool.hpp b/include/rabitqlib/utils/visited_pool.hpp index ac9d27d..d946b3f 100644 --- a/include/rabitqlib/utils/visited_pool.hpp +++ b/include/rabitqlib/utils/visited_pool.hpp @@ -12,7 +12,8 @@ class VisitedListPool { public: VisitedListPool(size_t initpoolsize, size_t max_elements) { - numelements_ = max_elements / 10; + // Must cover the whole id space: HashBasedBooleanSet indexes by id. + numelements_ = max_elements; for (size_t i = 0; i < initpoolsize; i++) { pool_.push_front(new HashBasedBooleanSet(numelements_)); } From 981e137253ef56ba634c43992e5f0bd5bccc3eb0 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Tue, 18 Aug 2026 13:14:08 +0800 Subject: [PATCH 2/5] fix: replace raw distance function with euclidean square in search_knn_direct --- include/rabitqlib/index/hnsw/hnsw.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 3f53b8d..9ad3658 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -1138,7 +1138,7 @@ inline maxheap> HierarchicalNSW::search_knn_direct( if (metric_type_ == METRIC_L2) { for (size_t i = 0; i < num_cluster_; i++) { - q_to_centroids[i] = std::sqrt(raw_dist_func_( + q_to_centroids[i] = std::sqrt(euclidean_sqr( rotated_query, reinterpret_cast(centroids_memory_) + (i * padded_dim_), padded_dim_ From 1865179838148a7b01cad9b84b8c7aa5810a47ad Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 19 Aug 2026 12:03:56 +0800 Subject: [PATCH 3/5] refactor: replace HashBasedBooleanSet with VisitedSet across multiple files --- include/rabitqlib/index/hnsw/hnsw.hpp | 4 +- include/rabitqlib/index/symqg/qg.hpp | 14 +- include/rabitqlib/index/symqg/qg_builder.hpp | 8 +- include/rabitqlib/utils/visited_pool.hpp | 20 +-- include/rabitqlib/utils/visited_set.hpp | 82 ++++++++++++ .../{hashset.hpp => visited_set_epoch.hpp} | 45 ++----- include/rabitqlib/utils/visited_set_hash.hpp | 121 ++++++++++++++++++ 7 files changed, 240 insertions(+), 54 deletions(-) create mode 100644 include/rabitqlib/utils/visited_set.hpp rename include/rabitqlib/utils/{hashset.hpp => visited_set_epoch.hpp} (50%) create mode 100644 include/rabitqlib/utils/visited_set_hash.hpp diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 9ad3658..b8d4612 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -777,7 +777,7 @@ inline void HierarchicalNSW::add_point( inline maxheap> HierarchicalNSW::search_base_layer( PID ep_id, PID cur_c, int layer ) { - HashBasedBooleanSet* vl = visited_list_pool_->get_free_vislist(); + VisitedSet* vl = visited_list_pool_->get_free_vislist(); maxheap> top_candidates; minheap> candidate_set; @@ -1225,7 +1225,7 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( [[maybe_unused]] const float* query, BoundedKNN& boundedKNN ) { - HashBasedBooleanSet* vl = visited_list_pool_->get_free_vislist(); + VisitedSet* vl = visited_list_pool_->get_free_vislist(); // Use our bounded priority queue instead of the maxheap. buffer::SearchBuffer candidate_set(ef); diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 3ad41dc..7d00eec 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -20,7 +20,7 @@ #include "rabitqlib/quantization/rabitq.hpp" #include "rabitqlib/utils/array.hpp" #include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/hashset.hpp" +#include "rabitqlib/utils/visited_set.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/rotator.hpp" @@ -98,20 +98,20 @@ class QuantizedGraph { PID, size_t, std::vector>&, - HashBasedBooleanSet&, + VisitedSet&, const std::vector& ) const; void update_qg(PID, const std::vector>&); - void update_results(buffer::SearchBuffer&, HashBasedBooleanSet&, const T*); + void update_results(buffer::SearchBuffer&, VisitedSet&, const T*); void scan_neighbors( const BatchQuery&, PID, T*, buffer::SearchBuffer&, - HashBasedBooleanSet&, + VisitedSet&, size_t ) const; @@ -352,7 +352,7 @@ void QuantizedGraph::scan_neighbors( PID data_id, T* est_dist, buffer::SearchBuffer& search_pool, - HashBasedBooleanSet& vis, + VisitedSet& vis, size_t cur_degree ) const { const auto* batch_data = get_batch_data(data_id); @@ -378,7 +378,7 @@ void QuantizedGraph::scan_neighbors( template inline void QuantizedGraph::update_results( - buffer::SearchBuffer& result_pool, HashBasedBooleanSet& vis, const T* query + buffer::SearchBuffer& result_pool, VisitedSet& vis, const T* query ) { if (result_pool.is_full()) { return; @@ -433,7 +433,7 @@ inline void QuantizedGraph::find_candidates( PID cur_id, size_t search_ef, std::vector>& results, - HashBasedBooleanSet& vis, + VisitedSet& vis, const std::vector& degrees ) const { const T* query = get_vector(cur_id); diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index a520553..2ae3e1f 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -11,7 +11,7 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/index/symqg/qg.hpp" -#include "rabitqlib/utils/hashset.hpp" +#include "rabitqlib/utils/visited_set.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -38,7 +38,7 @@ class QGBuilder { 300; // max number of recorded pruned candidates std::vector new_neighbors_; // new neighbors for current iteration std::vector pruned_neighbors_; // recorded pruned neighbors - std::vector visited_list_; // list of visited hash set + std::vector visited_list_; // per-thread visited sets std::vector degrees_; // record degree of qg void random_init(); void search_new_neighbors(bool refine); @@ -67,7 +67,7 @@ class QGBuilder { , pruned_neighbors_(qg_.num_vertices()) , visited_list_( num_threads_, - HashBasedBooleanSet(num_nodes_) + VisitedSet(num_nodes_) ) , degrees_(qg_.num_vertices(), degree_bound_) { omp_set_num_threads(static_cast(num_threads_)); @@ -236,7 +236,7 @@ inline void QGBuilder::search_new_neighbors(bool refine) { PID cur_id = i; auto tid = omp_get_thread_num(); CandidateList candidates; - HashBasedBooleanSet& vis = visited_list_[tid]; + VisitedSet& vis = visited_list_[tid]; candidates.reserve(2 * kMaxCandidatePoolSize); vis.clear(); qg_.find_candidates(cur_id, ef_build_, candidates, vis, degrees_); diff --git a/include/rabitqlib/utils/visited_pool.hpp b/include/rabitqlib/utils/visited_pool.hpp index d946b3f..8b2efff 100644 --- a/include/rabitqlib/utils/visited_pool.hpp +++ b/include/rabitqlib/utils/visited_pool.hpp @@ -2,49 +2,49 @@ #include #include -#include "rabitqlib/utils/hashset.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib { class VisitedListPool { - std::deque pool_; + std::deque pool_; std::mutex poolguard_; size_t numelements_; public: + // max_elements is the id space; each implementation sizes itself from it. VisitedListPool(size_t initpoolsize, size_t max_elements) { - // Must cover the whole id space: HashBasedBooleanSet indexes by id. numelements_ = max_elements; for (size_t i = 0; i < initpoolsize; i++) { - pool_.push_front(new HashBasedBooleanSet(numelements_)); + pool_.push_front(new VisitedSet(numelements_)); } } - HashBasedBooleanSet* get_free_vislist() { - HashBasedBooleanSet* rez; + VisitedSet* get_free_vislist() { + VisitedSet* rez; { std::unique_lock lock(poolguard_); if (pool_.size() > 0) { rez = pool_.front(); pool_.pop_front(); } else { - rez = new HashBasedBooleanSet(numelements_); + rez = new VisitedSet(numelements_); } } rez->clear(); return rez; } - void release_vis_list(HashBasedBooleanSet* vl) { + void release_vis_list(VisitedSet* vl) { std::unique_lock lock(poolguard_); pool_.push_front(vl); } ~VisitedListPool() { while (pool_.size() > 0) { - HashBasedBooleanSet* rez = pool_.front(); + VisitedSet* rez = pool_.front(); pool_.pop_front(); ::delete rez; } } }; -} // namespace rabitqlib \ No newline at end of file +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set.hpp b/include/rabitqlib/utils/visited_set.hpp new file mode 100644 index 0000000..4bb9eef --- /dev/null +++ b/include/rabitqlib/utils/visited_set.hpp @@ -0,0 +1,82 @@ +#pragma once + +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/visited_set_epoch.hpp" +#include "rabitqlib/utils/visited_set_hash.hpp" + +namespace rabitqlib { +/** + * @brief Selects the visited-set implementation used across the library. + * + * A visited set records which vertices a single graph search has already + * touched. Every implementation is a plain class -- there is no virtual base -- + * so the choice is made at compile time and the search path keeps direct calls. + * An implementation must provide: + * + * explicit Impl(size_t num_elements); // num_elements is the ID SPACE + * void initialize(size_t num_elements); // (re)size, leaving the set empty + * void clear(); // forget every mark + * bool get(PID data_id) const; // has data_id been marked? + * void set(PID data_id); // mark data_id + * + * num_elements is always the id space -- the number of vertices, such that + * every id passed to get()/set() is < num_elements. An implementation that + * wants a smaller table derives it internally (HashBasedVisitedSet does). + * + * Available implementations: + * EpochBasedVisitedSet (visited_set_epoch.hpp) -- default. O(1) clear, no + * allocation while searching, 2 bytes per id resident. + * HashBasedVisitedSet (visited_set_hash.hpp) -- sublinear memory, but + * allocates on collisions and clear() walks the table. + * + * To switch the whole library, change this alias -- it is the single place the + * implementation is named. + */ +namespace detail { +/** + * @brief Compile-time enforcement of the interface described above. + * + * Each assert compares the member's type exactly, so an implementation that + * takes a PID& or forgets the const on get() fails here, naming the member, + * rather than at some call site or -- worse -- silently, by binding to a + * conversion. Instantiating the struct is what fires the asserts. + * + * This is a check, not a base class: the implementations stay unrelated + * concrete types with non-virtual get()/set() that inline into the search + * loops. A virtual base would enforce the same thing through the type system, + * but at the price of an indirect call per visited vertex. + */ +template +struct check_visited_set { + static_assert( + std::is_constructible::value, + "visited set must be constructible from a size_t element count" + ); + static_assert( + std::is_same::value, + "visited set must declare: void initialize(size_t)" + ); + static_assert( + std::is_same::value, + "visited set must declare: void clear()" + ); + static_assert( + std::is_same::value, + "visited set must declare: bool get(PID) const" + ); + static_assert( + std::is_same::value, + "visited set must declare: void set(PID)" + ); + static constexpr bool value = true; +}; +} // namespace detail + +static_assert(detail::check_visited_set::value, ""); +static_assert(detail::check_visited_set::value, ""); + +using VisitedSet = HashBasedBooleanSet; +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/hashset.hpp b/include/rabitqlib/utils/visited_set_epoch.hpp similarity index 50% rename from include/rabitqlib/utils/hashset.hpp rename to include/rabitqlib/utils/visited_set_epoch.hpp index dd04364..ed3fb1e 100644 --- a/include/rabitqlib/utils/hashset.hpp +++ b/include/rabitqlib/utils/visited_set_epoch.hpp @@ -1,21 +1,3 @@ -// This code is modified based on NGT from Yahoo Japan -// https://github.com/yahoojapan/NGT -// -// Copyright (C) 2015 Yahoo Japan Corporation -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - #pragma once #include @@ -34,31 +16,32 @@ namespace rabitqlib { * clear() therefore just bumps the generation -- O(1), and only the rare * 16-bit wraparound pays for an actual refill. * - * This replaces a direct-mapped table with an std::unordered_set overflow. That + * This replaces HashBasedVisitedSet (visited_set_hash.hpp), whose direct-mapped * table was sized at max_elements/10 rounded down to a power of two, so a * single ef=100 search collided on nearly every visit and allocated an * unordered_set node each time: measured at ~1200 mallocs per query, 98% of all * allocation on the search path. Removing them was worth ~1.6x on HNSW search. + * The trade is memory: two bytes per id in the whole id space, resident for the + * lifetime of the set. See visited_set.hpp for how one of the two is selected. * - * The constructor argument is now the ID SPACE (number of vertices), not a - * bucket-count hint -- ids are used to index directly, so it must cover every - * id that will be passed to get()/set(). + * The constructor argument is the ID SPACE (number of vertices) -- ids are used + * to index directly, so it must cover every id passed to get()/set(). */ -class HashBasedBooleanSet { +class EpochBasedVisitedSet { private: std::vector> stamp_; uint16_t cur_ = 0; public: - HashBasedBooleanSet() = default; - ~HashBasedBooleanSet() = default; + EpochBasedVisitedSet() = default; + ~EpochBasedVisitedSet() = default; - HashBasedBooleanSet(const HashBasedBooleanSet&) = default; - HashBasedBooleanSet& operator=(const HashBasedBooleanSet&) = default; - HashBasedBooleanSet(HashBasedBooleanSet&&) noexcept = default; - HashBasedBooleanSet& operator=(HashBasedBooleanSet&&) noexcept = default; + EpochBasedVisitedSet(const EpochBasedVisitedSet&) = default; + EpochBasedVisitedSet& operator=(const EpochBasedVisitedSet&) = default; + EpochBasedVisitedSet(EpochBasedVisitedSet&&) noexcept = default; + EpochBasedVisitedSet& operator=(EpochBasedVisitedSet&&) noexcept = default; - explicit HashBasedBooleanSet(size_t num_elements) { initialize(num_elements); } + explicit EpochBasedVisitedSet(size_t num_elements) { initialize(num_elements); } void initialize(size_t num_elements) { stamp_.assign(num_elements, 0); @@ -75,7 +58,7 @@ class HashBasedBooleanSet { } } - // get if data_id is in the hashset + // get if data_id is in the visited set [[nodiscard]] bool get(PID data_id) const { return stamp_[data_id] == cur_; } void set(PID data_id) { stamp_[data_id] = cur_; } diff --git a/include/rabitqlib/utils/visited_set_hash.hpp b/include/rabitqlib/utils/visited_set_hash.hpp new file mode 100644 index 0000000..ed47434 --- /dev/null +++ b/include/rabitqlib/utils/visited_set_hash.hpp @@ -0,0 +1,121 @@ +// This code is modified based on NGT from Yahoo Japan +// https://github.com/yahoojapan/NGT +// +// Copyright (C) 2015 Yahoo Japan Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +#pragma once + +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/memory.hpp" + +namespace rabitqlib { +/** + * @brief Set of visited vertices, backed by a direct-mapped table plus an + * std::unordered_set for collisions. + * + * Memory is sublinear in the id space, so this is the choice when the id space + * is far larger than the number of vertices one search touches. The price is a + * heap allocation per colliding visit and an O(table) clear(). + * + * See EpochBasedVisitedSet in visited_set_epoch.hpp for the constant-time-clear + * alternative, and visited_set.hpp for how one of the two is selected. + * + * The constructor argument is the ID SPACE (number of vertices); the table is + * sized down from it internally, so callers pass the same number they would + * pass to any other visited set. + */ +class HashBasedVisitedSet { + private: + // The id space is divided by this before the table is sized, since only a + // small fraction of the vertices is visited by a single search. + static constexpr size_t kSizeDivisor = 10; + + size_t table_size_ = 0; + PID mask_ = 0; + std::vector> table_; + std::unordered_set stl_hash_; + + [[nodiscard]] auto hash1(const PID value) const { return value & mask_; } + + void initialize_table(const size_t table_size) { + table_size_ = table_size; + mask_ = static_cast(table_size_ - 1); + const PID check_val = hash1(static_cast(table_size)); + if (check_val != 0) { + std::cerr << "[WARN] table size is not 2^N : " << table_size << '\n'; + } + + table_ = std::vector>(table_size); + std::fill(table_.begin(), table_.end(), kPidMax); + stl_hash_.clear(); + } + + public: + HashBasedVisitedSet() = default; + ~HashBasedVisitedSet() = default; + + HashBasedVisitedSet(const HashBasedVisitedSet&) = default; + HashBasedVisitedSet& operator=(const HashBasedVisitedSet&) = default; + HashBasedVisitedSet(HashBasedVisitedSet&&) noexcept = default; + HashBasedVisitedSet& operator=(HashBasedVisitedSet&&) noexcept = default; + + explicit HashBasedVisitedSet(size_t num_elements) { initialize(num_elements); } + + void initialize(size_t num_elements) { + size_t size = num_elements / kSizeDivisor; + size_t bit_size = 0; + size_t bit = size; + while (bit != 0) { + bit_size++; + bit >>= 1; + } + size_t bucket_size = static_cast(0x1) << ((bit_size + 4) / 2 + 3); + initialize_table(bucket_size); + } + + void clear() { + std::fill(table_.begin(), table_.end(), kPidMax); + stl_hash_.clear(); + } + + // get if data_id is in the visited set + [[nodiscard]] bool get(PID data_id) const { + PID val = this->table_[hash1(data_id)]; + if (val == data_id) { + return true; + } + return (val != kPidMax && stl_hash_.find(data_id) != stl_hash_.end()); + } + + void set(PID data_id) { + PID& val = table_[hash1(data_id)]; + if (val == data_id) { + return; + } + if (val == kPidMax) { + val = data_id; + } else { + stl_hash_.emplace(data_id); + } + } +}; +} // namespace rabitqlib From 1f048ca4b930494415e1e6dda05779b04be1c6a1 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Wed, 19 Aug 2026 12:25:23 +0800 Subject: [PATCH 4/5] refactor: update VisitedSet initialization to include size hint for better memory management --- include/rabitqlib/index/symqg/qg_builder.hpp | 2 +- include/rabitqlib/utils/visited_pool.hpp | 5 +- include/rabitqlib/utils/visited_set.hpp | 52 +++---------------- include/rabitqlib/utils/visited_set_epoch.hpp | 28 ++-------- include/rabitqlib/utils/visited_set_hash.hpp | 33 ++++-------- 5 files changed, 25 insertions(+), 95 deletions(-) diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 2ae3e1f..26f467e 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -67,7 +67,7 @@ class QGBuilder { , pruned_neighbors_(qg_.num_vertices()) , visited_list_( num_threads_, - VisitedSet(num_nodes_) + VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) ) , degrees_(qg_.num_vertices(), degree_bound_) { omp_set_num_threads(static_cast(num_threads_)); diff --git a/include/rabitqlib/utils/visited_pool.hpp b/include/rabitqlib/utils/visited_pool.hpp index 8b2efff..2926d3d 100644 --- a/include/rabitqlib/utils/visited_pool.hpp +++ b/include/rabitqlib/utils/visited_pool.hpp @@ -11,11 +11,10 @@ class VisitedListPool { size_t numelements_; public: - // max_elements is the id space; each implementation sizes itself from it. VisitedListPool(size_t initpoolsize, size_t max_elements) { numelements_ = max_elements; for (size_t i = 0; i < initpoolsize; i++) { - pool_.push_front(new VisitedSet(numelements_)); + pool_.push_front(new VisitedSet(numelements_, numelements_ / 10)); } } @@ -27,7 +26,7 @@ class VisitedListPool { rez = pool_.front(); pool_.pop_front(); } else { - rez = new VisitedSet(numelements_); + rez = new VisitedSet(numelements_, numelements_ / 10); } } rez->clear(); diff --git a/include/rabitqlib/utils/visited_set.hpp b/include/rabitqlib/utils/visited_set.hpp index 4bb9eef..8e72520 100644 --- a/include/rabitqlib/utils/visited_set.hpp +++ b/include/rabitqlib/utils/visited_set.hpp @@ -8,56 +8,20 @@ #include "rabitqlib/utils/visited_set_hash.hpp" namespace rabitqlib { -/** - * @brief Selects the visited-set implementation used across the library. - * - * A visited set records which vertices a single graph search has already - * touched. Every implementation is a plain class -- there is no virtual base -- - * so the choice is made at compile time and the search path keeps direct calls. - * An implementation must provide: - * - * explicit Impl(size_t num_elements); // num_elements is the ID SPACE - * void initialize(size_t num_elements); // (re)size, leaving the set empty - * void clear(); // forget every mark - * bool get(PID data_id) const; // has data_id been marked? - * void set(PID data_id); // mark data_id - * - * num_elements is always the id space -- the number of vertices, such that - * every id passed to get()/set() is < num_elements. An implementation that - * wants a smaller table derives it internally (HashBasedVisitedSet does). - * - * Available implementations: - * EpochBasedVisitedSet (visited_set_epoch.hpp) -- default. O(1) clear, no - * allocation while searching, 2 bytes per id resident. - * HashBasedVisitedSet (visited_set_hash.hpp) -- sublinear memory, but - * allocates on collisions and clear() walks the table. - * - * To switch the whole library, change this alias -- it is the single place the - * implementation is named. - */ namespace detail { -/** - * @brief Compile-time enforcement of the interface described above. - * - * Each assert compares the member's type exactly, so an implementation that - * takes a PID& or forgets the const on get() fails here, naming the member, - * rather than at some call site or -- worse -- silently, by binding to a - * conversion. Instantiating the struct is what fires the asserts. - * - * This is a check, not a base class: the implementations stay unrelated - * concrete types with non-virtual get()/set() that inline into the search - * loops. A virtual base would enforce the same thing through the type system, - * but at the price of an indirect call per visited vertex. - */ template struct check_visited_set { static_assert( std::is_constructible::value, - "visited set must be constructible from a size_t element count" + "visited set must be constructible from a size_t id space" ); static_assert( - std::is_same::value, - "visited set must declare: void initialize(size_t)" + std::is_constructible::value, + "visited set must be constructible from (num_elements, size_hint)" + ); + static_assert( + std::is_same::value, + "visited set must declare: void initialize(size_t num_elements, size_t size_hint)" ); static_assert( std::is_same::value, @@ -78,5 +42,5 @@ struct check_visited_set { static_assert(detail::check_visited_set::value, ""); static_assert(detail::check_visited_set::value, ""); -using VisitedSet = HashBasedBooleanSet; +using VisitedSet = HashBasedVisitedSet; } // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set_epoch.hpp b/include/rabitqlib/utils/visited_set_epoch.hpp index ed3fb1e..04152a9 100644 --- a/include/rabitqlib/utils/visited_set_epoch.hpp +++ b/include/rabitqlib/utils/visited_set_epoch.hpp @@ -8,25 +8,6 @@ #include "rabitqlib/utils/memory.hpp" namespace rabitqlib { -/** - * @brief Set of visited vertices, backed by a generation-stamped array. - * - * One byte-pair per vertex holds the generation in which that vertex was last - * marked; a vertex is "visited" iff its stamp equals the current generation. - * clear() therefore just bumps the generation -- O(1), and only the rare - * 16-bit wraparound pays for an actual refill. - * - * This replaces HashBasedVisitedSet (visited_set_hash.hpp), whose direct-mapped - * table was sized at max_elements/10 rounded down to a power of two, so a - * single ef=100 search collided on nearly every visit and allocated an - * unordered_set node each time: measured at ~1200 mallocs per query, 98% of all - * allocation on the search path. Removing them was worth ~1.6x on HNSW search. - * The trade is memory: two bytes per id in the whole id space, resident for the - * lifetime of the set. See visited_set.hpp for how one of the two is selected. - * - * The constructor argument is the ID SPACE (number of vertices) -- ids are used - * to index directly, so it must cover every id passed to get()/set(). - */ class EpochBasedVisitedSet { private: std::vector> stamp_; @@ -41,13 +22,13 @@ class EpochBasedVisitedSet { EpochBasedVisitedSet(EpochBasedVisitedSet&&) noexcept = default; EpochBasedVisitedSet& operator=(EpochBasedVisitedSet&&) noexcept = default; - explicit EpochBasedVisitedSet(size_t num_elements) { initialize(num_elements); } + explicit EpochBasedVisitedSet(size_t num_elements, size_t = 0) { + initialize(num_elements); + } - void initialize(size_t num_elements) { + void initialize(size_t num_elements, size_t = 0) { stamp_.assign(num_elements, 0); cur_ = 0; - // Leaves the set usable without an explicit clear(): every stamp is 0 - // while the live generation is 1, so nothing reads as visited. clear(); } @@ -58,7 +39,6 @@ class EpochBasedVisitedSet { } } - // get if data_id is in the visited set [[nodiscard]] bool get(PID data_id) const { return stamp_[data_id] == cur_; } void set(PID data_id) { stamp_[data_id] = cur_; } diff --git a/include/rabitqlib/utils/visited_set_hash.hpp b/include/rabitqlib/utils/visited_set_hash.hpp index ed47434..f9f3722 100644 --- a/include/rabitqlib/utils/visited_set_hash.hpp +++ b/include/rabitqlib/utils/visited_set_hash.hpp @@ -18,9 +18,11 @@ #pragma once +#include #include #include #include +#include #include #include @@ -28,27 +30,11 @@ #include "rabitqlib/utils/memory.hpp" namespace rabitqlib { -/** - * @brief Set of visited vertices, backed by a direct-mapped table plus an - * std::unordered_set for collisions. - * - * Memory is sublinear in the id space, so this is the choice when the id space - * is far larger than the number of vertices one search touches. The price is a - * heap allocation per colliding visit and an O(table) clear(). - * - * See EpochBasedVisitedSet in visited_set_epoch.hpp for the constant-time-clear - * alternative, and visited_set.hpp for how one of the two is selected. - * - * The constructor argument is the ID SPACE (number of vertices); the table is - * sized down from it internally, so callers pass the same number they would - * pass to any other visited set. - */ class HashBasedVisitedSet { - private: - // The id space is divided by this before the table is sized, since only a - // small fraction of the vertices is visited by a single search. - static constexpr size_t kSizeDivisor = 10; + public: + static constexpr size_t kNoSizeHint = std::numeric_limits::max(); + private: size_t table_size_ = 0; PID mask_ = 0; std::vector> table_; @@ -78,10 +64,12 @@ class HashBasedVisitedSet { HashBasedVisitedSet(HashBasedVisitedSet&&) noexcept = default; HashBasedVisitedSet& operator=(HashBasedVisitedSet&&) noexcept = default; - explicit HashBasedVisitedSet(size_t num_elements) { initialize(num_elements); } + explicit HashBasedVisitedSet(size_t num_elements, size_t size_hint = kNoSizeHint) { + initialize(num_elements, size_hint); + } - void initialize(size_t num_elements) { - size_t size = num_elements / kSizeDivisor; + void initialize(size_t num_elements, size_t size_hint = kNoSizeHint) { + size_t size = std::min(num_elements, size_hint); size_t bit_size = 0; size_t bit = size; while (bit != 0) { @@ -97,7 +85,6 @@ class HashBasedVisitedSet { stl_hash_.clear(); } - // get if data_id is in the visited set [[nodiscard]] bool get(PID data_id) const { PID val = this->table_[hash1(data_id)]; if (val == data_id) { From 192a5223f7b62888dc36333fc761a528ce78562b Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 22 Aug 2026 13:24:10 +0800 Subject: [PATCH 5/5] fix: preserve visited set compatibility --- include/rabitqlib/index/hnsw/hnsw.hpp | 2 +- include/rabitqlib/index/symqg/qg.hpp | 19 +-- include/rabitqlib/index/symqg/qg_builder.hpp | 8 +- include/rabitqlib/utils/hashset.hpp | 8 + include/rabitqlib/utils/visited_set.hpp | 8 +- .../unit/rabitqlib/utils/visited_set_test.cpp | 138 ++++++++++++++++++ 6 files changed, 160 insertions(+), 23 deletions(-) create mode 100644 include/rabitqlib/utils/hashset.hpp create mode 100644 tests/unit/rabitqlib/utils/visited_set_test.cpp diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index b8d4612..7e290cc 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -1138,7 +1138,7 @@ inline maxheap> HierarchicalNSW::search_knn_direct( if (metric_type_ == METRIC_L2) { for (size_t i = 0; i < num_cluster_; i++) { - q_to_centroids[i] = std::sqrt(euclidean_sqr( + q_to_centroids[i] = std::sqrt(raw_dist_func_( rotated_query, reinterpret_cast(centroids_memory_) + (i * padded_dim_), padded_dim_ diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 7d00eec..9f91910 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -20,12 +20,12 @@ #include "rabitqlib/quantization/rabitq.hpp" #include "rabitqlib/utils/array.hpp" #include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/visited_set.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/visited_pool.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { @@ -94,25 +94,16 @@ class QuantizedGraph { ); } - void find_candidates( - PID, - size_t, - std::vector>&, - VisitedSet&, - const std::vector& - ) const; + void + find_candidates(PID, size_t, std::vector>&, VisitedSet&, const std::vector&) + const; void update_qg(PID, const std::vector>&); void update_results(buffer::SearchBuffer&, VisitedSet&, const T*); void scan_neighbors( - const BatchQuery&, - PID, - T*, - buffer::SearchBuffer&, - VisitedSet&, - size_t + const BatchQuery&, PID, T*, buffer::SearchBuffer&, VisitedSet&, size_t ) const; public: diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 26f467e..27ffef3 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -11,9 +11,9 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/index/symqg/qg.hpp" -#include "rabitqlib/utils/visited_set.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar @@ -37,9 +37,9 @@ class QGBuilder { static constexpr size_t kMaxPrunedSize = 300; // max number of recorded pruned candidates std::vector new_neighbors_; // new neighbors for current iteration - std::vector pruned_neighbors_; // recorded pruned neighbors - std::vector visited_list_; // per-thread visited sets - std::vector degrees_; // record degree of qg + std::vector pruned_neighbors_; // recorded pruned neighbors + std::vector visited_list_; // per-thread visited sets + std::vector degrees_; // record degree of qg void random_init(); void search_new_neighbors(bool refine); void heuristic_prune(PID, CandidateList&, CandidateList&, bool); diff --git a/include/rabitqlib/utils/hashset.hpp b/include/rabitqlib/utils/hashset.hpp new file mode 100644 index 0000000..2633417 --- /dev/null +++ b/include/rabitqlib/utils/hashset.hpp @@ -0,0 +1,8 @@ +// This compatibility header preserves the original visited-set API. +#pragma once + +#include "rabitqlib/utils/visited_set_hash.hpp" + +namespace rabitqlib { +using HashBasedBooleanSet = HashBasedVisitedSet; +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set.hpp b/include/rabitqlib/utils/visited_set.hpp index 8e72520..6ca6c14 100644 --- a/include/rabitqlib/utils/visited_set.hpp +++ b/include/rabitqlib/utils/visited_set.hpp @@ -10,7 +10,7 @@ namespace rabitqlib { namespace detail { template -struct check_visited_set { +struct CheckVisitedSet { static_assert( std::is_constructible::value, "visited set must be constructible from a size_t id space" @@ -35,12 +35,12 @@ struct check_visited_set { std::is_same::value, "visited set must declare: void set(PID)" ); - static constexpr bool value = true; + static constexpr bool kValue = true; }; } // namespace detail -static_assert(detail::check_visited_set::value, ""); -static_assert(detail::check_visited_set::value, ""); +static_assert(detail::CheckVisitedSet::kValue, ""); +static_assert(detail::CheckVisitedSet::kValue, ""); using VisitedSet = HashBasedVisitedSet; } // namespace rabitqlib diff --git a/tests/unit/rabitqlib/utils/visited_set_test.cpp b/tests/unit/rabitqlib/utils/visited_set_test.cpp new file mode 100644 index 0000000..61573b0 --- /dev/null +++ b/tests/unit/rabitqlib/utils/visited_set_test.cpp @@ -0,0 +1,138 @@ +#include "rabitqlib/utils/visited_set.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/utils/hashset.hpp" + +namespace { + +using rabitqlib::EpochBasedVisitedSet; +using rabitqlib::HashBasedBooleanSet; +using rabitqlib::HashBasedVisitedSet; +using rabitqlib::PID; +using rabitqlib::VisitedSet; + +static_assert( + std::is_same_v, + "HashBasedVisitedSet must remain the default visited-set implementation" +); +static_assert( + std::is_same_v, + "HashBasedBooleanSet must remain available as a compatibility alias" +); + +template +void expect_basic_visited_set_behavior() { + constexpr size_t kNumElements = 1024; + Set visited(kNumElements, 8); + + EXPECT_FALSE(visited.get(0)); + EXPECT_FALSE(visited.get(kNumElements - 1)); + + visited.set(0); + visited.set(17); + visited.set(kNumElements - 1); + + EXPECT_TRUE(visited.get(0)); + EXPECT_TRUE(visited.get(17)); + EXPECT_TRUE(visited.get(kNumElements - 1)); + + visited.set(17); + EXPECT_TRUE(visited.get(17)); + + visited.clear(); + EXPECT_FALSE(visited.get(0)); + EXPECT_FALSE(visited.get(17)); + EXPECT_FALSE(visited.get(kNumElements - 1)); + + visited.set(23); + EXPECT_TRUE(visited.get(23)); +} + +TEST(VisitedSet, HashBackendSatisfiesContract) { + expect_basic_visited_set_behavior(); +} + +TEST(VisitedSet, EpochBackendSatisfiesContract) { + expect_basic_visited_set_behavior(); +} + +TEST(VisitedSet, BackendsRemainBehaviorallyEquivalent) { + constexpr size_t kNumElements = 4096; + constexpr size_t kNumOperations = 20000; + + HashBasedVisitedSet hash_visited(kNumElements, 16); + EpochBasedVisitedSet epoch_visited(kNumElements); + std::vector expected(kNumElements, false); + std::mt19937 generator(42); + std::uniform_int_distribution id_distribution(0, kNumElements - 1); + + for (size_t operation = 0; operation < kNumOperations; ++operation) { + if (operation % 257 == 0) { + hash_visited.clear(); + epoch_visited.clear(); + std::fill(expected.begin(), expected.end(), false); + } else { + const PID id = id_distribution(generator); + if (operation % 3 == 0) { + EXPECT_EQ(hash_visited.get(id), expected[id]); + EXPECT_EQ(epoch_visited.get(id), expected[id]); + } else { + hash_visited.set(id); + epoch_visited.set(id); + expected[id] = true; + } + } + } + + for (PID id = 0; id < kNumElements; ++id) { + EXPECT_EQ(hash_visited.get(id), expected[id]) << "id=" << id; + EXPECT_EQ(epoch_visited.get(id), expected[id]) << "id=" << id; + } +} + +TEST(HashBasedVisitedSet, HandlesCollidingIds) { + HashBasedVisitedSet visited(1024, 1); + const std::vector colliding_ids{0, 32, 64, 96, 128, 1023}; + + for (PID id : colliding_ids) { + visited.set(id); + } + for (PID id : colliding_ids) { + EXPECT_TRUE(visited.get(id)) << "id=" << id; + } + + visited.clear(); + for (PID id : colliding_ids) { + EXPECT_FALSE(visited.get(id)) << "id=" << id; + } +} + +TEST(EpochBasedVisitedSet, ClearsStaleValuesWhenEpochWraps) { + EpochBasedVisitedSet visited(4); + visited.set(0); + + // Construction starts at epoch 1. Advance to the last uint16_t epoch. + for (size_t epoch = 0; epoch < UINT16_MAX - 1; ++epoch) { + visited.clear(); + } + visited.set(1); + EXPECT_TRUE(visited.get(1)); + + // The next clear wraps the counter and must reset all stored stamps. + visited.clear(); + EXPECT_FALSE(visited.get(0)); + EXPECT_FALSE(visited.get(1)); + + visited.set(3); + EXPECT_TRUE(visited.get(3)); +} + +} // namespace