diff --git a/src/impl/searcher/CMakeLists.txt b/src/impl/searcher/CMakeLists.txt index b3459adf70..61d443ac2e 100644 --- a/src/impl/searcher/CMakeLists.txt +++ b/src/impl/searcher/CMakeLists.txt @@ -19,6 +19,7 @@ set (SEARCHER_SRC basic_searcher.h parallel_searcher.cpp parallel_searcher.h + searcher_utils.h ) add_library (searcher OBJECT ${SEARCHER_SRC}) diff --git a/src/impl/searcher/basic_searcher.cpp b/src/impl/searcher/basic_searcher.cpp index c56244f30a..209ed07d1d 100644 --- a/src/impl/searcher/basic_searcher.cpp +++ b/src/impl/searcher/basic_searcher.cpp @@ -22,6 +22,7 @@ #include "algorithm/inner_index_interface.h" #include "datacell/flatten_interface.h" #include "impl/heap/standard_heap.h" +#include "impl/searcher/searcher_utils.h" #include "utils/filter_search_skip_strategy.h" #include "vsag/allocator.h" @@ -150,10 +151,11 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, if (iter_ctx->CheckPoint(cur_inner_id)) { lower_bound = std::max(lower_bound, cur_dist); flatten->Query(&cur_dist, computer, &cur_inner_id, 1, ctx); - if (cur_dist > inner_search_param.min_distance + THRESHOLD_ERROR) { + if (is_result_distance_eligible(cur_dist) and + cur_dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(cur_dist, cur_inner_id); } - candidate_set->Push(cur_dist, cur_inner_id); + candidate_set->Push(traversal_priority(cur_dist), cur_inner_id); if constexpr (mode == InnerSearchMode::RANGE_SEARCH) { if (cur_dist > inner_search_param.radius and not top_candidates->Empty()) { top_candidates->Pop(); @@ -164,19 +166,21 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, } } else { flatten->Query(&dist, computer, &ep, 1, ctx); - if ((not is_id_allowed || is_id_allowed->CheckValid(ep)) and + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(ep)) and !(dist <= inner_search_param.min_distance + THRESHOLD_ERROR)) { top_candidates->Push(dist, ep); lower_bound = top_candidates->Top().first; } - candidate_set->Push(-dist, ep); + candidate_set->Push(traversal_priority(dist), ep); vl->Set(ep); if (inner_search_param.consider_duplicate and label_table != nullptr and label_table->CompressDuplicateData()) { const auto& duplicate_ids = label_table->GetDuplicateId(ep); for (const auto& item : duplicate_ids) { - if ((not is_id_allowed || is_id_allowed->CheckValid(item)) and + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(item)) and iter_ctx->CheckPoint(item) and dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, item); @@ -230,14 +234,15 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, for (uint32_t i = 0; i < count_no_visited; i++) { dist = line_dists[i]; - if (top_candidates->Size() < ef || lower_bound > dist || + if (not std::isfinite(dist) || top_candidates->Size() < ef || lower_bound > dist || (mode == RANGE_SEARCH && dist <= inner_search_param.radius)) { if (!iter_ctx->CheckPoint(to_be_visited_id[i])) { continue; } - candidate_set->Push(-dist, to_be_visited_id[i]); + candidate_set->Push(traversal_priority(dist), to_be_visited_id[i]); flatten->Prefetch(candidate_set->Top().second); - if ((not is_id_allowed || is_id_allowed->CheckValid(to_be_visited_id[i])) && + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(to_be_visited_id[i])) && dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, to_be_visited_id[i]); } @@ -246,7 +251,8 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, label_table->CompressDuplicateData()) { const auto& duplicate_ids = label_table->GetDuplicateId(to_be_visited_id[i]); for (const auto& item : duplicate_ids) { - if ((not is_id_allowed || is_id_allowed->CheckValid(item)) and + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(item)) and iter_ctx->CheckPoint(item)) { top_candidates->Push(dist, item); } @@ -338,7 +344,8 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, flatten->Query(&dist, computer, &ep, 1, ctx); ++dist_cmp; - if (check_func(ep) && !(dist <= inner_search_param.min_distance + THRESHOLD_ERROR)) { + if (is_result_distance_eligible(dist) and check_func(ep) && + !(dist <= inner_search_param.min_distance + THRESHOLD_ERROR)) { top_candidates->Push(dist, ep); lower_bound = top_candidates->Top().first; } @@ -347,14 +354,15 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, top_candidates->Pop(); } } - candidate_set->Push(-dist, ep); + candidate_set->Push(traversal_priority(dist), ep); vl->Set(ep); if (inner_search_param.consider_duplicate and label_table != nullptr and label_table->CompressDuplicateData()) { const auto& duplicate_ids = label_table->GetDuplicateId(ep); for (const auto& item : duplicate_ids) { - if (check_func(item) && dist > inner_search_param.min_distance + THRESHOLD_ERROR) { + if (is_result_distance_eligible(dist) and check_func(item) && + dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, item); } } @@ -408,11 +416,11 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, for (uint32_t i = 0; i < count_no_visited; i++) { dist = line_dists[i]; - if (top_candidates->Size() < ef || lower_bound > dist || + if (not std::isfinite(dist) || top_candidates->Size() < ef || lower_bound > dist || (mode == RANGE_SEARCH && dist <= inner_search_param.radius)) { - candidate_set->Push(-dist, to_be_visited_id[i]); + candidate_set->Push(traversal_priority(dist), to_be_visited_id[i]); // flatten->Prefetch(candidate_set->Top().second); - if (check_func(to_be_visited_id[i]) && + if (is_result_distance_eligible(dist) and check_func(to_be_visited_id[i]) && dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, to_be_visited_id[i]); } @@ -420,7 +428,7 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, label_table->CompressDuplicateData()) { const auto& duplicate_ids = label_table->GetDuplicateId(to_be_visited_id[i]); for (const auto& item : duplicate_ids) { - if (check_func(item) && + if (is_result_distance_eligible(dist) and check_func(item) && dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, item); } diff --git a/src/impl/searcher/basic_searcher_test.cpp b/src/impl/searcher/basic_searcher_test.cpp index 77920dfae7..3362cda0b5 100644 --- a/src/impl/searcher/basic_searcher_test.cpp +++ b/src/impl/searcher/basic_searcher_test.cpp @@ -15,6 +15,8 @@ #include "basic_searcher.h" +#include +#include #include #include "algorithm/inner_index_interface.h" @@ -24,6 +26,57 @@ using namespace vsag; +TEST_CASE("BasicSearcher traverses through a non-finite-distance bridge", + "[ut][BasicSearcher][nonfinite]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + IndexCommonParam common; + common.dim_ = 1; + common.allocator_ = allocator; + common.metric_ = MetricType::METRIC_TYPE_L2SQR; + + constexpr const char* param_temp = R"({{"type": "{}"}})"; + auto quantizer_param = QuantizerParameter::GetQuantizerParameterByJson( + JsonType::Parse(fmt::format(param_temp, "fp32"))); + auto io_param = + IOParameter::GetIOParameterByJson(JsonType::Parse(fmt::format(param_temp, "memory_io"))); + auto flatten = + std::make_shared, MemoryIO>>( + quantizer_param, io_param, common); + flatten->SetQuantizer( + std::make_shared>(1, allocator.get())); + flatten->SetIO(std::make_unique(allocator.get())); + + const auto bridge_value = + GENERATE(std::numeric_limits::max(), std::numeric_limits::quiet_NaN()); + std::vector vectors = { + 10.0F, bridge_value, bridge_value, bridge_value, bridge_value, 1.0F}; + std::vector ids = {0, 1, 2, 3, 4, 5}; + flatten->Train(vectors.data(), ids.size()); + flatten->BatchInsertVector(vectors.data(), ids.size(), ids.data()); + + auto graph = std::make_shared( + std::vector>{{1}, {2}, {3}, {4}, {5}, {}}); + auto pool = std::make_shared(1, allocator.get(), ids.size(), allocator.get()); + InnerSearchParam param; + param.ep = 0; + param.ef = 2; + param.topk = 2; + float query = 0.0F; + auto visited = pool->TakeOne(); + QueryContext* ctx = nullptr; + auto result = + BasicSearcher(common).Search(graph, flatten, visited, &query, param, LabelTablePtr{}, ctx); + pool->ReturnOne(visited); + + bool found_target = false; + while (not result->Empty()) { + REQUIRE_FALSE(std::isnan(result->Top().first)); + found_target = found_target or result->Top().second == 5; + result->Pop(); + } + REQUIRE(found_target); +} + TEST_CASE("Basic Usage for GraphDataCell (adapter of hnsw)", "[ut][GraphDataCell]") { uint32_t M = 32; uint32_t data_size = 1000; diff --git a/src/impl/searcher/parallel_searcher.cpp b/src/impl/searcher/parallel_searcher.cpp index 5d18e735d7..c54baf76a2 100644 --- a/src/impl/searcher/parallel_searcher.cpp +++ b/src/impl/searcher/parallel_searcher.cpp @@ -15,11 +15,13 @@ #include "parallel_searcher.h" +#include #include #include #include "datacell/flatten_interface.h" #include "impl/heap/standard_heap.h" +#include "impl/searcher/searcher_utils.h" #include "utils/filter_search_skip_strategy.h" #include "utils/spsc_queue.h" @@ -139,7 +141,8 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, inner_search_param.skip_ratio); flatten->Query(&dist, computer, &ep, 1, ctx); - if (not is_id_allowed || is_id_allowed->CheckValid(ep)) { + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(ep))) { top_candidates->Push(dist, ep); lower_bound = top_candidates->Top().first; } @@ -151,14 +154,15 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, if (dist < THRESHOLD_ERROR) { inner_search_param.duplicate_id = ep; } - candidate_set->Push(-dist, ep); + candidate_set->Push(traversal_priority(dist), ep); vl->Set(ep); if (inner_search_param.consider_duplicate && label_table && label_table->CompressDuplicateData()) { const auto& duplicate_ids = label_table->GetDuplicateId(ep); for (const auto& item : duplicate_ids) { - if (not is_id_allowed || is_id_allowed->CheckValid(item)) { + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(item))) { top_candidates->Push(dist, item); } } @@ -192,8 +196,10 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, } }; + std::vector> worker_tasks; + worker_tasks.reserve(num_threads); for (uint64_t i = 0; i < num_threads; i++) { - pool->GeneralEnqueue(task, i); + worker_tasks.emplace_back(pool->GeneralEnqueue(task, i)); } while (not candidate_set->Empty()) { @@ -264,10 +270,11 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, if (dist < THRESHOLD_ERROR) { inner_search_param.duplicate_id = to_be_visited_id[i]; } - if (top_candidates->Size() < ef || lower_bound > dist || + if (not std::isfinite(dist) || top_candidates->Size() < ef || lower_bound > dist || (mode == RANGE_SEARCH && dist <= inner_search_param.radius)) { - candidate_set->Push(-dist, to_be_visited_id[i]); - if ((not is_id_allowed || is_id_allowed->CheckValid(to_be_visited_id[i])) && + candidate_set->Push(traversal_priority(dist), to_be_visited_id[i]); + if (is_result_distance_eligible(dist) and + (not is_id_allowed || is_id_allowed->CheckValid(to_be_visited_id[i])) && dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, to_be_visited_id[i]); } @@ -275,7 +282,8 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, label_table->CompressDuplicateData()) { const auto& duplicate_ids = label_table->GetDuplicateId(to_be_visited_id[i]); for (const auto& item : duplicate_ids) { - if (dist > inner_search_param.min_distance + THRESHOLD_ERROR) { + if (is_result_distance_eligible(dist) and + dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, item); } } @@ -313,6 +321,9 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, for (uint64_t i = 0; i < num_threads; i++) { queues[i].Push({nullptr, nullptr, 0}); } + for (auto& worker_task : worker_tasks) { + worker_task.wait(); + } return top_candidates; } diff --git a/src/impl/searcher/parallel_searcher_test.cpp b/src/impl/searcher/parallel_searcher_test.cpp index be7f668bd1..2f4275a9d9 100644 --- a/src/impl/searcher/parallel_searcher_test.cpp +++ b/src/impl/searcher/parallel_searcher_test.cpp @@ -15,6 +15,10 @@ #include "parallel_searcher.h" +#include +#include +#include + #include "searcher_test.h" using namespace vsag; @@ -175,4 +179,55 @@ TEST_CASE("Parallel search with HNSW", "[ut][ParallelSearcher]") { } } } -} \ No newline at end of file +} + +TEST_CASE("ParallelSearcher traverses through a non-finite-distance bridge", + "[ut][ParallelSearcher][nonfinite]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + IndexCommonParam common; + common.dim_ = 1; + common.allocator_ = allocator; + common.metric_ = MetricType::METRIC_TYPE_L2SQR; + + constexpr const char* param_temp = R"({{"type": "{}"}})"; + auto quantizer_param = QuantizerParameter::GetQuantizerParameterByJson( + JsonType::Parse(fmt::format(param_temp, "fp32"))); + auto io_param = + IOParameter::GetIOParameterByJson(JsonType::Parse(fmt::format(param_temp, "memory_io"))); + auto flatten = + std::make_shared, MemoryIO>>( + quantizer_param, io_param, common); + flatten->SetQuantizer( + std::make_shared>(1, allocator.get())); + flatten->SetIO(std::make_unique(allocator.get())); + + const auto bridge_value = + GENERATE(std::numeric_limits::max(), std::numeric_limits::quiet_NaN()); + std::vector vectors = { + 10.0F, bridge_value, bridge_value, bridge_value, bridge_value, 1.0F}; + std::vector ids = {0, 1, 2, 3, 4, 5}; + flatten->Train(vectors.data(), ids.size()); + flatten->BatchInsertVector(vectors.data(), ids.size(), ids.data()); + + auto graph = std::make_shared( + std::vector>{{1}, {2}, {3}, {4}, {5}, {}}); + auto pool = std::make_shared(1, allocator.get(), ids.size(), allocator.get()); + InnerSearchParam param; + param.ep = 0; + param.ef = 2; + param.topk = 2; + param.parallel_search_thread_count = 2; + float query = 0.0F; + auto visited = pool->TakeOne(); + auto result = ParallelSearcher(common, SafeThreadPool::FactoryDefaultThreadPool()) + .Search(graph, flatten, visited, &query, param); + pool->ReturnOne(visited); + + bool found_target = false; + while (not result->Empty()) { + REQUIRE_FALSE(std::isnan(result->Top().first)); + found_target = found_target or result->Top().second == 5; + result->Pop(); + } + REQUIRE(found_target); +} diff --git a/src/impl/searcher/searcher_utils.h b/src/impl/searcher/searcher_utils.h new file mode 100644 index 0000000000..aa9dea54b6 --- /dev/null +++ b/src/impl/searcher/searcher_utils.h @@ -0,0 +1,34 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace vsag { + +inline bool +is_result_distance_eligible(float distance) { + return not std::isnan(distance); +} + +inline float +traversal_priority(float distance) { + // Non-finite nodes can still connect the entry point to valid candidates. Give them a + // deterministic priority so NaN does not violate the heap's ordering requirements. + return std::isfinite(distance) ? -distance : std::numeric_limits::max(); +} + +} // namespace vsag