diff --git a/docs/docs/en/src/api/index_class.md b/docs/docs/en/src/api/index_class.md index eb24fd7938..a2a58a2aa6 100644 --- a/docs/docs/en/src/api/index_class.md +++ b/docs/docs/en/src/api/index_class.md @@ -89,6 +89,7 @@ Bit flags for [`GetDataByIdsWithFlag`](#getdatabyidswithflag), combined with bit | `DATA_FLAG_EXTRA_INFO` | `0x10` | extra info blobs | | `DATA_FLAG_ATTRIBUTE` | `0x20` | attributes | | `DATA_FLAG_ID` | `0x40` | ids | +| `DATA_FLAG_PATH` | `0x80` | Pyramid paths retained with `store_paths` | ### `WriteFuncType` @@ -244,8 +245,8 @@ See [Graph Enhancement](../advanced/enhance_graph.md). | `GetMinAndMaxId` | `tl::expected, Error> GetMinAndMaxId() const` | Smallest and largest ids in the index. | | `GetExtraInfoByIds` | `tl::expected GetExtraInfoByIds(const int64_t* ids, int64_t count, char* extra_infos) const` | Copies extra-info blobs for `ids` into a caller-provided buffer. | | `GetRawVectorByIds` | `tl::expected GetRawVectorByIds(const int64_t* ids, int64_t count, Allocator* specified_allocator = nullptr) const` | Returns stored vectors. Values are *close to* the originals but not guaranteed bit-identical (quantization/precision). | -| `GetDataByIds` | `tl::expected GetDataByIds(const int64_t* ids, int64_t count) const` | Returns all stored data (vectors, attributes, extra info) for `ids`. | -| `GetDataByIdsWithFlag` | `tl::expected GetDataByIdsWithFlag(const int64_t* ids, int64_t count, uint64_t selected_data_flag) const` | Like `GetDataByIds` but selects fields via [`DATA_FLAG_*`](#data-selection-flags). | +| `GetDataByIds` | `tl::expected GetDataByIds(const int64_t* ids, int64_t count) const` | Returns the implementation's default stored fields for `ids`; optional fields may require explicit selection. | +| `GetDataByIdsWithFlag` | `tl::expected GetDataByIdsWithFlag(const int64_t* ids, int64_t count, uint64_t selected_data_flag) const` | Selects supported fields via [`DATA_FLAG_*`](#data-selection-flags). Pyramid paths require `store_paths: true` and `DATA_FLAG_PATH`. | | `GetIndexDetailInfos` | `tl::expected, Error> GetIndexDetailInfos() const` | Lists the introspectable detail fields. See [`IndexDetailInfo`](types.md#index-detail-info). | | `GetDetailDataByName` | `tl::expected GetDetailDataByName(const std::string& name, IndexDetailInfo& info) const` | Fetches one detail-data payload by name. | diff --git a/docs/docs/en/src/indexes/pyramid.md b/docs/docs/en/src/indexes/pyramid.md index a3ae651466..382c2ac655 100644 --- a/docs/docs/en/src/indexes/pyramid.md +++ b/docs/docs/en/src/indexes/pyramid.md @@ -103,6 +103,7 @@ Build-time parameters live under `index_param`. | `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. | +| `store_paths` | bool | `false` | Top-level switch that preserves the original paths supplied to `Build` and `Add` so `GetDataByIdsWithFlag` can return them when `DATA_FLAG_PATH` is selected. It applies to every configured hierarchy and cannot be overridden per hierarchy. | | `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. | @@ -236,6 +237,29 @@ base->NumElements(n) index->Build(base); ``` +### Retrieving paths by ID + +Set the top-level build parameter `store_paths` to `true` to retain the original paths for ID-based +retrieval. Select `DATA_FLAG_PATH` in `GetDataByIdsWithFlag`; the returned paths follow the requested +ID order. Use `GetPaths()` for the default unnamed hierarchy and `GetPaths(hierarchy_name)` for a +named hierarchy: + +```cpp +int64_t requested_ids[] = {product_id_b, product_id_a}; +auto data = index->GetDataByIdsWithFlag( + requested_ids, 2, DATA_FLAG_ID | DATA_FLAG_PATH).value(); + +const std::string* site_paths = data->GetPaths("site"); +const std::string* category_paths = data->GetPaths("category"); +``` + +`GetDataByIds` and `GetDataByIdsWithFlag` calls without `DATA_FLAG_PATH` do not attach path arrays. +Selecting `DATA_FLAG_PATH` while `store_paths` is `false` returns an invalid-argument error. When +path storage is enabled, a hierarchy is included only if every requested ID has a recorded path in +that hierarchy. If even one requested ID was built or added without that hierarchy's path, its +getter returns `nullptr`; other hierarchies whose requested paths are complete are still returned. +In single-hierarchy mode, the same completeness rule applies to `GetPaths()`. + ### Searching a specific hierarchy Specify which hierarchy to search via `"hierarchies"` in the search parameters. @@ -282,8 +306,11 @@ auto result = index->RangeSearch( ### Serialize & Deserialize -Multi-hierarchy indexes serialize and deserialize transparently. The serialized -format includes all hierarchy names and their graph structures: +Multi-hierarchy indexes serialize and deserialize transparently. The serialized format includes +all hierarchy names and their graph structures. With `store_paths: true`, both regular and +streaming serialization also persist the retained original paths, making them available through +`GetDataByIdsWithFlag` after deserialization. With the default `false`, the graph hierarchy is +persisted but the original per-ID paths are not: ```cpp // Serialize diff --git a/docs/docs/en/src/resources/index_parameters.md b/docs/docs/en/src/resources/index_parameters.md index 7c49c578a9..eea5429b11 100644 --- a/docs/docs/en/src/resources/index_parameters.md +++ b/docs/docs/en/src/resources/index_parameters.md @@ -145,11 +145,18 @@ Pyramid build parameters also live under `index_param`: "index_param": { "base_quantization_type": "sq8", "max_degree": 24, - "ef_construction": 300 + "ef_construction": 300, + "store_paths": true } } ``` +`store_paths` is a top-level Pyramid build parameter and defaults to `false`. Enable it when +`GetDataByIdsWithFlag` must return the original default or named-hierarchy paths with +`DATA_FLAG_PATH`; see the +[Pyramid parameter table](../indexes/pyramid.md#build-parameters) for its completeness and +persistence semantics. + MRLE with split RaBitQ uses `base_quantization_type: "tq"`, `tq_chain: "mrle, rabitq"`, `mrle_dim`, and the `rabitq_bits_per_dim_base`/`rabitq_bits_per_dim_precise` pair. It automatically reorders from the diff --git a/docs/docs/zh/src/api/index_class.md b/docs/docs/zh/src/api/index_class.md index 3f9a2b3900..0c4f3b1736 100644 --- a/docs/docs/zh/src/api/index_class.md +++ b/docs/docs/zh/src/api/index_class.md @@ -87,6 +87,7 @@ struct Index::Checkpoint { | `DATA_FLAG_EXTRA_INFO` | `0x10` | extra info 数据块 | | `DATA_FLAG_ATTRIBUTE` | `0x20` | 属性 | | `DATA_FLAG_ID` | `0x40` | id | +| `DATA_FLAG_PATH` | `0x80` | Pyramid 通过 `store_paths` 保留的路径 | ### `WriteFuncType` @@ -239,8 +240,8 @@ RangeSearch(const DatasetPtr& query, float radius, const std::string& parameters | `GetMinAndMaxId` | `tl::expected, Error> GetMinAndMaxId() const` | 索引中最小与最大的 id。 | | `GetExtraInfoByIds` | `tl::expected GetExtraInfoByIds(const int64_t* ids, int64_t count, char* extra_infos) const` | 把 `ids` 的 extra-info 数据块拷贝到调用方提供的缓冲区。 | | `GetRawVectorByIds` | `tl::expected GetRawVectorByIds(const int64_t* ids, int64_t count, Allocator* specified_allocator = nullptr) const` | 返回已存向量。其值*接近*原始值，但不保证逐位一致（量化/精度）。 | -| `GetDataByIds` | `tl::expected GetDataByIds(const int64_t* ids, int64_t count) const` | 返回 `ids` 的全部已存数据（向量、属性、extra info）。 | -| `GetDataByIdsWithFlag` | `tl::expected GetDataByIdsWithFlag(const int64_t* ids, int64_t count, uint64_t selected_data_flag) const` | 类似 `GetDataByIds`，但通过 [`DATA_FLAG_*`](#数据选择标志) 选择字段。 | +| `GetDataByIds` | `tl::expected GetDataByIds(const int64_t* ids, int64_t count) const` | 返回实现默认提供的已存字段；可选字段可能需要显式选择。 | +| `GetDataByIdsWithFlag` | `tl::expected GetDataByIdsWithFlag(const int64_t* ids, int64_t count, uint64_t selected_data_flag) const` | 通过 [`DATA_FLAG_*`](#数据选择标志) 选择支持的字段。Pyramid 路径需要同时设置 `store_paths: true` 和 `DATA_FLAG_PATH`。 | | `GetIndexDetailInfos` | `tl::expected, Error> GetIndexDetailInfos() const` | 列出可自省的细节字段。见 [`IndexDetailInfo`](types.md#索引细节信息)。 | | `GetDetailDataByName` | `tl::expected GetDetailDataByName(const std::string& name, IndexDetailInfo& info) const` | 按名称获取一份细节数据负载。 | diff --git a/docs/docs/zh/src/indexes/pyramid.md b/docs/docs/zh/src/indexes/pyramid.md index 2aa06455cd..d3575a0107 100644 --- a/docs/docs/zh/src/indexes/pyramid.md +++ b/docs/docs/zh/src/indexes/pyramid.md @@ -98,6 +98,7 @@ auto result = index->KnnSearch( | `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 距离计算 | +| `store_paths` | bool | `false` | 顶层开关；保留传给 `Build` 和 `Add` 的原始路径，使 `GetDataByIdsWithFlag` 在选择 `DATA_FLAG_PATH` 时可以返回它们。该开关对所有已配置的 hierarchy 生效，不支持按 hierarchy 覆盖 | | `index_min_size` | int | `0` | 子索引的最小规模；小于该值的分区会退化为线性扫描 | | `support_duplicate` | bool | `false` | 是否允许重复 ID | | `build_thread_count` | int | `1` | 构建阶段并发线程数 | @@ -226,6 +227,29 @@ base->NumElements(n) index->Build(base); ``` +### 按 ID 取回路径 + +将顶层构建参数 `store_paths` 设为 `true` 后，Pyramid 会保留原始路径，供按 ID +取回。在 `GetDataByIdsWithFlag` 中选择 `DATA_FLAG_PATH` 后，返回路径的顺序与请求 ID +的顺序一致。默认匿名 hierarchy 使用 `GetPaths()`，命名 hierarchy 使用 +`GetPaths(hierarchy_name)`： + +```cpp +int64_t requested_ids[] = {product_id_b, product_id_a}; +auto data = index->GetDataByIdsWithFlag( + requested_ids, 2, DATA_FLAG_ID | DATA_FLAG_PATH).value(); + +const std::string* site_paths = data->GetPaths("site"); +const std::string* category_paths = data->GetPaths("category"); +``` + +`GetDataByIds` 以及未选择 `DATA_FLAG_PATH` 的 `GetDataByIdsWithFlag` 都不会附带路径 +数组。`store_paths` 为 `false` 时选择 `DATA_FLAG_PATH` 会返回参数错误。开启路径存储 +后，只有当所有请求 ID 在某个 hierarchy 中都有已记录的路径，结果才会包含该 +hierarchy。只要其中一个 ID 在构建或追加时没有提供该 hierarchy 的路径，对应 getter +就返回 `nullptr`；其他路径完整的 hierarchy 仍会正常返回。单 hierarchy 模式下， +`GetPaths()` 遵循相同的完整性规则。 + ### 检索指定层级 通过检索参数中的 `"hierarchies"` 指定要检索的层级。查询 Dataset 也需要在对应的 @@ -271,7 +295,10 @@ auto result = index->RangeSearch( ### 序列化与反序列化 -多层级索引的序列化和反序列化完全透明。序列化格式包含所有层级名称及其图结构： +多 hierarchy 索引的序列化和反序列化完全透明。序列化格式包含所有 hierarchy +名称及其图结构。当 `store_paths: true` 时，常规序列化和 streaming 序列化还会持久化 +已保留的原始路径，因此反序列化后仍可通过 `GetDataByIdsWithFlag` 获取。使用默认值 +`false` 时，图 hierarchy 会被持久化，但不会保留按 ID 索引的原始路径： ```cpp // 序列化 diff --git a/docs/docs/zh/src/resources/index_parameters.md b/docs/docs/zh/src/resources/index_parameters.md index 5dceb785ef..986697674e 100644 --- a/docs/docs/zh/src/resources/index_parameters.md +++ b/docs/docs/zh/src/resources/index_parameters.md @@ -138,11 +138,16 @@ Pyramid 构建参数同样放在 `index_param` 下： "index_param": { "base_quantization_type": "sq8", "max_degree": 24, - "ef_construction": 300 + "ef_construction": 300, + "store_paths": true } } ``` +`store_paths` 是 Pyramid 顶层构建参数，默认值为 `false`。需要通过 +`GetDataByIdsWithFlag` 和 `DATA_FLAG_PATH` 返回默认或命名 hierarchy 的原始路径时应启用它；路径完整性与持久化语义见 +[Pyramid 参数表](../indexes/pyramid.md#构建参数)。 + MRLE 与 split RaBitQ 组合使用 `base_quantization_type: "tq"`、 `tq_chain: "mrle, rabitq"`、`mrle_dim`，以及 `rabitq_bits_per_dim_base`/`rabitq_bits_per_dim_precise`。该配置会自动从 split base codes diff --git a/include/vsag/constants.h b/include/vsag/constants.h index 2f37981446..5db63ddf4b 100644 --- a/include/vsag/constants.h +++ b/include/vsag/constants.h @@ -120,6 +120,7 @@ extern const char* const PYRAMID_PARAMETER_HIERARCHY_OP; extern const char* const PYRAMID_NO_BUILD_LEVELS; extern const char* const PYRAMID_HIERARCHIES; extern const char* const PYRAMID_INDEX_MIN_SIZE; +extern const char* const PYRAMID_STORE_PATHS; extern const char PART_SLASH; extern const char PART_BAR; diff --git a/include/vsag/index.h b/include/vsag/index.h index 608e5bc8b7..ad519cb7a3 100644 --- a/include/vsag/index.h +++ b/include/vsag/index.h @@ -72,6 +72,7 @@ enum class IndexType { #define DATA_FLAG_EXTRA_INFO 0x10 #define DATA_FLAG_ATTRIBUTE 0x20 #define DATA_FLAG_ID 0x40 +#define DATA_FLAG_PATH 0x80 using OffsetType = uint64_t; using SizeType = uint64_t; @@ -700,7 +701,7 @@ class Index { } /** - * @brief Retrieve all data associated with vectors identified by given IDs. + * @brief Retrieve selected data associated with vectors identified by given IDs. * * This method fetches data stored with the vectors in the index * (e.g., attributes, labels, or extra infos). @@ -709,7 +710,7 @@ class Index { * @param count Number of IDs in the 'ids' array. * @param selected_data_flag selected data flag, set with DATA_FLAG_* * @return tl::expected - * - On success: A DatasetPtr containing the extra data, attribute and vector + * - On success: A DatasetPtr containing the selected supported fields * - On failure: An error object (e.g., invalid ID, out of memory). * @note The default base-class implementation returns tl::unexpected(ErrorType::UNSUPPORTED_INDEX_OPERATION) If the index implementation does not support this operation * (default behavior for base class). @@ -756,7 +757,7 @@ class Index { } /** - * @brief Retrieve all data associated with vectors identified by given IDs. + * @brief Retrieve the default data fields associated with vectors identified by given IDs. * * This method fetches data stored with the vectors in the index * (e.g., attributes, labels, or extra infos). @@ -764,11 +765,11 @@ class Index { * @param ids Array of vector IDs for which extra information is requested. * @param count Number of IDs in the 'ids' array. * @return tl::expected - * - On success: A DatasetPtr containing the extra data, attribute and vector + * - On success: A DatasetPtr containing the implementation's default fields * - On failure: An error object (e.g., invalid ID, out of memory). * @note The default base-class implementation returns tl::unexpected(ErrorType::UNSUPPORTED_INDEX_OPERATION) If the index implementation does not support this operation * (default behavior for base class). - * @note The default implementation returns all data which in current index + * @note Optional fields may require explicit selection through GetDataByIdsWithFlag. */ [[nodiscard]] virtual tl::expected GetDataByIds(const int64_t* ids, int64_t count) const { diff --git a/src/algorithm/inner_index_interface.cpp b/src/algorithm/inner_index_interface.cpp index 28e4930a29..1d712087b5 100644 --- a/src/algorithm/inner_index_interface.cpp +++ b/src/algorithm/inner_index_interface.cpp @@ -811,12 +811,22 @@ DatasetPtr InnerIndexInterface::GetDataByIdsWithFlag(const int64_t* ids, int64_t count, uint64_t selected_data_flag) const { - auto* inner_ids = - reinterpret_cast(this->allocator_->Allocate(count * sizeof(InnerIdType))); + Vector inner_ids(allocator_); + return this->get_data_by_ids_with_flag(ids, count, selected_data_flag, inner_ids); +} + +DatasetPtr +InnerIndexInterface::get_data_by_ids_with_flag(const int64_t* ids, + int64_t count, + uint64_t selected_data_flag, + Vector& inner_ids) const { + CHECK_ARGUMENT(count >= 0, "count must not be negative"); + inner_ids.clear(); + inner_ids.reserve(static_cast(count)); { std::shared_lock lock(this->label_lookup_mutex_); for (int64_t i = 0; i < count; ++i) { - inner_ids[i] = this->label_table_->GetIdByLabel(ids[i]); + inner_ids.emplace_back(this->label_table_->GetIdByLabel(ids[i])); } } auto thread_count = static_cast(this->build_thread_count_); @@ -832,8 +842,7 @@ InnerIndexInterface::GetDataByIdsWithFlag(const int64_t* ids, dataset->Float32Vectors(fp32_data); auto get_vec_func = [&](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; ++i) { - auto inner_id = this->label_table_->GetIdByLabel(ids[i]); - this->GetVectorByInnerId(inner_id, fp32_data + i * this->dim_); + this->GetVectorByInnerId(inner_ids[i], fp32_data + i * this->dim_); } }; if (this->thread_pool_ != nullptr) { @@ -859,8 +868,7 @@ InnerIndexInterface::GetDataByIdsWithFlag(const int64_t* ids, dataset->AttributeSets(attribute_data); auto get_attr_func = [&](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; ++i) { - auto inner_id = this->label_table_->GetIdByLabel(ids[i]); - this->GetAttributeSetByInnerId(inner_id, attribute_data + i); + this->GetAttributeSetByInnerId(inner_ids[i], attribute_data + i); } }; if (this->thread_pool_ != nullptr) { @@ -887,8 +895,8 @@ InnerIndexInterface::GetDataByIdsWithFlag(const int64_t* ids, dataset->ExtraInfos(extra_info)->ExtraInfoSize(static_cast(extra_info_size_)); auto get_extra_info_func = [&](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; ++i) { - auto inner_id = this->label_table_->GetIdByLabel(ids[i]); - this->extra_infos_->GetExtraInfoById(inner_id, extra_info + i * extra_info_size_); + this->extra_infos_->GetExtraInfoById(inner_ids[i], + extra_info + i * extra_info_size_); } }; if (this->thread_pool_ != nullptr) { @@ -912,7 +920,6 @@ InnerIndexInterface::GetDataByIdsWithFlag(const int64_t* ids, memcpy(new_ids, ids, count * sizeof(int64_t)); dataset->Ids(new_ids); } - this->allocator_->Deallocate(inner_ids); return dataset; } diff --git a/src/algorithm/inner_index_interface.h b/src/algorithm/inner_index_interface.h index 38bd7dc1ce..8e0ebdc2a3 100644 --- a/src/algorithm/inner_index_interface.h +++ b/src/algorithm/inner_index_interface.h @@ -552,6 +552,12 @@ class InnerIndexInterface { } protected: + DatasetPtr + get_data_by_ids_with_flag(const int64_t* ids, + int64_t count, + uint64_t selected_data_flag, + Vector& inner_ids) const; + virtual MetadataPtr collect_streaming_header() const; diff --git a/src/algorithm/pyramid/CMakeLists.txt b/src/algorithm/pyramid/CMakeLists.txt index fe91f0a3c3..18216d0483 100644 --- a/src/algorithm/pyramid/CMakeLists.txt +++ b/src/algorithm/pyramid/CMakeLists.txt @@ -18,6 +18,8 @@ set (PYRAMID_SRCS pyramid.cpp pyramid_build_cache.cpp pyramid_index_node.cpp + pyramid_path_store.cpp + pyramid_paths.cpp pyramid_zparameters.cpp ) diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index dd4e416493..21722dff0f 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -149,6 +149,18 @@ Pyramid::build_by_odescent(const DatasetPtr& base) { if (raw_vector_ != nullptr) { raw_vector_->BatchInsertVector(data_vectors, data_num); } + if (store_paths_) { + for (const auto& [hierarchy_name, hierarchy] : hierarchies_) { + const auto* paths = base->GetPaths(hierarchy_name); + if (paths != nullptr) { + auto writer = hierarchy->path_store->AcquireWriter(); + writer.Prepare(static_cast(data_num), static_cast(data_num)); + for (uint64_t offset = 0; offset < static_cast(data_num); ++offset) { + writer.Insert(static_cast(offset), paths[offset]); + } + } + } + } auto codes = has_precise_reorder() ? precise_codes_ : base_codes_; if (thread_pool_ != nullptr && hierarchies_.size() > 1) { @@ -628,6 +640,10 @@ Pyramid::Serialize(StreamWriter& writer) const { hierarchies_.at("")->root->Serialize(writer); } + if (store_paths_) { + serialize_paths(writer); + } + // serialize footer (introduced since v0.15) JsonType basic_info; basic_info["max_capacity"].SetInt(max_capacity_); @@ -715,6 +731,13 @@ Pyramid::collect_streaming_header() const { hierarchy_tag, StreamSerializationBlockCurrentVersion(hierarchy_tag), StreamSerializationTagCritical(hierarchy_tag)); + if (store_paths_) { + const auto path_tag = static_cast(StreamSerializationTag::PYRAMID_PATHS); + AppendStreamingManifestBlock(manifest, + path_tag, + StreamSerializationBlockCurrentVersion(path_tag), + StreamSerializationTagCritical(path_tag)); + } metadata->Set("block_manifest", manifest); metadata->SetEmptyIndex(this->GetNumElements() == 0); return metadata; @@ -767,6 +790,13 @@ Pyramid::serialize_streaming_body(StreamWriter& writer) const { hierarchy_tag, StreamSerializationTagCritical(hierarchy_tag), [this](StreamWriter& w) { this->serialize_hierarchies(w); }); + if (store_paths_) { + const auto path_tag = static_cast(StreamSerializationTag::PYRAMID_PATHS); + WriteStreamingBlock( + writer, path_tag, StreamSerializationTagCritical(path_tag), [this](StreamWriter& w) { + this->serialize_paths(w); + }); + } } void @@ -829,6 +859,7 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) bool loaded_precise_codes = false; bool loaded_raw_vector = false; bool loaded_hierarchies = false; + bool loaded_paths = false; while (true) { auto block_header = StreamBlockHeader::Read(reader); @@ -895,6 +926,20 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) }); loaded_hierarchies = true; break; + case StreamSerializationTag::PYRAMID_PATHS: + if (loaded_paths) { + throw VsagException(ErrorType::READ_ERROR, + "duplicate Pyramid streaming paths block"); + } + if (not loaded_base_codes) { + throw VsagException(ErrorType::READ_ERROR, + "Pyramid streaming paths block precedes base codes"); + } + ReadSeekableBlockPayload(block_reader, block_header, [this](StreamReader& block) { + this->deserialize_paths(block, static_cast(this->cur_element_count_)); + }); + loaded_paths = true; + break; default: if (block_header.IsCritical()) { throw VsagException( @@ -925,6 +970,10 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) throw VsagException(ErrorType::READ_ERROR, "Pyramid streaming serialization raw vector block is missing"); } + if (store_paths_ && !loaded_paths) { + throw VsagException(ErrorType::READ_ERROR, + "Pyramid streaming serialization paths block is missing"); + } resize(max_capacity); this->current_memory_usage_ = static_cast(this->CalSerializeSize()); @@ -938,6 +987,13 @@ Pyramid::Deserialize(StreamReader& reader) { throw VsagException(ErrorType::READ_ERROR, "failed to read index footer"); } auto max_capacity = basic_info["max_capacity"].GetInt(); + auto param_json = JsonType::Parse(basic_info[INDEX_PARAM].GetString()); + const bool serialized_store_paths = param_json.Contains(PYRAMID_STORE_PATHS_KEY) && + param_json[PYRAMID_STORE_PATHS_KEY].GetBool(); + if (serialized_store_paths != store_paths_) { + throw VsagException(ErrorType::READ_ERROR, + "serialized Pyramid store_paths does not match config"); + } BufferStreamReader buffer_reader( &reader, std::numeric_limits::max(), this->allocator_); @@ -955,7 +1011,6 @@ Pyramid::Deserialize(StreamReader& reader) { } cur_element_count_ = base_codes_->TotalCount(); - auto param_json = JsonType::Parse(basic_info[INDEX_PARAM].GetString()); if (param_json.Contains(PYRAMID_HIERARCHIES)) { uint64_t hierarchy_count = 0; StreamReader::ReadObj(buffer_reader, hierarchy_count); @@ -978,6 +1033,10 @@ Pyramid::Deserialize(StreamReader& reader) { h_iter->second->root->Deserialize(buffer_reader); } + if (store_paths_) { + deserialize_paths(buffer_reader, static_cast(cur_element_count_)); + } + resize(max_capacity); this->current_memory_usage_ = this->CalSerializeSize(); } @@ -1016,6 +1075,7 @@ Pyramid::Add(const DatasetPtr& base) { const auto* source_ids = base->GetSourceID(); std::vector failed_ids; Vector data_biases(allocator_); + Vector accepted_inner_ids(allocator_); int64_t local_cur_element_count = 0; { std::lock_guard lock(cur_element_count_mutex_); @@ -1042,6 +1102,9 @@ Pyramid::Add(const DatasetPtr& base) { } data_biases.reserve(data_num); + if (store_paths_) { + accepted_inner_ids.reserve(data_num); + } for (int64_t i = 0; i < data_num; ++i) { if (not label_table_->CheckLabel(data_ids[i])) { const auto inner_id = @@ -1051,6 +1114,9 @@ Pyramid::Add(const DatasetPtr& base) { label_table_->InsertSourceId(inner_id, source_ids[i]); } data_biases.push_back(i); + if (store_paths_) { + accepted_inner_ids.push_back(inner_id); + } } else { logger::warn("Label {} already exists, skip adding.", data_ids[i]); failed_ids.push_back(data_ids[i]); @@ -1133,6 +1199,14 @@ Pyramid::Add(const DatasetPtr& base) { for (const auto& [hname, h_ptr] : hierarchies_) { const auto* hpath = base->GetPaths(hname); if (hpath != nullptr) { + if (store_paths_ and not accepted_inner_ids.empty()) { + auto writer = h_ptr->path_store->AcquireWriter(); + writer.Prepare(static_cast(accepted_inner_ids.back()) + 1, + data_biases.size()); + for (uint64_t offset = 0; offset < data_biases.size(); ++offset) { + writer.Insert(accepted_inner_ids[offset], hpath[data_biases[offset]]); + } + } add_to_hierarchy(*h_ptr, data_vectors, hpath, data_biases, local_cur_element_count); } } @@ -1202,6 +1276,9 @@ Pyramid::InitFeatures() { IndexFeature::SUPPORT_EXPORT_MODEL, IndexFeature::SUPPORT_GET_MEMORY_USAGE, }); + if (store_paths_) { + this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_GET_DATA_BY_IDS); + } if (has_raw_vector_) { this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_GET_RAW_VECTOR_BY_IDS); } @@ -1287,7 +1364,8 @@ static const std::string HGRAPH_PARAMS_TEMPLATE = "{EF_CONSTRUCTION_KEY}": 400, "{NO_BUILD_LEVELS}":[], "{INDEX_MIN_SIZE}": 0, - "{SUPPORT_DUPLICATE}": false + "{SUPPORT_DUPLICATE}": false, + "{PYRAMID_STORE_PATHS_KEY}": false })"; ParamPtr @@ -1334,6 +1412,7 @@ Pyramid::CheckAndMappingExternalParam(const JsonType& external_param, {PYRAMID_NO_BUILD_LEVELS, {NO_BUILD_LEVELS}}, {PYRAMID_HIERARCHIES, {PYRAMID_HIERARCHIES}}, {PYRAMID_PERSIST_SOURCE_ID, {PYRAMID_PERSIST_SOURCE_ID_KEY}}, + {PYRAMID_STORE_PATHS, {PYRAMID_STORE_PATHS_KEY}}, {PYRAMID_BASE_PQ_DIM, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, PRODUCT_QUANTIZATION_DIM_KEY}}, {PYRAMID_BASE_FILE_PATH, {BASE_CODES_KEY, IO_PARAMS_KEY, IO_FILE_PATH_KEY}}, @@ -2053,6 +2132,13 @@ Pyramid::build_with_cache(const DatasetPtr& base) { for (const auto& [hname, h_ptr] : hierarchies_) { const auto* hpath = base->GetPaths(hname); if (hpath != nullptr) { + if (store_paths_) { + auto writer = h_ptr->path_store->AcquireWriter(); + writer.Prepare(static_cast(data_num), static_cast(data_num)); + for (uint64_t offset = 0; offset < static_cast(data_num); ++offset) { + writer.Insert(static_cast(offset), hpath[offset]); + } + } populate_path_tree(*h_ptr, hpath, data_num); } } diff --git a/src/algorithm/pyramid/pyramid.h b/src/algorithm/pyramid/pyramid.h index 1a43a5567e..7896057f6e 100644 --- a/src/algorithm/pyramid/pyramid.h +++ b/src/algorithm/pyramid/pyramid.h @@ -17,10 +17,12 @@ #include #include +#include #include #include "algorithm/inner_index_interface.h" #include "algorithm/pyramid/pyramid_build_cache.h" +#include "algorithm/pyramid/pyramid_path_store.h" #include "datacell/graph_interface.h" #include "datacell/sparse_graph_datacell_parameter.h" #include "impl/allocator/safe_allocator.h" @@ -67,6 +69,7 @@ class Pyramid : public InnerIndexInterface { RABITQ_SPLIT_DATA_CELL), support_duplicate_(pyramid_param->support_duplicate), persist_source_id_(pyramid_param->persist_source_id), + store_paths_(pyramid_param->store_paths), cache_(std::make_unique(common_param.allocator_.get())) { base_codes_ = FlattenInterface::MakeInstance(pyramid_param->base_codes_param, common_param); if (pyramid_param->has_hierarchies) { @@ -85,6 +88,9 @@ class Pyramid : public InnerIndexInterface { h_param.no_build_levels.end()); h->ef_construction = h_param.ef_construction; h->alpha = h_param.alpha; + if (store_paths_) { + h->path_store = std::make_unique(allocator_); + } hierarchies_.insert({h_param.name, std::move(h)}); } } else { @@ -95,6 +101,9 @@ class Pyramid : public InnerIndexInterface { pyramid_param->no_build_levels.end()); h->ef_construction = pyramid_param->ef_construction; h->alpha = pyramid_param->alpha; + if (store_paths_) { + h->path_store = std::make_unique(allocator_); + } hierarchies_.insert({"", std::move(h)}); } points_mutex_ = std::make_shared(max_capacity_, allocator_); @@ -142,6 +151,11 @@ class Pyramid : public InnerIndexInterface { bool calculate_precise_distance = true, int64_t topk = -1) const override; + DatasetPtr + GetDataByIdsWithFlag(const int64_t* ids, + int64_t count, + uint64_t selected_data_flag) const override; + void Deserialize(StreamReader& reader) override; @@ -241,6 +255,12 @@ class Pyramid : public InnerIndexInterface { void deserialize_hierarchies(StreamReader& reader, const JsonType& basic_info); + void + serialize_paths(StreamWriter& writer) const; + + void + deserialize_paths(StreamReader& reader, uint64_t max_count); + // RAII guard that returns the VisitedList to the pool on scope exit, // ensuring no leak if the search throws. class VisitedListGuard { @@ -277,6 +297,7 @@ class Pyramid : public InnerIndexInterface { Vector no_build_levels; // depths where graph build is skipped uint64_t ef_construction{400}; // expansion factor during graph build float alpha{1.2F}; // Relative Neighborhood Graph pruning coefficient + std::unique_ptr path_store{nullptr}; Hierarchy(const std::string& n, std::unique_ptr r, Allocator* alloc) : name(n), root(std::move(r)), no_build_levels(alloc) { @@ -423,6 +444,7 @@ class Pyramid : public InnerIndexInterface { uint32_t index_min_size_{0}; // min node size before graph is built bool persist_source_id_{false}; // whether to persist source_id in serialization + bool store_paths_{false}; // whether to retain paths for ID-based retrieval std::unique_ptr cache_{nullptr}; // per-graph caches for warm-start build diff --git a/src/algorithm/pyramid/pyramid_path_store.cpp b/src/algorithm/pyramid/pyramid_path_store.cpp new file mode 100644 index 0000000000..9b462dbd6f --- /dev/null +++ b/src/algorithm/pyramid/pyramid_path_store.cpp @@ -0,0 +1,293 @@ +// 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 "pyramid_path_store.h" + +#include +#include +#include + +#include "common.h" +#include "storage/stream_reader.h" +#include "storage/stream_writer.h" +#include "vsag_exception.h" + +namespace vsag { + +namespace { + +constexpr uint64_t MISSING_PATH_OFFSET = std::numeric_limits::max(); + +template +void +read_pyramid_path_vector(StreamReader& reader, Vector& values, uint64_t max_count) { + uint64_t size = 0; + StreamReader::ReadObj(reader, size); + if (size > max_count) { + throw VsagException(ErrorType::READ_ERROR, "corrupted Pyramid path vector size"); + } + + const auto cursor = reader.GetCursor(); + const auto reader_length = reader.Length(); + if (cursor > reader_length || size > (reader_length - cursor) / sizeof(T)) { + throw VsagException(ErrorType::READ_ERROR, + "corrupted Pyramid path vector exceeds remaining payload"); + } + if (size > values.max_size()) { + throw VsagException(ErrorType::READ_ERROR, "Pyramid path vector is too large"); + } + + values.resize(size); + if (size > 0) { + reader.Read(reinterpret_cast(values.data()), size * sizeof(T)); + } +} + +bool +are_pyramid_path_rows_valid(const Vector& offsets, + const Vector& counts, + uint64_t path_count) { + if (offsets.size() != counts.size()) { + return false; + } + for (uint64_t slot = 0; slot < offsets.size(); ++slot) { + const auto path_offset = offsets[slot]; + const auto row_path_count = static_cast(counts[slot]); + if (path_offset == MISSING_PATH_OFFSET) { + if (row_path_count != 0) { + return false; + } + continue; + } + if (path_offset > path_count || row_path_count > path_count - path_offset) { + return false; + } + } + return true; +} + +void +reserve_paths_geometrically(Vector& paths, uint64_t required_capacity) { + const auto current_capacity = static_cast(paths.capacity()); + if (required_capacity <= current_capacity) { + return; + } + + const auto max_capacity = static_cast(paths.max_size()); + auto target_capacity = std::max(current_capacity, 1); + if (target_capacity <= max_capacity / 2) { + target_capacity *= 2; + } else { + target_capacity = max_capacity; + } + paths.reserve(std::max(required_capacity, target_capacity)); +} + +} // namespace + +std::string +ReadPyramidPathString(StreamReader& reader) { + uint64_t length = 0; + StreamReader::ReadObj(reader, length); + const auto cursor = reader.GetCursor(); + const auto reader_length = reader.Length(); + if (cursor > reader_length || length > reader_length - cursor) { + throw VsagException(ErrorType::READ_ERROR, "corrupted Pyramid path string length"); + } + if (length > std::string{}.max_size()) { + throw VsagException(ErrorType::READ_ERROR, "Pyramid path string is too large"); + } + std::string value(length, '\0'); + if (length > 0) { + reader.Read(value.data(), length); + } + return value; +} + +void +PyramidPathStore::Writer::Prepare(uint64_t slot_count, uint64_t additional_path_count) { + constexpr uint64_t max_slot_count = + static_cast(std::numeric_limits::max()) + 1; + CHECK_ARGUMENT(slot_count <= max_slot_count, "Pyramid path slot count is too large"); + CHECK_ARGUMENT(slot_count <= store_.offsets_.max_size(), + "Pyramid path slot count is too large"); + CHECK_ARGUMENT(slot_count <= store_.counts_.max_size(), "Pyramid path slot count is too large"); + + const auto path_count = static_cast(store_.paths_.size()); + CHECK_ARGUMENT(additional_path_count <= store_.paths_.max_size() - path_count, + "Pyramid path count is too large"); + + const auto old_offset_count = store_.offsets_.size(); + const auto old_count_count = store_.counts_.size(); + try { + if (slot_count > old_offset_count) { + store_.offsets_.resize(slot_count, MISSING_PATH_OFFSET); + } + if (slot_count > old_count_count) { + store_.counts_.resize(slot_count, 0); + } + reserve_paths_geometrically(store_.paths_, path_count + additional_path_count); + } catch (...) { + store_.offsets_.resize(old_offset_count); + store_.counts_.resize(old_count_count); + throw; + } +} + +void +PyramidPathStore::Writer::Insert(InnerIdType inner_id, const std::string& path) { + Insert(inner_id, &path, 1); +} + +void +PyramidPathStore::Writer::Insert(InnerIdType inner_id, + const std::string* paths, + uint64_t path_count) { + CHECK_ARGUMENT(path_count <= std::numeric_limits::max(), + "too many Pyramid paths for one inner id"); + if (path_count > 0) { + CHECK_ARGUMENT(paths != nullptr, "Pyramid paths must not be null"); + } + const auto slot = static_cast(inner_id); + if (slot < store_.offsets_.size()) { + CHECK_ARGUMENT(store_.offsets_[slot] == MISSING_PATH_OFFSET, + "inner id already has Pyramid paths"); + } + + const auto old_path_count = static_cast(store_.paths_.size()); + const auto old_offset_count = store_.offsets_.size(); + const auto old_count_count = store_.counts_.size(); + try { + Prepare(slot + 1, path_count); + for (uint64_t offset = 0; offset < path_count; ++offset) { + store_.paths_.emplace_back(paths[offset]); + } + store_.offsets_[slot] = old_path_count; + store_.counts_[slot] = static_cast(path_count); + } catch (...) { + store_.paths_.resize(old_path_count); + store_.offsets_.resize(old_offset_count); + store_.counts_.resize(old_count_count); + throw; + } +} + +PyramidPathStore::Writer +PyramidPathStore::AcquireWriter() { + return Writer(*this); +} + +bool +PyramidPathStore::GetPaths(const Vector& inner_ids, std::string* paths) const { + std::shared_lock lock(mutex_); + if (offsets_.size() != counts_.size()) { + return false; + } + for (uint64_t offset = 0; offset < inner_ids.size(); ++offset) { + const auto inner_id = static_cast(inner_ids[offset]); + if (inner_id >= offsets_.size() || counts_[inner_id] != 1 || + offsets_[inner_id] >= paths_.size()) { + return false; + } + paths[offset] = paths_[offsets_[inner_id]]; + } + return true; +} + +bool +PyramidPathStore::GetPathRows(const Vector& inner_ids, + std::vector>& path_rows) const { + std::shared_lock lock(mutex_); + if (offsets_.size() != counts_.size()) { + return false; + } + path_rows.clear(); + path_rows.reserve(inner_ids.size()); + for (const auto inner_id : inner_ids) { + const auto slot = static_cast(inner_id); + if (slot >= offsets_.size()) { + return false; + } + const auto path_offset = offsets_[slot]; + const auto path_count = static_cast(counts_[slot]); + if (path_offset == MISSING_PATH_OFFSET || path_offset > paths_.size() || + path_count > paths_.size() - path_offset) { + return false; + } + const auto begin_offset = static_cast::difference_type>(path_offset); + const auto end_offset = + static_cast::difference_type>(path_offset + path_count); + path_rows.emplace_back(paths_.begin() + begin_offset, paths_.begin() + end_offset); + } + return true; +} + +void +PyramidPathStore::Serialize(StreamWriter& writer) const { + std::shared_lock lock(mutex_); + if (offsets_.size() != counts_.size()) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "Pyramid path store has inconsistent slot arrays"); + } + StreamWriter::WriteVector(writer, offsets_); + StreamWriter::WriteVector(writer, counts_); + StreamWriter::WriteObj(writer, static_cast(paths_.size())); + for (const auto& path : paths_) { + StreamWriter::WriteString(writer, path); + } +} + +void +PyramidPathStore::Deserialize(StreamReader& reader, uint64_t max_count) { + constexpr uint64_t max_slot_count = + static_cast(std::numeric_limits::max()) + 1; + const auto max_slot_vector_count = std::min(max_count, max_slot_count); + + Vector restored_offsets(allocator_); + Vector restored_counts(allocator_); + read_pyramid_path_vector(reader, restored_offsets, max_slot_vector_count); + read_pyramid_path_vector(reader, restored_counts, restored_offsets.size()); + if (restored_offsets.size() != restored_counts.size()) { + throw VsagException(ErrorType::READ_ERROR, + "corrupted Pyramid path slot vector sizes do not match"); + } + + uint64_t path_count = 0; + StreamReader::ReadObj(reader, path_count); + Vector restored_paths(allocator_); + const auto path_cursor = reader.GetCursor(); + const auto reader_length = reader.Length(); + if (path_cursor > reader_length || + path_count > (reader_length - path_cursor) / sizeof(uint64_t)) { + throw VsagException(ErrorType::READ_ERROR, + "corrupted Pyramid path count exceeds remaining payload"); + } + if (path_count > restored_paths.max_size()) { + throw VsagException(ErrorType::READ_ERROR, "Pyramid path count is too large"); + } + if (!are_pyramid_path_rows_valid(restored_offsets, restored_counts, path_count)) { + throw VsagException(ErrorType::READ_ERROR, "corrupted Pyramid path row range"); + } + restored_paths.reserve(path_count); + for (uint64_t offset = 0; offset < path_count; ++offset) { + restored_paths.emplace_back(ReadPyramidPathString(reader)); + } + + std::unique_lock lock(mutex_); + offsets_.swap(restored_offsets); + counts_.swap(restored_counts); + paths_.swap(restored_paths); +} + +} // namespace vsag diff --git a/src/algorithm/pyramid/pyramid_path_store.h b/src/algorithm/pyramid/pyramid_path_store.h new file mode 100644 index 0000000000..8f58c8f6bc --- /dev/null +++ b/src/algorithm/pyramid/pyramid_path_store.h @@ -0,0 +1,94 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "typing.h" + +namespace vsag { + +class StreamReader; +class StreamWriter; + +// Reads a serialized path string only after verifying that its payload fits in the reader. +std::string +ReadPyramidPathString(StreamReader& reader); + +class PyramidPathStore { +public: + class Writer { + public: + Writer(const Writer&) = delete; + Writer& + operator=(const Writer&) = delete; + Writer(Writer&&) = delete; + Writer& + operator=(Writer&&) = delete; + + // Prepares storage for one batch. Existing rows are never removed. + void + Prepare(uint64_t slot_count, uint64_t additional_path_count); + + void + Insert(InnerIdType inner_id, const std::string& path); + + void + Insert(InnerIdType inner_id, const std::string* paths, uint64_t path_count); + + private: + friend class PyramidPathStore; + + explicit Writer(PyramidPathStore& store) : store_(store), lock_(store.mutex_) { + } + + PyramidPathStore& store_; + std::unique_lock lock_; + }; + + explicit PyramidPathStore(Allocator* allocator) + : offsets_(allocator), counts_(allocator), paths_(allocator), allocator_(allocator) { + } + + // Holds the store's exclusive lock for the lifetime of the returned writer. + [[nodiscard]] Writer + AcquireWriter(); + + [[nodiscard]] bool + GetPaths(const Vector& inner_ids, std::string* paths) const; + + [[nodiscard]] bool + GetPathRows(const Vector& inner_ids, + std::vector>& path_rows) const; + + void + Serialize(StreamWriter& writer) const; + + void + Deserialize(StreamReader& reader, uint64_t max_count); + +private: + mutable std::shared_mutex mutex_; + Vector offsets_; + Vector counts_; + Vector paths_; + Allocator* allocator_{nullptr}; +}; + +} // namespace vsag diff --git a/src/algorithm/pyramid/pyramid_path_store_test.cpp b/src/algorithm/pyramid/pyramid_path_store_test.cpp new file mode 100644 index 0000000000..1dfd0d80eb --- /dev/null +++ b/src/algorithm/pyramid/pyramid_path_store_test.cpp @@ -0,0 +1,268 @@ +// 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 "pyramid_path_store.h" + +#include +#include +#include +#include +#include +#include + +#include "impl/allocator/safe_allocator.h" +#include "storage/stream_reader.h" +#include "storage/stream_writer.h" +#include "unittest.h" + +TEST_CASE("PyramidPathStore inserts and reorders paths", "[ut][pyramid][path_store]") { + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::PyramidPathStore store(allocator.get()); + const std::array source = {"root/a", "", "root/c", "root/d"}; + + { + auto writer = store.AcquireWriter(); + for (uint64_t slot = 0; slot < source.size(); ++slot) { + writer.Insert(static_cast(slot), source[slot]); + } + } + vsag::Vector inner_ids(allocator.get()); + inner_ids.push_back(2); + inner_ids.push_back(0); + inner_ids.push_back(1); + std::array restored; + + REQUIRE(store.GetPaths(inner_ids, restored.data())); + REQUIRE(restored == std::array{"root/c", "root/a", ""}); + + { + auto writer = store.AcquireWriter(); + REQUIRE_THROWS(writer.Insert(0, "duplicate")); + } +} + +TEST_CASE("PyramidPathStore stores zero, one, and multiple paths", "[ut][pyramid][path_store]") { + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::PyramidPathStore store(allocator.get()); + const std::array multiple_paths = {"root/a", "root/b", "root/a"}; + { + auto writer = store.AcquireWriter(); + writer.Prepare(6, 4); + writer.Insert(5, multiple_paths.data(), multiple_paths.size()); + writer.Insert(2, nullptr, 0); + writer.Insert(4, "single"); + } + vsag::Vector inner_ids(allocator.get()); + inner_ids.push_back(5); + inner_ids.push_back(2); + inner_ids.push_back(4); + std::vector> restored_rows; + REQUIRE(store.GetPathRows(inner_ids, restored_rows)); + REQUIRE(restored_rows == + std::vector>{{"root/a", "root/b", "root/a"}, {}, {"single"}}); + + std::array restored_single; + vsag::Vector single_id(allocator.get()); + single_id.push_back(4); + REQUIRE(store.GetPaths(single_id, restored_single.data())); + REQUIRE(restored_single[0] == "single"); + + vsag::Vector multiple_id(allocator.get()); + multiple_id.push_back(5); + REQUIRE_FALSE(store.GetPaths(multiple_id, restored_single.data())); + + vsag::Vector empty_id(allocator.get()); + empty_id.push_back(2); + REQUIRE_FALSE(store.GetPaths(empty_id, restored_single.data())); + + vsag::Vector hole_ids(allocator.get()); + hole_ids.push_back(3); + REQUIRE_FALSE(store.GetPaths(hole_ids, restored_single.data())); + REQUIRE_FALSE(store.GetPathRows(hole_ids, restored_rows)); + + vsag::Vector out_of_range_ids(allocator.get()); + out_of_range_ids.push_back(6); + REQUIRE_FALSE(store.GetPaths(out_of_range_ids, restored_single.data())); + + { + auto writer = store.AcquireWriter(); + REQUIRE_THROWS(writer.Insert(2, "duplicate")); + REQUIRE_THROWS(writer.Insert(1, nullptr, 1)); + const std::string path = "too-many"; + REQUIRE_THROWS(writer.Insert( + 1, &path, static_cast(std::numeric_limits::max()) + 1)); + writer.Insert(1, "valid-after-rejected-insert"); + } +} + +TEST_CASE("PyramidPathStore preparation does not shrink slots", "[ut][pyramid][path_store]") { + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + vsag::PyramidPathStore store(allocator.get()); + { + auto writer = store.AcquireWriter(); + writer.Prepare(8, 2); + writer.Insert(7, "last"); + } + { + auto writer = store.AcquireWriter(); + writer.Prepare(3, 1); + writer.Insert(2, "earlier"); + } + vsag::Vector present_ids(allocator.get()); + present_ids.push_back(7); + present_ids.push_back(2); + std::array restored; + REQUIRE(store.GetPaths(present_ids, restored.data())); + REQUIRE(restored == std::array{"last", "earlier"}); + + vsag::Vector hole_ids(allocator.get()); + hole_ids.push_back(3); + REQUIRE_FALSE(store.GetPaths(hole_ids, restored.data())); +} + +TEST_CASE("PyramidPathStore serializes direct dense fields", "[ut][pyramid][path_store]") { + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + const std::array multiple_paths = {"root/a", "root/b", "root/a"}; + + vsag::PyramidPathStore store(allocator.get()); + { + auto writer = store.AcquireWriter(); + writer.Insert(0, multiple_paths.data(), multiple_paths.size()); + writer.Insert(2, nullptr, 0); + writer.Insert(4, "single"); + } + + std::stringstream stream; + vsag::IOStreamWriter writer(stream); + store.Serialize(writer); + + std::stringstream expected_stream; + vsag::IOStreamWriter expected_writer(expected_stream); + const std::vector expected_offsets = { + 0, std::numeric_limits::max(), 3, std::numeric_limits::max(), 3}; + const std::vector expected_counts = {3, 0, 0, 0, 1}; + vsag::StreamWriter::WriteVector(expected_writer, expected_offsets); + vsag::StreamWriter::WriteVector(expected_writer, expected_counts); + vsag::StreamWriter::WriteObj(expected_writer, uint64_t{4}); + for (const auto& path : multiple_paths) { + vsag::StreamWriter::WriteString(expected_writer, path); + } + vsag::StreamWriter::WriteString(expected_writer, "single"); + REQUIRE(stream.str() == expected_stream.str()); + + vsag::PyramidPathStore restored_store(allocator.get()); + vsag::IOStreamReader reader(stream); + restored_store.Deserialize(reader, 5); + vsag::Vector inner_ids(allocator.get()); + inner_ids.push_back(0); + inner_ids.push_back(2); + inner_ids.push_back(4); + std::vector> restored_rows; + REQUIRE(restored_store.GetPathRows(inner_ids, restored_rows)); + REQUIRE(restored_rows == + std::vector>{{"root/a", "root/b", "root/a"}, {}, {"single"}}); + + vsag::Vector hole_ids(allocator.get()); + hole_ids.push_back(1); + REQUIRE_FALSE(restored_store.GetPathRows(hole_ids, restored_rows)); +} + +TEST_CASE("PyramidPathStore rejects malformed serialization", "[ut][pyramid][path_store]") { + auto allocator = vsag::SafeAllocator::FactoryDefaultAllocator(); + const auto write_payload = [](std::stringstream& stream, + const std::vector& offsets, + const std::vector& counts, + const std::vector& paths) { + vsag::IOStreamWriter writer(stream); + vsag::StreamWriter::WriteVector(writer, offsets); + vsag::StreamWriter::WriteVector(writer, counts); + vsag::StreamWriter::WriteObj(writer, static_cast(paths.size())); + for (const auto& path : paths) { + vsag::StreamWriter::WriteString(writer, path); + } + }; + + SECTION("vector size exceeds maximum") { + std::stringstream stream; + vsag::IOStreamWriter writer(stream); + vsag::StreamWriter::WriteObj(writer, uint64_t{3}); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 2)); + } + + SECTION("vector size exceeds remaining payload") { + std::stringstream stream; + vsag::IOStreamWriter writer(stream); + vsag::StreamWriter::WriteObj(writer, uint64_t{2}); + vsag::StreamWriter::WriteObj(writer, uint64_t{0}); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 2)); + } + + SECTION("slot vector sizes do not match") { + std::stringstream stream; + write_payload(stream, {0}, {}, {"path"}); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 1)); + } + + SECTION("hole has a nonzero path count") { + std::stringstream stream; + write_payload(stream, {std::numeric_limits::max()}, {1}, {}); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 1)); + } + + SECTION("path range is invalid") { + std::stringstream stream; + write_payload(stream, {1}, {1}, {"path"}); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 1)); + } + + SECTION("path count exceeds remaining payload") { + std::stringstream stream; + vsag::IOStreamWriter writer(stream); + vsag::StreamWriter::WriteVector(writer, std::vector{0}); + vsag::StreamWriter::WriteVector(writer, std::vector{1}); + vsag::StreamWriter::WriteObj(writer, uint64_t{1}); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 1)); + } + + SECTION("path string exceeds remaining payload") { + std::stringstream stream; + vsag::IOStreamWriter writer(stream); + vsag::StreamWriter::WriteVector(writer, std::vector{0}); + vsag::StreamWriter::WriteVector(writer, std::vector{1}); + vsag::StreamWriter::WriteObj(writer, uint64_t{1}); + vsag::StreamWriter::WriteObj(writer, std::numeric_limits::max()); + + vsag::PyramidPathStore store(allocator.get()); + vsag::IOStreamReader reader(stream); + REQUIRE_THROWS(store.Deserialize(reader, 1)); + } +} diff --git a/src/algorithm/pyramid/pyramid_paths.cpp b/src/algorithm/pyramid/pyramid_paths.cpp new file mode 100644 index 0000000000..7e366fa741 --- /dev/null +++ b/src/algorithm/pyramid/pyramid_paths.cpp @@ -0,0 +1,95 @@ +// 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 "pyramid.h" +#include "pyramid_path_store.h" +#include "storage/stream_reader.h" +#include "storage/stream_writer.h" +#include "vsag_exception.h" + +namespace vsag { + +void +Pyramid::serialize_paths(StreamWriter& writer) const { + StreamWriter::WriteObj(writer, static_cast(hierarchies_.size())); + for (const auto& [hierarchy_name, hierarchy] : hierarchies_) { + if (hierarchy->path_store == nullptr) { + throw VsagException(ErrorType::INTERNAL_ERROR, "Pyramid path store is missing"); + } + StreamWriter::WriteString(writer, hierarchy_name); + hierarchy->path_store->Serialize(writer); + } +} + +void +Pyramid::deserialize_paths(StreamReader& reader, uint64_t max_count) { + uint64_t hierarchy_count = 0; + StreamReader::ReadObj(reader, hierarchy_count); + if (hierarchy_count != hierarchies_.size()) { + throw VsagException(ErrorType::READ_ERROR, "corrupted Pyramid path hierarchy count"); + } + + for (uint64_t offset = 0; offset < hierarchy_count; ++offset) { + auto hierarchy_name = ReadPyramidPathString(reader); + const auto hierarchy = hierarchies_.find(hierarchy_name); + if (hierarchy == hierarchies_.end()) { + throw VsagException(ErrorType::READ_ERROR, + fmt::format("unknown Pyramid path hierarchy '{}'", hierarchy_name)); + } + if (hierarchy->second->path_store == nullptr) { + throw VsagException(ErrorType::READ_ERROR, "Pyramid path storage is disabled"); + } + hierarchy->second->path_store->Deserialize(reader, max_count); + } +} + +DatasetPtr +Pyramid::GetDataByIdsWithFlag(const int64_t* ids, + int64_t count, + uint64_t selected_data_flag) const { + const bool wants_paths = (selected_data_flag & DATA_FLAG_PATH) != 0U; + if (wants_paths) { + CHECK_ARGUMENT(store_paths_, + "DATA_FLAG_PATH requires store_paths=true in the Pyramid build parameters"); + } + + Vector inner_ids(allocator_); + auto result = get_data_by_ids_with_flag(ids, count, selected_data_flag, inner_ids); + if (not wants_paths) { + return result; + } + + for (const auto& [hierarchy_name, hierarchy] : hierarchies_) { + if (hierarchy->path_store == nullptr) { + throw VsagException(ErrorType::INTERNAL_ERROR, "Pyramid path store is missing"); + } + auto paths = std::make_unique(static_cast(count)); + if (not hierarchy->path_store->GetPaths(inner_ids, paths.get())) { + continue; + } + if (hierarchy_name.empty()) { + result->Paths(paths.get()); + } else { + result->Paths(hierarchy_name, paths.get()); + } + paths.release(); + } + return result; +} + +} // namespace vsag diff --git a/src/algorithm/pyramid/pyramid_zparameters.cpp b/src/algorithm/pyramid/pyramid_zparameters.cpp index 954b3dd697..fcaa4c308e 100644 --- a/src/algorithm/pyramid/pyramid_zparameters.cpp +++ b/src/algorithm/pyramid/pyramid_zparameters.cpp @@ -193,6 +193,10 @@ PyramidParameters::FromJson(const JsonType& json) { this->persist_source_id = json[PYRAMID_PERSIST_SOURCE_ID_KEY].GetBool(); } + if (json.Contains(PYRAMID_STORE_PATHS_KEY)) { + this->store_paths = json[PYRAMID_STORE_PATHS_KEY].GetBool(); + } + if (json.Contains(SUPPORT_DUPLICATE)) { this->support_duplicate = json[SUPPORT_DUPLICATE].GetBool(); } @@ -243,6 +247,7 @@ PyramidParameters::ToJson() const { json[INDEX_MIN_SIZE].SetInt(index_min_size); json[SUPPORT_DUPLICATE].SetBool(support_duplicate); json[PYRAMID_PERSIST_SOURCE_ID_KEY].SetBool(persist_source_id); + json[PYRAMID_STORE_PATHS_KEY].SetBool(store_paths); if (this->use_reorder && this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) { json[PRECISE_CODES_KEY].SetJson(precise_codes_param->ToJson()); } @@ -312,6 +317,7 @@ PyramidParameters::CheckCompatibility(const ParamPtr& other) const { } CHECK_FIELD_EQ(*this, *p, index_min_size); CHECK_FIELD_EQ(*this, *p, support_duplicate); + CHECK_FIELD_EQ(*this, *p, store_paths); return true; } diff --git a/src/algorithm/pyramid/pyramid_zparameters.h b/src/algorithm/pyramid/pyramid_zparameters.h index a40d67f3e9..d2d8fb268c 100644 --- a/src/algorithm/pyramid/pyramid_zparameters.h +++ b/src/algorithm/pyramid/pyramid_zparameters.h @@ -82,6 +82,7 @@ struct PyramidParameters : public InnerIndexParameter { bool support_duplicate{false}; bool has_hierarchies{false}; bool persist_source_id{false}; + bool store_paths{false}; }; class PyramidSearchParameters : public IndexSearchParameter { diff --git a/src/algorithm/pyramid/pyramid_zparameters_test.cpp b/src/algorithm/pyramid/pyramid_zparameters_test.cpp index b907a4e855..f8687aecda 100644 --- a/src/algorithm/pyramid/pyramid_zparameters_test.cpp +++ b/src/algorithm/pyramid/pyramid_zparameters_test.cpp @@ -159,6 +159,22 @@ TEST_CASE("Pyramid persist_source_id parameter", "[ut][PyramidParameters]") { REQUIRE(param->ToJson()["persist_source_id"].GetBool()); } +TEST_CASE("Pyramid store_paths parameter", "[ut][PyramidParameters]") { + PyramidDefaultParam default_param; + auto param_json = vsag::JsonType::Parse(generate_pyramid(default_param)); + auto param = std::make_shared(); + param->FromJson(param_json); + + REQUIRE_FALSE(param->store_paths); + REQUIRE_FALSE(param->ToJson()["store_paths"].GetBool()); + + param_json["store_paths"].SetBool(true); + param->FromJson(param_json); + + REQUIRE(param->store_paths); + REQUIRE(param->ToJson()["store_paths"].GetBool()); +} + TEST_CASE("Pyramid Hierarchy Parameters Test", "[ut][PyramidParameters][hierarchy]") { SECTION("parse string and object hierarchy definitions") { auto param = ParsePyramidWithHierarchies( @@ -296,6 +312,18 @@ TEST_CASE("Pyramid Parameters CheckCompatibility", "[ut][PyramidParameter][Check TEST_COMPATIBILITY_CASE("different index min size", index_min_size, 500, 1500, false); TEST_COMPATIBILITY_CASE("different support duplicate", support_duplicate, false, true, false); + SECTION("different store_paths") { + PyramidDefaultParam default_param; + auto param1 = std::make_shared(); + auto param2 = std::make_shared(); + param1->FromString(generate_pyramid(default_param)); + param2->FromString(generate_pyramid(default_param)); + param2->store_paths = true; + + REQUIRE_FALSE(param1->CheckCompatibility(param2)); + REQUIRE_FALSE(param2->CheckCompatibility(param1)); + } + SECTION("same hierarchies in different order") { auto param1 = ParsePyramidWithHierarchies( nlohmann::json::array({{{"name", "site"}, {"no_build_levels", {0, 1}}}, @@ -401,7 +429,7 @@ TEST_CASE("Pyramid Search Hierarchy Parameters Test", } } -TEST_CASE("Pyramid maps support_duplicate to graph parameter", "[ut][PyramidParameters]") { +TEST_CASE("Pyramid maps external parameters", "[ut][PyramidParameters]") { auto param = vsag::JsonType::Parse(R"({ "base_quantization_type": "fp32", "base_io_type": "memory_io", @@ -416,6 +444,7 @@ TEST_CASE("Pyramid maps support_duplicate to graph parameter", "[ut][PyramidPara "index_min_size": 0, "support_duplicate": true, "persist_source_id": true, + "store_paths": true, "hierarchies": [ "site", {"name": "taxonomy", "no_build_levels": [0, 2]} @@ -432,6 +461,7 @@ TEST_CASE("Pyramid maps support_duplicate to graph parameter", "[ut][PyramidPara REQUIRE(typed_param != nullptr); REQUIRE(typed_param->support_duplicate); REQUIRE(typed_param->persist_source_id); + REQUIRE(typed_param->store_paths); REQUIRE(typed_param->graph_param->support_duplicate_); REQUIRE(typed_param->has_hierarchies); REQUIRE(typed_param->hierarchies.size() == 2); diff --git a/src/constants.cpp b/src/constants.cpp index b4f396fab2..090a53a14d 100644 --- a/src/constants.cpp +++ b/src/constants.cpp @@ -253,6 +253,7 @@ const char* const PYRAMID_NO_BUILD_LEVELS = "no_build_levels"; // so either can evolve independently without coupling the other context const char* const PYRAMID_HIERARCHIES = "hierarchies"; const char* const PYRAMID_INDEX_MIN_SIZE = "index_min_size"; +const char* const PYRAMID_STORE_PATHS = "store_paths"; const char* const GNO_IMI_FIRST_ORDER_BUCKETS_COUNT = "first_order_buckets_count"; const char* const GNO_IMI_SECOND_ORDER_BUCKETS_COUNT = "second_order_buckets_count"; diff --git a/src/inner_string_params.h b/src/inner_string_params.h index 30408601bd..8d519897df 100644 --- a/src/inner_string_params.h +++ b/src/inner_string_params.h @@ -60,6 +60,7 @@ const char* const HGRAPH_MCI_INCREMENTAL_JOIN_RATIO_THRESHOLD_KEY = const char* const HGRAPH_MCI_INCREMENTAL_ADDED_MCT_KEY = "mci_incremental_added_mct"; const char* const HGRAPH_MCI_INCREMENTAL_CLIQUE_MAX_KEY = "mci_incremental_clique_max"; const char* const PYRAMID_PERSIST_SOURCE_ID_KEY = "persist_source_id"; +const char* const PYRAMID_STORE_PATHS_KEY = "store_paths"; const char* const LABEL_REMAP_TYPE_VALUE_ROBIN = "robin"; const char* const LABEL_REMAP_TYPE_VALUE_PG = "pg"; const char* const GRAPH_KEY = "graph"; @@ -310,6 +311,7 @@ const std::unordered_map DEFAULT_MAP = { {"HOLD_MOLDS", HOLD_MOLDS}, {"HGRAPH_PERSIST_SOURCE_ID_KEY", HGRAPH_PERSIST_SOURCE_ID_KEY}, {"PYRAMID_PERSIST_SOURCE_ID_KEY", PYRAMID_PERSIST_SOURCE_ID_KEY}, + {"PYRAMID_STORE_PATHS_KEY", PYRAMID_STORE_PATHS_KEY}, {"IVF_PARTITION_STRATEGY_TYPE_GNO_IMI", IVF_PARTITION_STRATEGY_TYPE_GNO_IMI}, {"STORE_RAW_VECTOR_KEY", STORE_RAW_VECTOR_KEY}, {"RAW_VECTOR_KEY", RAW_VECTOR_KEY}, diff --git a/src/storage/serialization_tags.h b/src/storage/serialization_tags.h index 7828cd03f9..ae4abc25e6 100644 --- a/src/storage/serialization_tags.h +++ b/src/storage/serialization_tags.h @@ -42,6 +42,7 @@ enum class StreamSerializationTag : uint32_t { IVF_BUCKET_GRAPH = 16, IVF_PRECISE_BUCKET = 17, CONJUGATE_GRAPH = 18, + PYRAMID_PATHS = 19, }; inline const char* @@ -85,6 +86,8 @@ StreamSerializationTagName(uint32_t tag) { return "ivf_precise_bucket"; case StreamSerializationTag::CONJUGATE_GRAPH: return "conjugate_graph"; + case StreamSerializationTag::PYRAMID_PATHS: + return "pyramid_paths"; } return "unknown"; } @@ -114,6 +117,10 @@ StreamSerializationTagCritical(uint32_t tag) { case StreamSerializationTag::IVF_BUCKET_GRAPH: case StreamSerializationTag::CONJUGATE_GRAPH: return false; + case StreamSerializationTag::PYRAMID_PATHS: + // Unknown readers may skip this optional capability; readers configured with + // store_paths validate the block presence separately. + return false; } return false; } @@ -140,6 +147,7 @@ StreamSerializationBlockCurrentVersion(uint32_t tag) { case StreamSerializationTag::PYRAMID_HIERARCHIES: case StreamSerializationTag::CODE_SLOT_MAP: case StreamSerializationTag::IVF_PRECISE_BUCKET: + case StreamSerializationTag::PYRAMID_PATHS: return kStreamSerializationBlockVersionV1; case StreamSerializationTag::IVF_BUCKET_GRAPH: case StreamSerializationTag::CONJUGATE_GRAPH: diff --git a/tests/test_pyramid.cpp b/tests/test_pyramid.cpp index 3c8ce9b47e..f3258b9fa7 100644 --- a/tests/test_pyramid.cpp +++ b/tests/test_pyramid.cpp @@ -2098,7 +2098,8 @@ TEST_CASE("Pyramid ExportCache + ImportCache + Build acceleration smoke test", "ef_construction": 16, "no_build_levels": [0, 1, 2], "index_min_size": 28, - "persist_source_id": true + "persist_source_id": true, + "store_paths": true } } )"; @@ -2152,6 +2153,10 @@ TEST_CASE("Pyramid ExportCache + ImportCache + Build acceleration smoke test", auto warmed_build = warmed->Build(make_dataset()); REQUIRE(warmed_build.has_value()); REQUIRE(warmed->GetNumElements() == TEST_COUNT); + auto restored_data = warmed->GetDataByIdsWithFlag(ids.data(), 1, DATA_FLAG_ID | DATA_FLAG_PATH); + REQUIRE(restored_data.has_value()); + REQUIRE(restored_data.value()->GetPaths() != nullptr); + REQUIRE(restored_data.value()->GetPaths()[0] == "a/b/c"); auto warm_stats = vsag::JsonType::Parse(warmed->GetStats()); REQUIRE(warm_stats["build_cache_hit_nodes"].GetInt() > 0); std::vector query_vec(TEST_DIM); diff --git a/tests/test_pyramid_path_serialization.cpp b/tests/test_pyramid_path_serialization.cpp new file mode 100644 index 0000000000..dcf8b3a9b5 --- /dev/null +++ b/tests/test_pyramid_path_serialization.cpp @@ -0,0 +1,266 @@ +// 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 "storage/serialization_tags.h" +#include "storage/streaming_serialization_test_utils.h" +#include "vsag/vsag.h" + +namespace { + +constexpr int64_t PYRAMID_PATH_TEST_DIM = 2; + +struct PyramidPathRow { + int64_t id; + std::array vector; + std::string site_path; + std::string category_path; +}; + +std::string +MakePyramidPathParameters(bool store_paths) { + nlohmann::json index_param = { + {"base_quantization_type", "fp32"}, + {"max_degree", 4}, + {"ef_construction", 8}, + {"alpha", 1.2}, + {"graph_type", "nsw"}, + {"no_build_levels", {0, 1, 2, 3}}, + {"use_reorder", false}, + {"index_min_size", 100}, + {"hierarchies", {"site", "category"}}, + }; + if (store_paths) { + index_param["store_paths"] = true; + } + return nlohmann::json({{"dtype", "float32"}, + {"metric_type", "l2"}, + {"dim", PYRAMID_PATH_TEST_DIM}, + {"index_param", std::move(index_param)}}) + .dump(); +} + +vsag::IndexPtr +MakePyramidPathIndex(bool store_paths) { + auto result = vsag::Factory::CreateIndex("pyramid", MakePyramidPathParameters(store_paths)); + REQUIRE(result.has_value()); + return result.value(); +} + +vsag::DatasetPtr +MakePyramidPathDataset(const std::vector& rows) { + auto* vectors = new float[rows.size() * PYRAMID_PATH_TEST_DIM]; + auto* ids = new int64_t[rows.size()]; + auto* site_paths = new std::string[rows.size()]; + auto* category_paths = new std::string[rows.size()]; + for (uint64_t offset = 0; offset < rows.size(); ++offset) { + ids[offset] = rows[offset].id; + std::copy(rows[offset].vector.begin(), + rows[offset].vector.end(), + vectors + offset * PYRAMID_PATH_TEST_DIM); + site_paths[offset] = rows[offset].site_path; + category_paths[offset] = rows[offset].category_path; + } + return vsag::Dataset::Make() + ->NumElements(static_cast(rows.size())) + ->Dim(PYRAMID_PATH_TEST_DIM) + ->Ids(ids) + ->Float32Vectors(vectors) + ->Paths("site", site_paths) + ->Paths("category", category_paths) + ->Owner(true); +} + +void +RequirePyramidPaths(const vsag::IndexPtr& index, + const std::vector& ids, + const std::vector& expected_site_paths, + const std::vector& expected_category_paths) { + REQUIRE(ids.size() == expected_site_paths.size()); + REQUIRE(ids.size() == expected_category_paths.size()); + auto result = index->GetDataByIdsWithFlag( + ids.data(), static_cast(ids.size()), DATA_FLAG_ID | DATA_FLAG_PATH); + REQUIRE(result.has_value()); + const auto data = result.value(); + REQUIRE(data->GetPaths() == nullptr); + const auto* site_paths = data->GetPaths("site"); + const auto* category_paths = data->GetPaths("category"); + REQUIRE(site_paths != nullptr); + REQUIRE(category_paths != nullptr); + for (uint64_t offset = 0; offset < ids.size(); ++offset) { + REQUIRE(data->GetIds()[offset] == ids[offset]); + REQUIRE(site_paths[offset] == expected_site_paths[offset]); + REQUIRE(category_paths[offset] == expected_category_paths[offset]); + } +} + +bool +HasPyramidPathsBlock(const std::string& bytes) { + const auto blocks = vsag::test::ParseStreamingBlocks(bytes); + return std::any_of(blocks.begin(), blocks.end(), [](const auto& block) { + return block.tag == static_cast(vsag::StreamSerializationTag::PYRAMID_PATHS); + }); +} + +const std::vector INITIAL_ROWS = { + {101, {1.0F, 0.0F}, "us/news", "media/news"}, + {205, {0.0F, 1.0F}, "uk/sports", "media/sports"}, +}; + +} // namespace + +TEST_CASE("Pyramid BinarySet serialization preserves paths and supports Add", + "[ft][pyramid][paths][serialization]") { + auto index = MakePyramidPathIndex(true); + REQUIRE(index->Build(MakePyramidPathDataset(INITIAL_ROWS)).has_value()); + + auto binary_set = index->Serialize(); + REQUIRE(binary_set.has_value()); + auto restored = MakePyramidPathIndex(true); + REQUIRE(restored->Deserialize(binary_set.value()).has_value()); + RequirePyramidPaths( + restored, {205, 101}, {"uk/sports", "us/news"}, {"media/sports", "media/news"}); + + const std::vector added_rows = { + {309, {0.5F, 0.5F}, "jp/technology", "media/technology"}, + }; + auto add_result = restored->Add(MakePyramidPathDataset(added_rows)); + REQUIRE(add_result.has_value()); + REQUIRE(add_result.value().empty()); + RequirePyramidPaths(restored, + {309, 101, 205}, + {"jp/technology", "us/news", "uk/sports"}, + {"media/technology", "media/news", "media/sports"}); +} + +TEST_CASE("Pyramid streaming serialization preserves required paths block", + "[ft][pyramid][paths][serialization][streaming]") { + auto index = MakePyramidPathIndex(true); + REQUIRE(index->Build(MakePyramidPathDataset(INITIAL_ROWS)).has_value()); + + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + const auto bytes = stream.str(); + REQUIRE(HasPyramidPathsBlock(bytes)); + + auto restored = MakePyramidPathIndex(true); + std::stringstream deserialize_stream(bytes); + REQUIRE(restored->DeserializeStreaming(deserialize_stream).has_value()); + RequirePyramidPaths( + restored, {205, 101}, {"uk/sports", "us/news"}, {"media/sports", "media/news"}); + + std::stringstream load_stream(bytes); + auto loaded = vsag::Index::Load(load_stream, "{}"); + REQUIRE(loaded.has_value()); + RequirePyramidPaths( + loaded.value(), {101, 205}, {"us/news", "uk/sports"}, {"media/news", "media/sports"}); + + const auto missing_paths = + vsag::test::EraseStreamingBlock(bytes, vsag::StreamSerializationTag::PYRAMID_PATHS); + auto missing_paths_target = MakePyramidPathIndex(true); + std::stringstream missing_deserialize_stream(missing_paths); + REQUIRE_FALSE( + missing_paths_target->DeserializeStreaming(missing_deserialize_stream).has_value()); + + std::stringstream missing_load_stream(missing_paths); + REQUIRE_FALSE(vsag::Index::Load(missing_load_stream, "{}").has_value()); +} + +TEST_CASE("Pyramid default streaming serialization omits paths block", + "[ft][pyramid][paths][serialization][streaming]") { + auto index = MakePyramidPathIndex(false); + REQUIRE(index->Build(MakePyramidPathDataset(INITIAL_ROWS)).has_value()); + + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + REQUIRE_FALSE(HasPyramidPathsBlock(stream.str())); +} + +TEST_CASE("Pyramid default BinarySet roundtrip does not return paths", + "[ft][pyramid][paths][serialization]") { + auto index = MakePyramidPathIndex(false); + REQUIRE(index->Build(MakePyramidPathDataset(INITIAL_ROWS)).has_value()); + + auto binary_set = index->Serialize(); + REQUIRE(binary_set.has_value()); + auto restored = MakePyramidPathIndex(false); + REQUIRE(restored->Deserialize(binary_set.value()).has_value()); + + const std::array ids = {205, 101}; + auto data = restored->GetDataByIds(ids.data(), static_cast(ids.size())); + REQUIRE(data.has_value()); + REQUIRE(data.value()->GetPaths("site") == nullptr); + REQUIRE(data.value()->GetPaths("category") == nullptr); +} + +TEST_CASE("Pyramid BinarySet path storage configuration must match", + "[ft][pyramid][paths][serialization]") { + SECTION("stored paths require an enabled reader") { + auto index = MakePyramidPathIndex(true); + REQUIRE(index->Build(MakePyramidPathDataset(INITIAL_ROWS)).has_value()); + auto binary_set = index->Serialize(); + REQUIRE(binary_set.has_value()); + + auto restored = MakePyramidPathIndex(false); + REQUIRE_FALSE(restored->Deserialize(binary_set.value()).has_value()); + } + + SECTION("a reader requiring paths rejects an index without them") { + auto index = MakePyramidPathIndex(false); + REQUIRE(index->Build(MakePyramidPathDataset(INITIAL_ROWS)).has_value()); + auto binary_set = index->Serialize(); + REQUIRE(binary_set.has_value()); + + auto restored = MakePyramidPathIndex(true); + REQUIRE_FALSE(restored->Deserialize(binary_set.value()).has_value()); + } +} + +TEST_CASE("Pyramid path storage supports empty serialization roundtrips", + "[ft][pyramid][paths][serialization][streaming]") { + SECTION("Clone") { + auto index = MakePyramidPathIndex(true); + auto clone = index->Clone(); + REQUIRE(clone.has_value()); + REQUIRE(clone.value()->GetNumElements() == 0); + } + + SECTION("BinarySet") { + auto index = MakePyramidPathIndex(true); + auto binary_set = index->Serialize(); + REQUIRE(binary_set.has_value()); + + auto restored = MakePyramidPathIndex(true); + REQUIRE(restored->Deserialize(binary_set.value()).has_value()); + } + + SECTION("streaming") { + auto index = MakePyramidPathIndex(true); + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + + auto restored = MakePyramidPathIndex(true); + std::stringstream deserialize_stream(stream.str()); + REQUIRE(restored->DeserializeStreaming(deserialize_stream).has_value()); + } +} diff --git a/tests/test_pyramid_paths.cpp b/tests/test_pyramid_paths.cpp new file mode 100644 index 0000000000..4a67372eec --- /dev/null +++ b/tests/test_pyramid_paths.cpp @@ -0,0 +1,252 @@ +// 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 "vsag/vsag.h" + +namespace { + +using HierarchyPaths = std::vector>>; + +std::string +MakePyramidParameters(const std::string& graph_type, + bool include_store_paths, + bool named_hierarchies = false) { + nlohmann::json index_param = { + {"max_degree", 4}, + {"alpha", 1.2}, + {"ef_construction", 8}, + {"graph_type", graph_type}, + {"graph_iter_turn", 1}, + {"neighbor_sample_rate", 0.2}, + {"base_quantization_type", "fp32"}, + {"use_reorder", false}, + {"index_min_size", 100}, + {"no_build_levels", {0, 1, 2, 3}}, + }; + if (include_store_paths) { + index_param["store_paths"] = true; + } + if (named_hierarchies) { + index_param["hierarchies"] = {"site", "category"}; + } + return nlohmann::json({{"dtype", "float32"}, + {"metric_type", "l2"}, + {"dim", 2}, + {"index_param", std::move(index_param)}}) + .dump(); +} + +vsag::IndexPtr +MakePyramidIndex(const std::string& graph_type, + bool include_store_paths, + bool named_hierarchies = false) { + auto index = vsag::Factory::CreateIndex( + "pyramid", MakePyramidParameters(graph_type, include_store_paths, named_hierarchies)); + REQUIRE(index.has_value()); + return index.value(); +} + +vsag::DatasetPtr +MakeDataset(const std::vector& ids, const HierarchyPaths& hierarchy_paths) { + auto* vectors = new float[ids.size() * 2]; + auto* raw_ids = new int64_t[ids.size()]; + for (uint64_t offset = 0; offset < ids.size(); ++offset) { + vectors[offset * 2] = static_cast(offset); + vectors[offset * 2 + 1] = static_cast(offset + 1); + raw_ids[offset] = ids[offset]; + } + + auto dataset = vsag::Dataset::Make(); + dataset->NumElements(static_cast(ids.size())) + ->Dim(2) + ->Float32Vectors(vectors) + ->Ids(raw_ids) + ->Owner(true); + for (const auto& [hierarchy_name, paths] : hierarchy_paths) { + REQUIRE(paths.size() == ids.size()); + auto* raw_paths = new std::string[paths.size()]; + for (uint64_t offset = 0; offset < paths.size(); ++offset) { + raw_paths[offset] = paths[offset]; + } + dataset->Paths(hierarchy_name, raw_paths); + } + return dataset; +} + +vsag::DatasetPtr +GetData(const vsag::IndexPtr& index, const std::vector& ids) { + auto result = index->GetDataByIds(ids.data(), static_cast(ids.size())); + REQUIRE(result.has_value()); + return result.value(); +} + +vsag::DatasetPtr +GetDataWithFlag(const vsag::IndexPtr& index, + const std::vector& ids, + uint64_t selected_data_flag) { + auto result = index->GetDataByIdsWithFlag( + ids.data(), static_cast(ids.size()), selected_data_flag); + REQUIRE(result.has_value()); + return result.value(); +} + +void +RequirePaths(const std::string* actual, const std::vector& expected) { + REQUIRE(actual != nullptr); + for (uint64_t offset = 0; offset < expected.size(); ++offset) { + REQUIRE(actual[offset] == expected[offset]); + } +} + +} // namespace + +TEST_CASE("Pyramid does not retain paths by default", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", false); + // CheckFeature discovers the opt-in retrieval capability. The inherited legacy call remains + // callable for compatibility and omits paths while path storage is disabled. + REQUIRE_FALSE(index->CheckFeature(vsag::IndexFeature::SUPPORT_GET_DATA_BY_IDS)); + auto base = MakeDataset({10, 20}, {{"", {"a", "b"}}}); + + REQUIRE(index->Build(base).has_value()); + auto data = GetData(index, {20, 10}); + + REQUIRE(data->GetPaths() == nullptr); +} + +TEST_CASE("Pyramid retains unnamed paths for NSW Build", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", true); + auto base = MakeDataset({101, 42, 900}, {{"", {"alpha", "", "beta/gamma"}}}); + + REQUIRE(index->Build(base).has_value()); + auto data = GetDataWithFlag(index, {900, 101, 42}, DATA_FLAG_PATH); + + RequirePaths(data->GetPaths(), {"beta/gamma", "alpha", ""}); + REQUIRE(data->GetIds() == nullptr); +} + +TEST_CASE("Pyramid retains unnamed paths for ODescent Build", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("odescent", true); + REQUIRE(index->CheckFeature(vsag::IndexFeature::SUPPORT_GET_DATA_BY_IDS)); + auto base = MakeDataset({7, 3, 11}, {{"", {"one", "two", "three"}}}); + + REQUIRE(index->Build(base).has_value()); + auto data = GetDataWithFlag(index, {3, 11, 7}, DATA_FLAG_ID | DATA_FLAG_PATH); + + RequirePaths(data->GetPaths(), {"two", "three", "one"}); + REQUIRE(data->GetIds() != nullptr); + REQUIRE(data->GetIds()[0] == 3); + REQUIRE(data->GetIds()[1] == 11); + REQUIRE(data->GetIds()[2] == 7); +} + +TEST_CASE("Pyramid Clone retains paths", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", true); + REQUIRE(index->Build(MakeDataset({4, 8}, {{"", {"left", "right"}}})).has_value()); + + auto clone = index->Clone(); + REQUIRE(clone.has_value()); + RequirePaths(GetDataWithFlag(clone.value(), {8, 4}, DATA_FLAG_PATH)->GetPaths(), + {"right", "left"}); +} + +TEST_CASE("Pyramid Add stores paths for allocated inner IDs", "[ft][pyramid][paths]") { + const auto verify = [](const std::string& graph_type) { + auto index = MakePyramidIndex(graph_type, true); + REQUIRE(index->Build(MakeDataset({10}, {{"", {"original"}}})).has_value()); + + auto add_result = + index->Add(MakeDataset({20, 10, 30}, {{"", {"path-20", "duplicate-path", "path-30"}}})); + REQUIRE(add_result.has_value()); + REQUIRE(add_result.value() == std::vector{10}); + + auto data = GetDataWithFlag(index, {20, 30, 10}, DATA_FLAG_PATH); + RequirePaths(data->GetPaths(), {"path-20", "path-30", "original"}); + }; + + verify("nsw"); + verify("odescent"); +} + +TEST_CASE("Pyramid concurrent Add keeps paths aligned with IDs", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", true); + REQUIRE(index->Build(MakeDataset({10}, {{"", {"original"}}})).has_value()); + auto first = MakeDataset({20, 21}, {{"", {"first/20", "first/21"}}}); + auto second = MakeDataset({30, 31}, {{"", {"second/30", "second/31"}}}); + + auto first_result = + std::async(std::launch::async, [index, first]() { return index->Add(first); }); + auto second_result = + std::async(std::launch::async, [index, second]() { return index->Add(second); }); + REQUIRE(first_result.get().has_value()); + REQUIRE(second_result.get().has_value()); + + auto data = GetDataWithFlag(index, {31, 20, 30, 21}, DATA_FLAG_PATH); + RequirePaths(data->GetPaths(), {"second/31", "first/20", "second/30", "first/21"}); +} + +TEST_CASE("Pyramid returns only complete named hierarchy paths", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", true, true); + auto base = MakeDataset( + {1, 2}, {{"site", {"us/ca", "uk/london"}}, {"category", {"tech/ai", "science/physics"}}}); + REQUIRE(index->Build(base).has_value()); + REQUIRE(index->Add(MakeDataset({3}, {{"site", {"jp/tokyo"}}})).has_value()); + + auto complete = GetDataWithFlag(index, {2, 1}, DATA_FLAG_PATH); + REQUIRE(complete->GetPaths() == nullptr); + RequirePaths(complete->GetPaths("site"), {"uk/london", "us/ca"}); + RequirePaths(complete->GetPaths("category"), {"science/physics", "tech/ai"}); + + auto partial = GetDataWithFlag(index, {1, 3}, DATA_FLAG_PATH); + RequirePaths(partial->GetPaths("site"), {"us/ca", "jp/tokyo"}); + REQUIRE(partial->GetPaths("category") == nullptr); +} + +TEST_CASE("Pyramid returns paths only when selected", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", true); + REQUIRE(index->Build(MakeDataset({1, 2}, {{"", {"a", "b"}}})).has_value()); + + auto all_standard_data = GetData(index, {2, 1}); + REQUIRE(all_standard_data->GetIds() != nullptr); + REQUIRE(all_standard_data->GetPaths() == nullptr); + + auto ids_only = GetDataWithFlag(index, {2, 1}, DATA_FLAG_ID); + REQUIRE(ids_only->GetIds() != nullptr); + REQUIRE(ids_only->GetPaths() == nullptr); + + auto paths_only = GetDataWithFlag(index, {2, 1}, DATA_FLAG_PATH); + REQUIRE(paths_only->GetIds() == nullptr); + RequirePaths(paths_only->GetPaths(), {"b", "a"}); +} + +TEST_CASE("Pyramid rejects selected paths when storage is disabled", "[ft][pyramid][paths]") { + auto index = MakePyramidIndex("nsw", false); + REQUIRE(index->Build(MakeDataset({1}, {{"", {"a"}}})).has_value()); + + const int64_t id = 1; + auto ids_only = index->GetDataByIdsWithFlag(&id, 1, DATA_FLAG_ID); + REQUIRE(ids_only.has_value()); + REQUIRE(ids_only.value()->GetIds()[0] == id); + + auto result = index->GetDataByIdsWithFlag(&id, 1, DATA_FLAG_PATH); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::INVALID_ARGUMENT); +}