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;
 
@@ -2222,6 +2233,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"
@@ -313,13 +314,54 @@ 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) {
+            if (duplicate_count_ > label_table_.size()) {
+                throw VsagException(ErrorType::INVALID_BINARY,
+                                    fmt::format("duplicate group count {} exceeds label count {}",
+                                                duplicate_count_,
+                                                label_table_.size()));
+            }
+
+            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);
-                duplicate_records_[id] = allocator_->New<DuplicateRecord>(allocator_);
+                if (id >= label_table_.size()) {
+                    throw VsagException(ErrorType::INVALID_BINARY,
+                                        fmt::format("duplicate head id {} exceeds label count {}",
+                                                    id,
+                                                    label_table_.size()));
+                }
+                if (not assigned_ids.insert(id).second) {
+                    throw VsagException(
+                        ErrorType::INVALID_BINARY,
+                        fmt::format("id {} belongs to multiple duplicate groups", id));
+                }
+
                 Vector<InnerIdType> id_list(allocator_);
                 StreamReader::ReadVector(reader, id_list);
+                for (const auto duplicate_id : id_list) {
+                    if (duplicate_id >= label_table_.size()) {
+                        throw VsagException(
+                            ErrorType::INVALID_BINARY,
+                            fmt::format("duplicate member id {} exceeds label count {}",
+                                        duplicate_id,
+                                        label_table_.size()));
+                    }
+                    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));
+            }
+
+            duplicate_records_.resize(label_table_.size(), nullptr);
+            for (auto& [id, id_list] : duplicate_groups) {
+                duplicate_records_[id] = allocator_->New<DuplicateRecord>(allocator_);
                 for (const auto& duplicate_id : id_list) {
                     duplicate_records_[id]->duplicate_ids.insert(duplicate_id);
                 }

src/impl/label_table_test.cpp

@@ -1,11 +1,45 @@
 #include "label_table.h"
 
 #include <catch2/catch_test_macros.hpp>
+#include <sstream>
+#include <utility>
+#include <vector>
 
 #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<InnerIdType, std::vector<InnerIdType>>;
+
+std::stringstream
+CreateSerializedLabelTable(const std::vector<LabelType>& labels,
+                           const std::vector<DuplicateGroup>& 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;
+}
+
+void
+DeserializeLabelTable(LabelTable& label_table, std::stringstream& stream) {
+    IOStreamReader reader(stream);
+    label_table.Deserialize(reader);
+}
+
+}  // namespace
+
 TEST_CASE("LabelTable Supports Configurable Remap Implementation", "[ut][LabelTable]") {
     auto allocator = std::make_shared<DefaultAllocator>();
 
@@ -27,3 +61,75 @@ 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<DefaultAllocator>();
+    LabelTable label_table(allocator.get(), true, true);
+    auto stream = CreateSerializedLabelTable({10, 11, 12, 13, 14}, {{0, {1, 2}}, {3, {4}}});
+
+    DeserializeLabelTable(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 rejects out-of-range duplicate members",
+          "[ut][LabelTable][duplicate][invalid-member]") {
+    auto allocator = std::make_shared<DefaultAllocator>();
+    LabelTable label_table(allocator.get(), true, true);
+    auto stream = CreateSerializedLabelTable({10, 11, 12}, {{0, {3}}});
+
+    REQUIRE_THROWS_AS(DeserializeLabelTable(label_table, stream), VsagException);
+}
+
+TEST_CASE("LabelTable rejects out-of-range duplicate heads",
+          "[ut][LabelTable][duplicate][invalid-head]") {
+    auto allocator = std::make_shared<DefaultAllocator>();
+    LabelTable label_table(allocator.get(), true, true);
+    auto stream = CreateSerializedLabelTable({10, 11, 12}, {{3, {1}}});
+
+    REQUIRE_THROWS_AS(DeserializeLabelTable(label_table, stream), VsagException);
+}
+
+TEST_CASE("LabelTable rejects overlapping duplicate groups",
+          "[ut][LabelTable][duplicate][overlap]") {
+    auto allocator = std::make_shared<DefaultAllocator>();
+
+    const auto require_invalid = [&](const std::vector<DuplicateGroup>& groups) {
+        LabelTable label_table(allocator.get(), true, true);
+        auto stream = CreateSerializedLabelTable({10, 11, 12, 13, 14, 15}, groups);
+        REQUIRE_THROWS_AS(DeserializeLabelTable(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}}});
+    }
+}