From a3041ff2d3a24ee8b6b97db500236ed079c4503a Mon Sep 17 00:00:00 2001 From: jac0626 Date: Wed, 22 Jul 2026 12:06:57 +0800 Subject: [PATCH 1/5] feat: support pyramid raw vector retrieval Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- .../docs/en/src/advanced/new_serialization.md | 1 + docs/docs/en/src/indexes/pyramid.md | 3 + .../docs/zh/src/advanced/new_serialization.md | 1 + docs/docs/zh/src/indexes/pyramid.md | 3 + ..._feature_compare_hgraph_pyramid_memory.cpp | 347 ++++++++++++++++++ examples/cpp/CMakeLists.txt | 3 + examples/cpp/README.md | 1 + src/algorithm/pyramid/pyramid.cpp | 180 +++++++++ src/algorithm/pyramid/pyramid.h | 18 +- src/algorithm/pyramid/pyramid_test.cpp | 289 +++++++++++++++ src/algorithm/pyramid/pyramid_zparameters.cpp | 27 ++ src/quantization/fp32_quantizer_parameter.cpp | 1 + .../fp32_quantizer_parameter_test.cpp | 4 +- 13 files changed, 876 insertions(+), 2 deletions(-) create mode 100644 examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp diff --git a/docs/docs/en/src/advanced/new_serialization.md b/docs/docs/en/src/advanced/new_serialization.md index f4389d24cd..a6dedec5ca 100644 --- a/docs/docs/en/src/advanced/new_serialization.md +++ b/docs/docs/en/src/advanced/new_serialization.md @@ -207,6 +207,7 @@ Pyramid writes these streaming blocks in order: | `label_table` | external labels and label remap | yes | | `base_codes` | base flatten codes used by graph search | yes | | `high_precision_codes` | precise reorder codes when reorder is enabled | conditional | +| `raw_vector` | separate FP32 vectors when `store_raw_vector` requires dedicated storage | conditional | | `pyramid_hierarchies` | hierarchy names and graph roots | yes | `DeserializeStreaming` restores the full in-memory Pyramid index. `Index::Load` can create the diff --git a/docs/docs/en/src/indexes/pyramid.md b/docs/docs/en/src/indexes/pyramid.md index 35ce1fc4d8..d00e4a863a 100644 --- a/docs/docs/en/src/indexes/pyramid.md +++ b/docs/docs/en/src/indexes/pyramid.md @@ -98,6 +98,9 @@ Build-time parameters live under `index_param`. | `fast_encode_rabitq_rounds` | int | `6` | Fast RaBitQ refinement rounds in `[1, 32]`. | | `base_io_type` / `precise_io_type` | string | `"block_memory_io"` | Base and reorder storage backends; `uring_io` is available in builds with liburing. | | `base_file_path` / `precise_file_path` | string | — | Required for disk-backed storage such as `buffer_io`, `async_io`, `uring_io`, or `mmap_io`. | +| `store_raw_vector` | bool | `false` | Preserve an FP32 copy for `GetRawVectorByIds` and precise distance-by-id calculations. Pyramid reuses an existing in-memory FP32 base/reorder copy when possible; otherwise it creates separate raw-vector storage. | +| `raw_vector_io_type` | string | `"block_memory_io"` | IO backend for separate raw-vector storage. | +| `raw_vector_file_path` | string | `"./default_file_path"` | File path used by file-backed raw-vector IO. | | `index_min_size` | int | `0` | Minimum sub-index size; smaller groups fall back to scan. | | `support_duplicate` | bool | `false` | Allow duplicate ids. | | `build_thread_count` | int | `1` | Threads used for parallel build. | diff --git a/docs/docs/zh/src/advanced/new_serialization.md b/docs/docs/zh/src/advanced/new_serialization.md index 1eba4cc7eb..5683cbc240 100644 --- a/docs/docs/zh/src/advanced/new_serialization.md +++ b/docs/docs/zh/src/advanced/new_serialization.md @@ -187,6 +187,7 @@ Pyramid 按顺序写入以下 streaming blocks： | `label_table` | 外部 label 和 label remap | 是 | | `base_codes` | 图搜索使用的 base flatten codes | 是 | | `high_precision_codes` | reorder 开启时的精排 codes | 条件必需 | +| `raw_vector` | `store_raw_vector` 需要独立存储时的 FP32 向量 | 条件必需 | | `pyramid_hierarchies` | hierarchy 名称和 graph roots | 是 | `DeserializeStreaming` 会恢复完整的内存 Pyramid 索引。`Index::Load` 可以直接从 streaming metadata diff --git a/docs/docs/zh/src/indexes/pyramid.md b/docs/docs/zh/src/indexes/pyramid.md index 9dd1a39d32..89f04e61ee 100644 --- a/docs/docs/zh/src/indexes/pyramid.md +++ b/docs/docs/zh/src/indexes/pyramid.md @@ -93,6 +93,9 @@ auto result = index->KnnSearch( | `fast_encode_rabitq_rounds` | int | `6` | RaBitQ 快速编码的微调轮数，范围 `[1, 32]` | | `base_io_type` / `precise_io_type` | string | `"block_memory_io"` | 底层与精排存储后端；以 liburing 构建时可用 `uring_io` | | `base_file_path` / `precise_file_path` | string | — | `buffer_io`、`async_io`、`uring_io`、`mmap_io` 等磁盘存储必须设置 | +| `store_raw_vector` | bool | `false` | 保留 FP32 向量，供 `GetRawVectorByIds` 和精确的按 ID 距离计算使用。若已有内存 FP32 base/reorder 副本，Pyramid 会直接复用；否则创建独立 raw-vector 存储。 | +| `raw_vector_io_type` | string | `"block_memory_io"` | 独立 raw-vector 存储使用的 IO 后端 | +| `raw_vector_file_path` | string | `"./default_file_path"` | 文件型 raw-vector IO 使用的文件路径 | | `index_min_size` | int | `0` | 子索引的最小规模；小于该值的分区会退化为线性扫描 | | `support_duplicate` | bool | `false` | 是否允许重复 ID | | `build_thread_count` | int | `1` | 构建阶段并发线程数 | diff --git a/examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp b/examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp new file mode 100644 index 0000000000..37f9a9c148 --- /dev/null +++ b/examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp @@ -0,0 +1,347 @@ +// 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 +#include +#include + +#if defined(__APPLE__) +#include +#elif defined(__linux__) +#include +#endif + +namespace { + +constexpr uint64_t DIM = 128; +constexpr uint64_t LEAF_COUNT = 9; +constexpr uint64_t DEFAULT_DATASET_MIB = 400; +constexpr uint64_t BYTES_PER_MIB = 1024 * 1024; + +struct Measurement { + std::string name; + uint64_t total_rss_bytes; + uint64_t tracked_bytes; + uint64_t reported_bytes; + std::vector> reported_detail; +}; + +class TrackingAllocator : public vsag::Allocator { +public: + std::string + Name() override { + return "memory-comparison-allocator"; + } + + void* + Allocate(uint64_t size) override { + auto* pointer = std::malloc(size); + if (pointer == nullptr) { + return nullptr; + } + std::lock_guard lock(mutex_); + allocations_[pointer] = size; + current_bytes_ += size; + peak_bytes_ = std::max(peak_bytes_, current_bytes_); + return pointer; + } + + void + Deallocate(void* pointer) override { + if (pointer == nullptr) { + return; + } + { + std::lock_guard lock(mutex_); + auto allocation = allocations_.find(pointer); + if (allocation != allocations_.end()) { + current_bytes_ -= allocation->second; + allocations_.erase(allocation); + } + } + std::free(pointer); + } + + void* + Reallocate(void* pointer, uint64_t size) override { + std::lock_guard lock(mutex_); + uint64_t old_size = 0; + auto allocation = allocations_.find(pointer); + if (allocation != allocations_.end()) { + old_size = allocation->second; + } + auto* new_pointer = std::realloc(pointer, size); + if (new_pointer == nullptr) { + return nullptr; + } + allocations_.erase(pointer); + allocations_[new_pointer] = size; + current_bytes_ = current_bytes_ - old_size + size; + peak_bytes_ = std::max(peak_bytes_, current_bytes_); + return new_pointer; + } + + uint64_t + CurrentBytes() const { + std::lock_guard lock(mutex_); + return current_bytes_; + } + +private: + mutable std::mutex mutex_; + std::unordered_map allocations_; + uint64_t current_bytes_{0}; + uint64_t peak_bytes_{0}; +}; + +double +to_mib(uint64_t bytes) { + return static_cast(bytes) / static_cast(BYTES_PER_MIB); +} + +uint64_t +get_current_rss() { +#if defined(__APPLE__) + mach_task_basic_info_data_t info{}; + mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; + if (task_info( + mach_task_self(), MACH_TASK_BASIC_INFO, reinterpret_cast(&info), &count) != + KERN_SUCCESS) { + return 0; + } + return info.resident_size; +#elif defined(__linux__) + std::ifstream statm("/proc/self/statm"); + uint64_t total_pages = 0; + uint64_t resident_pages = 0; + statm >> total_pages >> resident_pages; + return resident_pages * static_cast(sysconf(_SC_PAGESIZE)); +#else + return 0; +#endif +} + +Measurement +build_and_measure(const std::string& name, + const std::string& parameters, + const vsag::DatasetPtr& base) { + TrackingAllocator allocator; + const uint64_t baseline_rss = get_current_rss(); + vsag::Resource resource(&allocator, nullptr); + vsag::Engine engine(&resource); + const uint64_t baseline_bytes = allocator.CurrentBytes(); + auto create_result = engine.CreateIndex(name, parameters); + if (not create_result.has_value()) { + std::cerr << "Failed to create " << name << ": " << create_result.error().message + << std::endl; + std::exit(EXIT_FAILURE); + } + + auto index = create_result.value(); + auto build_result = index->Build(base); + if (not build_result.has_value()) { + std::cerr << "Failed to build " << name << ": " << build_result.error().message + << std::endl; + std::exit(EXIT_FAILURE); + } + + const uint64_t built_rss = get_current_rss(); + Measurement result{name, + built_rss > baseline_rss ? built_rss - baseline_rss : 0, + allocator.CurrentBytes() - baseline_bytes, + index->GetMemoryUsage(), + {}}; + const auto detail = index->GetMemoryUsageDetail(); + result.reported_detail.assign(detail.begin(), detail.end()); + std::sort(result.reported_detail.begin(), + result.reported_detail.end(), + [](const auto& lhs, const auto& rhs) { return lhs.second > rhs.second; }); + index.reset(); + engine.Shutdown(); + return result; +} + +} // namespace + +int +main(int argc, char** argv) { + // Keep internal block pools from reserving large default chunks that dominate small trial runs. + vsag::Options::Instance().set_block_size_limit(2 * BYTES_PER_MIB); + if (argc < 2 || (std::string(argv[1]) != "hgraph" && std::string(argv[1]) != "pyramid")) { + std::cerr << "Usage:\n " << argv[0] << " [dataset MiB]\n " << argv[0] + << " synthetic \n " << argv[0] + << " fbin " << std::endl; + return EXIT_FAILURE; + } + const std::string index_name = argv[1]; + uint64_t dim = DIM; + uint64_t num_vectors = 0; + std::vector vectors; + if (argc == 4 && std::string(argv[2]) == "fbin") { + std::ifstream input(argv[3], std::ios::binary); + uint32_t file_num_vectors = 0; + uint32_t file_dim = 0; + input.read(reinterpret_cast(&file_num_vectors), sizeof(file_num_vectors)); + input.read(reinterpret_cast(&file_dim), sizeof(file_dim)); + if (not input || file_num_vectors == 0 || file_dim == 0) { + std::cerr << "Invalid fbin header: " << argv[3] << std::endl; + return EXIT_FAILURE; + } + num_vectors = file_num_vectors; + dim = file_dim; + vectors.resize(num_vectors * dim); + input.read(reinterpret_cast(vectors.data()), + static_cast(vectors.size() * sizeof(float))); + if (not input) { + std::cerr << "Incomplete fbin data: " << argv[3] << std::endl; + return EXIT_FAILURE; + } + } else { + if (argc == 5 && std::string(argv[2]) == "synthetic") { + num_vectors = std::stoull(argv[3]); + dim = std::stoull(argv[4]); + } else { + const uint64_t dataset_mib = argc > 2 ? std::stoull(argv[2]) : DEFAULT_DATASET_MIB; + const uint64_t vector_bytes = dataset_mib * BYTES_PER_MIB; + num_vectors = vector_bytes / (dim * sizeof(float)); + } + vectors.resize(num_vectors * dim); + for (uint64_t i = 0; i < num_vectors; ++i) { + for (uint64_t j = 0; j < dim; ++j) { + uint64_t value = (i * dim + j) + 0x9E3779B97F4A7C15ULL; + value = (value ^ (value >> 30U)) * 0xBF58476D1CE4E5B9ULL; + value = (value ^ (value >> 27U)) * 0x94D049BB133111EBULL; + value ^= value >> 31U; + vectors[i * dim + j] = + static_cast(value & 0xFFFFFFU) / static_cast(0x1000000U); + } + } + } + const uint64_t actual_vector_bytes = vectors.size() * sizeof(float); + + std::vector ids(num_vectors); + std::vector paths(num_vectors); + const std::array leaf_paths = { + "leaf-0", "leaf-1", "leaf-2", "leaf-3", "leaf-4", "leaf-5", "leaf-6", "leaf-7", "leaf-8"}; + + for (uint64_t i = 0; i < num_vectors; ++i) { + ids[i] = static_cast(i); + paths[i] = leaf_paths[i % LEAF_COUNT]; + } + + auto hgraph_base = vsag::Dataset::Make(); + hgraph_base->NumElements(static_cast(num_vectors)) + ->Dim(static_cast(dim)) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); + + auto pyramid_base = vsag::Dataset::Make(); + pyramid_base->NumElements(static_cast(num_vectors)) + ->Dim(static_cast(dim)) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Paths(paths.data()) + ->Owner(false); + + std::string hgraph_parameters = R"( + { + "dtype": "float32", + "metric_type": "l2", + "dim": 128, + "index_param": { + "base_quantization_type": "fp32", + "max_degree": 64, + "ef_construction": 200, + "alpha": 1.2, + "use_reorder": false, + "use_reverse_edges": true, + "support_force_remove": true, + "store_raw_vector": true, + "base_io_type": "block_memory_io", + "build_thread_count": 1 + } + } + )"; + + std::string pyramid_parameters = R"( + { + "dtype": "float32", + "metric_type": "l2", + "dim": 128, + "index_param": { + "base_quantization_type": "fp32", + "max_degree": 32, + "ef_construction": 400, + "alpha": 1.2, + "graph_type": "nsw", + "no_build_levels": [0], + "use_reorder": false, + "store_raw_vector": true, + "index_min_size": 0, + "build_thread_count": 1 + } + } + )"; + + const std::string default_dim = "\"dim\": 128"; + const std::string actual_dim = "\"dim\": " + std::to_string(dim); + hgraph_parameters.replace(hgraph_parameters.find(default_dim), default_dim.size(), actual_dim); + pyramid_parameters.replace( + pyramid_parameters.find(default_dim), default_dim.size(), actual_dim); + + std::cout << "Dataset: " << num_vectors << " x " << dim << " float32 vectors (" << std::fixed + << std::setprecision(2) << to_mib(actual_vector_bytes) << " MiB)" << std::endl; + if (index_name == "pyramid") { + std::cout << "Pyramid layout: one root (graph disabled) + " << LEAF_COUNT << " leaf graphs" + << std::endl; + } + + const auto& parameters = index_name == "hgraph" ? hgraph_parameters : pyramid_parameters; + const auto& base = index_name == "hgraph" ? hgraph_base : pyramid_base; + const auto measurement = build_and_measure(index_name, parameters, base); + + std::cout << "\n" << measurement.name << " memory after Build():\n"; + std::cout << " Total process RSS increase: " << to_mib(measurement.total_rss_bytes) + << " MiB\n"; + std::cout << " VSAG allocator subset: " << to_mib(measurement.tracked_bytes) << " MiB\n"; + std::cout << " Index::GetMemoryUsage(): " << to_mib(measurement.reported_bytes) << " MiB\n"; + std::cout << " Total RSS bytes/vector: " + << static_cast(measurement.total_rss_bytes) / static_cast(num_vectors) + << "\n"; + + if (not measurement.reported_detail.empty()) { + uint64_t detail_total = 0; + std::cout << "\n Index::GetMemoryUsageDetail():\n"; + for (const auto& [component, bytes] : measurement.reported_detail) { + detail_total += bytes; + std::cout << " " << std::left << std::setw(24) << component << std::right + << std::setw(10) << to_mib(bytes) << " MiB\n"; + } + std::cout << " " << std::left << std::setw(24) << "detail total" << std::right + << std::setw(10) << to_mib(detail_total) << " MiB\n"; + } + + return 0; +} diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 3793625a87..ab5b290a5d 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -157,3 +157,6 @@ target_link_libraries(325_feature_uring_io vsag) add_executable(324_feature_hgraph_mci_companion 324_feature_hgraph_mci_companion.cpp) target_link_libraries(324_feature_hgraph_mci_companion vsag) +add_executable(325_feature_compare_hgraph_pyramid_memory + 325_feature_compare_hgraph_pyramid_memory.cpp) +target_link_libraries(325_feature_compare_hgraph_pyramid_memory vsag) diff --git a/examples/cpp/README.md b/examples/cpp/README.md index 879bce936d..be5230e0d3 100644 --- a/examples/cpp/README.md +++ b/examples/cpp/README.md @@ -102,6 +102,7 @@ together when the directory is listed: | [`320_feature_extra_info.cpp`](320_feature_extra_info.cpp) | Attach per-vector extra info / payload. | | [`322_feature_hgraph_brute_force_threshold.cpp`](322_feature_hgraph_brute_force_threshold.cpp) | HGraph search-time `brute_force_threshold`: automatically switch to an exact scan under highly selective filters. | | [`324_feature_lazy_hgraph_extra_info.cpp`](324_feature_lazy_hgraph_extra_info.cpp) | LazyHGraph `extra_info` filtering across flat and graph phases. | +| [`325_feature_compare_hgraph_pyramid_memory.cpp`](325_feature_compare_hgraph_pyramid_memory.cpp) | Compare HGraph with a nine-leaf Pyramid on the same 400 MiB vector dataset. | ### Persistence (`4xx`) diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index 3011b2af1f..756f949ebb 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -168,6 +168,23 @@ IndexNode::Serialize(StreamWriter& writer) const { item.second->Serialize(writer); } } + +uint64_t +IndexNode::GetMemoryUsage() const { + std::shared_lock lock(mutex_); + uint64_t memory = sizeof(IndexNode) + ids_.capacity() * sizeof(InnerIdType); + memory += + children_.bucket_count() * (sizeof(decltype(children_)::value_type) + sizeof(uint32_t)); + for (const auto& [key, child] : children_) { + memory += key.capacity() + 1; + memory += child->GetMemoryUsage(); + } + if (graph_ != nullptr) { + memory += graph_->GetMemoryUsage(); + } + return memory; +} + void IndexNode::Init() { if (status_ == Status::NO_INDEX) { @@ -228,6 +245,9 @@ Pyramid::build_by_odescent(const DatasetPtr& base) { if (use_reorder_) { precise_codes_->BatchInsertVector(data_vectors, data_num); } + if (create_new_raw_vector_) { + raw_vector_->BatchInsertVector(data_vectors, data_num); + } auto codes = use_reorder_ ? precise_codes_ : base_codes_; if (thread_pool_ != nullptr && hierarchies_.size() > 1) { @@ -420,6 +440,44 @@ Pyramid::GetNumberRemoved() const { return delete_count_.load(); } +uint64_t +Pyramid::GetMemoryUsage() const { + auto detail = GetMemoryUsageDetail(); + uint64_t memory = sizeof(Pyramid); + for (const auto& [name, usage] : detail) { + (void)name; + memory += usage; + } + return memory; +} + +std::unordered_map +Pyramid::GetMemoryUsageDetail() const { + std::shared_lock lock(resize_mutex_); + std::unordered_map memory_usage; + memory_usage["points_mutex"] = points_mutex_ == nullptr ? 0 : points_mutex_->GetMemoryUsage(); + memory_usage["pool"] = pool_ == nullptr ? 0 : pool_->GetMemoryUsage(); + memory_usage["label_table"] = label_table_ == nullptr ? 0 : label_table_->GetMemoryUsage(); + memory_usage["base_codes"] = base_codes_ == nullptr ? 0 : base_codes_->GetMemoryUsage(); + if (use_reorder_ && precise_codes_ != nullptr) { + memory_usage["precise_codes"] = precise_codes_->GetMemoryUsage(); + } + if (create_new_raw_vector_ && raw_vector_ != nullptr) { + memory_usage["raw_vector"] = raw_vector_->GetMemoryUsage(); + } + + uint64_t hierarchy_memory = hierarchies_.bucket_count() * + (sizeof(decltype(hierarchies_)::value_type) + sizeof(uint32_t)); + for (const auto& [name, hierarchy] : hierarchies_) { + hierarchy_memory += name.capacity() + 1; + hierarchy_memory += sizeof(Hierarchy) + hierarchy->name.capacity() + 1; + hierarchy_memory += hierarchy->no_build_levels.capacity() * sizeof(int32_t); + hierarchy_memory += hierarchy->root->GetMemoryUsage(); + } + memory_usage["hierarchies"] = hierarchy_memory; + return memory_usage; +} + uint32_t Pyramid::Remove(const std::vector& ids, RemoveMode mode) { if (mode != RemoveMode::MARK_REMOVE) { @@ -438,6 +496,9 @@ Pyramid::Serialize(StreamWriter& writer) const { if (use_reorder_) { precise_codes_->Serialize(writer); } + if (create_new_raw_vector_) { + raw_vector_->Serialize(writer); + } auto pyramid_param = std::dynamic_pointer_cast(create_param_ptr_); if (pyramid_param && pyramid_param->has_hierarchies) { @@ -492,6 +553,13 @@ Pyramid::collect_streaming_header() const { StreamSerializationBlockCurrentVersion(tag), StreamSerializationTagCritical(tag)); } + if (this->create_new_raw_vector_) { + auto tag = static_cast(StreamSerializationTag::RAW_VECTOR); + AppendStreamingManifestBlock(manifest, + tag, + StreamSerializationBlockCurrentVersion(tag), + StreamSerializationTagCritical(tag)); + } AppendStreamingManifestBlock(manifest, hierarchy_tag, StreamSerializationBlockCurrentVersion(hierarchy_tag), @@ -537,6 +605,13 @@ Pyramid::serialize_streaming_body(StreamWriter& writer) const { this->precise_codes_->Serialize(w); }); } + if (this->create_new_raw_vector_) { + auto tag = static_cast(StreamSerializationTag::RAW_VECTOR); + WriteStreamingBlock( + writer, tag, StreamSerializationTagCritical(tag), [this](StreamWriter& w) { + this->raw_vector_->Serialize(w); + }); + } WriteStreamingBlock(writer, hierarchy_tag, StreamSerializationTagCritical(hierarchy_tag), @@ -601,6 +676,7 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) bool loaded_label_table = false; bool loaded_base_codes = false; bool loaded_precise_codes = false; + bool loaded_raw_vector = false; bool loaded_hierarchies = false; while (true) { @@ -652,6 +728,15 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) loaded_precise_codes = true; } break; + case StreamSerializationTag::RAW_VECTOR: + if (this->create_new_raw_vector_) { + ReadSeekableBlockPayload( + block_reader, block_header, [this](StreamReader& block) { + this->raw_vector_->Deserialize(block); + }); + loaded_raw_vector = true; + } + break; case StreamSerializationTag::PYRAMID_HIERARCHIES: ReadSeekableBlockPayload( block_reader, block_header, [this, &basic_info](StreamReader& block) { @@ -685,6 +770,10 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) throw VsagException(ErrorType::READ_ERROR, "Pyramid streaming serialization precise codes block is missing"); } + if (this->create_new_raw_vector_ && !loaded_raw_vector) { + throw VsagException(ErrorType::READ_ERROR, + "Pyramid streaming serialization raw vector block is missing"); + } resize(max_capacity); this->current_memory_usage_ = static_cast(this->CalSerializeSize()); @@ -698,6 +787,17 @@ Pyramid::Deserialize(StreamReader& reader) { throw VsagException(ErrorType::READ_ERROR, "failed to read index footer"); } auto max_capacity = basic_info["max_capacity"].GetInt(); + if (basic_info.Contains(INDEX_PARAM)) { + auto index_param = std::make_shared(); + index_param->FromString(basic_info[INDEX_PARAM].GetString()); + if (not this->create_param_ptr_->CheckCompatibility(index_param)) { + auto message = fmt::format("Pyramid index parameter not match, current: {}, new: {}", + this->create_param_ptr_->ToString(), + index_param->ToString()); + logger::error(message); + throw VsagException(ErrorType::INVALID_ARGUMENT, message); + } + } BufferStreamReader buffer_reader( &reader, std::numeric_limits::max(), this->allocator_); @@ -709,6 +809,9 @@ Pyramid::Deserialize(StreamReader& reader) { if (use_reorder_) { precise_codes_->Deserialize(buffer_reader); } + if (create_new_raw_vector_) { + raw_vector_->Deserialize(buffer_reader); + } cur_element_count_ = base_codes_->TotalCount(); auto param_json = JsonType::Parse(basic_info[INDEX_PARAM].GetString()); @@ -787,6 +890,10 @@ Pyramid::Add(const DatasetPtr& base) { precise_codes_->InsertVector(data_vectors + dim_ * i, valid_id_count + local_cur_element_count); } + if (create_new_raw_vector_) { + raw_vector_->InsertVector(data_vectors + dim_ * i, + valid_id_count + local_cur_element_count); + } valid_id_count++; data_biases.push_back(i); } else { @@ -819,6 +926,9 @@ Pyramid::resize(int64_t new_max_capacity) { if (use_reorder_) { precise_codes_->Resize(new_max_capacity); } + if (create_new_raw_vector_) { + raw_vector_->Resize(new_max_capacity); + } points_mutex_->Resize(new_max_capacity); max_capacity_ = new_max_capacity; } @@ -867,6 +977,9 @@ Pyramid::InitFeatures() { IndexFeature::SUPPORT_EXPORT_MODEL, IndexFeature::SUPPORT_GET_MEMORY_USAGE, }); + if (raw_vector_ != nullptr) { + this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_GET_RAW_VECTOR_BY_IDS); + } this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_DELETE_BY_ID); } @@ -929,6 +1042,18 @@ static const std::string HGRAPH_PARAMS_TEMPLATE = "{HOLD_MOLDS}": false } }, + "{STORE_RAW_VECTOR_KEY}": false, + "{RAW_VECTOR_KEY}": { + "{IO_PARAMS_KEY}": { + "{TYPE_KEY}": "{IO_TYPE_VALUE_BLOCK_MEMORY_IO}", + "{IO_FILE_PATH_KEY}": "{DEFAULT_FILE_PATH_VALUE}" + }, + "{CODES_TYPE_KEY}": "flatten", + "{QUANTIZATION_PARAMS_KEY}": { + "{TYPE_KEY}": "{QUANTIZATION_TYPE_VALUE_FP32}", + "{HOLD_MOLDS}": true + } + }, "{BUILD_THREAD_COUNT_KEY}": 1, "{EF_CONSTRUCTION_KEY}": 400, "{NO_BUILD_LEVELS}":[], @@ -971,6 +1096,11 @@ Pyramid::CheckAndMappingExternalParam(const JsonType& external_param, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, PRODUCT_QUANTIZATION_DIM_KEY}}, {PYRAMID_BASE_FILE_PATH, {BASE_CODES_KEY, IO_PARAMS_KEY, IO_FILE_PATH_KEY}}, {PYRAMID_PRECISE_FILE_PATH, {PRECISE_CODES_KEY, IO_PARAMS_KEY, IO_FILE_PATH_KEY}}, + {STORE_RAW_VECTOR, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, HOLD_MOLDS}}, + {STORE_RAW_VECTOR, {PRECISE_CODES_KEY, QUANTIZATION_PARAMS_KEY, HOLD_MOLDS}}, + {STORE_RAW_VECTOR, {STORE_RAW_VECTOR_KEY}}, + {RAW_VECTOR_IO_TYPE, {RAW_VECTOR_KEY, IO_PARAMS_KEY, TYPE_KEY}}, + {RAW_VECTOR_FILE_PATH, {RAW_VECTOR_KEY, IO_PARAMS_KEY, IO_FILE_PATH_KEY}}, {ODESCENT_PARAMETER_BUILD_BLOCK_SIZE, {GRAPH_KEY, ODESCENT_PARAMETER_BUILD_BLOCK_SIZE}}, {ODESCENT_PARAMETER_MIN_IN_DEGREE, {GRAPH_KEY, ODESCENT_PARAMETER_MIN_IN_DEGREE}}, {ODESCENT_PARAMETER_GRAPH_ITER_TURN, {GRAPH_KEY, ODESCENT_PARAMETER_GRAPH_ITER_TURN}}, @@ -994,6 +1124,9 @@ Pyramid::Train(const DatasetPtr& base) { if (use_reorder_) { this->precise_codes_->Train(base->GetFloat32Vectors(), base->GetNumElements()); } + if (create_new_raw_vector_) { + this->raw_vector_->Train(base->GetFloat32Vectors(), base->GetNumElements()); + } } std::vector Pyramid::Build(const DatasetPtr& base) { @@ -1328,6 +1461,9 @@ Pyramid::CalcDistanceById(const float* query, int64_t id, bool calculate_precise if (use_reorder_ && calculate_precise_distance) { flat = this->precise_codes_; } + if (raw_vector_ != nullptr && calculate_precise_distance) { + flat = this->raw_vector_; + } return InnerIndexInterface::calc_distance_by_id(query, id, flat); } @@ -1350,6 +1486,9 @@ Pyramid::CalDistanceById(const float* query, if (use_reorder_ && calculate_precise_distance) { flat = this->precise_codes_; } + if (raw_vector_ != nullptr && calculate_precise_distance) { + flat = this->raw_vector_; + } std::vector validity; auto result = InnerIndexInterface::cal_distance_by_id(query, ids, count, flat, &validity); if (topk == -1) { @@ -1362,14 +1501,55 @@ void Pyramid::GetVectorByInnerId(InnerIdType inner_id, float* data) const { std::shared_lock lock(resize_mutex_); auto codes = (use_reorder_) ? precise_codes_ : base_codes_; + codes = raw_vector_ != nullptr ? raw_vector_ : codes; bool release = false; const auto* buffer = codes->GetCodesById(inner_id, release); + if (buffer == nullptr) { + throw VsagException(ErrorType::INTERNAL_ERROR, + fmt::format("failed to get vector by inner id {}", inner_id)); + } codes->Decode(buffer, data); if (release) { codes->Release(buffer); } } +void +Pyramid::check_and_init_raw_vector(const FlattenInterfaceParamPtr& raw_vector_param, + const IndexCommonParam& common_param) { + if (raw_vector_param == nullptr) { + return; + } + + raw_vector_ = FlattenInterface::MakeInstance(raw_vector_param, common_param); + if (base_codes_->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32 && + precise_codes_ == nullptr) { + create_new_raw_vector_ = true; + return; + } + if (base_codes_->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32 && + precise_codes_ != nullptr && + precise_codes_->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32) { + create_new_raw_vector_ = true; + return; + } + + auto io_type_name = raw_vector_param->io_parameter->GetTypeName(); + if (io_type_name != IO_TYPE_VALUE_BLOCK_MEMORY_IO && io_type_name != IO_TYPE_VALUE_MEMORY_IO) { + create_new_raw_vector_ = true; + return; + } + + if (base_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { + raw_vector_ = base_codes_; + return; + } + if (precise_codes_ != nullptr && + precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { + raw_vector_ = precise_codes_; + } +} + std::string Pyramid::GetStats() const { AnalyzerParam analyzer_param(allocator_); diff --git a/src/algorithm/pyramid/pyramid.h b/src/algorithm/pyramid/pyramid.h index 26f90f675b..5cb0c43841 100644 --- a/src/algorithm/pyramid/pyramid.h +++ b/src/algorithm/pyramid/pyramid.h @@ -92,6 +92,9 @@ class IndexNode { void Deserialize(StreamReader& reader); + [[nodiscard]] uint64_t + GetMemoryUsage() const; + friend class Pyramid; friend class PyramidAnalyzer; @@ -168,6 +171,7 @@ class Pyramid : public InnerIndexInterface { FlattenInterface::MakeInstance(pyramid_param->precise_codes_param, common_param); reorder_ = std::make_shared(precise_codes_, allocator_); } + check_and_init_raw_vector(pyramid_param->raw_vector_param, common_param); } explicit Pyramid(const ParamPtr& param, const IndexCommonParam& common_param) @@ -226,6 +230,12 @@ class Pyramid : public InnerIndexInterface { int64_t GetNumberRemoved() const override; + [[nodiscard]] uint64_t + GetMemoryUsage() const override; + + [[nodiscard]] std::unordered_map + GetMemoryUsageDetail() const override; + uint32_t Remove(const std::vector& ids, RemoveMode mode) override; @@ -286,6 +296,10 @@ class Pyramid : public InnerIndexInterface { void deserialize_hierarchies(StreamReader& reader, const JsonType& basic_info); + void + check_and_init_raw_vector(const FlattenInterfaceParamPtr& raw_vector_param, + const IndexCommonParam& common_param); + /// One named hierarchy with its own root IndexNode and build parameters. struct Hierarchy { std::string name; // hierarchy name (empty = default) @@ -367,6 +381,7 @@ class Pyramid : public InnerIndexInterface { UnorderedMap> hierarchies_; // named hierarchies FlattenInterfacePtr base_codes_{nullptr}; // coarse codes for graph build/search FlattenInterfacePtr precise_codes_{nullptr}; // precise codes for reorder (if enabled) + FlattenInterfacePtr raw_vector_{nullptr}; // raw FP32 vectors (if enabled) std::unique_ptr pool_ = nullptr; // pool of visited-lists for search MutexArrayPtr points_mutex_{nullptr}; // per-point locks for concurrent access @@ -385,7 +400,8 @@ class Pyramid : public InnerIndexInterface { 2021}; // random number generator for level promotion ReorderInterfacePtr reorder_{nullptr}; // reorder helper (if use_reorder_) - uint32_t index_min_size_{0}; // min node size before graph is built + uint32_t index_min_size_{0}; // min node size before graph is built + bool create_new_raw_vector_{false}; // whether raw_vector_ owns separate storage }; } // namespace vsag diff --git a/src/algorithm/pyramid/pyramid_test.cpp b/src/algorithm/pyramid/pyramid_test.cpp index 3e7ba18c9e..310a508353 100644 --- a/src/algorithm/pyramid/pyramid_test.cpp +++ b/src/algorithm/pyramid/pyramid_test.cpp @@ -15,11 +15,14 @@ #include "pyramid.h" +#include +#include #include #include "impl/allocator/safe_allocator.h" #include "index_common_param.h" #include "unittest.h" +#include "vsag/index.h" namespace { @@ -157,3 +160,289 @@ TEST_CASE("Pyramid promotes flat node at index minimum size", "[ut][pyramid]") { REQUIRE(result->GetIds()[0] == ids[i]); } } + +TEST_CASE("Pyramid stores and restores raw vectors", "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + constexpr int64_t count = 3; + std::array vectors = { + 0.0F, + 0.0F, + 0.0F, + 0.0F, + 0.123456F, + 0.234567F, + 0.345678F, + 0.456789F, + 1.0F, + 1.0F, + 1.0F, + 1.0F, + }; + std::array ids = {10, 11, 12}; + std::array paths = {"leaf", "leaf", "leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto pyramid_param = std::dynamic_pointer_cast(param); + + REQUIRE(pyramid_param != nullptr); + REQUIRE(pyramid_param->store_raw_vector); + REQUIRE(pyramid_param->raw_vector_param != nullptr); + REQUIRE(pyramid_param->raw_vector_param->quantizer_parameter->GetTypeName() == "fp32"); + + auto dataset = vsag::Dataset::Make() + ->NumElements(count) + ->Dim(dim) + ->Float32Vectors(vectors.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + auto index = std::make_shared(pyramid_param, common_param); + index->InitFeatures(); + REQUIRE(index->CheckFeature(vsag::IndexFeature::SUPPORT_GET_RAW_VECTOR_BY_IDS)); + REQUIRE(index->Build(dataset).empty()); + + std::array restored{}; + index->GetVectorByInnerId(1, restored.data()); + for (int64_t i = 0; i < dim; ++i) { + REQUIRE(restored[i] == vectors[dim + i]); + } + REQUIRE(index->CalcDistanceById(vectors.data() + dim, ids[1], true) == 0.0F); + + auto binary_set = index->vsag::InnerIndexInterface::Serialize(); + auto loaded = std::make_shared(pyramid_param, common_param); + loaded->vsag::InnerIndexInterface::Deserialize(binary_set); + restored.fill(0.0F); + loaded->GetVectorByInnerId(1, restored.data()); + for (int64_t i = 0; i < dim; ++i) { + REQUIRE(restored[i] == vectors[dim + i]); + } + + std::stringstream streaming_buffer; + index->SerializeStreaming(streaming_buffer); + auto streaming_loaded = std::make_shared(pyramid_param, common_param); + std::stringstream streaming_reader(streaming_buffer.str()); + streaming_loaded->DeserializeStreaming(streaming_reader); + restored.fill(0.0F); + streaming_loaded->GetVectorByInnerId(1, restored.data()); + for (int64_t i = 0; i < dim; ++i) { + REQUIRE(restored[i] == vectors[dim + i]); + } +} + +TEST_CASE("Pyramid stores raw vectors during ODescent build", "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + std::array vectors = { + 0.123456F, 0.234567F, 0.345678F, 0.456789F, 1.0F, 1.0F, 1.0F, 1.0F}; + std::array ids = {10, 11}; + std::array paths = {"leaf", "leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "store_raw_vector": true, + "graph_type": "odescent", + "max_degree": 4, + "no_build_levels": [0, 1] + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto index = std::make_shared(param, common_param); + auto dataset = vsag::Dataset::Make() + ->NumElements(2) + ->Dim(dim) + ->Float32Vectors(vectors.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + + REQUIRE(index->Build(dataset).empty()); + std::array restored{}; + index->GetVectorByInnerId(0, restored.data()); + for (int64_t i = 0; i < dim; ++i) { + REQUIRE(restored[i] == vectors[i]); + } +} + +TEST_CASE("Pyramid reuses FP32 codes as raw vectors", "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + std::array vector = {0.123456F, 0.234567F, 0.345678F, 0.456789F}; + std::array ids = {10}; + std::array paths = {"leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "fp32", + "use_reorder": true, + "precise_quantization_type": "sq8", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto index = std::make_shared(param, common_param); + auto dataset = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(vector.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + + REQUIRE(index->Build(dataset).empty()); + std::array restored{}; + index->GetVectorByInnerId(0, restored.data()); + REQUIRE(restored == vector); +} + +TEST_CASE("Pyramid streaming load preserves cosine raw vectors", "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + std::array vectors = {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F}; + std::array ids = {10, 11}; + std::array paths = {"leaf", "leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_COSINE; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "fp32", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto index = std::make_shared(param, common_param); + auto dataset = vsag::Dataset::Make() + ->NumElements(2) + ->Dim(dim) + ->Float32Vectors(vectors.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + + REQUIRE(index->Build(dataset).empty()); + std::stringstream stream; + index->SerializeStreaming(stream); + + std::stringstream load_stream(stream.str()); + auto loaded = vsag::Index::Load(load_stream, "{}"); + REQUIRE(loaded.has_value()); + auto raw = loaded.value()->GetRawVectorByIds(ids.data(), 2); + REQUIRE(raw.has_value()); + const auto* restored = raw.value()->GetFloat32Vectors(); + for (int64_t i = 0; i < dim * 2; ++i) { + REQUIRE(restored[i] == vectors[i]); + } +} + +TEST_CASE("Pyramid legacy deserialize validates raw vector config before reading", + "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + std::array vector = {1.0F, 2.0F, 3.0F, 4.0F}; + std::array ids = {10}; + std::array paths = {"leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto stored_external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto stored_param = + vsag::Pyramid::CheckAndMappingExternalParam(stored_external_param, common_param); + auto stored = std::make_shared(stored_param, common_param); + auto dataset = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(vector.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + REQUIRE(stored->Build(dataset).empty()); + auto binary_set = stored->vsag::InnerIndexInterface::Serialize(); + + auto target_external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "store_raw_vector": false, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto target_param = + vsag::Pyramid::CheckAndMappingExternalParam(target_external_param, common_param); + auto target = std::make_shared(target_param, common_param); + + REQUIRE_THROWS(target->vsag::InnerIndexInterface::Deserialize(binary_set)); + REQUIRE(target->GetNumElements() == 0); +} + +TEST_CASE("Pyramid reports live raw vector memory", "[ut][pyramid][raw_vector][memory]") { + constexpr int64_t dim = 4; + std::array vectors = { + 0.123456F, 0.234567F, 0.345678F, 0.456789F, 1.0F, 2.0F, 3.0F, 4.0F}; + std::array ids = {10, 11}; + std::array paths = {"leaf", "leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto index = std::make_shared(param, common_param); + auto dataset = vsag::Dataset::Make() + ->NumElements(2) + ->Dim(dim) + ->Float32Vectors(vectors.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + + REQUIRE(index->Build(dataset).empty()); + REQUIRE(index->GetMemoryUsage() > 0); + auto detail = index->GetMemoryUsageDetail(); + REQUIRE(detail.at("base_codes") > 0); + REQUIRE(detail.at("raw_vector") >= sizeof(float) * dim * 2); + REQUIRE(index->GetMemoryUsage() >= detail.at("raw_vector")); +} diff --git a/src/algorithm/pyramid/pyramid_zparameters.cpp b/src/algorithm/pyramid/pyramid_zparameters.cpp index 5bd0b3b27a..f51b63cf8c 100644 --- a/src/algorithm/pyramid/pyramid_zparameters.cpp +++ b/src/algorithm/pyramid/pyramid_zparameters.cpp @@ -180,6 +180,29 @@ PyramidParameters::FromJson(const JsonType& json) { this->precise_codes_param = CreateFlattenParam(json[PRECISE_CODES_KEY]); } + auto restore_legacy_raw_fp32_param = [this, &json](const FlattenInterfaceParamPtr& param, + const char* key) { + if (not this->store_raw_vector || param == nullptr || not json.Contains(key)) { + return; + } + const auto& flatten_json = json[key]; + if (not flatten_json.Contains(QUANTIZATION_PARAMS_KEY)) { + return; + } + const auto& quantization_json = flatten_json[QUANTIZATION_PARAMS_KEY]; + if (quantization_json.Contains(HOLD_MOLDS) || + param->quantizer_parameter->GetTypeName() != QUANTIZATION_TYPE_VALUE_FP32) { + return; + } + std::dynamic_pointer_cast(param->quantizer_parameter)->hold_molds = + true; + }; + restore_legacy_raw_fp32_param(this->base_codes_param, BASE_CODES_KEY); + if (this->use_reorder) { + restore_legacy_raw_fp32_param(this->precise_codes_param, PRECISE_CODES_KEY); + } + restore_legacy_raw_fp32_param(this->raw_vector_param, RAW_VECTOR_KEY); + if (json.Contains(INDEX_MIN_SIZE)) { this->index_min_size = json[INDEX_MIN_SIZE].GetInt(); } @@ -295,6 +318,10 @@ PyramidParameters::CheckCompatibility(const ParamPtr& other) const { if (this->use_reorder) { CHECK_SUB_PARAM(*this, *p, precise_codes_param); } + CHECK_FIELD_EQ(*this, *p, store_raw_vector); + if (this->store_raw_vector) { + CHECK_SUB_PARAM(*this, *p, raw_vector_param); + } CHECK_FIELD_EQ(*this, *p, index_min_size); CHECK_FIELD_EQ(*this, *p, support_duplicate); return true; diff --git a/src/quantization/fp32_quantizer_parameter.cpp b/src/quantization/fp32_quantizer_parameter.cpp index 2b4adc0464..ef82a1e73f 100644 --- a/src/quantization/fp32_quantizer_parameter.cpp +++ b/src/quantization/fp32_quantizer_parameter.cpp @@ -33,6 +33,7 @@ JsonType FP32QuantizerParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(QUANTIZATION_TYPE_VALUE_FP32); + json[HOLD_MOLDS].SetBool(hold_molds); return json; } } // namespace vsag diff --git a/src/quantization/fp32_quantizer_parameter_test.cpp b/src/quantization/fp32_quantizer_parameter_test.cpp index 53e0edc3b9..8b5264ace1 100644 --- a/src/quantization/fp32_quantizer_parameter_test.cpp +++ b/src/quantization/fp32_quantizer_parameter_test.cpp @@ -15,15 +15,17 @@ #include "fp32_quantizer_parameter.h" +#include "inner_string_params.h" #include "parameter_test.h" #include "unittest.h" using namespace vsag; TEST_CASE("FP32 Quantizer Parameter ToJson Test", "[ut][FP32QuantizerParameter]") { - std::string param_str = "{}"; + std::string param_str = R"({"hold_molds": true})"; auto param = std::make_shared(); JsonType param_json = JsonType::Parse(param_str); param->FromJson(param_json); ParameterTest::TestToJson(param); + REQUIRE(param->ToJson()[HOLD_MOLDS].GetBool()); } From 0629ca6a1641046aea9ecd78c4e1034bc0602f17 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Wed, 22 Jul 2026 14:10:03 +0800 Subject: [PATCH 2/5] chore: remove pyramid memory comparison example Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- ..._feature_compare_hgraph_pyramid_memory.cpp | 347 ------------------ examples/cpp/CMakeLists.txt | 4 - examples/cpp/README.md | 1 - 3 files changed, 352 deletions(-) delete mode 100644 examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp diff --git a/examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp b/examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp deleted file mode 100644 index 37f9a9c148..0000000000 --- a/examples/cpp/325_feature_compare_hgraph_pyramid_memory.cpp +++ /dev/null @@ -1,347 +0,0 @@ -// 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 -#include -#include - -#if defined(__APPLE__) -#include -#elif defined(__linux__) -#include -#endif - -namespace { - -constexpr uint64_t DIM = 128; -constexpr uint64_t LEAF_COUNT = 9; -constexpr uint64_t DEFAULT_DATASET_MIB = 400; -constexpr uint64_t BYTES_PER_MIB = 1024 * 1024; - -struct Measurement { - std::string name; - uint64_t total_rss_bytes; - uint64_t tracked_bytes; - uint64_t reported_bytes; - std::vector> reported_detail; -}; - -class TrackingAllocator : public vsag::Allocator { -public: - std::string - Name() override { - return "memory-comparison-allocator"; - } - - void* - Allocate(uint64_t size) override { - auto* pointer = std::malloc(size); - if (pointer == nullptr) { - return nullptr; - } - std::lock_guard lock(mutex_); - allocations_[pointer] = size; - current_bytes_ += size; - peak_bytes_ = std::max(peak_bytes_, current_bytes_); - return pointer; - } - - void - Deallocate(void* pointer) override { - if (pointer == nullptr) { - return; - } - { - std::lock_guard lock(mutex_); - auto allocation = allocations_.find(pointer); - if (allocation != allocations_.end()) { - current_bytes_ -= allocation->second; - allocations_.erase(allocation); - } - } - std::free(pointer); - } - - void* - Reallocate(void* pointer, uint64_t size) override { - std::lock_guard lock(mutex_); - uint64_t old_size = 0; - auto allocation = allocations_.find(pointer); - if (allocation != allocations_.end()) { - old_size = allocation->second; - } - auto* new_pointer = std::realloc(pointer, size); - if (new_pointer == nullptr) { - return nullptr; - } - allocations_.erase(pointer); - allocations_[new_pointer] = size; - current_bytes_ = current_bytes_ - old_size + size; - peak_bytes_ = std::max(peak_bytes_, current_bytes_); - return new_pointer; - } - - uint64_t - CurrentBytes() const { - std::lock_guard lock(mutex_); - return current_bytes_; - } - -private: - mutable std::mutex mutex_; - std::unordered_map allocations_; - uint64_t current_bytes_{0}; - uint64_t peak_bytes_{0}; -}; - -double -to_mib(uint64_t bytes) { - return static_cast(bytes) / static_cast(BYTES_PER_MIB); -} - -uint64_t -get_current_rss() { -#if defined(__APPLE__) - mach_task_basic_info_data_t info{}; - mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; - if (task_info( - mach_task_self(), MACH_TASK_BASIC_INFO, reinterpret_cast(&info), &count) != - KERN_SUCCESS) { - return 0; - } - return info.resident_size; -#elif defined(__linux__) - std::ifstream statm("/proc/self/statm"); - uint64_t total_pages = 0; - uint64_t resident_pages = 0; - statm >> total_pages >> resident_pages; - return resident_pages * static_cast(sysconf(_SC_PAGESIZE)); -#else - return 0; -#endif -} - -Measurement -build_and_measure(const std::string& name, - const std::string& parameters, - const vsag::DatasetPtr& base) { - TrackingAllocator allocator; - const uint64_t baseline_rss = get_current_rss(); - vsag::Resource resource(&allocator, nullptr); - vsag::Engine engine(&resource); - const uint64_t baseline_bytes = allocator.CurrentBytes(); - auto create_result = engine.CreateIndex(name, parameters); - if (not create_result.has_value()) { - std::cerr << "Failed to create " << name << ": " << create_result.error().message - << std::endl; - std::exit(EXIT_FAILURE); - } - - auto index = create_result.value(); - auto build_result = index->Build(base); - if (not build_result.has_value()) { - std::cerr << "Failed to build " << name << ": " << build_result.error().message - << std::endl; - std::exit(EXIT_FAILURE); - } - - const uint64_t built_rss = get_current_rss(); - Measurement result{name, - built_rss > baseline_rss ? built_rss - baseline_rss : 0, - allocator.CurrentBytes() - baseline_bytes, - index->GetMemoryUsage(), - {}}; - const auto detail = index->GetMemoryUsageDetail(); - result.reported_detail.assign(detail.begin(), detail.end()); - std::sort(result.reported_detail.begin(), - result.reported_detail.end(), - [](const auto& lhs, const auto& rhs) { return lhs.second > rhs.second; }); - index.reset(); - engine.Shutdown(); - return result; -} - -} // namespace - -int -main(int argc, char** argv) { - // Keep internal block pools from reserving large default chunks that dominate small trial runs. - vsag::Options::Instance().set_block_size_limit(2 * BYTES_PER_MIB); - if (argc < 2 || (std::string(argv[1]) != "hgraph" && std::string(argv[1]) != "pyramid")) { - std::cerr << "Usage:\n " << argv[0] << " [dataset MiB]\n " << argv[0] - << " synthetic \n " << argv[0] - << " fbin " << std::endl; - return EXIT_FAILURE; - } - const std::string index_name = argv[1]; - uint64_t dim = DIM; - uint64_t num_vectors = 0; - std::vector vectors; - if (argc == 4 && std::string(argv[2]) == "fbin") { - std::ifstream input(argv[3], std::ios::binary); - uint32_t file_num_vectors = 0; - uint32_t file_dim = 0; - input.read(reinterpret_cast(&file_num_vectors), sizeof(file_num_vectors)); - input.read(reinterpret_cast(&file_dim), sizeof(file_dim)); - if (not input || file_num_vectors == 0 || file_dim == 0) { - std::cerr << "Invalid fbin header: " << argv[3] << std::endl; - return EXIT_FAILURE; - } - num_vectors = file_num_vectors; - dim = file_dim; - vectors.resize(num_vectors * dim); - input.read(reinterpret_cast(vectors.data()), - static_cast(vectors.size() * sizeof(float))); - if (not input) { - std::cerr << "Incomplete fbin data: " << argv[3] << std::endl; - return EXIT_FAILURE; - } - } else { - if (argc == 5 && std::string(argv[2]) == "synthetic") { - num_vectors = std::stoull(argv[3]); - dim = std::stoull(argv[4]); - } else { - const uint64_t dataset_mib = argc > 2 ? std::stoull(argv[2]) : DEFAULT_DATASET_MIB; - const uint64_t vector_bytes = dataset_mib * BYTES_PER_MIB; - num_vectors = vector_bytes / (dim * sizeof(float)); - } - vectors.resize(num_vectors * dim); - for (uint64_t i = 0; i < num_vectors; ++i) { - for (uint64_t j = 0; j < dim; ++j) { - uint64_t value = (i * dim + j) + 0x9E3779B97F4A7C15ULL; - value = (value ^ (value >> 30U)) * 0xBF58476D1CE4E5B9ULL; - value = (value ^ (value >> 27U)) * 0x94D049BB133111EBULL; - value ^= value >> 31U; - vectors[i * dim + j] = - static_cast(value & 0xFFFFFFU) / static_cast(0x1000000U); - } - } - } - const uint64_t actual_vector_bytes = vectors.size() * sizeof(float); - - std::vector ids(num_vectors); - std::vector paths(num_vectors); - const std::array leaf_paths = { - "leaf-0", "leaf-1", "leaf-2", "leaf-3", "leaf-4", "leaf-5", "leaf-6", "leaf-7", "leaf-8"}; - - for (uint64_t i = 0; i < num_vectors; ++i) { - ids[i] = static_cast(i); - paths[i] = leaf_paths[i % LEAF_COUNT]; - } - - auto hgraph_base = vsag::Dataset::Make(); - hgraph_base->NumElements(static_cast(num_vectors)) - ->Dim(static_cast(dim)) - ->Ids(ids.data()) - ->Float32Vectors(vectors.data()) - ->Owner(false); - - auto pyramid_base = vsag::Dataset::Make(); - pyramid_base->NumElements(static_cast(num_vectors)) - ->Dim(static_cast(dim)) - ->Ids(ids.data()) - ->Float32Vectors(vectors.data()) - ->Paths(paths.data()) - ->Owner(false); - - std::string hgraph_parameters = R"( - { - "dtype": "float32", - "metric_type": "l2", - "dim": 128, - "index_param": { - "base_quantization_type": "fp32", - "max_degree": 64, - "ef_construction": 200, - "alpha": 1.2, - "use_reorder": false, - "use_reverse_edges": true, - "support_force_remove": true, - "store_raw_vector": true, - "base_io_type": "block_memory_io", - "build_thread_count": 1 - } - } - )"; - - std::string pyramid_parameters = R"( - { - "dtype": "float32", - "metric_type": "l2", - "dim": 128, - "index_param": { - "base_quantization_type": "fp32", - "max_degree": 32, - "ef_construction": 400, - "alpha": 1.2, - "graph_type": "nsw", - "no_build_levels": [0], - "use_reorder": false, - "store_raw_vector": true, - "index_min_size": 0, - "build_thread_count": 1 - } - } - )"; - - const std::string default_dim = "\"dim\": 128"; - const std::string actual_dim = "\"dim\": " + std::to_string(dim); - hgraph_parameters.replace(hgraph_parameters.find(default_dim), default_dim.size(), actual_dim); - pyramid_parameters.replace( - pyramid_parameters.find(default_dim), default_dim.size(), actual_dim); - - std::cout << "Dataset: " << num_vectors << " x " << dim << " float32 vectors (" << std::fixed - << std::setprecision(2) << to_mib(actual_vector_bytes) << " MiB)" << std::endl; - if (index_name == "pyramid") { - std::cout << "Pyramid layout: one root (graph disabled) + " << LEAF_COUNT << " leaf graphs" - << std::endl; - } - - const auto& parameters = index_name == "hgraph" ? hgraph_parameters : pyramid_parameters; - const auto& base = index_name == "hgraph" ? hgraph_base : pyramid_base; - const auto measurement = build_and_measure(index_name, parameters, base); - - std::cout << "\n" << measurement.name << " memory after Build():\n"; - std::cout << " Total process RSS increase: " << to_mib(measurement.total_rss_bytes) - << " MiB\n"; - std::cout << " VSAG allocator subset: " << to_mib(measurement.tracked_bytes) << " MiB\n"; - std::cout << " Index::GetMemoryUsage(): " << to_mib(measurement.reported_bytes) << " MiB\n"; - std::cout << " Total RSS bytes/vector: " - << static_cast(measurement.total_rss_bytes) / static_cast(num_vectors) - << "\n"; - - if (not measurement.reported_detail.empty()) { - uint64_t detail_total = 0; - std::cout << "\n Index::GetMemoryUsageDetail():\n"; - for (const auto& [component, bytes] : measurement.reported_detail) { - detail_total += bytes; - std::cout << " " << std::left << std::setw(24) << component << std::right - << std::setw(10) << to_mib(bytes) << " MiB\n"; - } - std::cout << " " << std::left << std::setw(24) << "detail total" << std::right - << std::setw(10) << to_mib(detail_total) << " MiB\n"; - } - - return 0; -} diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index ab5b290a5d..4745241985 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -154,9 +154,5 @@ target_link_libraries(324_feature_lazy_hgraph_extra_info vsag) add_executable(325_feature_uring_io 325_feature_uring_io.cpp) target_link_libraries(325_feature_uring_io vsag) - add_executable(324_feature_hgraph_mci_companion 324_feature_hgraph_mci_companion.cpp) target_link_libraries(324_feature_hgraph_mci_companion vsag) -add_executable(325_feature_compare_hgraph_pyramid_memory - 325_feature_compare_hgraph_pyramid_memory.cpp) -target_link_libraries(325_feature_compare_hgraph_pyramid_memory vsag) diff --git a/examples/cpp/README.md b/examples/cpp/README.md index be5230e0d3..879bce936d 100644 --- a/examples/cpp/README.md +++ b/examples/cpp/README.md @@ -102,7 +102,6 @@ together when the directory is listed: | [`320_feature_extra_info.cpp`](320_feature_extra_info.cpp) | Attach per-vector extra info / payload. | | [`322_feature_hgraph_brute_force_threshold.cpp`](322_feature_hgraph_brute_force_threshold.cpp) | HGraph search-time `brute_force_threshold`: automatically switch to an exact scan under highly selective filters. | | [`324_feature_lazy_hgraph_extra_info.cpp`](324_feature_lazy_hgraph_extra_info.cpp) | LazyHGraph `extra_info` filtering across flat and graph phases. | -| [`325_feature_compare_hgraph_pyramid_memory.cpp`](325_feature_compare_hgraph_pyramid_memory.cpp) | Compare HGraph with a nine-leaf Pyramid on the same 400 MiB vector dataset. | ### Persistence (`4xx`) From 78aa96c529d7039871d1e4be2d58e0c00933c301 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Thu, 23 Jul 2026 11:18:53 +0800 Subject: [PATCH 3/5] fix: harden pyramid raw vector storage Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- src/algorithm/pyramid/pyramid.cpp | 135 +++++++++--- src/algorithm/pyramid/pyramid.h | 8 +- src/algorithm/pyramid/pyramid_test.cpp | 194 ++++++++++++++++++ src/quantization/fp32_quantizer_parameter.cpp | 1 - .../fp32_quantizer_parameter_test.cpp | 4 +- 5 files changed, 310 insertions(+), 32 deletions(-) diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index 756f949ebb..99356daf63 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -34,6 +34,8 @@ namespace vsag { const static float RADIUS_EPSILON = 1.1F; +constexpr const char* PYRAMID_RAW_VECTOR_SIZE = "raw_vector_size"; +constexpr uint64_t RAW_VECTOR_COPY_BUFFER_SIZE = 1024 * 1024; std::vector split(const std::string& str, char delimiter) { @@ -453,7 +455,8 @@ Pyramid::GetMemoryUsage() const { std::unordered_map Pyramid::GetMemoryUsageDetail() const { - std::shared_lock lock(resize_mutex_); + std::lock_guard count_lock(cur_element_count_mutex_); + std::unique_lock resize_lock(resize_mutex_); std::unordered_map memory_usage; memory_usage["points_mutex"] = points_mutex_ == nullptr ? 0 : points_mutex_->GetMemoryUsage(); memory_usage["pool"] = pool_ == nullptr ? 0 : pool_->GetMemoryUsage(); @@ -496,8 +499,11 @@ Pyramid::Serialize(StreamWriter& writer) const { if (use_reorder_) { precise_codes_->Serialize(writer); } + uint64_t raw_vector_size = 0; if (create_new_raw_vector_) { + auto begin = writer.GetCursor(); raw_vector_->Serialize(writer); + raw_vector_size = writer.GetCursor() - begin; } auto pyramid_param = std::dynamic_pointer_cast(create_param_ptr_); @@ -515,6 +521,7 @@ Pyramid::Serialize(StreamWriter& writer) const { // serialize footer (introduced since v0.15) JsonType basic_info; basic_info["max_capacity"].SetInt(max_capacity_); + basic_info[PYRAMID_RAW_VECTOR_SIZE].SetUint64(raw_vector_size); basic_info[INDEX_PARAM].SetString(this->create_param_ptr_->ToString()); write_index_footer(writer, basic_info); } @@ -729,12 +736,12 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) } break; case StreamSerializationTag::RAW_VECTOR: + loaded_raw_vector = true; if (this->create_new_raw_vector_) { ReadSeekableBlockPayload( block_reader, block_header, [this](StreamReader& block) { this->raw_vector_->Deserialize(block); }); - loaded_raw_vector = true; } break; case StreamSerializationTag::PYRAMID_HIERARCHIES: @@ -770,12 +777,10 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) throw VsagException(ErrorType::READ_ERROR, "Pyramid streaming serialization precise codes block is missing"); } + resize(max_capacity); if (this->create_new_raw_vector_ && !loaded_raw_vector) { - throw VsagException(ErrorType::READ_ERROR, - "Pyramid streaming serialization raw vector block is missing"); + this->restore_raw_vector_from_source(); } - - resize(max_capacity); this->current_memory_usage_ = static_cast(this->CalSerializeSize()); } @@ -787,6 +792,10 @@ Pyramid::Deserialize(StreamReader& reader) { throw VsagException(ErrorType::READ_ERROR, "failed to read index footer"); } auto max_capacity = basic_info["max_capacity"].GetInt(); + uint64_t raw_vector_size = 0; + if (basic_info.Contains(PYRAMID_RAW_VECTOR_SIZE)) { + raw_vector_size = basic_info[PYRAMID_RAW_VECTOR_SIZE].GetUint64(); + } if (basic_info.Contains(INDEX_PARAM)) { auto index_param = std::make_shared(); index_param->FromString(basic_info[INDEX_PARAM].GetString()); @@ -809,8 +818,15 @@ Pyramid::Deserialize(StreamReader& reader) { if (use_reorder_) { precise_codes_->Deserialize(buffer_reader); } - if (create_new_raw_vector_) { - raw_vector_->Deserialize(buffer_reader); + if (raw_vector_size > 0) { + auto raw_vector_begin = buffer_reader.GetCursor(); + if (create_new_raw_vector_) { + raw_vector_->Deserialize(buffer_reader); + CHECK_ARGUMENT(buffer_reader.GetCursor() - raw_vector_begin == raw_vector_size, + "serialized Pyramid raw vector size mismatch"); + } else { + buffer_reader.Seek(raw_vector_begin + raw_vector_size); + } } cur_element_count_ = base_codes_->TotalCount(); @@ -838,6 +854,9 @@ Pyramid::Deserialize(StreamReader& reader) { } resize(max_capacity); + if (create_new_raw_vector_ && raw_vector_size == 0) { + restore_raw_vector_from_source(); + } this->current_memory_usage_ = this->CalSerializeSize(); } @@ -977,7 +996,7 @@ Pyramid::InitFeatures() { IndexFeature::SUPPORT_EXPORT_MODEL, IndexFeature::SUPPORT_GET_MEMORY_USAGE, }); - if (raw_vector_ != nullptr) { + if (has_raw_vector_) { this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_GET_RAW_VECTOR_BY_IDS); } @@ -1521,32 +1540,94 @@ Pyramid::check_and_init_raw_vector(const FlattenInterfaceParamPtr& raw_vector_pa return; } - raw_vector_ = FlattenInterface::MakeInstance(raw_vector_param, common_param); - if (base_codes_->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32 && - precise_codes_ == nullptr) { - create_new_raw_vector_ = true; - return; + auto io_type_name = raw_vector_param->io_parameter->GetTypeName(); + if (io_type_name == IO_TYPE_VALUE_BLOCK_MEMORY_IO || io_type_name == IO_TYPE_VALUE_MEMORY_IO) { + raw_vector_ = find_raw_vector_source(true); } - if (base_codes_->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32 && - precise_codes_ != nullptr && - precise_codes_->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32) { + if (raw_vector_ == nullptr) { + raw_vector_ = FlattenInterface::MakeInstance(raw_vector_param, common_param); create_new_raw_vector_ = true; - return; } - auto io_type_name = raw_vector_param->io_parameter->GetTypeName(); - if (io_type_name != IO_TYPE_VALUE_BLOCK_MEMORY_IO && io_type_name != IO_TYPE_VALUE_MEMORY_IO) { - create_new_raw_vector_ = true; - return; + CHECK_ARGUMENT(raw_vector_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32, + "Pyramid raw vector storage must use fp32 quantization"); + CHECK_ARGUMENT(metric_ != MetricType::METRIC_TYPE_COSINE || raw_vector_->HoldMolds(), + "Pyramid cosine raw vector storage must preserve vector molds"); + has_raw_vector_ = raw_vector_ != nullptr; +} + +FlattenInterfacePtr +Pyramid::find_raw_vector_source(bool require_in_memory) const { + auto is_usable = [this, require_in_memory](const FlattenInterfacePtr& candidate) { + if (candidate == nullptr || candidate->GetQuantizerName() != QUANTIZATION_TYPE_VALUE_FP32) { + return false; + } + if (require_in_memory && !candidate->InMemory()) { + return false; + } + return metric_ != MetricType::METRIC_TYPE_COSINE || candidate->HoldMolds(); + }; + + if (is_usable(base_codes_)) { + return base_codes_; } + if (is_usable(precise_codes_)) { + return precise_codes_; + } + return nullptr; +} - if (base_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { - raw_vector_ = base_codes_; +void +Pyramid::restore_raw_vector_from_source() { + auto source = find_raw_vector_source(false); + if (source == nullptr) { + throw VsagException(ErrorType::READ_ERROR, + "serialized Pyramid index does not contain restorable raw vectors"); + } + if (raw_vector_->TotalCount() != 0) { + throw VsagException(ErrorType::READ_ERROR, + "Pyramid raw vector storage must be empty before restoration"); + } + + auto total_count = static_cast(source->TotalCount()); + if (total_count == 0) { return; } - if (precise_codes_ != nullptr && - precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { - raw_vector_ = precise_codes_; + auto vector_size = static_cast(dim_) * sizeof(float); + auto batch_size = std::max(1, RAW_VECTOR_COPY_BUFFER_SIZE / vector_size); + Vector vectors(batch_size * static_cast(dim_), allocator_); + for (uint64_t begin = 0; begin < total_count; begin += batch_size) { + auto current_batch_size = std::min(batch_size, total_count - begin); + for (uint64_t offset = 0; offset < current_batch_size; ++offset) { + bool release = false; + const auto* codes = + source->GetCodesById(static_cast(begin + offset), release); + if (codes == nullptr) { + throw VsagException( + ErrorType::READ_ERROR, + fmt::format("failed to restore Pyramid raw vector {}", begin + offset)); + } + bool decoded = false; + try { + decoded = + source->Decode(codes, vectors.data() + offset * static_cast(dim_)); + } catch (...) { + if (release) { + source->Release(codes); + } + throw; + } + if (release) { + source->Release(codes); + } + if (not decoded) { + throw VsagException( + ErrorType::READ_ERROR, + fmt::format("failed to decode Pyramid raw vector {}", begin + offset)); + } + } + raw_vector_->BatchInsertVector(vectors.data(), + static_cast(current_batch_size)); } } diff --git a/src/algorithm/pyramid/pyramid.h b/src/algorithm/pyramid/pyramid.h index 5cb0c43841..22532da84f 100644 --- a/src/algorithm/pyramid/pyramid.h +++ b/src/algorithm/pyramid/pyramid.h @@ -300,6 +300,12 @@ class Pyramid : public InnerIndexInterface { check_and_init_raw_vector(const FlattenInterfaceParamPtr& raw_vector_param, const IndexCommonParam& common_param); + FlattenInterfacePtr + find_raw_vector_source(bool require_in_memory) const; + + void + restore_raw_vector_from_source(); + /// One named hierarchy with its own root IndexNode and build parameters. struct Hierarchy { std::string name; // hierarchy name (empty = default) @@ -392,7 +398,7 @@ class Pyramid : public InnerIndexInterface { bool support_duplicate_{false}; // whether to allow duplicate ids mutable std::shared_mutex resize_mutex_; // guards resize operations - std::mutex cur_element_count_mutex_; // guards cur_element_count_ updates + mutable std::mutex cur_element_count_mutex_; // guards cur_element_count_ updates std::string graph_type_{GRAPH_TYPE_VALUE_NSW}; // graph algorithm type std::mutex entry_point_mutex_; // guards entry-point selection diff --git a/src/algorithm/pyramid/pyramid_test.cpp b/src/algorithm/pyramid/pyramid_test.cpp index 310a508353..5ca6e37866 100644 --- a/src/algorithm/pyramid/pyramid_test.cpp +++ b/src/algorithm/pyramid/pyramid_test.cpp @@ -16,6 +16,7 @@ #include "pyramid.h" #include +#include #include #include @@ -446,3 +447,196 @@ TEST_CASE("Pyramid reports live raw vector memory", "[ut][pyramid][raw_vector][m REQUIRE(detail.at("raw_vector") >= sizeof(float) * dim * 2); REQUIRE(index->GetMemoryUsage() >= detail.at("raw_vector")); } + +TEST_CASE("Pyramid raw vector serialization handles IO storage changes", + "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + constexpr int64_t count = 2; + std::array vectors = { + 0.123456F, 0.234567F, 0.345678F, 0.456789F, 1.0F, 2.0F, 3.0F, 4.0F}; + std::array ids = {10, 11}; + std::array paths = {"leaf", "leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + fixtures::TempDir dir("pyramid_raw_vector_io"); + auto make_param = [&](bool use_dedicated_storage) { + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "fp32", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto io_type = use_dedicated_storage ? "buffer_io" : "block_memory_io"; + external_param["base_io_type"].SetString(io_type); + external_param["base_file_path"].SetString(dir.GenerateRandomFile(false)); + external_param["raw_vector_io_type"].SetString(io_type); + external_param["raw_vector_file_path"].SetString(dir.GenerateRandomFile(false)); + return vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + }; + auto dataset = vsag::Dataset::Make() + ->NumElements(count) + ->Dim(dim) + ->Float32Vectors(vectors.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + + auto round_trip = [&](bool producer_uses_dedicated_storage, bool streaming) { + auto producer = std::make_shared(make_param(producer_uses_dedicated_storage), + common_param); + REQUIRE(producer->Build(dataset).empty()); + + auto loaded = std::make_shared( + make_param(not producer_uses_dedicated_storage), common_param); + if (streaming) { + std::stringstream buffer; + producer->SerializeStreaming(buffer); + std::stringstream reader(buffer.str()); + loaded->DeserializeStreaming(reader); + } else { + auto binary_set = producer->vsag::InnerIndexInterface::Serialize(); + loaded->vsag::InnerIndexInterface::Deserialize(binary_set); + } + + auto restored = loaded->GetDataByIds(ids.data(), count); + REQUIRE(restored->GetFloat32Vectors() != nullptr); + for (int64_t i = 0; i < dim * count; ++i) { + REQUIRE(restored->GetFloat32Vectors()[i] == vectors[i]); + } + + auto memory_detail = loaded->GetMemoryUsageDetail(); + if (producer_uses_dedicated_storage) { + REQUIRE(memory_detail.count("raw_vector") == 0); + } else { + REQUIRE(memory_detail.at("raw_vector") >= sizeof(float) * dim * count); + } + }; + + SECTION("legacy alias to dedicated") { + round_trip(false, false); + } + SECTION("legacy dedicated to alias") { + round_trip(true, false); + } + SECTION("streaming alias to dedicated") { + round_trip(false, true); + } + SECTION("streaming dedicated to alias") { + round_trip(true, true); + } +} + +TEST_CASE("Pyramid does not alias normalized FP32 codes as cosine raw vectors", + "[ut][pyramid][raw_vector]") { + constexpr int64_t dim = 4; + std::array vector = {1.0F, 2.0F, 3.0F, 4.0F}; + std::array ids = {10}; + std::array paths = {"leaf"}; + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_COSINE; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "fp32", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8, + "no_build_levels": [0, 1] + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto pyramid_param = std::dynamic_pointer_cast(param); + REQUIRE(pyramid_param != nullptr); + vsag::JsonType no_molds; + no_molds["hold_molds"].SetBool(false); + pyramid_param->base_codes_param->quantizer_parameter->FromJson(no_molds); + + auto index = std::make_shared(pyramid_param, common_param); + auto dataset = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(vector.data()) + ->Ids(ids.data()) + ->Paths(paths.data()) + ->Owner(false); + REQUIRE(index->Build(dataset).empty()); + + auto memory_detail = index->GetMemoryUsageDetail(); + REQUIRE(memory_detail.at("raw_vector") >= sizeof(float) * dim); + auto restored = index->GetDataByIds(ids.data(), 1); + REQUIRE(restored->GetFloat32Vectors() != nullptr); + for (int64_t i = 0; i < dim; ++i) { + REQUIRE(restored->GetFloat32Vectors()[i] == vector[i]); + } +} + +TEST_CASE("Pyramid reports memory while adding vectors", "[ut][pyramid][raw_vector][memory]") { + constexpr int64_t dim = 4; + constexpr int64_t add_count = 256; + std::array initial_vector = {0.0F, 0.0F, 0.0F, 0.0F}; + std::array initial_id = {1}; + std::array initial_path = {"leaf"}; + std::vector vectors(dim * add_count, 1.0F); + std::vector ids(add_count); + std::vector paths(add_count, "leaf"); + for (int64_t i = 0; i < add_count; ++i) { + ids[i] = i + 2; + } + + vsag::IndexCommonParam common_param; + common_param.dim_ = dim; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + common_param.metric_ = vsag::MetricType::METRIC_TYPE_L2SQR; + common_param.allocator_ = vsag::SafeAllocator::FactoryDefaultAllocator(); + auto external_param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "sq8", + "store_raw_vector": true, + "max_degree": 4, + "ef_construction": 8 + })"); + auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); + auto index = std::make_shared(param, common_param); + auto initial_dataset = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(initial_vector.data()) + ->Ids(initial_id.data()) + ->Paths(initial_path.data()) + ->Owner(false); + REQUIRE(index->Build(initial_dataset).empty()); + + std::promise add_started; + auto started = add_started.get_future(); + auto add_result = std::async(std::launch::async, [&]() { + add_started.set_value(); + for (int64_t i = 0; i < add_count; ++i) { + auto add_dataset = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(vectors.data() + i * dim) + ->Ids(ids.data() + i) + ->Paths(paths.data() + i) + ->Owner(false); + if (not index->Add(add_dataset).empty()) { + return false; + } + } + return true; + }); + started.wait(); + for (int64_t i = 0; i < 32; ++i) { + auto detail = index->GetMemoryUsageDetail(); + REQUIRE(detail.at("label_table") > 0); + REQUIRE(detail.at("raw_vector") > 0); + } + REQUIRE(add_result.get()); + REQUIRE(index->GetNumElements() == add_count + 1); +} diff --git a/src/quantization/fp32_quantizer_parameter.cpp b/src/quantization/fp32_quantizer_parameter.cpp index ef82a1e73f..2b4adc0464 100644 --- a/src/quantization/fp32_quantizer_parameter.cpp +++ b/src/quantization/fp32_quantizer_parameter.cpp @@ -33,7 +33,6 @@ JsonType FP32QuantizerParameter::ToJson() const { JsonType json; json[TYPE_KEY].SetString(QUANTIZATION_TYPE_VALUE_FP32); - json[HOLD_MOLDS].SetBool(hold_molds); return json; } } // namespace vsag diff --git a/src/quantization/fp32_quantizer_parameter_test.cpp b/src/quantization/fp32_quantizer_parameter_test.cpp index 8b5264ace1..53e0edc3b9 100644 --- a/src/quantization/fp32_quantizer_parameter_test.cpp +++ b/src/quantization/fp32_quantizer_parameter_test.cpp @@ -15,17 +15,15 @@ #include "fp32_quantizer_parameter.h" -#include "inner_string_params.h" #include "parameter_test.h" #include "unittest.h" using namespace vsag; TEST_CASE("FP32 Quantizer Parameter ToJson Test", "[ut][FP32QuantizerParameter]") { - std::string param_str = R"({"hold_molds": true})"; + std::string param_str = "{}"; auto param = std::make_shared(); JsonType param_json = JsonType::Parse(param_str); param->FromJson(param_json); ParameterTest::TestToJson(param); - REQUIRE(param->ToJson()[HOLD_MOLDS].GetBool()); } From a64fd2cd5cdc2f99685175189ed368b8f406b11d Mon Sep 17 00:00:00 2001 From: jc543239 Date: Thu, 23 Jul 2026 06:00:42 +0000 Subject: [PATCH 4/5] fix: resolve pyramid clang-tidy errors Signed-off-by: jc543239 Assisted-by: Codex:gpt-5 --- src/algorithm/pyramid/pyramid.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index 99356daf63..ef16ec0073 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -35,7 +35,7 @@ namespace vsag { const static float RADIUS_EPSILON = 1.1F; constexpr const char* PYRAMID_RAW_VECTOR_SIZE = "raw_vector_size"; -constexpr uint64_t RAW_VECTOR_COPY_BUFFER_SIZE = 1024 * 1024; +constexpr uint64_t RAW_VECTOR_COPY_BUFFER_SIZE = uint64_t{1024} * 1024; std::vector split(const std::string& str, char delimiter) { @@ -1551,7 +1551,9 @@ Pyramid::check_and_init_raw_vector(const FlattenInterfaceParamPtr& raw_vector_pa CHECK_ARGUMENT(raw_vector_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32, "Pyramid raw vector storage must use fp32 quantization"); - CHECK_ARGUMENT(metric_ != MetricType::METRIC_TYPE_COSINE || raw_vector_->HoldMolds(), + const bool valid_cosine_raw_vector = + metric_ != MetricType::METRIC_TYPE_COSINE || raw_vector_->HoldMolds(); + CHECK_ARGUMENT(valid_cosine_raw_vector, "Pyramid cosine raw vector storage must preserve vector molds"); has_raw_vector_ = raw_vector_ != nullptr; } From d8dc386a0ade6eb74faf3a8f7b3732516f69a98d Mon Sep 17 00:00:00 2001 From: jac0626 Date: Thu, 30 Jul 2026 15:48:52 +0800 Subject: [PATCH 5/5] fix: address pyramid raw vector review feedback Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- examples/cpp/CMakeLists.txt | 1 + src/algorithm/pyramid/pyramid.cpp | 16 +++++------ src/algorithm/pyramid/pyramid_test.cpp | 27 +++++++++++++++++++ src/algorithm/pyramid/pyramid_zparameters.cpp | 6 +++-- src/factory/factory.cpp | 5 ++-- 5 files changed, 43 insertions(+), 12 deletions(-) diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 4745241985..3793625a87 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -154,5 +154,6 @@ target_link_libraries(324_feature_lazy_hgraph_extra_info vsag) add_executable(325_feature_uring_io 325_feature_uring_io.cpp) target_link_libraries(325_feature_uring_io vsag) + add_executable(324_feature_hgraph_mci_companion 324_feature_hgraph_mci_companion.cpp) target_link_libraries(324_feature_hgraph_mci_companion vsag) diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index ef16ec0073..12cfee430e 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -241,7 +241,9 @@ Pyramid::build_by_odescent(const DatasetPtr& base) { const auto* data_ids = base->GetIds(); resize(data_num); - std::memcpy(label_table_->label_table_.data(), data_ids, sizeof(LabelType) * data_num); + for (InnerIdType inner_id = 0; inner_id < data_num; ++inner_id) { + label_table_->Insert(inner_id, data_ids[inner_id]); + } base_codes_->BatchInsertVector(data_vectors, data_num); if (use_reorder_) { @@ -902,17 +904,15 @@ Pyramid::Add(const DatasetPtr& base) { int64_t valid_id_count = 0; for (int64_t i = 0; i < data_num; ++i) { if (not label_table_->CheckLabel(data_ids[i])) { - label_table_->Insert(valid_id_count + local_cur_element_count, data_ids[i]); - base_codes_->InsertVector(data_vectors + dim_ * i, - valid_id_count + local_cur_element_count); + const auto inner_id = valid_id_count + local_cur_element_count; + base_codes_->InsertVector(data_vectors + dim_ * i, inner_id); if (use_reorder_) { - precise_codes_->InsertVector(data_vectors + dim_ * i, - valid_id_count + local_cur_element_count); + precise_codes_->InsertVector(data_vectors + dim_ * i, inner_id); } if (create_new_raw_vector_) { - raw_vector_->InsertVector(data_vectors + dim_ * i, - valid_id_count + local_cur_element_count); + raw_vector_->InsertVector(data_vectors + dim_ * i, inner_id); } + label_table_->Insert(inner_id, data_ids[i]); valid_id_count++; data_biases.push_back(i); } else { diff --git a/src/algorithm/pyramid/pyramid_test.cpp b/src/algorithm/pyramid/pyramid_test.cpp index 5ca6e37866..41df77af43 100644 --- a/src/algorithm/pyramid/pyramid_test.cpp +++ b/src/algorithm/pyramid/pyramid_test.cpp @@ -16,6 +16,7 @@ #include "pyramid.h" #include +#include #include #include #include @@ -279,6 +280,11 @@ TEST_CASE("Pyramid stores raw vectors during ODescent build", "[ut][pyramid][raw for (int64_t i = 0; i < dim; ++i) { REQUIRE(restored[i] == vectors[i]); } + auto raw = index->GetDataByIds(ids.data(), 2); + REQUIRE(raw->GetFloat32Vectors() != nullptr); + for (int64_t i = 0; i < dim * 2; ++i) { + REQUIRE(raw->GetFloat32Vectors()[i] == vectors[i]); + } } TEST_CASE("Pyramid reuses FP32 codes as raw vectors", "[ut][pyramid][raw_vector]") { @@ -530,6 +536,27 @@ TEST_CASE("Pyramid raw vector serialization handles IO storage changes", SECTION("streaming dedicated to alias") { round_trip(true, true); } + SECTION("streaming load applies storage overrides") { + auto producer = std::make_shared(make_param(true), common_param); + REQUIRE(producer->Build(dataset).empty()); + std::stringstream buffer; + producer->SerializeStreaming(buffer); + producer.reset(); + std::filesystem::remove_all(dir.path); + + std::stringstream reader(buffer.str()); + auto loaded = vsag::Index::Load(reader, + R"({ + "base_io_type": "block_memory_io", + "raw_vector_io_type": "block_memory_io" + })"); + REQUIRE(loaded.has_value()); + auto restored = loaded.value()->GetRawVectorByIds(ids.data(), count); + REQUIRE(restored.has_value()); + for (int64_t i = 0; i < dim * count; ++i) { + REQUIRE(restored.value()->GetFloat32Vectors()[i] == vectors[i]); + } + } } TEST_CASE("Pyramid does not alias normalized FP32 codes as cosine raw vectors", diff --git a/src/algorithm/pyramid/pyramid_zparameters.cpp b/src/algorithm/pyramid/pyramid_zparameters.cpp index f51b63cf8c..11d29342ee 100644 --- a/src/algorithm/pyramid/pyramid_zparameters.cpp +++ b/src/algorithm/pyramid/pyramid_zparameters.cpp @@ -194,8 +194,10 @@ PyramidParameters::FromJson(const JsonType& json) { param->quantizer_parameter->GetTypeName() != QUANTIZATION_TYPE_VALUE_FP32) { return; } - std::dynamic_pointer_cast(param->quantizer_parameter)->hold_molds = - true; + auto fp32_param = + std::dynamic_pointer_cast(param->quantizer_parameter); + CHECK_ARGUMENT(fp32_param != nullptr, "invalid fp32 quantizer parameter"); + fp32_param->hold_molds = true; }; restore_legacy_raw_fp32_param(this->base_codes_param, BASE_CODES_KEY); if (this->use_reorder) { diff --git a/src/factory/factory.cpp b/src/factory/factory.cpp index 0ffe1608ac..9a9c02949c 100644 --- a/src/factory/factory.cpp +++ b/src/factory/factory.cpp @@ -122,7 +122,7 @@ set_streaming_io_override(JsonType& index_param, } void -apply_hgraph_streaming_load_parameters(JsonType& index_param, const std::string& parameters) { +apply_flatten_streaming_load_parameters(JsonType& index_param, const std::string& parameters) { auto load_json = JsonType::Parse(parameters.empty() ? "{}" : parameters); if (load_json.Contains(HGRAPH_BASE_IO_TYPE)) { require_string_load_parameter(load_json, HGRAPH_BASE_IO_TYPE); @@ -241,13 +241,14 @@ create_streaming_index_from_metadata(const MetadataPtr& metadata, return create_streaming_index(index_param, common_param); } if (index_name == INDEX_HGRAPH) { - apply_hgraph_streaming_load_parameters(index_param, parameters); + apply_flatten_streaming_load_parameters(index_param, parameters); return create_streaming_index(index_param, common_param); } if (index_name == INDEX_IVF) { return create_streaming_index(index_param, common_param); } if (index_name == INDEX_PYRAMID) { + apply_flatten_streaming_load_parameters(index_param, parameters); return create_streaming_index(index_param, common_param); } if (index_name == INDEX_SINDI) {