diff --git a/docs/docs/en/src/indexes/hgraph.md b/docs/docs/en/src/indexes/hgraph.md index 91f27d31df..d0761362ef 100644 --- a/docs/docs/en/src/indexes/hgraph.md +++ b/docs/docs/en/src/indexes/hgraph.md @@ -90,7 +90,7 @@ most users need; the exhaustive list is in [Index Parameters](../resources/index | `rabitq_fused_datacell` | bool | `false` | Fuse the bottom HGraph node and RaBitQ split codes into one in-memory record. Requires L2/IP, flat in-memory graph storage, RaBitQ x+y split codes with x in `[1, 4]`, and the other constraints described in [RaBitQ x+y split](../quantization/rabitq_split.md). | | `train_sample_count` | int | `65536` | Maximum number of vectors sampled for quantizer training; must be at least `512` when set explicitly. | | `build_thread_count` | int | `100` | Threads used to parallelise build | -| `support_duplicate` | bool | `false` | Enable duplicate-ID detection on insert | +| `support_duplicate` | bool | `false` | Group duplicate vectors while retaining their individual labels | | `deduplicate_storage` | bool | `false` | Share vector storage between duplicates; requires `support_duplicate: true` | | `duplicate_distance_threshold` | float | `0.0` | Duplicate-detection distance threshold. When greater than `0`, deduplicate by the nearest candidate distance; when `0`, compare codes with the nearest candidate first, then check the remaining returned build candidates until a match is found | | `support_remove` | bool | `false` | Enable graph delete-tracking metadata used by mark-remove recovery paths | @@ -122,8 +122,8 @@ full `Build`. The PiPNN builder accepts dense `float32` input with `metric_type` `"ip"`, or `"cosine"`. It builds the bottom graph from the original build vectors, then reuses HGraph's route layers, storage, search, filtering, reordering, incremental `Add`, removal, and serialization paths. The persistent base storage may use a supported quantizer such as `sq8`, -including RaBitQ with SQ8 reorder. Cache-assisted build and deduplicated vector storage are not -supported with PiPNN. The `pipnn_*` parameters in the table above tune this builder; +including RaBitQ with SQ8 reorder. Cache-assisted build is not supported with PiPNN. +The `pipnn_*` parameters in the table above tune this builder; `ef_construction` does not. The existing `alpha` parameter controls PiPNN's final robust pruning. `build_thread_count` parallelizes vector preparation, partitioning, candidate generation, and @@ -146,8 +146,14 @@ The policy and step are serialized. Loading an enabled index requires matching p Set both `support_duplicate: true` and `deduplicate_storage: true` to let duplicate vectors share one physical code slot while retaining their individual labels. This option -currently supports only dense-vector HGraph indexes using `graph_type: "nsw"`; it is not -available for `graph_type: "odescent"` or `graph_type: "pipnn"`. +supports dense-vector HGraph indexes using `graph_type: "nsw"` or `"pipnn"`; +`graph_type: "odescent"` is not supported. + +For PiPNN, initial `Build` groups byte-identical FP32 input rows and encodes each +representative once. All labels, attributes, and extra information remain independent. +This also works with adaptive alpha. `duplicate_distance_threshold` does not affect +PiPNN batch grouping; later incremental `Add` retains its existing threshold/code-based +duplicate detection. Quantizer training still uses the original training sample. The following operations and configurations are not supported while storage deduplication is enabled: diff --git a/docs/docs/zh/src/indexes/hgraph.md b/docs/docs/zh/src/indexes/hgraph.md index e94d0875db..f5019a6387 100644 --- a/docs/docs/zh/src/indexes/hgraph.md +++ b/docs/docs/zh/src/indexes/hgraph.md @@ -84,7 +84,7 @@ auto result = index->KnnSearch( | `rabitq_fused_datacell` | bool | `false` | 将底层 HGraph 节点与 RaBitQ split 编码融合到一条内存记录中；要求 L2/IP、flat 内存图、x 在 `[1, 4]` 的 RaBitQ x+y split 编码，并满足 [RaBitQ x+y split](../quantization/rabitq_split.md) 中的其他约束 | | `train_sample_count` | int | `65536` | 量化器训练的最大采样向量数；显式配置时最小为 `512` | | `build_thread_count` | int | `100` | 构建阶段并发线程数 | -| `support_duplicate` | bool | `false` | 是否在插入时做重复 ID 检测 | +| `support_duplicate` | bool | `false` | 合并重复向量的图节点，并保留各自的标签 | | `deduplicate_storage` | bool | `false` | 让重复向量共享存储；需同时设置 `support_duplicate: true` | | `duplicate_distance_threshold` | float | `0.0` | 重复判定距离阈值。大于 `0` 时按最近候选的距离判重；等于 `0` 时先比较最近候选的编码，不匹配则继续检查构建搜索返回的其余候选，直到找到匹配 | | `support_remove` | bool | `false` | 是否启用 mark-remove 恢复路径所需的图删除追踪元数据 | @@ -113,8 +113,8 @@ auto result = index->KnnSearch( [PiPNN](https://arxiv.org/abs/2602.21247)。PiPNN 构建器接收 `dtype: "float32"`、 `metric_type: "l2"`、`"ip"` 或 `"cosine"` 的稠密向量输入，从原始构建向量生成底层图，并复用 HGraph 现有的路由层、向量存储、搜索、过滤、精排、增量 `Add`、删除和序列化路径。持久化底层 -存储可以使用 `sq8` 等受支持的量化器，包括 RaBitQ 配合 SQ8 精排。缓存辅助构建和向量存储去重 -暂不支持 PiPNN。上表的 `pipnn_*` 参数用于调节该构建器，`ef_construction` 不生效； +存储可以使用 `sq8` 等受支持的量化器，包括 RaBitQ 配合 SQ8 精排。缓存辅助构建暂不支持 PiPNN。 +上表的 `pipnn_*` 参数用于调节该构建器，`ef_construction` 不生效； 现有的 `alpha` 参数控制 PiPNN 的最终 robust pruning。 `build_thread_count` 会并行化向量预处理、分区、候选边生成和最终剪枝。性能测试时应固定 @@ -136,8 +136,13 @@ PiPNN，并通过查询验证每一份生成的图。 ### 向量存储去重 同时设置 `support_duplicate: true` 和 `deduplicate_storage: true` 后，重复向量会共享 -同一个物理编码槽位，但仍保留各自的标签。该选项目前仅支持使用 `graph_type: "nsw"` 的 -稠密向量 HGraph 索引；`graph_type: "odescent"` 和 `graph_type: "pipnn"` 不支持。 +同一个物理编码槽位，但仍保留各自的标签。该选项支持使用 `graph_type: "nsw"` 或 `"pipnn"` 的 +稠密向量 HGraph 索引；`graph_type: "odescent"` 不支持。 + +PiPNN 初始 `Build` 按 FP32 原始输入行的字节相等关系分组，每个代表向量只编码一次， +各 label 的属性和 extra_info 仍独立保存，也可同时启用动态 alpha。 +`duplicate_distance_threshold` 不参与 PiPNN 批量分组；后续增量 `Add` 沿用现有的 +距离阈值或编码判重规则。量化器仍使用原有训练样本进行训练。 启用存储去重后，暂不支持以下操作和配置： diff --git a/src/algorithm/hgraph/hgraph.cpp b/src/algorithm/hgraph/hgraph.cpp index ba14da2d50..07b944b665 100644 --- a/src/algorithm/hgraph/hgraph.cpp +++ b/src/algorithm/hgraph/hgraph.cpp @@ -96,9 +96,10 @@ HGraph::HGraph(const HGraphParameterPtr& hgraph_param, const vsag::IndexCommonPa "HGraph PiPNN only supports dense float32 indexes"); } } - if (this->deduplicate_storage_ && this->graph_type_ != GRAPH_TYPE_VALUE_NSW) { + if (this->deduplicate_storage_ && this->graph_type_ != GRAPH_TYPE_VALUE_NSW && + this->graph_type_ != GRAPH_TYPE_VALUE_PIPNN) { throw VsagException(ErrorType::INVALID_ARGUMENT, - "HGraph deduplicate_storage only supports nsw graph"); + "HGraph deduplicate_storage only supports nsw or pipnn graph"); } this->persist_source_id_ = hgraph_param->persist_source_id; if (this->using_dedup_storage()) { @@ -149,7 +150,6 @@ HGraph::HGraph(const HGraphParameterPtr& hgraph_param, const vsag::IndexCommonPa mult_ = 1 / log(1.0 * static_cast(this->bottom_graph_->MaximumDegree())); this->resize_increase_count_bit_ = hgraph_param->resize_increase_count_bit; - init_resize_bit_and_reorder(); this->parallel_searcher_ = std::make_shared(common_param, thread_pool_, neighbors_mutex_); @@ -170,11 +170,18 @@ HGraph::HGraph(const HGraphParameterPtr& hgraph_param, const vsag::IndexCommonPa this->high_precise_codes_ = MakeCodeSlotFlattenAdapter( this->high_precise_codes_, this->code_slot_map_, allocator_, &this->total_count_); } + if (this->raw_vector_ != nullptr and not this->create_new_raw_vector_) { + // Rebind an existing FP32 raw alias to the wrapped base/reorder codes. + // false refreshes the alias/flags without allocating another raw flatten. + this->check_and_init_raw_vector(hgraph_param->raw_vector_param, common_param, false); + } if (this->create_new_raw_vector_ && this->raw_vector_ != nullptr) { this->raw_vector_ = MakeCodeSlotFlattenAdapter( this->raw_vector_, this->code_slot_map_, allocator_, &this->total_count_); } } + // Reorder retains its flatten pointer, so initialize it after slot wrapping. + init_resize_bit_and_reorder(); resize(bottom_graph_->max_capacity_); } @@ -793,6 +800,11 @@ void HGraph::init_resize_bit_and_reorder() { if (use_reorder_) { auto reorder_codes = this->get_reorder_codes(); + if (this->using_dedup_storage() and + GetCodeSlotPhysicalFlatten(reorder_codes) == reorder_codes) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "storage deduplication requires wrapped reorder codes"); + } auto fused_graph = reorder_codes.get() == basic_flatten_codes_.get() ? rabitq_fused_datacell_ : nullptr; reorder_ = std::make_shared(reorder_codes, allocator_, fused_graph); diff --git a/src/algorithm/hgraph/hgraph.h b/src/algorithm/hgraph/hgraph.h index 5c86bbb316..6c2d75896a 100644 --- a/src/algorithm/hgraph/hgraph.h +++ b/src/algorithm/hgraph/hgraph.h @@ -81,6 +81,8 @@ struct RaBitQFusedCodeView; */ class HGraph : public InnerIndexInterface { public: + friend struct HGraphPiPNNStorageTestAccess; + static ParamPtr CheckAndMappingExternalParam(const JsonType& external_param, const IndexCommonParam& common_param); @@ -469,6 +471,12 @@ class HGraph : public InnerIndexInterface { std::vector build_by_batch_graph(const DatasetPtr& data, bool use_pipnn); + /// Encode each PiPNN representative once, then bind all logical IDs to its slot. + void + materialize_pipnn_storage(const DatasetPtr& data, + const Vector& inner_ids, + const Vector& valid_indices); + /// Write codes for inner_id into the persistent flatten storage. void insert_persistent_codes(const void* data, diff --git a/src/algorithm/hgraph/hgraph_add_test.cpp b/src/algorithm/hgraph/hgraph_add_test.cpp index cd9bc3dd1b..b10abcf96e 100644 --- a/src/algorithm/hgraph/hgraph_add_test.cpp +++ b/src/algorithm/hgraph/hgraph_add_test.cpp @@ -1272,6 +1272,7 @@ TEST_CASE("HGraph deduplicate_storage batch Add supports internal parallel add", constexpr uint64_t build_threads = 4; auto common_param = MakeCommonParam(dim, build_threads); auto hgraph_json = MakeFp32HGraphJson(true, static_cast(build_threads)); + hgraph_json["graph_type"].SetString(GENERATE("nsw", "pipnn")); auto index = MakeHGraphIndex(hgraph_json, common_param); std::vector seed_vectors = { @@ -1419,6 +1420,7 @@ TEST_CASE("HGraph deduplicate_storage concurrent Add keeps visible duplicates se constexpr int64_t dim = 2; auto common_param = MakeCommonParam(dim); auto hgraph_json = MakeFp32HGraphJson(); + hgraph_json["graph_type"].SetString(GENERATE("nsw", "pipnn")); auto index = MakeHGraphIndex(hgraph_json, common_param); std::vector seed_vectors = { diff --git a/src/algorithm/hgraph/hgraph_build.cpp b/src/algorithm/hgraph/hgraph_build.cpp index eda541c081..ee9747fb9a 100644 --- a/src/algorithm/hgraph/hgraph_build.cpp +++ b/src/algorithm/hgraph/hgraph_build.cpp @@ -321,9 +321,9 @@ HGraph::build_by_batch_graph(const DatasetPtr& data, bool use_pipnn) { batch_insert_flatten_codes(temporary_sq8_build_data); temporary_codes_batched = true; } - // Every code page of the batch is written at this point (or by the per-row loop below when - // rows were skipped), so the batch can be published exactly like a sequence of Add() calls. - this->PublishThroughTotalCount(); + if (not this->using_dedup_storage()) { + this->PublishThroughTotalCount(); + } Vector> deferred_code_ids(allocator_); Vector pipnn_vectors(allocator_); if (use_pipnn) { @@ -341,7 +341,7 @@ HGraph::build_by_batch_graph(const DatasetPtr& data, bool use_pipnn) { } if (defer_persistent_codes) { deferred_code_ids.emplace_back(inner_id, i); - } else if (not persistent_codes_batched) { + } else if (not persistent_codes_batched and not this->using_dedup_storage()) { this->insert_persistent_codes(get_data(data, i), inner_id); } if (temporary_sq8_build_data != nullptr and not temporary_codes_batched) { @@ -406,6 +406,9 @@ HGraph::build_by_batch_graph(const DatasetPtr& data, bool use_pipnn) { this->thread_pool_.get(), this->build_thread_count_); pipnn_builder.Build(bottom_graph_, inner_ids, pipnn_vectors); + if (this->using_dedup_storage() and not defer_persistent_codes) { + this->materialize_pipnn_storage(data, inner_ids, valid_indices); + } if (this->support_duplicate_) { entry_point_id_ = bottom_graph_->GetGroupId(entry_point_id_); for (auto& route_graph_id : route_graph_ids) { @@ -458,7 +461,9 @@ HGraph::build_by_batch_graph(const DatasetPtr& data, bool use_pipnn) { if (this->rabitq_fused_datacell_ == nullptr) { this->train_codes_with_dataset(this->sample_train_dataset(data)); } - if (all_rows_are_valid) { + if (this->using_dedup_storage()) { + this->materialize_pipnn_storage(data, inner_ids, valid_indices); + } else if (all_rows_are_valid) { batch_insert_persistent_codes(); } else { for (const auto& [inner_id, local_idx] : deferred_code_ids) { @@ -466,9 +471,92 @@ HGraph::build_by_batch_graph(const DatasetPtr& data, bool use_pipnn) { } } } + if (this->using_dedup_storage()) { + // All logical IDs now have complete codes and a published slot mapping. + this->PublishThroughTotalCount(); + } return failed_ids; } +void +HGraph::materialize_pipnn_storage(const DatasetPtr& data, + const Vector& inner_ids, + const Vector& valid_indices) { + struct Representative { + InnerIdType inner_id; + CodeSlotIdType slot; + int64_t row; + }; + Vector representatives(allocator_); + for (uint64_t i = 0; i < inner_ids.size(); ++i) { + const auto id = inner_ids[i]; + if (bottom_graph_->GetGroupId(id) == id) { + representatives.push_back({id, 0, valid_indices[i]}); + } + } + // Build is exclusive. Encode into provisional slots so a failed task does not consume + // slot IDs or change logical mappings. The allocated buffers can be reused after failure. + const auto first_slot = code_slot_map_->PhysicalCount(); + CHECK_ARGUMENT( + representatives.size() <= std::numeric_limits::max() - first_slot, + "PiPNN storage exceeds the code slot limit"); + for (uint64_t i = 0; i < representatives.size(); ++i) { + representatives[i].slot = first_slot + static_cast(i); + } + this->ensure_physical_code_capacity(first_slot + + static_cast(representatives.size())); + + // Capacity is fixed before workers write. Wait for every scheduled task before publishing + // mappings or letting representative rows go out of scope, even if scheduling/encoding fails. + const uint64_t workers = std::min( + representatives.size(), + thread_pool_ != nullptr ? std::max(1, build_thread_count_) : 1); + std::vector> futures; + futures.reserve(workers); + std::exception_ptr failure; + try { + for (uint64_t worker = 0; worker < workers; ++worker) { + const auto begin = representatives.size() * worker / workers; + const auto end = representatives.size() * (worker + 1) / workers; + auto encode = [this, &representatives, data, begin, end]() { + for (auto i = begin; i < end; ++i) { + const auto& representative = representatives[i]; + this->insert_persistent_codes_to_slot(get_data(data, representative.row), + representative.slot); + } + }; + if (workers > 1) { + futures.emplace_back(thread_pool_->GeneralEnqueue(std::move(encode))); + } else { + encode(); + } + } + } catch (...) { + failure = std::current_exception(); + } + for (auto& future : futures) { + try { + future.get(); + } catch (...) { + if (failure == nullptr) { + failure = std::current_exception(); + } + } + } + if (failure != nullptr) { + std::rethrow_exception(failure); + } + for (const auto& representative : representatives) { + code_slot_map_->PublishSlot(representative.inner_id, code_slot_map_->AllocateSlot()); + } + for (auto id : inner_ids) { + const auto group = bottom_graph_->GetGroupId(id); + if (group != id) { + code_slot_map_->PublishSlot(id, code_slot_map_->Resolve(group)); + } + } +} + std::vector HGraph::Add(const DatasetPtr& data) { if (graph_type_ == GRAPH_TYPE_VALUE_PIPNN) { diff --git a/src/algorithm/hgraph/hgraph_pipnn_storage_test.cpp b/src/algorithm/hgraph/hgraph_pipnn_storage_test.cpp new file mode 100644 index 0000000000..6456d3598f --- /dev/null +++ b/src/algorithm/hgraph/hgraph_pipnn_storage_test.cpp @@ -0,0 +1,333 @@ +// 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 "datacell/flatten_datacell.h" +#include "hgraph.h" +#include "impl/allocator/safe_allocator.h" +#include "impl/thread_pool/safe_thread_pool.h" +#include "index/index_impl.h" +#include "io/memory_io/memory_io.h" +#include "quantization/fp32_quantizer.h" +#include "unittest.h" +#include "vsag/options.h" + +namespace vsag { + +struct HGraphPiPNNStorageTestAccess { + static void + MaterializeWithCodes(HGraph& index, const DatasetPtr& data, FlattenInterfacePtr codes) { + auto original = index.basic_flatten_codes_; + index.basic_flatten_codes_ = std::move(codes); + Vector ids(index.allocator_); + Vector rows(index.allocator_); + for (int64_t row = 0; row < data->GetNumElements(); ++row) { + ids.push_back(static_cast(row)); + rows.push_back(row); + } + try { + index.materialize_pipnn_storage(data, ids, rows); + } catch (...) { + index.basic_flatten_codes_ = std::move(original); + throw; + } + index.basic_flatten_codes_ = std::move(original); + } +}; + +} // namespace vsag + +namespace { + +class FailingStorageFlatten + : public vsag::FlattenDataCell, + vsag::FixedLayout> { +public: + void + InsertVector(const void* vector, vsag::InnerIdType id) override { + ++attempts; + throw std::runtime_error("injected storage encoding failure"); + } + + std::atomic attempts{0}; +}; + +using Index = vsag::IndexImpl; + +class SmallStorageBlocks { +public: + SmallStorageBlocks() : previous_(vsag::Options::Instance().block_size_limit()) { + vsag::Options::Instance().set_block_size_limit(512ULL * 1024); + } + ~SmallStorageBlocks() { + vsag::Options::Instance().set_block_size_limit(previous_); + } + +private: + uint64_t previous_; +}; + +vsag::IndexCommonParam +make_common(int64_t dim, uint64_t threads = 1) { + vsag::IndexCommonParam common; + common.dim_ = dim; + common.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + common.thread_pool_ = vsag::SafeThreadPool::FactoryDefaultThreadPool(); + common.thread_pool_->SetPoolSize(threads); + return common; +} + +vsag::JsonType +make_parameters() { + return vsag::JsonType::Parse(R"({ + "base_quantization_type": "fp32", "graph_type": "pipnn", + "max_degree": 8, "ef_construction": 64, "build_thread_count": 1, + "support_duplicate": true, "deduplicate_storage": true + })"); +} + +vsag::DatasetPtr +make_data(std::vector& vectors, std::vector& labels, int64_t dim) { + return vsag::Dataset::Make() + ->NumElements(static_cast(labels.size())) + ->Dim(dim) + ->Ids(labels.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); +} + +void +require_counts(const std::shared_ptr& index, uint32_t logical, uint32_t physical) { + auto hgraph = std::dynamic_pointer_cast(index->GetInnerIndex()); + REQUIRE(hgraph != nullptr); + REQUIRE(index->GetNumElements() == logical); + REQUIRE(hgraph->GetCodeStorageCounts() == + std::pair{logical, physical}); +} + +void +require_labels(const std::shared_ptr& index, + float* vector, + int64_t dim, + const std::set& labels) { + auto query = + vsag::Dataset::Make()->NumElements(1)->Dim(dim)->Float32Vectors(vector)->Owner(false); + auto result = index->KnnSearch(query, + static_cast(labels.size()), + R"({"hgraph":{"ef_search":128,"brute_force_threshold":0.0}})"); + REQUIRE(result.has_value()); + REQUIRE(result->get()->GetDim() == static_cast(labels.size())); + const auto* ids = result->get()->GetIds(); + REQUIRE(std::set(ids, ids + labels.size()) == labels); +} + +} // namespace + +TEST_CASE("PiPNN storage dedup shares codes across quantizers and round trips", + "[ut][hgraph][pipnn][pipnn_storage][duplicate][serialize]") { + SmallStorageBlocks blocks; + constexpr int64_t dim = 64; + constexpr int64_t count = 128; + constexpr int64_t unique = 16; + const auto threads = GENERATE(1, 4); + const auto adaptive = GENERATE(false, true); + const auto* const encoding = GENERATE("fp32", "sq8_reorder", "rabitq_raw", "rabitq_deferred"); + CAPTURE(threads, adaptive, encoding); + auto common = make_common(dim, threads); + auto parameter = make_parameters(); + parameter["build_thread_count"].SetInt(threads); + parameter["pipnn_adaptive_pruning"].SetBool(adaptive); + const std::string kind = encoding; + parameter["store_raw_vector"].SetBool(kind != "rabitq_deferred"); + if (kind == "sq8_reorder") { + parameter["base_quantization_type"].SetString("sq8"); + parameter["use_reorder"].SetBool(true); + parameter["precise_quantization_type"].SetString("fp32"); + } else if (kind == "rabitq_raw" or kind == "rabitq_deferred") { + parameter["base_quantization_type"].SetString("rabitq"); + parameter["rabitq_bits_per_dim_base"].SetInt(1); + parameter["store_raw_vector"].SetBool(kind == "rabitq_raw"); + } + std::vector vectors(count * dim); + std::vector labels(count); + for (int64_t row = 0; row < count; ++row) { + labels[row] = 1000 + row * 7; + for (int64_t d = 0; d < dim; ++d) { + vectors[row * dim + d] = static_cast(((row % unique) * 13 + d * 7) % 71); + } + } + auto index = std::make_shared(parameter, common); + auto built = index->Build(make_data(vectors, labels, dim)); + REQUIRE(built.has_value()); + REQUIRE(built.value().empty()); + require_counts(index, count, unique); + + auto restored = std::make_shared(parameter, common); + SECTION("binary set") { + auto binary = index->Serialize(); + REQUIRE(binary.has_value()); + REQUIRE(restored->Deserialize(binary.value()).has_value()); + } + SECTION("stream") { + std::stringstream stream; + REQUIRE(index->Serialize(stream).has_value()); + stream.seekg(0); + REQUIRE(restored->Deserialize(stream).has_value()); + } + require_counts(restored, count, unique); + if (kind != "rabitq_deferred") { + auto raw = restored->GetRawVectorByIds(labels.data(), count, nullptr); + REQUIRE(raw.has_value()); + REQUIRE(std::equal(vectors.begin(), vectors.end(), raw.value()->GetFloat32Vectors())); + } + for (int64_t row = 0; row < count; ++row) { + auto distance = restored->CalcDistanceById(vectors.data(), labels[row]); + auto representative = restored->CalcDistanceById(vectors.data(), labels[row % unique]); + REQUIRE(distance.has_value()); + REQUIRE(representative.has_value()); + REQUIRE(distance.value() == representative.value()); + } + if (kind == "fp32" or kind == "sq8_reorder") { + std::set expected; + for (int64_t row = 0; row < count; row += unique) { + expected.insert(labels[row]); + } + require_labels(restored, vectors.data(), dim, expected); + } +} + +TEST_CASE("PiPNN storage dedup preserves filtered rows and extends after restore", + "[ut][hgraph][pipnn][pipnn_storage][duplicate][add]") { + SmallStorageBlocks blocks; + const auto first_add = GENERATE(false, true); + const auto adaptive = GENERATE(false, true); + auto common = make_common(2, 4); + auto parameter = make_parameters(); + parameter["build_thread_count"].SetInt(4); + parameter["pipnn_adaptive_pruning"].SetBool(adaptive); + auto index = std::make_shared(parameter, common); + std::vector vectors{0, 0, 1, 0, 99, 99, 0, 0, 1, 0, 3, 0}; + std::vector labels{90, 10, 90, 30, 50, 70}; + auto data = make_data(vectors, labels, 2); + auto built = first_add ? index->Add(data) : index->Build(data); + REQUIRE(built.has_value()); + REQUIRE(built.value() == std::vector{90}); + require_counts(index, 5, 3); + require_labels(index, vectors.data(), 2, {90, 30}); + require_labels(index, vectors.data() + 2, 2, {10, 50}); + + auto binary = index->Serialize(); + REQUIRE(binary.has_value()); + auto restored = std::make_shared(parameter, common); + REQUIRE(restored->Deserialize(binary.value()).has_value()); + std::vector added{1, 0, 7, 0}; + std::vector added_labels{110, 130}; + auto added_result = restored->Add(make_data(added, added_labels, 2)); + REQUIRE(added_result.has_value()); + REQUIRE(added_result.value().empty()); + require_counts(restored, 7, 4); + require_labels(restored, added.data(), 2, {10, 50, 110}); + require_labels(restored, added.data() + 2, 2, {130}); + + std::vector duplicate_label{150}; + std::vector duplicate{7, 0}; + REQUIRE(restored->Add(make_data(duplicate, duplicate_label, 2)).has_value()); + require_counts(restored, 8, 4); + require_labels(restored, duplicate.data(), 2, {130, 150}); +} + +TEST_CASE("PiPNN storage dedup handles one representative and unique-only input", + "[ut][hgraph][pipnn][pipnn_storage][duplicate]") { + SmallStorageBlocks blocks; + const auto count = GENERATE(1, 48); + const auto all_same = GENERATE(false, true); + auto common = make_common(4, 4); + auto parameter = make_parameters(); + parameter["build_thread_count"].SetInt(4); + std::vector vectors(static_cast(count) * 4, 1.0F); + std::vector labels(count); + for (int64_t row = 0; row < count; ++row) { + labels[row] = 100 + row; + if (not all_same) { + vectors[row * 4] = static_cast(row); + } + } + auto index = std::make_shared(parameter, common); + REQUIRE(index->Build(make_data(vectors, labels, 4)).has_value()); + require_counts(index, count, all_same ? 1 : count); + if (all_same) { + require_labels(index, vectors.data(), 4, std::set(labels.begin(), labels.end())); + } +} + +TEST_CASE("PiPNN storage dedup retains per-label extra info and filtering", + "[ut][hgraph][pipnn][pipnn_storage][duplicate][filter]") { + SmallStorageBlocks blocks; + auto common = make_common(2, 4); + common.extra_info_size_ = 2; + auto parameter = make_parameters(); + parameter["build_thread_count"].SetInt(4); + auto index = std::make_shared(parameter, common); + std::vector vectors{1, 0, 1, 0, 1, 0, 9, 0}; + std::vector labels{10, 20, 30, 40}; + std::vector extra{'a', 'A', 'b', 'B', 'c', 'C', 'd', 'D'}; + auto data = make_data(vectors, labels, 2)->ExtraInfoSize(2)->ExtraInfos(extra.data()); + REQUIRE(index->Build(data).has_value()); + require_counts(index, 4, 2); + auto binary = index->Serialize(); + REQUIRE(binary.has_value()); + auto restored = std::make_shared(parameter, common); + REQUIRE(restored->Deserialize(binary.value()).has_value()); + std::vector fetched(extra.size()); + REQUIRE(restored->GetExtraInfoByIds(labels.data(), 4, fetched.data()).has_value()); + REQUIRE(fetched == extra); + + auto query = + vsag::Dataset::Make()->NumElements(1)->Dim(2)->Float32Vectors(vectors.data())->Owner(false); + const std::function exclude = [](int64_t label) { return label != 20; }; + auto result = restored->KnnSearch(query, 1, R"({"hgraph":{"ef_search":64}})", exclude); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 1); + REQUIRE(result.value()->GetIds()[0] == 20); +} + +TEST_CASE("PiPNN storage failure drains workers without consuming slots", + "[ut][hgraph][pipnn][pipnn_storage][duplicate]") { + SmallStorageBlocks blocks; + const auto threads = GENERATE(1, 4); + auto parameter = make_parameters(); + parameter["build_thread_count"].SetInt(threads); + auto index = std::make_shared(parameter, make_common(2, threads)); + std::vector vectors{1, 0, 1, 0, 2, 0, 3, 0}; + std::vector labels{10, 20, 30, 40}; + auto data = make_data(vectors, labels, 2); + REQUIRE(index->Build(data).has_value()); + auto hgraph = std::dynamic_pointer_cast(index->GetInnerIndex()); + auto failing_codes = std::make_shared(); + REQUIRE_THROWS_AS( + vsag::HGraphPiPNNStorageTestAccess::MaterializeWithCodes(*hgraph, data, failing_codes), + std::runtime_error); + REQUIRE(failing_codes->attempts.load() == static_cast(threads == 1 ? 1 : 3)); + require_counts(index, 4, 3); + require_labels(index, vectors.data(), 2, {10, 20}); +} diff --git a/src/algorithm/hgraph/hgraph_pipnn_test.cpp b/src/algorithm/hgraph/hgraph_pipnn_test.cpp index 15de29e47c..154961d8af 100644 --- a/src/algorithm/hgraph/hgraph_pipnn_test.cpp +++ b/src/algorithm/hgraph/hgraph_pipnn_test.cpp @@ -323,7 +323,7 @@ TEST_CASE("HGraph validates PiPNN graph type and data boundary", "[ut][pipnn][hg auto parameter = MakePiPNNHGraphParam(); parameter["support_duplicate"].SetBool(true); parameter["deduplicate_storage"].SetBool(true); - REQUIRE_THROWS(MakePiPNNIndex(parameter, common_param)); + REQUIRE_NOTHROW(MakePiPNNIndex(parameter, common_param)); } SECTION("non-fp32 base codes") { auto parameter = MakePiPNNHGraphParam();