diff --git a/AGENTS.md b/AGENTS.md index 15c7dd9c..17a83263 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -174,6 +174,12 @@ Use `python_bindings/bindings_common.hpp` for shared conversions. Validate array index state, and parameter ranges before entering the core. `py::array::forcecast` permits copies; do not use it where callers expect in-place mutation or pointer identity. +Python index wrappers use `IndexAccess` for concurrent reads and fail-fast exclusive updates. +Acquire and destroy its guards with the GIL held; release the GIL only inside the guard's +lifetime, after extracting native pointers, and retain Python array owners until it is reacquired. +Keep search parameters local to each call. C++ callers remain responsible for synchronizing +index mutations. + ### Rotation and padded dimensions Indexes quantize and search in the rotated, padded domain. FHT/Kac, RaBitQKMeans, QGKMeans, diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 62bae2c8..a07bef15 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -69,6 +69,11 @@ std::vector>> HierarchicalNSW::search(const fl - **efSearch**: The size of the candidate set for searching HNSW base layer. - **thread_num**: Number of threads to use. Each query is processed by one thread. +The search window belongs to each call, so concurrent C++ searches can use +different `efSearch` values and separate outputs while the index is unchanged. +Python search releases the GIL. Conflicting updates raise `RuntimeError` +immediately; see the [Python concurrency contract](../quick_start.md#threading-and-file-paths). + We first pre-process the query: 1. Rotate the raw query vector. @@ -146,8 +151,9 @@ new_ids = index.add(vectors) # route to the nearest centro index.resize(index.max_elements + 100_000) ``` -It is not safe to call `add` or `resize` while another thread searches the same -index. +In C++, callers must prevent `add` or `resize` from overlapping with searches +on the same index. Python detects these conflicts and raises `RuntimeError` +immediately. `add` prepares the result IDs, quantized vectors, and link-list allocations before changing the graph or point count. A failure during preparation leaves those @@ -214,8 +220,9 @@ still walks through it and `add` may link new points to it, which keeps the grap connected, as in hnswlib. Search never returns it, so a query can get fewer than `k` results once points are removed: in Python the missing slots hold `kPidMax` (`2**32 - 1`) with an infinite distance. Removed points still count in -`num_points()` and cannot be restored. It is not safe to call `remove` while -another thread searches or adds to the same index. +`num_points()` and cannot be restored. In C++, callers must prevent `remove` +from overlapping with searches or updates on the same index. Python detects +these conflicts and raises `RuntimeError` immediately. #### Choosing `ef` after removals diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index bbbef31b..b2547254 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -244,6 +244,11 @@ need separate output buffers, and must not overlap index updates. Python's existing `index.search(queries, k, nprobe, num_threads=4)` uses this batch path with the same output and precision rules. +Python search releases the GIL. Concurrent searches may use independent +`nprobe`, `high_accuracy`, and `num_threads` values. Conflicting updates raise +`RuntimeError` immediately; see the +[Python concurrency contract](../quick_start.md#threading-and-file-paths). + ## Updating an Index Starting with 0.5.1, C++ `IVF` and Python `IvfIndex` support `add()` and `remove()` @@ -292,8 +297,9 @@ IDs and returns the number newly removed. Duplicate IDs count once. Python The stored point count (`max_elements()` in C++, `index.max_elements` in Python) increases after `add()` and includes removed points. Rebuilding with C++ -`construct` or Python `build` expects that many rows. Calls to `add()` or `remove()` -must not overlap with searches or other updates on the same index. +`construct` or Python `build` expects that many rows. In C++, callers must prevent +`add()` or `remove()` from overlapping with searches or other updates on the same +index. Python detects these conflicts and raises `RuntimeError` immediately. ### Cost of `add` diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index fc7acd54..137e7750 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -152,5 +152,24 @@ including empty batches. Both APIs default to one worker; `num_threads=0` selects the available logical CPU count, capped by the number of queries. Queries use the original dimension, including when the internal rotation pads it. +Python search releases the GIL and uses the `ef` passed to each call without +changing a shared default. Concurrent searches may use different windows. +A concurrent `build()` on the same index raises `RuntimeError` immediately; +see the [Python concurrency contract](../quick_start.md#threading-and-file-paths). + +The current source checkout also provides C++ +`search_batch_with_ef(queries, num_queries, k, ids, distances, ef, num_threads=1)` +for concurrent calls with independent windows: + +```cpp +qg.search_batch_with_ef( + queries.data(), num_queries, 10, batch_ids.data(), batch_distances.data(), 100, 4 +); +``` + +This method does not read or change the default set by `set_ef()`. The existing +`search()`, `search_batch()`, and `search_with_scratch()` retain that default. +All concurrent searches require separate outputs and an unchanged index. + See `sample/cpp/symqg_indexing.cpp`, `sample/cpp/symqg_querying.cpp`, and their Python counterparts for complete examples. diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index 7cd37f80..b086dbfa 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -153,6 +153,31 @@ counts, or [QGKMeans](clustering.md#qgkmeans) for graph assignment. ### Threading and file paths +The current source checkout releases the Python GIL during native index search, +construction, updates, and file I/O. The following contract applies to +`IvfIndex`, `HnswIndex`, and `SymqgIndex`: + +| Operation on the same index | Concurrent access | +| --- | --- | +| `search`, `search_batch`, `save`, and property reads | May run together | +| `build`, `add`, `remove`, and `resize`, where available | Require exclusive access | + +A conflicting call immediately raises +`RuntimeError("Index is busy: conflicting operation in progress")`. +It does not wait or queue an update. Access protection includes Python input +conversion and is released when the operation returns or raises an exception. +Different index objects can operate independently. Concurrent saves must use +different output paths, including any sidecar files. + +Search parameters such as `ef`, `nprobe`, precision, and `num_threads` belong to +each call. When using a Python thread pool, consider `num_threads=1` to avoid +starting multiple native worker pools. Keep input arrays and any shared backing +storage unchanged for the entire call; do not resize or modify them from another +thread. Returned result arrays own their storage. + +These access checks belong to the Python wrappers. C++ callers must synchronize +index updates themselves and use separate search output buffers. + `num_threads=0` selects the detected available logical CPU count. Positive values set an upper limit, capped at that count; small workloads may use fewer workers. Python index methods default to one thread; clustering defaults to `0`. diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index d25781be..0c9f198a 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -182,19 +182,19 @@ class HierarchicalNSW { private: friend struct HnswPruningTestAccess; friend maxheap> detail::search_knn_neon( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ); friend maxheap> detail::search_knn_generic( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ); friend maxheap> detail::search_knn_avx2( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ); friend maxheap> detail::search_knn_avx512_core( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ); friend maxheap> detail::search_knn_avx512_popcnt( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ); static constexpr PID kMaxLabelOperationLock = 65536; @@ -400,10 +400,10 @@ class HierarchicalNSW { get_full_est_direct(std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord&) const; - maxheap> search_knn(const float*, size_t); + maxheap> search_knn(const float*, size_t, size_t); template - maxheap> search_knn_direct(const float*, size_t); + maxheap> search_knn_direct(const float*, size_t, size_t); template void searchBaseLayerST_AdaptiveRerankOptDirect( @@ -1405,7 +1405,6 @@ inline void HierarchicalNSW::get_full_est_direct( inline std::vector>> HierarchicalNSW::search( const float* queries, size_t query_num, size_t TOPK, size_t efSearch, size_t thread_num ) { - set_ef(efSearch); std::vector>> results(query_num); rabitqlib::ivf::parallel_for( 0, @@ -1414,7 +1413,8 @@ inline std::vector>> HierarchicalNSW::search( [&](size_t idx, size_t /*threadId*/) { std::vector rotated_query(padded_dim_); this->rotator_->rotate(queries + (idx * dim_), rotated_query.data()); - maxheap> knn = search_knn(rotated_query.data(), TOPK); + maxheap> knn = + search_knn(rotated_query.data(), TOPK, std::max(efSearch, TOPK)); while (knn.size()) { results[idx].emplace_back(knn.top()); knn.pop(); @@ -1463,14 +1463,14 @@ inline size_t HierarchicalNSW::remove(const PID* labels, size_t n) { } inline maxheap> HierarchicalNSW::search_knn( - const float* rotated_query, size_t TOPK + const float* rotated_query, size_t TOPK, size_t ef ) { - return detail::search_knn(*this, rotated_query, TOPK); + return detail::search_knn(*this, rotated_query, TOPK, ef); } template inline maxheap> HierarchicalNSW::search_knn_direct( - const float* rotated_query, size_t TOPK + const float* rotated_query, size_t TOPK, size_t ef ) { maxheap> result; if (cur_element_count_ == 0) { @@ -1545,7 +1545,7 @@ inline maxheap> HierarchicalNSW::search_knn_direct( BoundedKNN boundedKnn(TOPK); searchBaseLayerST_AdaptiveRerankOptDirect( curr_obj, - std::max(ef_, TOPK), + std::max(ef, TOPK), TOPK, query_wrapper, q_to_centroids, diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 21e2f370..cc7460c1 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -124,7 +124,13 @@ class QuantizedGraph { return get_row_data(data_id); } - void validate_search(const float*, uint32_t, const uint32_t*, const float*) const; + void validate_search(const float*, uint32_t, const uint32_t*, const float*, size_t) + const; + + void search_impl(const float*, uint32_t, uint32_t*, float*, size_t); + + void + search_with_scratch_impl(const float*, uint32_t, uint32_t*, float*, float*, float*, float*, BatchQuery&, buffer::SearchBuffer&, buffer::SearchBuffer&, VisitedSet&, PID, const float*); const float* prepare_build_query( PID, @@ -256,6 +262,18 @@ class QuantizedGraph { size_t num_threads = 1 ); + // Like search_batch, but ef belongs to this call and never reads or changes + // the default set by set_ef(). Concurrent calls may use different windows. + void search_batch_with_ef( + const float* queries, + size_t num_queries, + uint32_t knn, + uint32_t* results, + float* dists, + size_t ef, + size_t num_threads = 1 + ); + /** * Search using caller-owned scratch buffers. * diff --git a/include/rabitqlib/simd/hnsw_dispatch.hpp b/include/rabitqlib/simd/hnsw_dispatch.hpp index f6674955..0cd23752 100644 --- a/include/rabitqlib/simd/hnsw_dispatch.hpp +++ b/include/rabitqlib/simd/hnsw_dispatch.hpp @@ -10,26 +10,26 @@ namespace rabitqlib::hnsw { class HierarchicalNSW; namespace detail { std::priority_queue> search_knn( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ef = 0 ); std::priority_queue> search_knn_neon( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ef = 0 ); std::priority_queue> search_knn_generic( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ef = 0 ); std::priority_queue> search_knn_avx2( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ef = 0 ); std::priority_queue> search_knn_avx512_core( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ef = 0 ); std::priority_queue> search_knn_avx512_popcnt( - HierarchicalNSW&, const float*, size_t + HierarchicalNSW&, const float*, size_t, size_t ef = 0 ); } // namespace detail diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 50312123..9a1f9bca 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -16,6 +16,48 @@ namespace py = pybind11; namespace rabitqlib::python_bindings { +// All guard acquisition/destruction happens with the GIL held. Native work may +// release the GIL inside the guard's lifetime, but must reacquire it first on +// exit (including exceptions). This makes conflicts fail immediately without +// blocking the GIL or introducing a lock-order dependency. +class IndexAccess { + private: + size_t readers_ = 0; + bool writing_ = false; + + public: + class Guard { + private: + IndexAccess& access_; + bool writing_; + + public: + Guard(IndexAccess& access, bool writing) : access_(access), writing_(writing) { + if (access.writing_ || (writing && access.readers_ != 0)) { + throw std::runtime_error("Index is busy: conflicting operation in progress" + ); + } + if (writing) { + access.writing_ = true; + } else { + ++access.readers_; + } + } + Guard(const Guard&) = delete; + Guard& operator=(const Guard&) = delete; + ~Guard() { + if (writing_) { + access_.writing_ = false; + } else { + --access_.readers_; + } + } + }; + + [[nodiscard]] Guard read() { return Guard(*this, false); } + [[nodiscard]] Guard write() { return Guard(*this, true); } +}; + inline rabitqlib::MetricType metric_from_string(const std::string& metric) { if (metric == "l2") { return rabitqlib::METRIC_L2; @@ -85,4 +127,4 @@ ensure_1d_integer_array(py::handle value, const char* name) { return ensure_1d_array(value, name); } -} // namespace rabitqlib::python_bindings \ No newline at end of file +} // namespace rabitqlib::python_bindings diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 703f17a9..1e506d25 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -51,6 +52,7 @@ class HnswIndex { size_t num_threads = 1, bool fast_quantization = false ) { + const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); auto centroids_array = ensure_2d_array(centroids, "centroids"); auto cluster_ids_array = @@ -94,15 +96,21 @@ class HnswIndex { cluster_ids_vec.size() * sizeof(rabitqlib::PID) ); - index_->construct( - num_clusters, - centroids_array.data(), - static_cast(data_array.shape(0)), - data_array.data(), - cluster_ids_vec.data(), - num_threads, - fast_quantization - ); + const auto* data_ptr = data_array.data(); + const auto* centroids_ptr = centroids_array.data(); + const auto rows = static_cast(data_array.shape(0)); + { + py::gil_scoped_release release; + index_->construct( + num_clusters, + centroids_ptr, + rows, + data_ptr, + cluster_ids_vec.data(), + num_threads, + fast_quantization + ); + } num_clusters_ = num_clusters; built_ = true; } @@ -113,6 +121,7 @@ class HnswIndex { size_t num_threads = 1, bool fast_quantization = false ) { + const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before add"); @@ -144,26 +153,33 @@ class HnswIndex { // 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 - ); + const auto* data_ptr = data_array.data(); + const auto labels = [&] { + py::gil_scoped_release release; + return index_->add( + data_ptr, + 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) { + const auto operation = access_.write(); if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before resize"); } + py::gil_scoped_release release; index_->resize(new_max_elements); max_elements_ = index_->max_elements(); } size_t remove(py::handle ids) { + const auto operation = access_.write(); auto ids_array = ensure_1d_integer_array(ids, "ids"); if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before remove"); @@ -179,12 +195,17 @@ class HnswIndex { } to_remove.push_back(static_cast(id)); } + py::gil_scoped_release release; return index_->remove(to_remove.data(), to_remove.size()); } - [[nodiscard]] size_t num_points() const { return built_ ? index_->num_points() : 0; } + [[nodiscard]] size_t num_points() const { + const auto operation = access_.read(); + return built_ ? index_->num_points() : 0; + } py::tuple search(py::handle queries, size_t k, size_t ef = 0, size_t num_threads = 1) { + const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before search"); @@ -212,13 +233,13 @@ class HnswIndex { auto ids_buf = ids.mutable_unchecked<2>(); auto dists_buf = dists.mutable_unchecked<2>(); - std::vector>> results = index_->search( - query_array.data(), - static_cast(query_array.shape(0)), - k, - ef, - num_threads - ); + const auto* queries_data = query_array.data(); + const auto nq = static_cast(query_array.shape(0)); + std::vector>> results; + { + py::gil_scoped_release release; + results = index_->search(queries_data, nq, k, ef, num_threads); + } for (py::ssize_t i = 0; i < static_cast(results.size()); ++i) { for (py::ssize_t j = 0; j < static_cast(std::min( @@ -235,16 +256,21 @@ class HnswIndex { } void save(const std::string& path) const { + const auto operation = access_.read(); if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before save"); } + py::gil_scoped_release release; index_->save(path.c_str()); } static HnswIndex load(const std::string& path) { HnswIndex wrapper; wrapper.index_ = std::make_unique(); - wrapper.index_->load(path.c_str()); + { + py::gil_scoped_release release; + wrapper.index_->load(path.c_str()); + } wrapper.dim_ = wrapper.index_->dimension(); wrapper.max_elements_ = wrapper.index_->max_elements(); wrapper.M_ = wrapper.index_->M(); @@ -256,12 +282,30 @@ class HnswIndex { return wrapper; } - [[nodiscard]] size_t dim() const { return dim_; } - [[nodiscard]] size_t max_elements() const { return max_elements_; } - [[nodiscard]] size_t nbits() const { return nbits_; } - [[nodiscard]] bool is_built() const { return built_; } - [[nodiscard]] size_t num_clusters() const { return num_clusters_; } - [[nodiscard]] std::string metric() const { return metric_to_string(metric_); } + [[nodiscard]] size_t dim() const { + const auto operation = access_.read(); + return dim_; + } + [[nodiscard]] size_t max_elements() const { + const auto operation = access_.read(); + return max_elements_; + } + [[nodiscard]] size_t nbits() const { + const auto operation = access_.read(); + return nbits_; + } + [[nodiscard]] bool is_built() const { + const auto operation = access_.read(); + return built_; + } + [[nodiscard]] size_t num_clusters() const { + const auto operation = access_.read(); + return num_clusters_; + } + [[nodiscard]] std::string metric() const { + const auto operation = access_.read(); + return metric_to_string(metric_); + } private: HnswIndex() = default; @@ -275,6 +319,7 @@ class HnswIndex { size_t random_seed_ = 100; size_t num_clusters_ = 0; bool built_ = false; + mutable IndexAccess access_; std::unique_ptr index_; }; diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 55bd8d3a..9fe992c8 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -87,6 +88,7 @@ class IvfIndex { size_t num_threads = 1, bool fast_quantization = false ) { + const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); auto centroids_array = ensure_2d_array(centroids, "centroids"); auto cluster_ids_array = @@ -114,13 +116,15 @@ class IvfIndex { ); } - index_->construct( - data_array.data(), - centroids_array.data(), - cluster_ids_array.data(), - fast_quantization, - num_threads - ); + const auto* data_ptr = data_array.data(); + const auto* centroids_ptr = centroids_array.data(); + const auto* cluster_ids_ptr = cluster_ids_array.data(); + { + py::gil_scoped_release release; + index_->construct( + data_ptr, centroids_ptr, cluster_ids_ptr, fast_quantization, num_threads + ); + } built_ = true; } @@ -130,6 +134,7 @@ class IvfIndex { size_t num_threads = 1, bool fast_quantization = false ) { + const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before add"); @@ -167,18 +172,23 @@ class IvfIndex { ids_data[i] = static_cast(first + i); } - index_->add( - data_array.data(), - rows, - assigned.empty() ? nullptr : assigned.data(), - fast_quantization, - num_threads - ); + const auto* data_ptr = data_array.data(); + { + py::gil_scoped_release release; + index_->add( + data_ptr, + rows, + assigned.empty() ? nullptr : assigned.data(), + fast_quantization, + num_threads + ); + } max_elements_ = index_->max_elements(); return ids; } size_t remove(py::handle ids) { + const auto operation = access_.write(); auto ids_array = ensure_1d_integer_array(ids, "ids"); if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before remove"); @@ -192,6 +202,7 @@ class IvfIndex { } to_remove.push_back(static_cast(id)); } + py::gil_scoped_release release; return index_->remove(to_remove.data(), to_remove.size()); } @@ -202,6 +213,7 @@ class IvfIndex { std::optional high_accuracy = std::nullopt, size_t num_threads = 1 ) { + const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before search"); @@ -228,31 +240,40 @@ class IvfIndex { dists_data, dists_data + dists.size(), std::numeric_limits::infinity() ); - index_->search_batch( - query_array.data(), - nq, - k, - nprobe, - ids_data, - dists_data, - high_accuracy, - num_threads - ); + const auto* queries_data = query_array.data(); + { + py::gil_scoped_release release; + index_->search_batch( + queries_data, + nq, + k, + nprobe, + ids_data, + dists_data, + high_accuracy, + num_threads + ); + } return py::make_tuple(ids, dists); } void save(const std::string& path) const { + const auto operation = access_.read(); if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before save"); } + py::gil_scoped_release release; index_->save(path.c_str()); } static IvfIndex load(const std::string& path) { IvfIndex wrapper; wrapper.index_ = std::make_unique(); - wrapper.index_->load(path.c_str()); + { + py::gil_scoped_release release; + wrapper.index_->load(path.c_str()); + } wrapper.dim_ = wrapper.index_->dimension(); wrapper.max_elements_ = wrapper.index_->max_elements(); wrapper.num_clusters_ = wrapper.index_->num_clusters(); @@ -262,15 +283,34 @@ class IvfIndex { return wrapper; } - [[nodiscard]] size_t dim() const { return dim_; } - [[nodiscard]] size_t max_elements() const { return max_elements_; } - [[nodiscard]] size_t num_clusters() const { return num_clusters_; } - [[nodiscard]] size_t nbits() const { return nbits_; } - [[nodiscard]] bool is_built() const { return built_; } + [[nodiscard]] size_t dim() const { + const auto operation = access_.read(); + return dim_; + } + [[nodiscard]] size_t max_elements() const { + const auto operation = access_.read(); + return max_elements_; + } + [[nodiscard]] size_t num_clusters() const { + const auto operation = access_.read(); + return num_clusters_; + } + [[nodiscard]] size_t nbits() const { + const auto operation = access_.read(); + return nbits_; + } + [[nodiscard]] bool is_built() const { + const auto operation = access_.read(); + return built_; + } [[nodiscard]] std::string initializer() const { + const auto operation = access_.read(); return initializer_to_string(index_->initializer_type()); } - [[nodiscard]] std::string metric() const { return metric_to_string(metric_); } + [[nodiscard]] std::string metric() const { + const auto operation = access_.read(); + return metric_to_string(metric_); + } private: IvfIndex() = default; @@ -281,6 +321,7 @@ class IvfIndex { size_t nbits_ = 0; rabitqlib::MetricType metric_ = rabitqlib::METRIC_L2; bool built_ = false; + mutable IndexAccess access_; std::unique_ptr index_; }; diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 1d463517..e1bbce45 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -50,6 +51,7 @@ class SymqgIndex { size_t num_threads = 1, const std::string& init = "pipnn" ) { + const auto operation = access_.write(); if (init != "random" && init != "pipnn") { throw std::invalid_argument("init must be 'random' or 'pipnn'"); } @@ -74,20 +76,25 @@ class SymqgIndex { const auto initialization = init == "pipnn" ? symqg::QGInitialization::PiPNN : symqg::QGInitialization::Random; - symqg::QGBuilder builder( - *index, - static_cast(ef_construction), - data_array.data(), - num_threads, - initialization - ); - builder.build(); + const auto* data_ptr = data_array.data(); + { + py::gil_scoped_release release; + symqg::QGBuilder builder( + *index, + static_cast(ef_construction), + data_ptr, + num_threads, + initialization + ); + builder.build(); + } index_ = std::move(index); num_points_ = num_points; built_ = true; } py::tuple search(py::handle queries, size_t k, size_t ef, size_t num_threads = 1) { + const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); if (!built_) { throw std::runtime_error("SymqgIndex must be built or loaded before search"); @@ -105,8 +112,6 @@ class SymqgIndex { throw std::invalid_argument("query dimension does not match index dim"); } - index_->set_ef(ef); - const size_t nq = static_cast(query_array.shape(0)); const auto shape = std::vector{ static_cast(nq), static_cast(k)}; @@ -114,29 +119,39 @@ class SymqgIndex { auto dists = py::array_t(shape); auto* ids_data = ids.mutable_data(); auto* dists_data = dists.mutable_data(); - index_->search_batch( - query_array.data(), - nq, - static_cast(k), - ids_data, - dists_data, - num_threads - ); + const auto* queries_data = query_array.data(); + { + py::gil_scoped_release release; + index_->search_batch_with_ef( + queries_data, + nq, + static_cast(k), + ids_data, + dists_data, + ef, + num_threads + ); + } return py::make_tuple(ids, dists); } void save(const std::string& path) const { + const auto operation = access_.read(); if (!built_) { throw std::runtime_error("SymqgIndex must be built before save"); } + py::gil_scoped_release release; index_->save(path.c_str()); } static SymqgIndex load(const std::string& path) { SymqgIndex wrapper; wrapper.index_ = std::make_unique>(); - wrapper.index_->load(path.c_str()); + { + py::gil_scoped_release release; + wrapper.index_->load(path.c_str()); + } wrapper.num_points_ = wrapper.index_->num_vertices(); wrapper.dim_ = wrapper.index_->dimension(); wrapper.max_degree_ = wrapper.index_->degree_bound(); @@ -146,12 +161,30 @@ class SymqgIndex { return wrapper; } - [[nodiscard]] size_t dim() const { return dim_; } - [[nodiscard]] size_t max_degree() const { return max_degree_; } - [[nodiscard]] size_t num_points() const { return num_points_; } - [[nodiscard]] bool is_built() const { return built_; } - [[nodiscard]] std::string metric() const { return metric_to_string(metric_); } - [[nodiscard]] size_t quantization_bits() const { return quantization_bits_; } + [[nodiscard]] size_t dim() const { + const auto operation = access_.read(); + return dim_; + } + [[nodiscard]] size_t max_degree() const { + const auto operation = access_.read(); + return max_degree_; + } + [[nodiscard]] size_t num_points() const { + const auto operation = access_.read(); + return num_points_; + } + [[nodiscard]] bool is_built() const { + const auto operation = access_.read(); + return built_; + } + [[nodiscard]] std::string metric() const { + const auto operation = access_.read(); + return metric_to_string(metric_); + } + [[nodiscard]] size_t quantization_bits() const { + const auto operation = access_.read(); + return quantization_bits_; + } private: SymqgIndex() = default; @@ -162,6 +195,7 @@ class SymqgIndex { rabitqlib::MetricType metric_ = rabitqlib::METRIC_L2; size_t quantization_bits_ = 0; bool built_ = false; + mutable IndexAccess access_; std::unique_ptr> index_; }; diff --git a/src/index/hnsw_search_avx2.cpp b/src/index/hnsw_search_avx2.cpp index e5b5873e..d1de0d63 100644 --- a/src/index/hnsw_search_avx2.cpp +++ b/src/index/hnsw_search_avx2.cpp @@ -28,9 +28,11 @@ struct HnswAvx2Kernel { }; maxheap> search_knn_avx2( - HierarchicalNSW& index, const float* rotated_query, size_t topk + HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct(rotated_query, topk); + return index.search_knn_direct( + rotated_query, topk, ef == 0 ? index.ef_ : ef + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_avx512_core.cpp b/src/index/hnsw_search_avx512_core.cpp index 25ff1825..a7647c76 100644 --- a/src/index/hnsw_search_avx512_core.cpp +++ b/src/index/hnsw_search_avx512_core.cpp @@ -29,9 +29,11 @@ struct HnswAvx512CoreKernel { }; maxheap> search_knn_avx512_core( - HierarchicalNSW& index, const float* rotated_query, size_t topk + HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct(rotated_query, topk); + return index.search_knn_direct( + rotated_query, topk, ef == 0 ? index.ef_ : ef + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_avx512_popcnt.cpp b/src/index/hnsw_search_avx512_popcnt.cpp index 8330642e..c32adc22 100644 --- a/src/index/hnsw_search_avx512_popcnt.cpp +++ b/src/index/hnsw_search_avx512_popcnt.cpp @@ -28,9 +28,11 @@ struct HnswAvx512PopcntKernel { }; maxheap> search_knn_avx512_popcnt( - HierarchicalNSW& index, const float* rotated_query, size_t topk + HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct(rotated_query, topk); + return index.search_knn_direct( + rotated_query, topk, ef == 0 ? index.ef_ : ef + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_generic.cpp b/src/index/hnsw_search_generic.cpp index 5afc59c0..a49bbf31 100644 --- a/src/index/hnsw_search_generic.cpp +++ b/src/index/hnsw_search_generic.cpp @@ -31,9 +31,11 @@ struct HnswGenericKernel { }; maxheap> search_knn_generic( - HierarchicalNSW& index, const float* rotated_query, size_t topk + HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct(rotated_query, topk); + return index.search_knn_direct( + rotated_query, topk, ef == 0 ? index.ef_ : ef + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_neon.cpp b/src/index/hnsw_search_neon.cpp index ad0ace15..0d8473f5 100644 --- a/src/index/hnsw_search_neon.cpp +++ b/src/index/hnsw_search_neon.cpp @@ -34,9 +34,11 @@ struct HnswNeonKernel { }; maxheap> search_knn_neon( - HierarchicalNSW& index, const float* rotated_query, size_t topk + HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct(rotated_query, topk); + return index.search_knn_direct( + rotated_query, topk, ef == 0 ? index.ef_ : ef + ); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 67a98655..b88628e4 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -414,12 +414,18 @@ void QuantizedGraph::search( uint32_t* __restrict__ results, float* __restrict__ dists ) { - validate_search(query, k, results, dists); + search_impl(query, k, results, dists, ef_); +} + +void QuantizedGraph::search_impl( + const float* query, uint32_t k, uint32_t* results, float* dists, size_t ef +) { + validate_search(query, k, results, dists, ef); std::vector rotated_query(padded_dim_); std::vector est_dist(degree_bound_); std::vector lut_float(padded_dim_ * 4); BatchQuery batch_query; - buffer::SearchBuffer search_pool(ef_); + buffer::SearchBuffer search_pool(ef); buffer::SearchBuffer result_pool(k); thread_local VisitedSet visited; thread_local size_t visited_size = 0; @@ -427,7 +433,7 @@ void QuantizedGraph::search( visited.initialize(num_points_, num_points_ / 10); visited_size = num_points_; } - search_with_scratch( + search_with_scratch_impl( query, k, results, @@ -438,7 +444,9 @@ void QuantizedGraph::search( batch_query, search_pool, result_pool, - visited + visited, + kPidMax, + nullptr ); } @@ -449,17 +457,29 @@ void QuantizedGraph::search_batch( uint32_t* results, float* dists, size_t num_threads +) { + search_batch_with_ef(queries, num_queries, knn, results, dists, ef_, num_threads); +} + +void QuantizedGraph::search_batch_with_ef( + const float* queries, + size_t num_queries, + uint32_t knn, + uint32_t* results, + float* dists, + size_t ef, + size_t num_threads ) { if (num_queries == 0) { return; } - validate_search(queries, knn, results, dists); + validate_search(queries, knn, results, dists, ef); const size_t max_values = std::numeric_limits::max() / sizeof(float); if (num_queries > max_values / dim_ || num_queries > max_values / knn) { throw std::length_error("QuantizedGraph query batch is too large"); } if (num_queries == 1) { - search(queries, knn, results, dists); + search_impl(queries, knn, results, dists, ef); return; } @@ -502,7 +522,7 @@ void QuantizedGraph::search_batch( // Allocate and release scratch on its worker, retaining it across queries. std::optional scratch; try { - scratch.emplace(padded_dim_, degree_bound_, ef_, knn, num_points_); + scratch.emplace(padded_dim_, degree_bound_, ef, knn, num_points_); } catch (...) { capture_error(); } #pragma omp for schedule(dynamic) for (std::ptrdiff_t query_index = 0; @@ -514,7 +534,7 @@ void QuantizedGraph::search_batch( const size_t i = static_cast(query_index); try { auto& local = *scratch; - search_with_scratch( + search_with_scratch_impl( queries + i * dim_, knn, results + i * knn, @@ -525,7 +545,9 @@ void QuantizedGraph::search_batch( local.batch_query, local.search_pool, local.result_pool, - local.visited + local.visited, + kPidMax, + nullptr ); } catch (...) { capture_error(); } } @@ -536,7 +558,7 @@ void QuantizedGraph::search_batch( } void QuantizedGraph::validate_search( - const float* query, uint32_t k, const uint32_t* results, const float* dists + const float* query, uint32_t k, const uint32_t* results, const float* dists, size_t ef ) const { if (!ready_ || rotator_ == nullptr) { throw std::logic_error("QuantizedGraph must be built or loaded before search"); @@ -547,7 +569,7 @@ void QuantizedGraph::validate_search( if (k == 0 || k > num_points_) { throw std::invalid_argument("QuantizedGraph k must be between 1 and num_points"); } - if (ef_ < k) { + if (ef < k) { throw std::invalid_argument("QuantizedGraph ef must be at least k"); } } @@ -881,7 +903,39 @@ void QuantizedGraph::search_with_scratch( PID hint, const float* k1xsumq ) { - validate_search(query, knn, results, dists); + validate_search(query, knn, results, dists, ef_); + search_with_scratch_impl( + query, + knn, + results, + dists, + rotated_query_scratch, + est_dist_scratch, + lut_float_scratch, + batch_query, + search_pool, + result_pool, + visited, + hint, + k1xsumq + ); +} + +void QuantizedGraph::search_with_scratch_impl( + const float* __restrict__ query, + uint32_t knn, + uint32_t* __restrict__ results, + float* __restrict__ dists, + float* __restrict__ rotated_query_scratch, + float* __restrict__ est_dist_scratch, + float* __restrict__ lut_float_scratch, + BatchQuery& batch_query, + buffer::SearchBuffer& search_pool, + buffer::SearchBuffer& result_pool, + VisitedSet& visited, + PID hint, + const float* k1xsumq +) { if (rotated_query_scratch == nullptr || est_dist_scratch == nullptr || lut_float_scratch == nullptr) { throw std::invalid_argument("QuantizedGraph search buffers must not be null"); diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 14630c2f..b841d38a 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -659,8 +659,9 @@ float warmup_ip_x0_q_512( namespace rabitqlib::hnsw::detail { namespace { -using SearchKnnFn = - std::priority_queue> (*)(HierarchicalNSW&, const float*, size_t); +using SearchKnnFn = std::priority_queue> (*)( + HierarchicalNSW&, const float*, size_t, size_t +); // The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. #if defined(__x86_64__) || defined(_M_X64) const SearchKnnFn kSearchKnnFn = cpu::has_avx512_popcnt() @@ -677,8 +678,8 @@ const SearchKnnFn kSearchKnnFn = #endif } // namespace std::priority_queue> search_knn( - HierarchicalNSW& index, const float* query, size_t topk + HierarchicalNSW& index, const float* query, size_t topk, size_t ef ) { - return kSearchKnnFn(index, query, topk); + return kSearchKnnFn(index, query, topk, ef); } } // namespace rabitqlib::hnsw::detail diff --git a/tests/python/test_concurrent_indexes.py b/tests/python/test_concurrent_indexes.py new file mode 100644 index 00000000..b1753bc7 --- /dev/null +++ b/tests/python/test_concurrent_indexes.py @@ -0,0 +1,201 @@ +"""Independent search parameters, GIL release, and fail-fast index access.""" + +import sys +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier, Event, Thread + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex + +BUSY = "Index is busy: conflicting operation in progress" +CASES = [ + ("ivf", 1), + ("ivf", 4), + ("ivf", 32), + ("ivf_hnsw", 4), + ("ivf_flat_rabitq", 4), + ("hnsw", 1), + ("hnsw", 9), + ("qg", 0), + ("qg", 4), + ("qg", 8), +] + + +def make_index(kind, bits=4, metric="l2"): + rng = np.random.default_rng(713) + data = rng.standard_normal((129, 65)).astype(np.float32) + labels = (np.arange(len(data)) % 5).astype(np.uint32) + centroids = np.stack([data[labels == i].mean(axis=0) for i in range(5)]) + if kind.startswith("ivf"): + initializer = kind.removeprefix("ivf_") if "_" in kind else "flat" + index = IvfIndex(65, len(data), 5, bits, metric, initializer) + elif kind == "hnsw": + index = HnswIndex(65, len(data) + 8, 8, 64, bits, metric, 42) + else: + index = SymqgIndex(65, 32, metric, bits) + + def build(vectors=data): + if kind == "qg": + index.build(vectors, 64, 1) + else: + index.build(vectors, centroids, labels, 1) + + build() + return index, data, build + + +def parameters(kind, wide): + if kind.startswith("ivf"): + return {"nprobe": 5 if wide else 1, "high_accuracy": wide} + return {"ef": 96 if wide else 5} + + +@pytest.mark.parametrize("kind,bits", CASES) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_concurrent_search_matches_serial_and_loaded(kind, bits, metric, tmp_path): + index, _, _ = make_index(kind, bits, metric) + # Exercise conversions and ownership of a temporary contiguous float32 copy. + queries = np.random.default_rng(18).standard_normal((17, 130))[:, ::2] + expected = [index.search(queries, 5, **parameters(kind, wide)) for wide in (0, 1)] + path = str(tmp_path / "concurrent.index") + index.save(path) + for current in (index, type(index).load(path)): + start = Barrier(4) + + def search(worker): + wide = worker % 2 + start.wait(timeout=10) + for count in (0, 1, 17, 17): + method = ( + current.search_batch + if kind == "qg" and worker >= 2 + else current.search + ) + actual = method( + queries[:count], + 5, + num_threads=1 + worker // 2, + **parameters(kind, wide), + ) + for result, reference in zip(actual, expected[wide], strict=True): + np.testing.assert_array_equal(result, reference[:count]) + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(search, range(4))) + + +class PausedArray: + """Hold an operation open deterministically during Python input conversion.""" + + def __init__(self, array): + self.array = array + self.entered = Event() + self.resume = Event() + + def __array__(self, dtype=None, copy=None): + self.entered.set() + if not self.resume.wait(timeout=15): + raise RuntimeError("test did not release input conversion") + return np.array(self.array, dtype=dtype, copy=True) + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) +def test_active_search_allows_reads_and_rejects_mutations(kind, tmp_path): + index, data, build = make_index(kind) + queries = PausedArray(data[:3]) + options = parameters(kind, True) + with ThreadPoolExecutor(max_workers=1) as pool: + running = pool.submit(index.search, queries, 5, **options) + try: + assert queries.entered.wait(timeout=10) + assert index.dim == 65 + index.search(data[:1], 5, **options) + index.save(str(tmp_path / "reading.index")) + mutations = [build] + if kind != "qg": + mutations += [lambda: index.add(data[:1]), lambda: index.remove([0])] + if kind == "hnsw": + mutations.append(lambda: index.resize(256)) + for mutate in mutations: + with pytest.raises(RuntimeError, match=f"^{BUSY}$"): + mutate() + finally: + queries.resume.set() + running.result(timeout=10) + # Failed conflicts must not leave the access state busy. + if kind == "hnsw": + index.resize(256) + else: + build() + index.search(data[:1], 5, **options) + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) +def test_active_mutation_rejects_reads_and_other_mutations(kind, tmp_path): + index, data, build = make_index(kind) + vectors = PausedArray(data[:1] if kind == "hnsw" else data) + with ThreadPoolExecutor(max_workers=1) as pool: + running = pool.submit(index.add if kind == "hnsw" else build, vectors) + try: + assert vectors.entered.wait(timeout=10) + calls = [ + lambda: index.search(data[:1], 5, **parameters(kind, True)), + lambda: index.save(str(tmp_path / "busy.index")), + lambda: index.dim, + lambda: index.is_built, + build, + ] + if kind != "qg": + calls += [lambda: index.add(data[:1]), lambda: index.remove([0])] + if kind == "hnsw": + calls.append(lambda: index.resize(256)) + for call in calls: + with pytest.raises(RuntimeError, match=f"^{BUSY}$"): + call() + other, other_data, _ = make_index(kind) + other.search(other_data[:1], 5, **parameters(kind, True)) + finally: + vectors.resume.set() + running.result(timeout=10) + assert index.is_built + index.search(data[:1], 5, **parameters(kind, True)) + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) +def test_search_releases_gil(kind): + index, data, _ = make_index(kind) + queries = np.tile(data, (32, 1)) + go, finished = Event(), Event() + observed = [] + + def python_worker(): + go.wait(timeout=10) + observed.append(not finished.is_set()) + + worker = Thread(target=python_worker) + worker.start() + previous = sys.getswitchinterval() + try: + # Prevent an interpreter timeslice between go.set() and the native call. + # The worker must run while that call has released the GIL. + sys.setswitchinterval(100) + go.set() + index.search(queries, 5, num_threads=1, **parameters(kind, True)) + finished.set() + finally: + sys.setswitchinterval(previous) + finished.set() + worker.join(timeout=10) + assert not worker.is_alive() + assert observed == [True] + + +def test_native_exception_reacquires_gil_and_releases_access(): + index, data, build = make_index("qg") + with pytest.raises(ValueError, match="SearchBuffer capacity is too large"): + index.search(data[:2], 5, ef=int(np.iinfo(np.uintp).max)) + build() + ids, distances = index.search(data[:2], 5, ef=64) + assert ids.shape == distances.shape == (2, 5) diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 4f5fbedd..03554393 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -1206,12 +1207,11 @@ TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { ); // Insertion reconstructs stored sign/extra-bit tails before linking. index.add(data.data() + (count - 4) * dim, 4, clusters.data(), true, 1); - index.search(query.data(), 1, 5, count, 1); const auto reference = - hnsw::detail::search_knn_generic(index, query.data(), 5); + hnsw::detail::search_knn_generic(index, query.data(), 5, count); for (auto search : backends) { auto expected = reference; - auto actual = search(index, query.data(), 5); + auto actual = search(index, query.data(), 5, count); ASSERT_EQ(actual.size(), expected.size()); while (!expected.empty()) { EXPECT_EQ(actual.top().second, expected.top().second); @@ -1228,3 +1228,43 @@ TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { } } } + +TEST(HnswSearchTest, ConcurrentCallsKeepIndependentSearchWindows) { + using namespace rabitqlib; + constexpr size_t kCount = 129, kDim = 65, kQueries = 13, k = 5; + std::mt19937 random(314); + std::normal_distribution value; + std::vector data(kCount * kDim), queries(kQueries * kDim), centroid(kDim, 0); + for (auto& x : data) + x = value(random); + for (auto& x : queries) + x = value(random); + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1U, 4U, 9U}) { + SCOPED_TRACE(metric); + SCOPED_TRACE(bits); + std::vector clusters(kCount, 0); + hnsw::HierarchicalNSW index(kCount, kDim, bits, 8, 64, 42, metric); + index.construct( + 1, centroid.data(), kCount, data.data(), clusters.data(), 1, true + ); + const auto narrow = index.search(queries.data(), kQueries, k, k, 1); + const auto wide = index.search(queries.data(), kQueries, k, 96, 1); + std::promise start; + const auto ready = start.get_future().share(); + const auto run = [&](size_t ef, size_t threads, const auto& expected) { + ready.wait(); + for (size_t i = 0; i < 16; ++i) { + if (index.search(queries.data(), kQueries, k, ef, threads) != expected) + return false; + } + return true; + }; + auto first = std::async(std::launch::async, [&] { return run(k, 1, narrow); }); + auto second = std::async(std::launch::async, [&] { return run(96, 2, wide); }); + start.set_value(); + EXPECT_TRUE(first.get()); + EXPECT_TRUE(second.get()); + } + } +} diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index d309861d..2b5274d1 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -1888,3 +1888,94 @@ TEST(QuantizedGraphSearchTest, ScratchSearchRejectsIncompleteResults) { })); } } // namespace rabitqlib::symqg + +namespace rabitqlib::symqg { +TEST(QGSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { + constexpr size_t kCount = 129, kDim = 65, kQueries = 13, k = 5; + std::mt19937 random(314); + std::normal_distribution value; + std::vector data(kCount * kDim), queries(kQueries * kDim); + for (auto& x : data) + x = value(random); + for (auto& x : queries) + x = value(random); + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(metric); + SCOPED_TRACE(bits); + QuantizedGraph graph( + kCount, kDim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGBuilder builder(graph, 64, data.data(), 1); + builder.build(); + std::vector narrow(kQueries * k), wide(narrow.size()); + std::vector narrow_distances(narrow.size()), wide_distances(wide.size()); + graph.set_ef(k); + graph.search_batch( + queries.data(), kQueries, k, narrow.data(), narrow_distances.data() + ); + graph.set_ef(96); + graph.search_batch( + queries.data(), kQueries, k, wide.data(), wide_distances.data() + ); + graph.set_ef(1); + std::promise start; + const auto ready = start.get_future().share(); + const auto run = [&](size_t ef, + size_t threads, + const auto& expected, + const auto& expected_distances) { + std::vector actual(kQueries * k); + std::vector distances(actual.size()); + ready.wait(); + for (size_t i = 0; i < 16; ++i) { + graph.search_batch_with_ef( + queries.data(), + kQueries, + k, + actual.data(), + distances.data(), + ef, + threads + ); + if (actual != expected || distances != expected_distances) + return false; + graph.search_batch_with_ef( + queries.data(), 1, k, actual.data(), distances.data(), ef, threads + ); + if (!std::equal(actual.begin(), actual.begin() + k, expected.begin()) || + !std::equal( + distances.begin(), + distances.begin() + k, + expected_distances.begin() + )) + return false; + } + return true; + }; + auto first = std::async(std::launch::async, [&] { + return run(k, 1, narrow, narrow_distances); + }); + auto second = std::async(std::launch::async, [&] { + return run(96, 2, wide, wide_distances); + }); + start.set_value(); + EXPECT_TRUE(first.get()); + EXPECT_TRUE(second.get()); + EXPECT_THROW( + graph.search_batch( + queries.data(), kQueries, k, narrow.data(), narrow_distances.data() + ), + std::invalid_argument + ); + EXPECT_THROW( + graph.search_batch_with_ef( + queries.data(), kQueries, k, narrow.data(), narrow_distances.data(), 0 + ), + std::invalid_argument + ); + EXPECT_NO_THROW(graph.search_batch_with_ef(nullptr, 0, 0, nullptr, nullptr, 0)); + } + } +} +} // namespace rabitqlib::symqg diff --git a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp index 86ec044c..3757dcbd 100644 --- a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp +++ b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp @@ -104,11 +104,10 @@ TEST(NeonHnsw, SearchMatchesGenericOnSameGraphForBothMetricsAndAllWidths) { value = distribution(random); } for (size_t ef : {10U, 64U}) { - // Configure ef through the public search API. - index.search(query.data(), 1, 5, ef, 1); auto expected = - hnsw::detail::search_knn_generic(index, query.data(), 5); - auto actual = hnsw::detail::search_knn_neon(index, query.data(), 5); + hnsw::detail::search_knn_generic(index, query.data(), 5, ef); + auto actual = + hnsw::detail::search_knn_neon(index, query.data(), 5, ef); ASSERT_EQ(expected.size(), actual.size()); while (!expected.empty()) { EXPECT_EQ(actual.top().second, expected.top().second);