diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index f3e363ab..69ba7050 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -98,3 +98,81 @@ 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, + size_t num_threads = 1 +); + +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**, **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 C++, initialize the index through its parameterized constructor or `load` + before calling `resize`. + +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. + +`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 +`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. + +### 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. diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 5c778f46..4cb1d218 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -61,11 +61,50 @@ 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_; } + [[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*); void construct(size_t, const float*, size_t, const float*, PID*, size_t, bool); + + /** + * @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. + * 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 + * @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 + ); + + /** + * @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. + * Requires initialization through the parameterized constructor or `load`. + * + * @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 ); @@ -124,6 +163,7 @@ class HierarchicalNSW { }; private: + friend struct HnswPruningTestAccess; friend maxheap> detail::search_knn_neon( HierarchicalNSW&, const float*, size_t ); @@ -360,6 +400,14 @@ class HierarchicalNSW { 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, 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); }; inline HierarchicalNSW::HierarchicalNSW( @@ -1520,3 +1568,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..025c0297 --- /dev/null +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -0,0 +1,627 @@ +// 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 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 +#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 +) { + std::unique_ptr 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.get()); + vl.release(); + 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(cur_c, selected_neighbor); + 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; +} + +// 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, std::unique_ptr& link_list +) { + detail::quant_query().clear(); + + std::unique_lock lock_el(link_list_locks_[cur_c]); + + std::unique_lock templock(global_); + int maxlevelcopy = maxlevel_; + if (curlevel <= maxlevelcopy) { + templock.unlock(); + } + PID curr_obj = enterpoint_node_; + const PID previous_entry_point = curr_obj; + + linkLists_[cur_c] = link_list.release(); + element_levels_[cur_c] = curlevel; + + 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); + } + } else { + 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"); + } + 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, size_t num_threads +) { + 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) { + 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"); + } + 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); + 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_); + 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) { + const float dist = raw_dist_func_( + rotated.data(), centroids + (c * padded_dim_), padded_dim_ + ); + if (dist < best) { + best = dist; + assigned[i] = static_cast(c); + } + } + }); + } + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + // 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); + } + + // 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"); + } + 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) { + 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(labels[i], levels[i], link_lists[i]); + }); + + return labels; +} + +} // namespace rabitqlib::hnsw 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/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 559237dd..5a3e1cf2 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,64 @@ class HnswIndex { built_ = true; } + py::array_t add( + 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_) { + 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, + num_threads + ); + 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 +294,20 @@ 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("num_threads") = 1, + 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) 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/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/python/test_hnsw.py b/tests/python/test_hnsw.py index f6959699..82ee4194 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -179,3 +179,136 @@ 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) diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 459bfad4..732f3aed 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))), @@ -77,6 +138,484 @@ 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, 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, 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); + 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) { + 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]); + } +} + +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; + 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); + 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); + } + 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); + + 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); 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