fix(hgraph): backport duplicate handling to v0.18

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.
+means that the search uses an ef_search value of 200 and returns at most two additional IDs from
+each duplicate group.

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;

src/algorithm/hgraph.cpp

@@ -933,10 +933,20 @@ HGraph::KnnSearch(const DatasetPtr& query,
     auto search_result = DistanceHeap::MakeInstanceBySize<true, false>(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<int64_t>(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;
     }
 
@@ -1474,6 +1488,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<uint64_t>(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 +2241,7 @@ HGraph::SearchWithRequest(const SearchRequest& request) const {
             search_param.topk, static_cast<int64_t>(static_cast<float>(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<Timer>();
         search_param.time_cost->SetThreshold(params.timeout_ms);

src/algorithm/hgraph_parameter.cpp

@@ -15,6 +15,8 @@
 
 #include "hgraph_parameter.h"
 
+#include <nlohmann/json.hpp>
+
 #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<uint64_t>() <=
+                               static_cast<uint64_t>(std::numeric_limits<int64_t>::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;
 }

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<float>::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;

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",

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";

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<float>::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;
     }

src/impl/label_table.h

@@ -25,6 +25,7 @@
 #include <mutex>
 #include <tuple>
 #include <utility>
+#include <vector>
 
 #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<InnerIdType> 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<InnerIdType> 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<StreamReader> reader) {
         StreamReader::ReadVector(reader, label_table_);
@@ -312,25 +318,107 @@ 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<InnerIdType>(reader, id);
-                duplicate_records_[id] = allocator_->New<DuplicateRecord>(allocator_);
-                Vector<InnerIdType> 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_);
         }
         this->total_count_.store(label_table_.size());
     }
 
+    void
+    DeserializeDuplicateRecords(lvalue_or_rvalue<StreamReader> 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<std::pair<InnerIdType, Vector<InnerIdType>>> duplicate_groups;
+        duplicate_groups.reserve(duplicate_count);
+        UnorderedSet<InnerIdType> assigned_ids(allocator_);
+        for (uint64_t i = 0; i < duplicate_count; ++i) {
+            InnerIdType id;
+            StreamReader::ReadObj<InnerIdType>(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<uint64_t>(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<InnerIdType> id_list(allocator_);
+            id_list.resize(member_count);
+            reader->Read(reinterpret_cast<char*>(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<DuplicateRecord*> restored_records(label_table_.size(), nullptr, allocator_);
+        try {
+            for (auto& [id, id_list] : duplicate_groups) {
+                restored_records[id] = allocator_->New<DuplicateRecord>(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_) {