diff --git a/docs/hgraph.md b/docs/hgraph.md index 5b01dd1dcb..8c3df9881a 100644 --- a/docs/hgraph.md +++ b/docs/hgraph.md @@ -225,10 +225,21 @@ means that the index uses PQ quantization with 64 subspaces, enables reordering - **Optional Values**: 1 to INT_MAX - **Default Value**: Must be provided (no default value) +### max_duplicates_per_group +- **Parameter Type**: int +- **Parameter Description**: Maximum number of additional duplicate IDs returned for each + duplicate group during KNN graph search +- **Optional Values**: -1 (unlimited), 0 (do not expand duplicate IDs), or any positive integer +- **Default Value**: -1 +- **Notes**: Only applies when `support_duplicate` is enabled; range search keeps its existing + expansion behavior + ## Examples for Search Parameter String ```json "hgraph": { - "ef_search": 200 + "ef_search": 200, + "max_duplicates_per_group": 2 } ``` -means that the search will use an ef_search value of 200 to control the search quality and performance trade-off. \ No newline at end of file +means that the search uses an ef_search value of 200 and returns at most two additional IDs from +each duplicate group. diff --git a/include/vsag/constants.h b/include/vsag/constants.h index 13ce467085..1777849610 100644 --- a/include/vsag/constants.h +++ b/include/vsag/constants.h @@ -184,6 +184,7 @@ extern const char* const HGRAPH_PRECISE_IO_TYPE; extern const char* const HGRAPH_PRECISE_FILE_PATH; extern const char* const HGRAPH_PARAMETER_EF_RUNTIME; extern const char* const HGRAPH_PARAMETER_HOPS_LIMIT; +extern const char* const HGRAPH_PARAMETER_MAX_DUPLICATES_PER_GROUP; extern const char* const HGRAPH_EXTRA_INFO_SIZE; extern const char* const HGRAPH_SUPPORT_DUPLICATE; extern const char* const HGRAPH_DUPLICATE_DISTANCE_THRESHOLD; diff --git a/src/algorithm/hgraph.cpp b/src/algorithm/hgraph.cpp index 283db02396..da9a911b40 100644 --- a/src/algorithm/hgraph.cpp +++ b/src/algorithm/hgraph.cpp @@ -933,10 +933,20 @@ HGraph::KnnSearch(const DatasetPtr& query, auto search_result = DistanceHeap::MakeInstanceBySize(ctx.alloc, k); const auto* query_data = get_data(query); if (is_last_filter) { + if (const auto* pending_duplicates = iter_filter_ctx->GetPendingDuplicates(); + pending_duplicates != nullptr) { + for (const auto& [pending_id, pending_dist] : *pending_duplicates) { + if (iter_filter_ctx->CheckPoint(pending_id)) { + search_result->Push(pending_dist, pending_id); + } + } + } while (!iter_filter_ctx->Empty()) { uint32_t cur_inner_id = iter_filter_ctx->GetTopID(); float cur_dist = iter_filter_ctx->GetTopDist(); - search_result->Push(cur_dist, cur_inner_id); + if (not iter_filter_ctx->IsPendingDuplicate(cur_inner_id)) { + search_result->Push(cur_dist, cur_inner_id); + } iter_filter_ctx->PopDiscard(); } } else { @@ -964,6 +974,7 @@ HGraph::KnnSearch(const DatasetPtr& query, search_param.is_inner_id_allowed = ft; search_param.topk = static_cast(search_param.ef); search_param.consider_duplicate = this->label_table_->CompressDuplicateData(); + search_param.max_duplicates_per_group = params.max_duplicates_per_group; search_param.parallel_search_thread_count = params.parallel_search_thread_count; search_param.min_distance = params.min_distance; @@ -1410,6 +1421,9 @@ HGraph::Serialize(StreamWriter& writer) const { if (this->use_attribute_filter_ and this->attr_filter_index_ != nullptr) { this->attr_filter_index_->Serialize(writer); } + if (this->label_table_->CompressDuplicateData()) { + this->label_table_->SerializeDuplicateRecords(writer); + } return; } @@ -1442,10 +1456,53 @@ HGraph::Serialize(StreamWriter& writer) const { footer->Write(writer); } +void +HGraph::Deserialize(std::istream& in_stream) { + if (not this->use_old_serial_format_) { + InnerIndexInterface::Deserialize(in_stream); + return; + } + + try { + uint64_t cursor = 0; + auto read_func = [&](uint64_t offset, uint64_t size, void* data) { + if (offset != cursor) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "v0.14 sequential stream does not support seek"); + } + in_stream.read(static_cast(data), static_cast(size)); + if (in_stream.gcount() != static_cast(size)) { + throw VsagException( + ErrorType::READ_ERROR, + fmt::format("Attempted to read: {} bytes. Remaining content size: {} bytes.", + size, + in_stream.gcount())); + } + cursor += size; + }; + ReadFuncStreamReader reader(read_func, 0, 0); + this->deserialize(reader, true); + if (reader.GetCursor() != cursor) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "v0.14 sequential stream does not support seek"); + } + } catch (const std::bad_alloc& e) { + throw VsagException(ErrorType::NO_ENOUGH_MEMORY, "failed to Deserialize: ", e.what()); + } +} + void HGraph::Deserialize(StreamReader& reader) { - // try to deserialize footer (only in new version) - auto footer = Footer::Parse(reader); + this->deserialize(reader, false); +} + +void +HGraph::deserialize(StreamReader& reader, bool force_v0_14) { + FooterPtr footer = nullptr; + if (not force_v0_14) { + // try to deserialize footer (only in new version) + footer = Footer::Parse(reader); + } if (footer == nullptr) { // old format, DON'T EDIT, remove in the future logger::debug("parse with v0.14 version format"); @@ -1474,6 +1531,11 @@ HGraph::Deserialize(StreamReader& reader) { if (this->use_attribute_filter_ and this->attr_filter_index_ != nullptr) { this->attr_filter_index_->Deserialize(reader); } + if (this->label_table_->CompressDuplicateData()) { + const auto logical_element_count = + static_cast(this->label_table_->GetTotalCount()); + this->label_table_->DeserializeDuplicateRecords(reader, logical_element_count); + } } else { // create like `else if ( ver in [v0.15, v0.17] )` here if need in the future logger::debug("parse with new version format"); @@ -2222,6 +2284,7 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { search_param.topk, static_cast(static_cast(k) * params.topk_factor)); } search_param.consider_duplicate = true; + search_param.max_duplicates_per_group = params.max_duplicates_per_group; if (params.enable_time_record) { search_param.time_cost = std::make_shared(); search_param.time_cost->SetThreshold(params.timeout_ms); diff --git a/src/algorithm/hgraph.h b/src/algorithm/hgraph.h index eb11c7a9fe..6075579ad1 100644 --- a/src/algorithm/hgraph.h +++ b/src/algorithm/hgraph.h @@ -15,6 +15,7 @@ #pragma once +#include #include #include #include @@ -79,6 +80,9 @@ class HGraph : public InnerIndexInterface { DatasetPtr CalDistanceById(const float* query, const int64_t* ids, int64_t count) const override; + void + Deserialize(std::istream& in_stream) override; + void Deserialize(StreamReader& reader) override; @@ -285,6 +289,9 @@ class HGraph : public InnerIndexInterface { void deserialize_label_info(StreamReader& reader) const; + void + deserialize(StreamReader& reader, bool force_v0_14); + // used in version [0.12.*, 0.14.*] void serialize_basic_info_v0_14(StreamWriter& writer) const; diff --git a/src/algorithm/hgraph_parameter.cpp b/src/algorithm/hgraph_parameter.cpp index e78a663750..03ebc64a3a 100644 --- a/src/algorithm/hgraph_parameter.cpp +++ b/src/algorithm/hgraph_parameter.cpp @@ -15,6 +15,8 @@ #include "hgraph_parameter.h" +#include + #include "datacell/extra_info_datacell_parameter.h" #include "datacell/flatten_datacell_parameter.h" #include "datacell/graph_datacell_parameter.h" @@ -223,6 +225,22 @@ HGraphSearchParameters::FromJson(const std::string& json_string) { if (params[INDEX_TYPE_HGRAPH].Contains("min_distance")) { obj.min_distance = params[INDEX_TYPE_HGRAPH]["min_distance"].GetFloat(); } + if (params[INDEX_TYPE_HGRAPH].Contains(HGRAPH_PARAMETER_MAX_DUPLICATES_PER_GROUP)) { + const auto& max_duplicates = + params[INDEX_TYPE_HGRAPH][HGRAPH_PARAMETER_MAX_DUPLICATES_PER_GROUP]; + CHECK_ARGUMENT(max_duplicates.IsNumberInteger(), + "max_duplicates_per_group must be an integer"); + const auto* max_duplicates_json = max_duplicates.GetInnerJson(); + if (max_duplicates_json->is_number_unsigned()) { + CHECK_ARGUMENT(max_duplicates_json->get() <= + static_cast(std::numeric_limits::max()), + "max_duplicates_per_group exceeds int64_t range"); + } + obj.max_duplicates_per_group = max_duplicates.GetInt(); + CHECK_ARGUMENT(obj.max_duplicates_per_group >= -1, + fmt::format("max_duplicates_per_group({}) must be >= -1", + obj.max_duplicates_per_group)); + } return obj; } diff --git a/src/algorithm/hgraph_parameter.h b/src/algorithm/hgraph_parameter.h index 213ba8b01b..94f30f8057 100644 --- a/src/algorithm/hgraph_parameter.h +++ b/src/algorithm/hgraph_parameter.h @@ -82,6 +82,9 @@ class HGraphSearchParameters : public IndexSearchParameter { bool use_reorder{false}; bool use_extra_info_filter{false}; float min_distance{std::numeric_limits::lowest()}; + // Maximum additional duplicate IDs returned per group during KNN graph search. + // -1 means unlimited; 0 disables duplicate expansion. + int64_t max_duplicates_per_group{-1}; private: HGraphSearchParameters() = default; diff --git a/src/algorithm/hgraph_parameter_test.cpp b/src/algorithm/hgraph_parameter_test.cpp index 949256f127..aed5342ff2 100644 --- a/src/algorithm/hgraph_parameter_test.cpp +++ b/src/algorithm/hgraph_parameter_test.cpp @@ -204,6 +204,24 @@ TEST_CASE("HGraph maps support_duplicate to graph parameter", "[ut][HGraphParame REQUIRE(typed_param->duplicate_distance_threshold == 0.25F); } +TEST_CASE("HGraph Search Parameters validate max_duplicates_per_group", + "[ut][HGraphSearchParameters][duplicate]") { + REQUIRE(vsag::HGraphSearchParameters::FromJson(R"({"hgraph": {"ef_search": 32}})") + .max_duplicates_per_group == -1); + REQUIRE(vsag::HGraphSearchParameters::FromJson( + R"({"hgraph": {"ef_search": 32, "max_duplicates_per_group": 0}})") + .max_duplicates_per_group == 0); + REQUIRE(vsag::HGraphSearchParameters::FromJson( + R"({"hgraph": {"ef_search": 32, "max_duplicates_per_group": 2}})") + .max_duplicates_per_group == 2); + REQUIRE_THROWS(vsag::HGraphSearchParameters::FromJson( + R"({"hgraph": {"ef_search": 32, "max_duplicates_per_group": -2}})")); + REQUIRE_THROWS(vsag::HGraphSearchParameters::FromJson( + R"({"hgraph": {"ef_search": 32, "max_duplicates_per_group": 1.5}})")); + REQUIRE_THROWS(vsag::HGraphSearchParameters::FromJson( + R"({"hgraph": {"ef_search": 32, "max_duplicates_per_group": 9223372036854775808}})")); +} + TEST_CASE("HGraph maps label_remap_type to inner index parameter", "[ut][HGraphParameter]") { auto param = vsag::JsonType::Parse(R"({ "base_quantization_type": "fp32", diff --git a/src/analyzer/hgraph_analyzer.cpp b/src/analyzer/hgraph_analyzer.cpp index e45225c6db..4669611ea2 100644 --- a/src/analyzer/hgraph_analyzer.cpp +++ b/src/analyzer/hgraph_analyzer.cpp @@ -508,6 +508,9 @@ HGraphAnalyzer::GetDegreeDistribution() { Vector in_degree(this->total_count_, allocator_); Vector out_degree(this->total_count_, allocator_); for (InnerIdType i = 0; i < this->total_count_; ++i) { + if (is_duplicate_ids_[i]) { + continue; + } Vector neighbors(allocator_); hgraph_->bottom_graph_->GetNeighbors(i, neighbors); out_degree[i] = neighbors.size(); diff --git a/src/constants.cpp b/src/constants.cpp index 4729adbe5c..42207c426c 100644 --- a/src/constants.cpp +++ b/src/constants.cpp @@ -164,6 +164,7 @@ const char* const HGRAPH_PRECISE_IO_TYPE = "precise_io_type"; const char* const HGRAPH_PRECISE_FILE_PATH = "precise_file_path"; const char* const HGRAPH_PARAMETER_EF_RUNTIME = "ef_search"; const char* const HGRAPH_PARAMETER_HOPS_LIMIT = "hops_limit"; +const char* const HGRAPH_PARAMETER_MAX_DUPLICATES_PER_GROUP = "max_duplicates_per_group"; const char* const HGRAPH_EXTRA_INFO_SIZE = "extra_info_size"; const char* const HGRAPH_SUPPORT_DUPLICATE = "support_duplicate"; const char* const HGRAPH_DUPLICATE_DISTANCE_THRESHOLD = "duplicate_distance_threshold"; diff --git a/src/impl/inner_search_param.h b/src/impl/inner_search_param.h index 664cd401e5..181e636e77 100644 --- a/src/impl/inner_search_param.h +++ b/src/impl/inner_search_param.h @@ -64,6 +64,7 @@ class InnerSearchParam { // use in search process with duplicate ids bool consider_duplicate{false}; + int64_t max_duplicates_per_group{-1}; // skip results with dist <= min_distance (for search iterator) float min_distance{std::numeric_limits::lowest()}; @@ -87,6 +88,8 @@ class InnerSearchParam { factor = other.factor; first_order_scan_ratio = other.first_order_scan_ratio; parallel_search_thread_count = other.parallel_search_thread_count; + consider_duplicate = other.consider_duplicate; + max_duplicates_per_group = other.max_duplicates_per_group; } return *this; } diff --git a/src/impl/label_table.h b/src/impl/label_table.h index 938f8d1fec..835cf2de26 100644 --- a/src/impl/label_table.h +++ b/src/impl/label_table.h @@ -25,6 +25,7 @@ #include #include #include +#include #include "storage/stream_reader.h" #include "storage/stream_writer.h" @@ -284,23 +285,28 @@ class LabelTable { Serialize(StreamWriter& writer) const { StreamWriter::WriteVector(writer, label_table_); if (compress_duplicate_data_) { - StreamWriter::WriteObj(writer, duplicate_count_); - for (InnerIdType i = 0; i < label_table_.size(); ++i) { - if (duplicate_records_[i] != nullptr) { - StreamWriter::WriteObj(writer, i); - Vector id_list(allocator_); - for (const auto& duplicate_id : duplicate_records_[i]->duplicate_ids) { - id_list.push_back(duplicate_id); - } - StreamWriter::WriteVector(writer, id_list); - } - } + this->SerializeDuplicateRecords(writer); } if (support_tombstone_) { StreamWriter::WriteObj(writer, deleted_ids_); } } + void + SerializeDuplicateRecords(StreamWriter& writer) const { + StreamWriter::WriteObj(writer, duplicate_count_); + for (InnerIdType i = 0; i < label_table_.size(); ++i) { + if (duplicate_records_[i] != nullptr) { + StreamWriter::WriteObj(writer, i); + Vector id_list(allocator_); + for (const auto& duplicate_id : duplicate_records_[i]->duplicate_ids) { + id_list.push_back(duplicate_id); + } + StreamWriter::WriteVector(writer, id_list); + } + } + } + void Deserialize(lvalue_or_rvalue reader) { StreamReader::ReadVector(reader, label_table_); @@ -312,18 +318,7 @@ class LabelTable { } } if (compress_duplicate_data_) { - StreamReader::ReadObj(reader, duplicate_count_); - duplicate_records_.resize(label_table_.size(), nullptr); - for (InnerIdType i = 0; i < duplicate_count_; ++i) { - InnerIdType id; - StreamReader::ReadObj(reader, id); - duplicate_records_[id] = allocator_->New(allocator_); - Vector id_list(allocator_); - StreamReader::ReadVector(reader, id_list); - for (const auto& duplicate_id : id_list) { - duplicate_records_[id]->duplicate_ids.insert(duplicate_id); - } - } + this->DeserializeDuplicateRecords(reader, label_table_.size()); } if (support_tombstone_) { StreamReader::ReadObj(reader, deleted_ids_); @@ -331,6 +326,99 @@ class LabelTable { this->total_count_.store(label_table_.size()); } + void + DeserializeDuplicateRecords(lvalue_or_rvalue reader, + uint64_t logical_element_count) { + if (logical_element_count > label_table_.size()) { + throw VsagException(ErrorType::INVALID_BINARY, + fmt::format("logical element count {} exceeds label capacity {}", + logical_element_count, + label_table_.size())); + } + + uint64_t duplicate_count = 0; + StreamReader::ReadObj(reader, duplicate_count); + if (duplicate_count > logical_element_count / 2) { + throw VsagException( + ErrorType::INVALID_BINARY, + fmt::format("duplicate group count {} exceeds logical element limit {}", + duplicate_count, + logical_element_count / 2)); + } + + std::vector>> duplicate_groups; + duplicate_groups.reserve(duplicate_count); + UnorderedSet assigned_ids(allocator_); + for (uint64_t i = 0; i < duplicate_count; ++i) { + InnerIdType id; + StreamReader::ReadObj(reader, id); + if (id >= logical_element_count) { + throw VsagException(ErrorType::INVALID_BINARY, + fmt::format("duplicate head id {} exceeds logical element " + "count {}", + id, + logical_element_count)); + } + if (not assigned_ids.insert(id).second) { + throw VsagException(ErrorType::INVALID_BINARY, + fmt::format("id {} belongs to multiple duplicate groups", id)); + } + + uint64_t member_count = 0; + StreamReader::ReadObj(reader, member_count); + const auto assigned_count = static_cast(assigned_ids.size()); + if (member_count == 0 or member_count > logical_element_count - assigned_count) { + throw VsagException( + ErrorType::INVALID_BINARY, + fmt::format("duplicate member count {} exceeds remaining logical element " + "count {}", + member_count, + logical_element_count - assigned_count)); + } + + Vector id_list(allocator_); + id_list.resize(member_count); + reader->Read(reinterpret_cast(id_list.data()), + member_count * sizeof(InnerIdType)); + for (const auto duplicate_id : id_list) { + if (duplicate_id >= logical_element_count) { + throw VsagException( + ErrorType::INVALID_BINARY, + fmt::format("duplicate member id {} exceeds logical element count {}", + duplicate_id, + logical_element_count)); + } + if (not assigned_ids.insert(duplicate_id).second) { + throw VsagException( + ErrorType::INVALID_BINARY, + fmt::format("id {} belongs to multiple duplicate groups", duplicate_id)); + } + } + duplicate_groups.emplace_back(id, std::move(id_list)); + } + + Vector restored_records(label_table_.size(), nullptr, allocator_); + try { + for (auto& [id, id_list] : duplicate_groups) { + restored_records[id] = allocator_->New(allocator_); + for (const auto& duplicate_id : id_list) { + restored_records[id]->duplicate_ids.insert(duplicate_id); + } + } + } catch (...) { + for (auto* record : restored_records) { + allocator_->Delete(record); + } + throw; + } + + for (auto* record : duplicate_records_) { + allocator_->Delete(record); + } + duplicate_records_ = std::move(restored_records); + duplicate_count_ = duplicate_count; + } + void Resize(uint64_t new_size) { if (new_size < total_count_) { diff --git a/src/impl/label_table_test.cpp b/src/impl/label_table_test.cpp index 1da273f05b..bae37db471 100644 --- a/src/impl/label_table_test.cpp +++ b/src/impl/label_table_test.cpp @@ -1,11 +1,105 @@ #include "label_table.h" #include +#include +#include +#include +#include +#include +#include #include "impl/allocator/default_allocator.h" +#include "storage/stream_reader.h" +#include "storage/stream_writer.h" +#include "vsag_exception.h" using namespace vsag; +namespace { + +using DuplicateGroup = std::pair>; + +class ForwardOnlyStreamReader final : public StreamReader { +public: + explicit ForwardOnlyStreamReader(std::string data) : data_(std::move(data)) { + } + + void + Read(char* data, uint64_t size) override { + if (cursor_ > data_.size() or size > data_.size() - cursor_) { + throw std::runtime_error("read exceeds stream boundary"); + } + std::memcpy(data, data_.data() + cursor_, size); + cursor_ += size; + } + + void + Seek(uint64_t cursor) override { + (void)cursor; + throw std::runtime_error("seek is not supported"); + } + + [[nodiscard]] uint64_t + GetCursor() const override { + return cursor_; + } + + [[nodiscard]] uint64_t + Length() override { + throw std::runtime_error("length is not supported"); + } + +private: + std::string data_; + uint64_t cursor_{0}; +}; + +std::stringstream +create_serialized_label_table(const std::vector& labels, + const std::vector& duplicate_groups) { + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + IOStreamWriter writer(stream); + StreamWriter::WriteVector(writer, labels); + const uint64_t duplicate_count = duplicate_groups.size(); + StreamWriter::WriteObj(writer, duplicate_count); + for (const auto& [head_id, duplicate_ids] : duplicate_groups) { + StreamWriter::WriteObj(writer, head_id); + StreamWriter::WriteVector(writer, duplicate_ids); + } + stream.seekg(0); + return stream; +} + +std::stringstream +create_serialized_duplicate_records(const std::vector& duplicate_groups) { + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + IOStreamWriter writer(stream); + const uint64_t duplicate_count = duplicate_groups.size(); + StreamWriter::WriteObj(writer, duplicate_count); + for (const auto& [head_id, duplicate_ids] : duplicate_groups) { + StreamWriter::WriteObj(writer, head_id); + StreamWriter::WriteVector(writer, duplicate_ids); + } + stream.seekg(0); + return stream; +} + +void +deserialize_label_table(LabelTable& label_table, std::stringstream& stream) { + IOStreamReader reader(stream); + label_table.Deserialize(reader); +} + +void +deserialize_duplicate_records(LabelTable& label_table, + std::stringstream& stream, + uint64_t logical_element_count) { + IOStreamReader reader(stream); + label_table.DeserializeDuplicateRecords(reader, logical_element_count); +} + +} // namespace + TEST_CASE("LabelTable Supports Configurable Remap Implementation", "[ut][LabelTable]") { auto allocator = std::make_shared(); @@ -27,3 +121,168 @@ TEST_CASE("LabelTable Supports Configurable Remap Implementation", "[ut][LabelTa REQUIRE(label_table.GetIdByLabel(100) == 0); } } + +TEST_CASE("LabelTable deserializes duplicate groups", "[ut][LabelTable][duplicate]") { + auto allocator = std::make_shared(); + LabelTable label_table(allocator.get(), true, true); + auto stream = create_serialized_label_table({10, 11, 12, 13, 14}, {{0, {1, 2}}, {3, {4}}}); + + deserialize_label_table(label_table, stream); + + const auto first_group = label_table.GetDuplicateId(0); + REQUIRE(first_group.size() == 2); + REQUIRE(first_group.contains(1)); + REQUIRE(first_group.contains(2)); + const auto second_group = label_table.GetDuplicateId(3); + REQUIRE(second_group.size() == 1); + REQUIRE(second_group.contains(4)); + REQUIRE(label_table.GetDuplicateId(1).empty()); + REQUIRE(label_table.GetDuplicateId(2).empty()); + REQUIRE(label_table.GetDuplicateId(4).empty()); +} + +TEST_CASE("LabelTable reads duplicate groups sequentially", "[ut][LabelTable][duplicate]") { + auto allocator = std::make_shared(); + LabelTable source(allocator.get(), true, true); + for (InnerIdType id = 0; id < 5; ++id) { + source.Insert(id, 10 + id); + } + source.Resize(5); + source.SetDuplicateId(0, 1); + source.SetDuplicateId(0, 2); + source.SetDuplicateId(3, 4); + + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + IOStreamWriter writer(stream); + source.Serialize(writer); + + ForwardOnlyStreamReader reader(stream.str()); + LabelTable restored(allocator.get(), true, true); + REQUIRE_NOTHROW(restored.Deserialize(reader)); + REQUIRE(reader.GetCursor() == stream.str().size()); + + const auto first_group = restored.GetDuplicateId(0); + REQUIRE(first_group.size() == 2); + REQUIRE(first_group.contains(1)); + REQUIRE(first_group.contains(2)); + const auto second_group = restored.GetDuplicateId(3); + REQUIRE(second_group.size() == 1); + REQUIRE(second_group.contains(4)); +} + +TEST_CASE("LabelTable rejects out-of-range duplicate members", + "[ut][LabelTable][duplicate][invalid-member]") { + auto allocator = std::make_shared(); + LabelTable label_table(allocator.get(), true, true); + auto stream = create_serialized_label_table({10, 11, 12}, {{0, {3}}}); + + REQUIRE_THROWS_AS(deserialize_label_table(label_table, stream), VsagException); +} + +TEST_CASE("LabelTable rejects out-of-range duplicate heads", + "[ut][LabelTable][duplicate][invalid-head]") { + auto allocator = std::make_shared(); + LabelTable label_table(allocator.get(), true, true); + auto stream = create_serialized_label_table({10, 11, 12}, {{3, {1}}}); + + REQUIRE_THROWS_AS(deserialize_label_table(label_table, stream), VsagException); +} + +TEST_CASE("LabelTable rejects overlapping duplicate groups", + "[ut][LabelTable][duplicate][overlap]") { + auto allocator = std::make_shared(); + + const auto require_invalid = [&](const std::vector& groups) { + LabelTable label_table(allocator.get(), true, true); + auto stream = create_serialized_label_table({10, 11, 12, 13, 14, 15}, groups); + REQUIRE_THROWS_AS(deserialize_label_table(label_table, stream), VsagException); + }; + + SECTION("head is repeated") { + require_invalid({{0, {1}}, {0, {2}}}); + } + + SECTION("member is repeated in one group") { + require_invalid({{0, {1, 1}}}); + } + + SECTION("member is shared by groups") { + require_invalid({{0, {1}}, {2, {1}}}); + } + + SECTION("member later becomes a head") { + require_invalid({{0, {1}}, {1, {2}}}); + } + + SECTION("head later becomes a member") { + require_invalid({{0, {1}}, {2, {0}}}); + } + + SECTION("head contains itself") { + require_invalid({{0, {0}}}); + } +} + +TEST_CASE("LabelTable validates duplicate records against logical count", + "[ut][LabelTable][duplicate][logical-count]") { + auto allocator = std::make_shared(); + LabelTable label_table(allocator.get(), true, true); + for (InnerIdType id = 0; id < 3; ++id) { + label_table.Insert(id, 10 + id); + } + label_table.Resize(1024); + + SECTION("head within capacity but outside logical elements") { + auto stream = create_serialized_duplicate_records({{3, {1}}}); + REQUIRE_THROWS_AS(deserialize_duplicate_records(label_table, stream, 3), VsagException); + } + + SECTION("member within capacity but outside logical elements") { + auto stream = create_serialized_duplicate_records({{0, {3}}}); + REQUIRE_THROWS_AS(deserialize_duplicate_records(label_table, stream, 3), VsagException); + } +} + +TEST_CASE("LabelTable validates duplicate record counts before allocation", + "[ut][LabelTable][duplicate][invalid-count]") { + auto allocator = std::make_shared(); + LabelTable label_table(allocator.get(), true, true); + for (InnerIdType id = 0; id < 5; ++id) { + label_table.Insert(id, 10 + id); + } + label_table.Resize(1024); + + SECTION("group count") { + auto stream = create_serialized_duplicate_records({{0, {1}}, {2, {3}}, {4, {}}}); + REQUIRE_THROWS_AS(deserialize_duplicate_records(label_table, stream, 5), VsagException); + } + + SECTION("empty group") { + auto stream = create_serialized_duplicate_records({{0, {}}}); + REQUIRE_THROWS_AS(deserialize_duplicate_records(label_table, stream, 5), VsagException); + } + + SECTION("member count exceeds remaining logical elements") { + auto stream = create_serialized_duplicate_records({{0, {1, 2}}, {3, {4, 4}}}); + REQUIRE_THROWS_AS(deserialize_duplicate_records(label_table, stream, 5), VsagException); + } +} + +TEST_CASE("LabelTable keeps duplicate records unchanged after invalid input", + "[ut][LabelTable][duplicate][atomic-publish]") { + auto allocator = std::make_shared(); + LabelTable label_table(allocator.get(), true, true); + for (InnerIdType id = 0; id < 4; ++id) { + label_table.Insert(id, 10 + id); + } + label_table.Resize(1024); + label_table.SetDuplicateId(0, 1); + + auto stream = create_serialized_duplicate_records({{2, {4}}}); + REQUIRE_THROWS_AS(deserialize_duplicate_records(label_table, stream, 4), VsagException); + + REQUIRE(label_table.duplicate_count_ == 1); + const auto duplicate_ids = label_table.GetDuplicateId(0); + REQUIRE(duplicate_ids.size() == 1); + REQUIRE(duplicate_ids.contains(1)); +} diff --git a/src/impl/searcher/basic_searcher.cpp b/src/impl/searcher/basic_searcher.cpp index c56244f30a..b7650e59c0 100644 --- a/src/impl/searcher/basic_searcher.cpp +++ b/src/impl/searcher/basic_searcher.cpp @@ -139,13 +139,66 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, : 1.0F, inner_search_param.skip_ratio); + auto add_pending_duplicates = [&](float duplicate_dist, InnerIdType group_head_id) { + if constexpr (mode != InnerSearchMode::KNN_SEARCH) { + return; + } + if (not inner_search_param.consider_duplicate or + inner_search_param.max_duplicates_per_group == 0 or label_table == nullptr or + not label_table->CompressDuplicateData() or + duplicate_dist <= inner_search_param.min_distance + THRESHOLD_ERROR) { + return; + } + + int64_t duplicate_count = 0; + for (const auto duplicate_id : label_table->GetDuplicateId(group_head_id)) { + if (inner_search_param.max_duplicates_per_group >= 0 and + duplicate_count >= inner_search_param.max_duplicates_per_group) { + break; + } + if (is_id_allowed != nullptr and not is_id_allowed->CheckValid(duplicate_id)) { + continue; + } + + ++duplicate_count; + if (not iter_ctx->CheckPoint(duplicate_id) or + not iter_ctx->AddPendingDuplicate(duplicate_dist, duplicate_id)) { + continue; + } + top_candidates->Push(duplicate_dist, duplicate_id); + } + }; + + auto shrink_top_candidates = [&](uint64_t limit) { + while (top_candidates->Size() > limit) { + const auto current = top_candidates->Top(); + if (iter_ctx->CheckPoint(current.second)) { + iter_ctx->AddDiscardNode(current.first, current.second); + } + top_candidates->Pop(); + } + }; + if (!iter_ctx->IsFirstUsed()) { - if (iter_ctx->Empty()) { + if (iter_ctx->Empty() and iter_ctx->GetPendingDuplicateElementNum() == 0) { return top_candidates; } + + if (const auto* pending_duplicates = iter_ctx->GetPendingDuplicates(); + pending_duplicates != nullptr) { + for (const auto& [pending_id, pending_dist] : *pending_duplicates) { + if (iter_ctx->CheckPoint(pending_id)) { + top_candidates->Push(pending_dist, pending_id); + } + } + } while (!iter_ctx->Empty()) { uint32_t cur_inner_id = iter_ctx->GetTopID(); float cur_dist = iter_ctx->GetTopDist(); + if (iter_ctx->IsPendingDuplicate(cur_inner_id)) { + iter_ctx->PopDiscard(); + continue; + } vl->Set(cur_inner_id); if (iter_ctx->CheckPoint(cur_inner_id)) { lower_bound = std::max(lower_bound, cur_dist); @@ -153,7 +206,8 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, if (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); + add_pending_duplicates(cur_dist, cur_inner_id); + candidate_set->Push(-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(); @@ -162,39 +216,27 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, } iter_ctx->PopDiscard(); } + if constexpr (mode == KNN_SEARCH) { + shrink_top_candidates(ef); + } + if (not top_candidates->Empty()) { + lower_bound = top_candidates->Top().first; + } } else { flatten->Query(&dist, computer, &ep, 1, ctx); if ((not is_id_allowed || is_id_allowed->CheckValid(ep)) and !(dist <= inner_search_param.min_distance + THRESHOLD_ERROR)) { top_candidates->Push(dist, ep); + } + add_pending_duplicates(dist, ep); + if constexpr (mode == KNN_SEARCH) { + shrink_top_candidates(ef); + } + if (not top_candidates->Empty()) { lower_bound = top_candidates->Top().first; } candidate_set->Push(-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 - iter_ctx->CheckPoint(item) and - dist > inner_search_param.min_distance + THRESHOLD_ERROR) { - top_candidates->Push(dist, item); - } - } - if constexpr (mode == KNN_SEARCH) { - if (top_candidates->Size() > ef) { - if (iter_ctx->CheckPoint(top_candidates->Top().second)) { - auto cur_node_pair = top_candidates->Top(); - iter_ctx->AddDiscardNode(cur_node_pair.first, cur_node_pair.second); - } - top_candidates->Pop(); - } - } - if (not top_candidates->Empty()) { - lower_bound = top_candidates->Top().first; - } - } } while (not candidate_set->Empty()) { @@ -242,25 +284,10 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, top_candidates->Push(dist, to_be_visited_id[i]); } - if (inner_search_param.consider_duplicate and label_table != nullptr and - 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 - iter_ctx->CheckPoint(item)) { - top_candidates->Push(dist, item); - } - } - } + add_pending_duplicates(dist, to_be_visited_id[i]); if constexpr (mode == KNN_SEARCH) { - if (top_candidates->Size() > ef) { - if (iter_ctx->CheckPoint(top_candidates->Top().second)) { - auto cur_node_pair = top_candidates->Top(); - iter_ctx->AddDiscardNode(cur_node_pair.first, cur_node_pair.second); - } - top_candidates->Pop(); - } + shrink_top_candidates(ef); } if (not top_candidates->Empty()) { @@ -271,13 +298,7 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, } if constexpr (mode == KNN_SEARCH) { - while (top_candidates->Size() > inner_search_param.topk) { - auto cur_node_pair = top_candidates->Top(); - if (iter_ctx->CheckPoint(cur_node_pair.second)) { - iter_ctx->AddDiscardNode(cur_node_pair.first, cur_node_pair.second); - } - top_candidates->Pop(); - } + shrink_top_candidates(inner_search_param.topk); } return top_candidates; @@ -336,6 +357,33 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, (attr_ft == nullptr or attr_ft->CheckValid(id)); }; + auto add_duplicate_results = [&](float duplicate_dist, InnerIdType group_head_id) { + if (not inner_search_param.consider_duplicate or label_table == nullptr or + not label_table->CompressDuplicateData() or + duplicate_dist <= inner_search_param.min_distance + THRESHOLD_ERROR) { + return; + } + if constexpr (mode == InnerSearchMode::KNN_SEARCH) { + if (inner_search_param.max_duplicates_per_group == 0) { + return; + } + } + + int64_t duplicate_count = 0; + for (const auto duplicate_id : label_table->GetDuplicateId(group_head_id)) { + if constexpr (mode == InnerSearchMode::KNN_SEARCH) { + if (inner_search_param.max_duplicates_per_group >= 0 and + duplicate_count >= inner_search_param.max_duplicates_per_group) { + break; + } + } + if (check_func(duplicate_id)) { + top_candidates->Push(duplicate_dist, duplicate_id); + ++duplicate_count; + } + } + }; + flatten->Query(&dist, computer, &ep, 1, ctx); ++dist_cmp; if (check_func(ep) && !(dist <= inner_search_param.min_distance + THRESHOLD_ERROR)) { @@ -350,23 +398,15 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, candidate_set->Push(-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) { - top_candidates->Push(dist, item); - } - } - if constexpr (mode == KNN_SEARCH) { - if (top_candidates->Size() > ef) { - top_candidates->Pop(); - } - } - if (not top_candidates->Empty()) { - lower_bound = top_candidates->Top().first; + add_duplicate_results(dist, ep); + if constexpr (mode == KNN_SEARCH) { + while (top_candidates->Size() > ef) { + top_candidates->Pop(); } } + if (not top_candidates->Empty()) { + lower_bound = top_candidates->Top().first; + } while (not candidate_set->Empty()) { ++hops; @@ -416,19 +456,10 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, to_be_visited_id[i]); } - if (inner_search_param.consider_duplicate and label_table != nullptr and - 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) && - dist > inner_search_param.min_distance + THRESHOLD_ERROR) { - top_candidates->Push(dist, item); - } - } - } + add_duplicate_results(dist, to_be_visited_id[i]); if constexpr (mode == KNN_SEARCH) { - if (top_candidates->Size() > ef) { + while (top_candidates->Size() > ef) { top_candidates->Pop(); } } diff --git a/src/impl/searcher/parallel_searcher.cpp b/src/impl/searcher/parallel_searcher.cpp index 5d18e735d7..67ca6934da 100644 --- a/src/impl/searcher/parallel_searcher.cpp +++ b/src/impl/searcher/parallel_searcher.cpp @@ -15,6 +15,7 @@ #include "parallel_searcher.h" +#include #include #include @@ -138,8 +139,34 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, : 1.0F, inner_search_param.skip_ratio); + auto add_knn_duplicate_results = [&](float duplicate_dist, InnerIdType group_head_id) { + if (not inner_search_param.consider_duplicate or label_table == nullptr or + not label_table->CompressDuplicateData() or + inner_search_param.max_duplicates_per_group == 0 or + duplicate_dist <= inner_search_param.min_distance + THRESHOLD_ERROR) { + return; + } + + int64_t duplicate_count = 0; + for (const auto duplicate_id : label_table->GetDuplicateId(group_head_id)) { + if (inner_search_param.max_duplicates_per_group >= 0 and + duplicate_count >= inner_search_param.max_duplicates_per_group) { + break; + } + if (is_id_allowed == nullptr or is_id_allowed->CheckValid(duplicate_id)) { + top_candidates->Push(duplicate_dist, duplicate_id); + ++duplicate_count; + } + } + }; + flatten->Query(&dist, computer, &ep, 1, ctx); - if (not is_id_allowed || is_id_allowed->CheckValid(ep)) { + bool entry_point_allowed = not is_id_allowed or is_id_allowed->CheckValid(ep); + if constexpr (mode == InnerSearchMode::KNN_SEARCH) { + entry_point_allowed = + entry_point_allowed and dist > inner_search_param.min_distance + THRESHOLD_ERROR; + } + if (entry_point_allowed) { top_candidates->Push(dist, ep); lower_bound = top_candidates->Top().first; } @@ -154,22 +181,23 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, candidate_set->Push(-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)) { - top_candidates->Push(dist, item); - } + if constexpr (mode == KNN_SEARCH) { + add_knn_duplicate_results(dist, ep); + while (top_candidates->Size() > ef) { + top_candidates->Pop(); } - if constexpr (mode == KNN_SEARCH) { - if (top_candidates->Size() > ef) { - top_candidates->Pop(); + } else if constexpr (mode == RANGE_SEARCH) { + if (inner_search_param.consider_duplicate and label_table != nullptr and + label_table->CompressDuplicateData()) { + for (const auto duplicate_id : label_table->GetDuplicateId(ep)) { + if (is_id_allowed == nullptr or is_id_allowed->CheckValid(duplicate_id)) { + top_candidates->Push(dist, duplicate_id); + } } } - if (not top_candidates->Empty()) { - lower_bound = top_candidates->Top().first; - } + } + if (not top_candidates->Empty()) { + lower_bound = top_candidates->Top().first; } auto num_threads = inner_search_param.parallel_search_thread_count - 1; @@ -192,8 +220,10 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, } }; + std::vector> futures; + futures.reserve(num_threads); for (uint64_t i = 0; i < num_threads; i++) { - pool->GeneralEnqueue(task, i); + futures.emplace_back(pool->GeneralEnqueue(task, i)); } while (not candidate_set->Empty()) { @@ -271,20 +301,21 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, dist > inner_search_param.min_distance + THRESHOLD_ERROR) { top_candidates->Push(dist, to_be_visited_id[i]); } - if (inner_search_param.consider_duplicate && label_table && - 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) { - top_candidates->Push(dist, item); - } - } - } - if constexpr (mode == KNN_SEARCH) { - if (top_candidates->Size() > ef) { + add_knn_duplicate_results(dist, to_be_visited_id[i]); + while (top_candidates->Size() > ef) { top_candidates->Pop(); } + } else if constexpr (mode == RANGE_SEARCH) { + if (inner_search_param.consider_duplicate and label_table != nullptr and + label_table->CompressDuplicateData()) { + for (const auto duplicate_id : + label_table->GetDuplicateId(to_be_visited_id[i])) { + if (dist > inner_search_param.min_distance + THRESHOLD_ERROR) { + top_candidates->Push(dist, duplicate_id); + } + } + } } if (not top_candidates->Empty()) { @@ -313,7 +344,9 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, for (uint64_t i = 0; i < num_threads; i++) { queues[i].Push({nullptr, nullptr, 0}); } - + for (auto& future : futures) { + future.get(); + } return top_candidates; } diff --git a/src/index/iterator_filter.cpp b/src/index/iterator_filter.cpp index 7b10c029e7..f52d4c0944 100644 --- a/src/index/iterator_filter.cpp +++ b/src/index/iterator_filter.cpp @@ -52,6 +52,11 @@ IteratorFilterContext::init(InnerIdType max_size, int64_t ef_search, Allocator* void IteratorFilterContext::AddDiscardNode(float dis, uint32_t inner_id) { + if (this->IsPendingDuplicate(inner_id)) { + // Pending duplicates retain the graph-search distance. Reorder may evict the same ID with + // a refined distance, which must not change the next graph search's lower bound. + return; + } if (discard_->size() >= ef_search_) { if (discard_->top().first > dis) { discard_->pop(); @@ -62,6 +67,29 @@ IteratorFilterContext::AddDiscardNode(float dis, uint32_t inner_id) { } } +bool +IteratorFilterContext::AddPendingDuplicate(float dis, InnerIdType inner_id) { + if (pending_duplicates_ == nullptr) { + pending_duplicates_ = std::make_unique>(allocator_); + } + return pending_duplicates_->try_emplace(inner_id, dis).second; +} + +bool +IteratorFilterContext::IsPendingDuplicate(InnerIdType inner_id) const { + return pending_duplicates_ != nullptr && pending_duplicates_->contains(inner_id); +} + +const UnorderedMap* +IteratorFilterContext::GetPendingDuplicates() const { + return pending_duplicates_.get(); +} + +uint64_t +IteratorFilterContext::GetPendingDuplicateElementNum() const { + return pending_duplicates_ == nullptr ? 0 : pending_duplicates_->size(); +} + uint32_t IteratorFilterContext::GetTopID() { return discard_->top().second; @@ -94,6 +122,9 @@ IteratorFilterContext::SetOFFFirstUsed() { void IteratorFilterContext::SetPoint(InnerIdType inner_id) { + if (pending_duplicates_ != nullptr) { + pending_duplicates_->erase(inner_id); + } if (inner_id >= max_size_) { return; } diff --git a/src/index/iterator_filter.h b/src/index/iterator_filter.h index 33e136d21d..218508acd9 100644 --- a/src/index/iterator_filter.h +++ b/src/index/iterator_filter.h @@ -15,6 +15,7 @@ #pragma once +#include #include #include "typing.h" @@ -36,6 +37,18 @@ class IteratorFilterContext : public IteratorContext { void AddDiscardNode(float dis, uint32_t inner_id); + bool + AddPendingDuplicate(float dis, InnerIdType inner_id); + + bool + IsPendingDuplicate(InnerIdType inner_id) const; + + const UnorderedMap* + GetPendingDuplicates() const; + + uint64_t + GetPendingDuplicateElementNum() const; + uint32_t GetTopID(); @@ -84,6 +97,7 @@ class IteratorFilterContext : public IteratorContext { Allocator* allocator_{nullptr}; uint8_t* list_{nullptr}; std::unique_ptr discard_; + std::unique_ptr> pending_duplicates_; }; }; // namespace vsag diff --git a/src/index/iterator_filter_test.cpp b/src/index/iterator_filter_test.cpp index 7bf12e660c..c3687df445 100644 --- a/src/index/iterator_filter_test.cpp +++ b/src/index/iterator_filter_test.cpp @@ -83,3 +83,25 @@ TEST_CASE("Iterator Context CheckPoint And SetPoint", "[ut][hnsw][filter]") { filter_context.SetPoint(3); REQUIRE_FALSE(filter_context.CheckPoint(3)); } + +TEST_CASE("Iterator Context keeps pending duplicates out of graph discard", + "[ut][hnsw][filter][duplicate]") { + auto allocator = std::make_shared(); + IteratorFilterContext filter_context; + REQUIRE(filter_context.init(100, 10, allocator.get()).has_value()); + + REQUIRE(filter_context.AddPendingDuplicate(0.25F, 7)); + REQUIRE_FALSE(filter_context.AddPendingDuplicate(0.75F, 7)); + REQUIRE(filter_context.IsPendingDuplicate(7)); + REQUIRE(filter_context.GetPendingDuplicateElementNum() == 1); + + filter_context.AddDiscardNode(0.5F, 7); + REQUIRE(filter_context.Empty()); + REQUIRE(filter_context.GetPendingDuplicateElementNum() == 1); + REQUIRE(filter_context.GetPendingDuplicates()->at(7) == 0.25F); + + filter_context.SetPoint(7); + REQUIRE_FALSE(filter_context.IsPendingDuplicate(7)); + REQUIRE(filter_context.GetPendingDuplicateElementNum() == 0); + REQUIRE_FALSE(filter_context.CheckPoint(7)); +} diff --git a/tests/test_hgraph.cpp b/tests/test_hgraph.cpp index 3a51208114..c66e84a432 100644 --- a/tests/test_hgraph.cpp +++ b/tests/test_hgraph.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include "algorithm/hgraph.h" @@ -358,6 +359,40 @@ HGraphTestIndex::TestMemoryUsageDetail(const IndexPtr& index) { } } // namespace fixtures +namespace { + +class ForwardOnlyStringBuffer final : public std::stringbuf { +public: + explicit ForwardOnlyStringBuffer(const std::string& data) + : std::stringbuf(data, std::ios::in | std::ios::binary) { + } + + [[nodiscard]] uint64_t + GetSeekCount() const { + return seek_count_; + } + +protected: + pos_type + seekoff(off_type, + std::ios_base::seekdir, + std::ios_base::openmode = std::ios_base::in | std::ios_base::out) override { + ++seek_count_; + return pos_type(off_type(-1)); + } + + pos_type + seekpos(pos_type, std::ios_base::openmode = std::ios_base::in | std::ios_base::out) override { + ++seek_count_; + return pos_type(off_type(-1)); + } + +private: + uint64_t seek_count_{0}; +}; + +} // namespace + TEST_CASE_PERSISTENT_FIXTURE(fixtures::HGraphTestIndex, "HGraph Factory Test With Exceptions", "[ft][hgraph]") { @@ -1635,6 +1670,142 @@ TEST_CASE("HGraph Deserialize Old Format With Duplicate Support", vsag::Options::Instance().set_block_size_limit(origin_size); } +TEST_CASE("HGraph old format supports forward-only streams", "[ft][hgraph][serialization][pr]") { + constexpr int64_t dim = 8; + const bool support_duplicate = GENERATE(false, true); + const auto build_param = fmt::format( + R"({{ + "dtype": "float32", + "metric_type": "l2", + "dim": {}, + "use_old_serial_format": true, + "index_param": {{ + "max_degree": 16, + "ef_construction": 100, + "base_quantization_type": "fp32", + "build_thread_count": 1, + "support_duplicate": {} + }} + }})", + dim, + support_duplicate); + + auto index_result = vsag::Factory::CreateIndex("hgraph", build_param); + REQUIRE(index_result.has_value()); + auto index = index_result.value(); + + std::vector ids{0, 1, 2, 3, 4, 5, 6, 7}; + std::vector vectors(ids.size() * dim, 0.0F); + for (uint64_t i = 4; i < ids.size(); ++i) { + std::fill(vectors.begin() + i * dim, vectors.begin() + (i + 1) * dim, i * 10.0F); + } + auto base = vsag::Dataset::Make(); + base->NumElements(ids.size()) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); + REQUIRE(index->Build(base).has_value()); + + std::ostringstream serialized(std::ios::out | std::ios::binary); + REQUIRE(index->Serialize(serialized).has_value()); + + auto restored_result = vsag::Factory::CreateIndex("hgraph", build_param); + REQUIRE(restored_result.has_value()); + ForwardOnlyStringBuffer buffer(serialized.str()); + std::istream stream(&buffer); + REQUIRE(restored_result.value()->Deserialize(stream).has_value()); + REQUIRE(buffer.GetSeekCount() == 0); + REQUIRE(restored_result.value()->GetNumElements() == ids.size()); + if (support_duplicate) { + auto impl = + std::dynamic_pointer_cast>(restored_result.value()); + REQUIRE(impl != nullptr); + auto hgraph = std::dynamic_pointer_cast(impl->GetInnerIndex()); + REQUIRE(hgraph != nullptr); + REQUIRE(hgraph->label_table_->duplicate_count_ == 1); + REQUIRE(hgraph->label_table_->duplicate_records_[0] != nullptr); + REQUIRE(hgraph->label_table_->duplicate_records_[0]->duplicate_ids.size() == 3); + REQUIRE_NOTHROW(static_cast(restored_result.value()->GetStats())); + } +} + +TEST_CASE("HGraph Old Format Preserves Duplicate Records", "[ft][hgraph][serialization][pr]") { + auto origin_size = vsag::Options::Instance().block_size_limit(); + vsag::Options::Instance().set_block_size_limit(1024 * 1024 * 2); + + constexpr int64_t dim = 8; + constexpr const char* build_param = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 8, + "use_old_serial_format": true, + "index_param": { + "max_degree": 16, + "ef_construction": 100, + "base_quantization_type": "fp32", + "build_thread_count": 1, + "support_duplicate": true + } + })"; + + auto index_result = vsag::Factory::CreateIndex("hgraph", build_param); + REQUIRE(index_result.has_value()); + auto index = index_result.value(); + + std::vector ids{0, 100, 101, 102, 4, 5, 6, 7}; + std::vector vectors(ids.size() * dim, 0.0F); + for (uint64_t i = 4; i < ids.size(); ++i) { + std::fill(vectors.begin() + i * dim, vectors.begin() + (i + 1) * dim, i * 10.0F); + } + auto base = vsag::Dataset::Make(); + base->NumElements(ids.size()) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); + REQUIRE(index->Build(base).has_value()); + + std::vector query_vector(dim, 0.0F); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector.data())->Owner(false); + constexpr const char* search_param = R"({"hgraph":{"ef_search":32}})"; + const auto search_duplicate_group = [&](const vsag::IndexPtr& target) { + auto result = target->KnnSearch(query, 4, search_param); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 4); + std::vector result_ids(result.value()->GetIds(), result.value()->GetIds() + 4); + std::sort(result_ids.begin(), result_ids.end()); + for (int64_t i = 0; i < 4; ++i) { + REQUIRE(result.value()->GetDistances()[i] == 0.0F); + } + return result_ids; + }; + const std::vector expected_ids{0, 100, 101, 102}; + REQUIRE(search_duplicate_group(index) == expected_ids); + + auto serialized = index->Serialize(); + REQUIRE(serialized.has_value()); + auto reloaded_result = vsag::Factory::CreateIndex("hgraph", build_param); + REQUIRE(reloaded_result.has_value()); + auto reloaded = reloaded_result.value(); + REQUIRE(reloaded->Deserialize(serialized.value()).has_value()); + + auto impl = std::dynamic_pointer_cast>(reloaded); + REQUIRE(impl != nullptr); + auto hgraph = std::dynamic_pointer_cast(impl->GetInnerIndex()); + REQUIRE(hgraph != nullptr); + REQUIRE(hgraph->label_table_->duplicate_count_ == 1); + const auto duplicate_ids = hgraph->label_table_->GetDuplicateId(0); + REQUIRE(duplicate_ids.size() == 3); + REQUIRE(duplicate_ids.contains(1)); + REQUIRE(duplicate_ids.contains(2)); + REQUIRE(duplicate_ids.contains(3)); + REQUIRE(search_duplicate_group(reloaded) == expected_ids); + + vsag::Options::Instance().set_block_size_limit(origin_size); +} + static void TestHGraphSearchWithDirtyVector(const fixtures::HGraphTestIndexPtr& test_index, const fixtures::HGraphResourcePtr& resource) { diff --git a/tests/test_hgraph_dedup_search_control.cpp b/tests/test_hgraph_dedup_search_control.cpp new file mode 100644 index 0000000000..f77134c3d4 --- /dev/null +++ b/tests/test_hgraph_dedup_search_control.cpp @@ -0,0 +1,560 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "vsag/vsag.h" + +namespace { + +constexpr int64_t DIM = 32; + +std::string +MakeBuildParam(bool support_duplicate, float duplicate_threshold = 0.0F) { + return fmt::format(R"({{ + "dtype": "float32", + "metric_type": "l2", + "dim": {}, + "index_param": {{ + "base_quantization_type": "fp32", + "graph_type": "nsw", + "max_degree": 24, + "ef_construction": 100, + "build_thread_count": 4, + "support_duplicate": {}, + "duplicate_distance_threshold": {} + }} + }})", + DIM, + support_duplicate ? "true" : "false", + duplicate_threshold); +} + +std::string +MakeSearchParam(int64_t ef_search = 100, + int64_t max_duplicates_per_group = -1, + int64_t parallelism = 1) { + return fmt::format(R"({{ + "hgraph": {{ + "ef_search": {}, + "parallelism": {}, + "max_duplicates_per_group": {} + }} + }})", + ef_search, + parallelism, + max_duplicates_per_group); +} + +struct TestVectors { + std::vector base; + std::vector base_ids; + std::vector duplicates; + std::vector duplicate_ids; + std::vector queries; +}; + +TestVectors +GenerateTestData(int64_t dim, int64_t base_count, int64_t duplicate_count, uint32_t seed = 42) { + TestVectors vectors; + std::mt19937 random(seed); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + + vectors.base.resize(base_count * dim); + vectors.base_ids.resize(base_count); + for (int64_t i = 0; i < base_count; ++i) { + for (int64_t d = 0; d < dim; ++d) { + vectors.base[i * dim + d] = distribution(random); + } + vectors.base_ids[i] = i; + } + + vectors.duplicates.resize(duplicate_count * dim); + vectors.duplicate_ids.resize(duplicate_count); + for (int64_t i = 0; i < duplicate_count; ++i) { + const int64_t source = i % base_count; + std::memcpy(vectors.duplicates.data() + i * dim, + vectors.base.data() + source * dim, + dim * sizeof(float)); + vectors.duplicate_ids[i] = base_count + i; + } + + vectors.queries.resize(10 * dim); + for (int64_t i = 0; i < 10; ++i) { + const int64_t source = i % base_count; + std::memcpy(vectors.queries.data() + i * dim, + vectors.base.data() + source * dim, + dim * sizeof(float)); + } + + return vectors; +} + +vsag::IndexPtr +BuildIndexWithDuplicates(const TestVectors& vectors, const std::string& build_param) { + auto index = vsag::Factory::CreateIndex("hgraph", build_param); + REQUIRE(index.has_value()); + + const auto base_count = static_cast(vectors.base_ids.size()); + auto base_dataset = vsag::Dataset::Make(); + base_dataset->NumElements(base_count) + ->Dim(DIM) + ->Float32Vectors(vectors.base.data()) + ->Ids(vectors.base_ids.data()) + ->Owner(false); + auto build_result = index.value()->Build(base_dataset); + REQUIRE(build_result.has_value()); + REQUIRE(build_result.value().empty()); + + const auto duplicate_count = static_cast(vectors.duplicate_ids.size()); + if (duplicate_count > 0) { + auto duplicate_dataset = vsag::Dataset::Make(); + duplicate_dataset->NumElements(duplicate_count) + ->Dim(DIM) + ->Float32Vectors(vectors.duplicates.data()) + ->Ids(vectors.duplicate_ids.data()) + ->Owner(false); + auto add_result = index.value()->Add(duplicate_dataset); + REQUIRE(add_result.has_value()); + REQUIRE(add_result.value().empty()); + } + + return index.value(); +} + +std::map +CountDuplicateIdsByGroup(const vsag::DatasetPtr& result, int64_t base_count) { + std::map group_duplicate_count; + const auto* ids = result->GetIds(); + for (int64_t i = 0; i < result->GetDim(); ++i) { + if (ids[i] >= base_count) { + const int64_t source = (ids[i] - base_count) % base_count; + ++group_duplicate_count[source]; + } + } + return group_duplicate_count; +} + +int64_t +CountDuplicateIds(const vsag::DatasetPtr& result, int64_t base_count) { + int64_t duplicate_count = 0; + const auto* ids = result->GetIds(); + for (int64_t i = 0; i < result->GetDim(); ++i) { + if (ids[i] >= base_count) { + ++duplicate_count; + } + } + return duplicate_count; +} + +struct IteratorItem { + int64_t id; + float distance; + int64_t page; +}; + +struct IteratorContextGuard { + vsag::IteratorContext*& context; + + ~IteratorContextGuard() { + delete context; + } +}; + +class AllowedIdFilter : public vsag::Filter { +public: + explicit AllowedIdFilter(std::set allowed_ids) : allowed_ids_(std::move(allowed_ids)) { + } + + [[nodiscard]] bool + CheckValid(int64_t id) const override { + return allowed_ids_.count(id) != 0; + } + +private: + std::set allowed_ids_; +}; + +std::vector +CollectIteratorResults(const vsag::IndexPtr& index, + const vsag::DatasetPtr& query, + const std::string& params, + int64_t page_size = 1, + vsag::FilterPtr filter = nullptr) { + vsag::IteratorContext* iterator_context = nullptr; + IteratorContextGuard guard{iterator_context}; + + std::vector items; + std::set seen_ids; + bool exhausted = false; + constexpr int64_t max_pages = 32; + for (int64_t page = 0; page < max_pages; ++page) { + auto result = index->KnnSearch(query, page_size, params, filter, iterator_context, false); + REQUIRE(result.has_value()); + + const auto count = result.value()->GetDim(); + if (count == 0) { + exhausted = true; + break; + } + REQUIRE(count <= page_size); + for (int64_t i = 0; i < count; ++i) { + const auto id = result.value()->GetIds()[i]; + REQUIRE(seen_ids.insert(id).second); + items.push_back({id, result.value()->GetDistances()[i], page}); + } + } + + REQUIRE(exhausted); + return items; +} + +} // namespace + +TEST_CASE("HGraph dedup search expands the entry-point group", + "[ft][hgraph][duplicate][search_control][entry_point]") { + constexpr int64_t entry_base_count = 1; + constexpr int64_t entry_duplicate_count = 3; + constexpr int64_t duplicate_limit = 2; + auto vectors = GenerateTestData(DIM, entry_base_count, entry_duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + for (const auto parallelism : {1, 2}) { + DYNAMIC_SECTION("parallelism=" << parallelism) { + auto result = index->KnnSearch( + query, 1 + duplicate_limit, MakeSearchParam(4, duplicate_limit, parallelism)); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 1 + duplicate_limit); + REQUIRE(CountDuplicateIds(result.value(), entry_base_count) == duplicate_limit); + + std::set ids(result.value()->GetIds(), + result.value()->GetIds() + result.value()->GetDim()); + REQUIRE(ids.count(0) == 1); + } + } + + auto default_limit = + index->KnnSearch(query, 4, R"({"hgraph":{"ef_search":4,"parallelism":2}})"); + REQUIRE(default_limit.has_value()); + REQUIRE(default_limit.value()->GetDim() == 4); + REQUIRE(CountDuplicateIds(default_limit.value(), entry_base_count) == entry_duplicate_count); +} + +TEST_CASE("HGraph dedup iterator expands the entry-point group across pages", + "[ft][hgraph][duplicate][search_control][iterator][entry_point]") { + constexpr int64_t entry_base_count = 1; + constexpr int64_t entry_duplicate_count = 3; + constexpr int64_t duplicate_limit = 2; + auto vectors = GenerateTestData(DIM, entry_base_count, entry_duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + const auto items = CollectIteratorResults(index, query, MakeSearchParam(4, duplicate_limit), 1); + REQUIRE(items.size() == static_cast(1 + duplicate_limit)); + + int64_t root_count = 0; + int64_t duplicate_count = 0; + std::set duplicate_pages; + for (const auto& item : items) { + if (item.id == 0) { + ++root_count; + } else { + ++duplicate_count; + duplicate_pages.insert(item.page); + } + } + REQUIRE(root_count == 1); + REQUIRE(duplicate_count == duplicate_limit); + REQUIRE(duplicate_pages.size() == static_cast(duplicate_limit)); +} + +TEST_CASE("HGraph dedup search limits duplicate expansion per group", + "[ft][hgraph][duplicate][search_control]") { + constexpr int64_t multi_base_count = 2; + constexpr int64_t multi_duplicate_count = 6; + auto vectors = GenerateTestData(DIM, multi_base_count, multi_duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + auto result = index->KnnSearch(query, 4, MakeSearchParam(8, 1)); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 4); + const auto group_duplicate_count = CountDuplicateIdsByGroup(result.value(), multi_base_count); + REQUIRE(group_duplicate_count.size() == 2); + for (int64_t group = 0; group < multi_base_count; ++group) { + REQUIRE(group_duplicate_count.at(group) == 1); + } +} + +TEST_CASE("HGraph dedup search supports zero and unlimited duplicate expansion", + "[ft][hgraph][duplicate][search_control]") { + constexpr int64_t base_count = 2; + constexpr int64_t duplicate_count = 6; + auto vectors = GenerateTestData(DIM, base_count, duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + auto limited = index->KnnSearch(query, 8, MakeSearchParam(8, 0)); + REQUIRE(limited.has_value()); + REQUIRE(limited.value()->GetDim() == base_count); + REQUIRE(CountDuplicateIds(limited.value(), base_count) == 0); + + auto unlimited = index->KnnSearch(query, 8, MakeSearchParam(8, -1)); + REQUIRE(unlimited.has_value()); + REQUIRE(unlimited.value()->GetDim() == base_count + duplicate_count); + REQUIRE(CountDuplicateIds(unlimited.value(), base_count) == duplicate_count); +} + +TEST_CASE("HGraph dedup search applies min_distance consistently at the entry point", + "[ft][hgraph][duplicate][search_control][entry_point][min_distance]") { + constexpr int64_t base_count = 1; + constexpr int64_t duplicate_count = 3; + auto vectors = GenerateTestData(DIM, base_count, duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + for (const auto parallelism : {1, 2}) { + DYNAMIC_SECTION("parallelism=" << parallelism) { + const auto params = fmt::format( + R"({{"hgraph":{{"ef_search":4,"parallelism":{},"min_distance":0}}}})", parallelism); + auto result = index->KnnSearch(query, 4, params); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 0); + } + } +} + +TEST_CASE("HGraph dedup search ignores duplicate limit when tracking is disabled", + "[ft][hgraph][duplicate][search_control]") { + constexpr int64_t base_count = 2; + constexpr int64_t duplicate_count = 6; + auto vectors = GenerateTestData(DIM, base_count, duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(false)); + + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + auto limited = index->KnnSearch(query, 8, MakeSearchParam(16, 0)); + auto unlimited = index->KnnSearch(query, 8, MakeSearchParam(16, -1)); + REQUIRE(limited.has_value()); + REQUIRE(unlimited.has_value()); + REQUIRE(limited.value()->GetDim() == base_count + duplicate_count); + REQUIRE(unlimited.value()->GetDim() == base_count + duplicate_count); + REQUIRE(std::set(limited.value()->GetIds(), + limited.value()->GetIds() + limited.value()->GetDim()) == + std::set(unlimited.value()->GetIds(), + unlimited.value()->GetIds() + unlimited.value()->GetDim())); +} + +TEST_CASE("HGraph dedup search applies filters before counting duplicate limits", + "[ft][hgraph][duplicate][search_control][filter]") { + constexpr int64_t base_count = 2; + auto vectors = GenerateTestData(DIM, base_count, 3 * base_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + + for (const auto parallelism : {1, 2}) { + DYNAMIC_SECTION("parallelism=" << parallelism) { + for (int64_t group = 0; group < base_count; ++group) { + for (int64_t member = 0; member < 3; ++member) { + auto query = vsag::Dataset::Make(); + query->NumElements(1) + ->Dim(DIM) + ->Float32Vectors(vectors.base.data() + group * DIM) + ->Owner(false); + const auto allowed_duplicate = base_count + group + member * base_count; + auto filter = + std::make_shared(std::set{allowed_duplicate}); + + auto result = + index->KnnSearch(query, 1, MakeSearchParam(8, 1, parallelism), filter); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 1); + REQUIRE(result.value()->GetIds()[0] == allowed_duplicate); + } + } + } + } +} + +TEST_CASE("HGraph duplicate limit does not change range-search expansion", + "[ft][hgraph][duplicate][search_control][range]") { + constexpr int64_t base_count = 1; + constexpr int64_t duplicate_count = 3; + auto vectors = GenerateTestData(DIM, base_count, duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + vsag::FilterPtr filter = nullptr; + for (const auto parallelism : {1, 2}) { + DYNAMIC_SECTION("parallelism=" << parallelism) { + auto result = + index->RangeSearch(query, 0.1F, MakeSearchParam(4, 0, parallelism), filter, -1); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == base_count + duplicate_count); + } + } +} + +TEST_CASE("HGraph dedup iterator keeps limited duplicates across pages", + "[ft][hgraph][duplicate][search_control][iterator]") { + constexpr int64_t iterator_base_count = 2; + constexpr int64_t iterator_duplicate_count = 6; + auto vectors = GenerateTestData(DIM, iterator_base_count, iterator_duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + const std::vector> cases{{0, 0}, {1, 1}, {2, 2}, {-1, 3}}; + for (const auto& [limit, expected_duplicate_count] : cases) { + DYNAMIC_SECTION("max_duplicates_per_group=" << limit) { + const auto items = CollectIteratorResults(index, query, MakeSearchParam(2, limit)); + std::map duplicate_counts; + std::map> duplicate_pages; + for (const auto& item : items) { + REQUIRE(item.distance >= 0.0F); + if (item.id < iterator_base_count) { + continue; + } + const auto group = (item.id - iterator_base_count) % iterator_base_count; + ++duplicate_counts[group]; + duplicate_pages[group].insert(item.page); + } + + if (expected_duplicate_count == 0) { + REQUIRE(duplicate_counts.empty()); + continue; + } + + bool found_full_group = false; + for (const auto& [group, count] : duplicate_counts) { + REQUIRE(count <= expected_duplicate_count); + if (count == expected_duplicate_count) { + found_full_group = true; + REQUIRE(duplicate_pages[group].size() == + static_cast(expected_duplicate_count)); + } + } + REQUIRE(found_full_group); + } + } +} + +TEST_CASE("HGraph dedup iterator filters duplicate groups before applying limits", + "[ft][hgraph][duplicate][search_control][iterator][filter]") { + constexpr int64_t base_count = 1; + constexpr int64_t duplicate_count = 3; + auto vectors = GenerateTestData(DIM, base_count, duplicate_count); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + for (const auto allowed_duplicate : vectors.duplicate_ids) { + DYNAMIC_SECTION("allowed_duplicate=" << allowed_duplicate) { + auto filter = std::make_shared(std::set{allowed_duplicate}); + const auto items = + CollectIteratorResults(index, query, MakeSearchParam(4, 1), 1, filter); + REQUIRE(items.size() == 1); + REQUIRE(items[0].id == allowed_duplicate); + } + } +} + +TEST_CASE("HGraph dedup iterator drains pending duplicates in last-search mode", + "[ft][hgraph][duplicate][search_control][iterator]") { + constexpr int64_t iterator_base_count = 2; + auto vectors = GenerateTestData(DIM, iterator_base_count, 6); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(true, 0.001F)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + vsag::IteratorContext* iterator_context = nullptr; + IteratorContextGuard guard{iterator_context}; + vsag::FilterPtr filter = nullptr; + const auto params = MakeSearchParam(2, 2); + auto first = index->KnnSearch(query, 1, params, filter, iterator_context, false); + REQUIRE(first.has_value()); + REQUIRE(first.value()->GetDim() == 1); + auto last = index->KnnSearch(query, 8, params, filter, iterator_context, true); + REQUIRE(last.has_value()); + REQUIRE(last.value()->GetDim() > 0); + + std::set result_ids; + std::map duplicate_counts; + const auto collect_result = [&](const vsag::DatasetPtr& result) { + for (int64_t i = 0; i < result->GetDim(); ++i) { + const auto id = result->GetIds()[i]; + REQUIRE(result_ids.insert(id).second); + if (id < iterator_base_count) { + continue; + } + const auto group = (id - iterator_base_count) % iterator_base_count; + ++duplicate_counts[group]; + } + }; + collect_result(first.value()); + collect_result(last.value()); + + bool found_full_group = false; + for (const auto& [group, count] : duplicate_counts) { + (void)group; + REQUIRE(count <= 2); + found_full_group = found_full_group or count == 2; + } + REQUIRE(found_full_group); + + auto exhausted = index->KnnSearch(query, 8, params, filter, iterator_context, true); + REQUIRE(exhausted.has_value()); + REQUIRE(exhausted.value()->GetDim() == 0); +} + +TEST_CASE("HGraph iterator ignores duplicate limit without duplicate tracking", + "[ft][hgraph][duplicate][search_control][iterator]") { + constexpr int64_t iterator_base_count = 2; + auto vectors = GenerateTestData(DIM, iterator_base_count, 6); + auto index = BuildIndexWithDuplicates(vectors, MakeBuildParam(false)); + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(DIM)->Float32Vectors(vectors.queries.data())->Owner(false); + + const auto limited = CollectIteratorResults(index, query, MakeSearchParam(16, 0), 2); + const auto unlimited = CollectIteratorResults(index, query, MakeSearchParam(16, -1), 2); + std::set limited_ids; + std::set unlimited_ids; + for (const auto& item : limited) { + limited_ids.insert(item.id); + } + for (const auto& item : unlimited) { + unlimited_ids.insert(item.id); + } + + REQUIRE(limited_ids == unlimited_ids); + REQUIRE(limited_ids.lower_bound(iterator_base_count) != limited_ids.end()); +}