diff --git a/docs/docs/en/src/advanced/quantization_transform.md b/docs/docs/en/src/advanced/quantization_transform.md index ea474bd598..36c7af1778 100644 --- a/docs/docs/en/src/advanced/quantization_transform.md +++ b/docs/docs/en/src/advanced/quantization_transform.md @@ -27,16 +27,11 @@ Quantizer)**. ## Quick start -`tq` is currently exposed as a **public, externally configurable** quantization type only by -**HGraph**. HGraph maps the top-level keys `tq_chain` and `rabitq_pca_dim` into the nested -`base_codes.quantization_params` JSON via its external-parameter mapping -(`src/algorithm/hgraph.cpp:370-385`). IVF, BruteForce, Pyramid and WARP all internally render -a `tq_chain` field into their inner JSON template, but none of them expose `tq_chain` (or any -other TQ parameter) in their external mapping today. `CheckAndMappingExternalParam` rejects -unknown external keys with `invalid config param` -(`src/utils/util_functions.cpp:50-53`), so passing `tq_chain` in the `index_param` JSON of -those indexes will fail at index construction. Configuring TQ on non-HGraph indexes -therefore requires code-side changes to add the external mapping. +`tq` is exposed as a public, externally configurable quantization type by **HGraph** and +**Pyramid**. Both map `tq_chain` and dimension parameters into +`base_codes.quantization_params`. The MRLE + RaBitQ x+y split combination is available in both +indexes and automatically reorders from the base split datacell. IVF, BruteForce and WARP do +not currently expose `tq_chain` through their external parameter mapping. ```cpp std::string params = R"({ @@ -85,6 +80,7 @@ Examples: | `"pca, rom, sq8_uniform"` | PCA reduction, random rotation, then 8-bit uniform — the example chain. | | `"pca, rom, rabitq"` | PCA + rotation feeding the RaBitQ binary quantizer. | | `"mrle, fp32"` | MRLE projection then store as fp32 (MRLE must be first). | +| `"mrle, rabitq"` | MRLE reduction followed by RaBitQ; with x+y split storage, filter and supplement codes are produced by the terminal RaBitQ. | Constraints (`transform_quantizer_parameter.cpp:33-45`): @@ -98,6 +94,9 @@ Constraints (`transform_quantizer_parameter.cpp:33-45`): when `is_transform_quantizer` is true (`src/datacell/flatten_interface.cpp:166`), so using any of those three as the terminal quantizer fails at index construction with an "unsupported quantization type" error. +- RaBitQ x+y split storage supports TQ only for the exact chain `"mrle, rabitq"`. The split + datacell reuses the terminal RaBitQ encoder and keeps RaBitQ internal FHT/ROM rotation + behavior unchanged. Other transformed split chains are rejected. - Any unrecognized transformer name raises `INVALID_ARGUMENT: invalid transformer name` (`transform_quantizer.h:225-227`). diff --git a/docs/docs/en/src/indexes/hgraph.md b/docs/docs/en/src/indexes/hgraph.md index 4ac81a8442..712d62bb48 100644 --- a/docs/docs/en/src/indexes/hgraph.md +++ b/docs/docs/en/src/indexes/hgraph.md @@ -71,7 +71,7 @@ most users need; the exhaustive list is in [Index Parameters](../resources/index | `use_reverse_edges` | bool | `false` | Track incoming neighbors for O(1) reverse-edge lookup. Roughly doubles edge storage and is unsupported with `graph_storage_type: "compressed"`. | | `label_remap_type` | string | `"pg"` | Label-to-inner-ID map implementation: `"pg"` or `"robin"`. Keep the same value when restoring or combining compatible indexes. | | `use_reorder` | bool | `false` | Keep a high-precision copy and re-rank after the coarse search | -| `reorder_source` | string | `"precise"` | Reorder from `"precise"` codes or directly from `"base"` codes. RaBitQ x+y split sets `"base"` automatically. | +| `reorder_source` | string | `"precise"` | Reorder from `"precise"` codes or directly from `"base"` codes. RaBitQ x+y split, including `tq_chain: "mrle, rabitq"`, sets `"base"` automatically. | | `precise_quantization_type` | string | `"fp32"` | Quantizer used for reordering (takes effect only with `use_reorder: true`) | | `base_pq_dim` | int | `1` | Number of PQ subspaces. When using `pq` / `pqfs`, set this explicitly instead of relying on the default. | | `mrle_dim` | int | `0` | Output dimension for an MRLE transform in `tq_chain`; allowed range `[0, dim]`, where `0` means the input dimension. | diff --git a/docs/docs/en/src/resources/index_parameters.md b/docs/docs/en/src/resources/index_parameters.md index c67ea9f140..b67d89803e 100644 --- a/docs/docs/en/src/resources/index_parameters.md +++ b/docs/docs/en/src/resources/index_parameters.md @@ -50,7 +50,7 @@ HGraph places its build parameters under the generic `index_param` key (see | `base_quantization_type` | `fp32` / `fp16` / `bf16` / `sq8` / `sq4` / `pq` | Quantization of the base storage — see the [Quantization chapter](../quantization/README.md) for all supported values | | `use_reverse_edges` | `false` | Track incoming neighbors for O(1) reverse-edge lookup; roughly doubles edge storage and is unsupported with compressed graph storage | | `label_remap_type` | `pg` | Label-map implementation: `pg` (default) or `robin` | -| `reorder_source` | `precise` | Reorder from the `precise` store or directly from `base`; RaBitQ x+y split selects `base` automatically | +| `reorder_source` | `precise` | Reorder from the `precise` store or directly from `base`; RaBitQ x+y split, including `tq_chain="mrle, rabitq"`, selects `base` automatically | | `persist_source_id` | `false` | Include HGraph source-ID metadata in serialization; useful when a restored index must later export a build cache | | `mrle_dim` | `0` | MRLE output dimension in `[0, dim]`; `0` means input dimension | | `fast_encode_rabitq` | `true` | Use fast multi-bit RaBitQ encoding; `false` restores the exact encoder | diff --git a/docs/docs/zh/src/advanced/quantization_transform.md b/docs/docs/zh/src/advanced/quantization_transform.md index 34cb3d4264..b6672f0c1b 100644 --- a/docs/docs/zh/src/advanced/quantization_transform.md +++ b/docs/docs/zh/src/advanced/quantization_transform.md @@ -20,13 +20,10 @@ ## 快速上手 -`tq` 目前作为**对外可配置**的量化类型，只有 **HGraph** 真正暴露了它。HGraph 通过外部参数映射把 -顶层键 `tq_chain` 和 `rabitq_pca_dim` 写到嵌套的 `base_codes.quantization_params` -（`src/algorithm/hgraph.cpp:370-385`）。IVF、BruteForce、Pyramid、WARP 虽然在内部 JSON 模板中 -也会渲染 `tq_chain` 字段，但它们的外部参数映射里**都没有** `tq_chain`（或其它 TQ 参数）。 -`CheckAndMappingExternalParam` 遇到未映射的外部键会直接抛 `invalid config param` -（`src/utils/util_functions.cpp:50-53`），因此在这些索引的 `index_param` JSON 中传 `tq_chain` -会在构建时报错。在非 HGraph 索引上启用 TQ 目前需要在代码侧补一条外部映射。 +`tq` 已作为对外可配置的量化类型暴露给 **HGraph** 和 **Pyramid**。两者都会把 +`tq_chain` 与降维参数映射到 `base_codes.quantization_params`。MRLE + RaBitQ x+y split +组合在两个索引中均可用，并自动从 base split datacell 重排。IVF、BruteForce 与 WARP +目前仍未通过外部参数映射暴露 `tq_chain`。 ```cpp std::string params = R"({ @@ -74,6 +71,7 @@ token 两侧的空白会被自动 trim | `"pca, rom, sq8_uniform"` | 先 PCA 降维，再随机旋转，再 8 位均匀量化 —— 即示例 501。 | | `"pca, rom, rabitq"` | PCA + 旋转后喂给 RaBitQ 二值量化器。 | | `"mrle, fp32"` | MRLE 投影再以 fp32 存储（MRLE 必须放在最前）。 | +| `"mrle, rabitq"` | 先做 MRLE 降维，再由 RaBitQ 编码；使用 x+y split 存储时，filter 与 supplement 编码仍由末端 RaBitQ 生成。 | 约束（`transform_quantizer_parameter.cpp:33-45`）： @@ -86,6 +84,8 @@ token 两侧的空白会被自动 trim `is_transform_quantizer=true` 时显式拒绝 `sparse` （`src/datacell/flatten_interface.cpp:166`），因此这三个不能用作 TQ 末端，否则会在构建索引时 以 "unsupported quantization type" 失败。 +- RaBitQ x+y split 存储仅支持精确链 `"mrle, rabitq"`。split datacell 复用末端 + RaBitQ 编码器，并保留 RaBitQ 内部的 FHT/ROM 随机旋转；其他带变换的 split 链会被拒绝。 - 未识别的变换名会抛 `INVALID_ARGUMENT: invalid transformer name` （`transform_quantizer.h:225-227`）。 diff --git a/docs/docs/zh/src/indexes/hgraph.md b/docs/docs/zh/src/indexes/hgraph.md index ad27f8c6d4..8939a24f0b 100644 --- a/docs/docs/zh/src/indexes/hgraph.md +++ b/docs/docs/zh/src/indexes/hgraph.md @@ -65,7 +65,7 @@ auto result = index->KnnSearch( | `use_reverse_edges` | bool | `false` | 跟踪入边，实现 O(1) 反向邻居查找；边存储约翻倍，且 `graph_storage_type: "compressed"` 不支持 | | `label_remap_type` | string | `"pg"` | label 到内部 ID 的 map 实现：`"pg"` 或 `"robin"`；恢复或组合兼容索引时应保持一致 | | `use_reorder` | bool | `false` | 是否额外保留一份高精度副本用于精排 | -| `reorder_source` | string | `"precise"` | 从 `"precise"` 编码或直接从 `"base"` 编码重排；RaBitQ x+y split 会自动设置为 `"base"` | +| `reorder_source` | string | `"precise"` | 从 `"precise"` 编码或直接从 `"base"` 编码重排；RaBitQ x+y split（包括 `tq_chain: "mrle, rabitq"`）会自动设置为 `"base"` | | `precise_quantization_type` | string | `"fp32"` | 精排使用的量化类型（仅在 `use_reorder: true` 时生效） | | `base_pq_dim` | int | `1` | PQ 子空间数（`pq` / `pqfs` 时必填） | | `mrle_dim` | int | `0` | `tq_chain` 中 MRLE 的输出维度，范围 `[0, dim]`；`0` 表示输入维度 | diff --git a/docs/docs/zh/src/resources/index_parameters.md b/docs/docs/zh/src/resources/index_parameters.md index c5580deac9..6945113379 100644 --- a/docs/docs/zh/src/resources/index_parameters.md +++ b/docs/docs/zh/src/resources/index_parameters.md @@ -47,7 +47,7 @@ HGraph 的构建参数使用通用的 `index_param` 键（参见 `examples/cpp/1 | `base_quantization_type` | `fp32` / `fp16` / `bf16` / `sq8` / `sq4` / `pq` | 主存储的量化策略 —— 支持的全部取值见[量化章节](../quantization/README.md) | | `use_reverse_edges` | `false` | 跟踪入边，实现 O(1) 反向邻居查找；边存储约翻倍，且压缩图存储不支持 | | `label_remap_type` | `pg` | label map 实现：默认 `pg`，或 `robin` | -| `reorder_source` | `precise` | 从 `precise` 存储或直接从 `base` 重排；RaBitQ x+y split 会自动选择 `base` | +| `reorder_source` | `precise` | 从 `precise` 存储或直接从 `base` 重排；RaBitQ x+y split（包括 `tq_chain="mrle, rabitq"`）会自动选择 `base` | | `persist_source_id` | `false` | 序列化 HGraph 时保留 Source ID 元数据；适用于恢复索引后继续导出构建缓存 | | `mrle_dim` | `0` | MRLE 输出维度，范围 `[0, dim]`；`0` 表示输入维度 | | `fast_encode_rabitq` | `true` | 使用多 bit RaBitQ 快速编码；设为 `false` 恢复精确编码器 | diff --git a/src/algorithm/hgraph/hgraph_param_mapping.cpp b/src/algorithm/hgraph/hgraph_param_mapping.cpp index efc7d47583..7c2559527b 100644 --- a/src/algorithm/hgraph/hgraph_param_mapping.cpp +++ b/src/algorithm/hgraph/hgraph_param_mapping.cpp @@ -17,78 +17,9 @@ #include "common.h" #include "hgraph.h" // IWYU pragma: keep #include "hgraph_parameter.h" -#include "quantization/rabitq_quantization/rabitq_quantizer_parameter.h" namespace vsag { -namespace { - -void -map_rabitq_split_param(const JsonType& external_json, JsonType& inner_json) { - if (not external_json.Contains(RABITQ_BITS_PER_DIM_PRECISE)) { - return; - } - - CHECK_ARGUMENT( - external_json.Contains(RABITQ_BITS_PER_DIM_BASE), - fmt::format("{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, RABITQ_BITS_PER_DIM_BASE)); - CHECK_ARGUMENT( - external_json.Contains(HGRAPH_BASE_QUANTIZATION_TYPE), - fmt::format("{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, HGRAPH_BASE_QUANTIZATION_TYPE)); - CHECK_ARGUMENT( - external_json.Contains(HGRAPH_PRECISE_QUANTIZATION_TYPE), - fmt::format( - "{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, HGRAPH_PRECISE_QUANTIZATION_TYPE)); - - const auto base_quantization_type = external_json[HGRAPH_BASE_QUANTIZATION_TYPE].GetString(); - const auto precise_quantization_type = - external_json[HGRAPH_PRECISE_QUANTIZATION_TYPE].GetString(); - CHECK_ARGUMENT(base_quantization_type == QUANTIZATION_TYPE_VALUE_RABITQ, - fmt::format("{} requires {}={}", - RABITQ_BITS_PER_DIM_PRECISE, - HGRAPH_BASE_QUANTIZATION_TYPE, - QUANTIZATION_TYPE_VALUE_RABITQ)); - CHECK_ARGUMENT(precise_quantization_type == QUANTIZATION_TYPE_VALUE_RABITQ, - fmt::format("{} requires {}={}", - RABITQ_BITS_PER_DIM_PRECISE, - HGRAPH_PRECISE_QUANTIZATION_TYPE, - QUANTIZATION_TYPE_VALUE_RABITQ)); - - const int64_t filter_bits = external_json[RABITQ_BITS_PER_DIM_BASE].GetInt(); - const int64_t supplement_bits = external_json[RABITQ_BITS_PER_DIM_PRECISE].GetInt(); - CHECK_ARGUMENT( - filter_bits >= 1, - fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_BASE, filter_bits)); - CHECK_ARGUMENT( - filter_bits <= 8, - fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_BASE, filter_bits)); - CHECK_ARGUMENT( - supplement_bits >= 1, - fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_PRECISE, supplement_bits)); - CHECK_ARGUMENT( - supplement_bits <= 8, - fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_PRECISE, supplement_bits)); - const int64_t total_bits = filter_bits + supplement_bits; - CHECK_ARGUMENT(total_bits <= 8, - fmt::format("{} + {} must be no greater than 8, got {}", - RABITQ_BITS_PER_DIM_BASE, - RABITQ_BITS_PER_DIM_PRECISE, - total_bits)); - - inner_json[REORDER_SOURCE_KEY].SetString(HGRAPH_REORDER_SOURCE_BASE); - inner_json[BASE_CODES_KEY][CODES_TYPE_KEY].SetString(RABITQ_SPLIT_CODES); - inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_VERSION_KEY].SetString( - RaBitQuantizerParameter::RABITQ_VERSION_SPLIT); - inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_QUERY_KEY] - .SetInt(32); - inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_FILTER_KEY] - .SetInt(filter_bits); - inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_BASE_KEY] - .SetInt(total_bits); -} - -} // namespace - JsonType HGraph::map_hgraph_param(const JsonType& hgraph_json) { static const ConstParamMap external_mapping = { @@ -613,7 +544,7 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { std::string str = format_map(hgraph_params_template, DEFAULT_MAP); auto inner_json = JsonType::Parse(str); mapping_external_param_to_inner(hgraph_json, external_mapping, inner_json); - map_rabitq_split_param(hgraph_json, inner_json); + MapRaBitQSplitParam(hgraph_json, inner_json); return inner_json; } @@ -628,16 +559,7 @@ HGraph::CheckAndMappingExternalParam(const JsonType& external_param, inner_json[RAW_VECTOR_KEY][CODES_TYPE_KEY].SetString(SPARSE_CODES); } - if (external_param.Contains(INDEX_MRLE_DIM)) { - CHECK_ARGUMENT(external_param[INDEX_MRLE_DIM].IsNumberInteger(), - fmt::format("mrle_dim must be an integer, got {}", - external_param[INDEX_MRLE_DIM].Dump())); - int64_t mrle_dim = external_param[INDEX_MRLE_DIM].GetInt(); - bool valid_mrle_dim = mrle_dim >= 0 and mrle_dim <= static_cast(common_param.dim_); - CHECK_ARGUMENT( - valid_mrle_dim, - fmt::format("mrle_dim({}) must be in range [0, {}]", mrle_dim, common_param.dim_)); - } + ValidateMRLEDim(external_param, common_param.dim_); auto hgraph_parameter = std::make_shared(); hgraph_parameter->data_type = common_param.data_type_; diff --git a/src/algorithm/hgraph/hgraph_parameter_test.cpp b/src/algorithm/hgraph/hgraph_parameter_test.cpp index 3ae16b13ca..3d10ccd89a 100644 --- a/src/algorithm/hgraph/hgraph_parameter_test.cpp +++ b/src/algorithm/hgraph/hgraph_parameter_test.cpp @@ -636,3 +636,35 @@ TEST_CASE("HGraph maps fast RaBitQ to base and precise quantizers", "[ut][HGraph REQUIRE_FALSE(precise_json["quantization_params"]["fast_encode_rabitq"].GetBool()); REQUIRE(precise_json["quantization_params"]["fast_encode_rabitq_rounds"].GetInt() == 9); } + +TEST_CASE("HGraph maps MRLE RaBitQ split to base reorder", "[ut][HGraphParameter][MRLE]") { + auto param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "tq", + "tq_chain": "mrle, rabitq", + "mrle_dim": 64, + "precise_quantization_type": "rabitq", + "rabitq_bits_per_dim_base": 3, + "rabitq_bits_per_dim_precise": 5, + "use_reorder": true + })"); + + vsag::IndexCommonParam common_param; + common_param.dim_ = 128; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + auto mapped = vsag::HGraph::CheckAndMappingExternalParam(param, common_param); + auto typed_param = std::dynamic_pointer_cast(mapped); + + REQUIRE(typed_param != nullptr); + REQUIRE(typed_param->reorder_source == std::string("base")); + REQUIRE(typed_param->precise_codes_param == nullptr); + const auto base_json = typed_param->base_codes_param->ToJson(); + REQUIRE(base_json["codes_type"].GetString() == std::string("rabitq_split")); + REQUIRE(base_json["quantization_params"]["type"].GetString() == std::string("tq")); + REQUIRE(base_json["quantization_params"]["tq_chain"].GetString() == std::string("mrle,rabitq")); + REQUIRE(base_json["quantization_params"]["mrle_dim"].GetInt() == 64); + REQUIRE(base_json["quantization_params"]["rabitq_bits_per_dim_filter"].GetInt() == 3); + REQUIRE(base_json["quantization_params"]["rabitq_bits_per_dim_base"].GetInt() == 8); + + param["tq_chain"].SetString("pca, rabitq"); + REQUIRE_THROWS(vsag::HGraph::CheckAndMappingExternalParam(param, common_param)); +} diff --git a/src/algorithm/inner_index_parameter.cpp b/src/algorithm/inner_index_parameter.cpp index c5f94730ee..a8568fd152 100644 --- a/src/algorithm/inner_index_parameter.cpp +++ b/src/algorithm/inner_index_parameter.cpp @@ -21,6 +21,8 @@ #include "datacell/flatten_datacell_parameter.h" #include "impl/logger/logger.h" #include "inner_string_params.h" +#include "quantization/rabitq_quantization/rabitq_quantizer_parameter.h" +#include "quantization/transform_quantization/transform_quantizer_parameter.h" #include "utils/param_compat_macros.h" #include "vsag/constants.h" @@ -48,6 +50,95 @@ dump_label_remap_type(LabelRemapType remap_type) -> const char* { } // namespace +void +MapRaBitQSplitParam(const JsonType& external_json, JsonType& inner_json) { + if (not external_json.Contains(RABITQ_BITS_PER_DIM_PRECISE)) { + return; + } + + CHECK_ARGUMENT( + external_json.Contains(RABITQ_BITS_PER_DIM_BASE), + fmt::format("{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, RABITQ_BITS_PER_DIM_BASE)); + CHECK_ARGUMENT( + external_json.Contains(HGRAPH_BASE_QUANTIZATION_TYPE), + fmt::format("{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, HGRAPH_BASE_QUANTIZATION_TYPE)); + CHECK_ARGUMENT( + external_json.Contains(HGRAPH_PRECISE_QUANTIZATION_TYPE), + fmt::format( + "{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, HGRAPH_PRECISE_QUANTIZATION_TYPE)); + + const auto base_type = external_json[HGRAPH_BASE_QUANTIZATION_TYPE].GetString(); + const auto precise_type = external_json[HGRAPH_PRECISE_QUANTIZATION_TYPE].GetString(); + const bool is_direct = base_type == QUANTIZATION_TYPE_VALUE_RABITQ; + const bool is_mrle = base_type == QUANTIZATION_TYPE_VALUE_TQ; + CHECK_ARGUMENT(is_direct or is_mrle, + fmt::format("{} requires {}={} or {}", + RABITQ_BITS_PER_DIM_PRECISE, + HGRAPH_BASE_QUANTIZATION_TYPE, + QUANTIZATION_TYPE_VALUE_RABITQ, + QUANTIZATION_TYPE_VALUE_TQ)); + CHECK_ARGUMENT(precise_type == QUANTIZATION_TYPE_VALUE_RABITQ, + fmt::format("{} requires {}={}", + RABITQ_BITS_PER_DIM_PRECISE, + HGRAPH_PRECISE_QUANTIZATION_TYPE, + QUANTIZATION_TYPE_VALUE_RABITQ)); + if (is_mrle) { + CHECK_ARGUMENT(external_json.Contains(INDEX_TQ_CHAIN), + fmt::format("{}={} requires {}", + HGRAPH_BASE_QUANTIZATION_TYPE, + QUANTIZATION_TYPE_VALUE_TQ, + INDEX_TQ_CHAIN)); + const auto chain = + TransformQuantizerParameter::SplitString(external_json[INDEX_TQ_CHAIN].GetString()); + const std::vector expected_chain{TRANSFORMER_TYPE_VALUE_MRLE, + QUANTIZATION_TYPE_VALUE_RABITQ}; + CHECK_ARGUMENT(chain == expected_chain, + "rabitq split transform quantizer requires tq_chain=\"mrle, rabitq\""); + } + + const int64_t filter_bits = external_json[RABITQ_BITS_PER_DIM_BASE].GetInt(); + const int64_t supplement_bits = external_json[RABITQ_BITS_PER_DIM_PRECISE].GetInt(); + const bool valid_filter_bits = filter_bits >= 1 and filter_bits <= 8; + CHECK_ARGUMENT( + valid_filter_bits, + fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_BASE, filter_bits)); + const bool valid_supplement_bits = supplement_bits >= 1 and supplement_bits <= 8; + CHECK_ARGUMENT( + valid_supplement_bits, + fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_PRECISE, supplement_bits)); + const int64_t total_bits = filter_bits + supplement_bits; + CHECK_ARGUMENT(total_bits <= 8, + fmt::format("{} + {} must be no greater than 8, got {}", + RABITQ_BITS_PER_DIM_BASE, + RABITQ_BITS_PER_DIM_PRECISE, + total_bits)); + + inner_json[REORDER_SOURCE_KEY].SetString(HGRAPH_REORDER_SOURCE_BASE); + inner_json[BASE_CODES_KEY][CODES_TYPE_KEY].SetString(RABITQ_SPLIT_CODES); + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_VERSION_KEY].SetString( + RaBitQuantizerParameter::RABITQ_VERSION_SPLIT); + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_QUERY_KEY] + .SetInt(32); + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_FILTER_KEY] + .SetInt(filter_bits); + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_BASE_KEY] + .SetInt(total_bits); +} + +void +ValidateMRLEDim(const JsonType& external_json, uint64_t dim) { + if (not external_json.Contains(INDEX_MRLE_DIM)) { + return; + } + CHECK_ARGUMENT( + external_json[INDEX_MRLE_DIM].IsNumberInteger(), + fmt::format("mrle_dim must be an integer, got {}", external_json[INDEX_MRLE_DIM].Dump())); + const int64_t mrle_dim = external_json[INDEX_MRLE_DIM].GetInt(); + const bool valid_mrle_dim = mrle_dim >= 0 and mrle_dim <= static_cast(dim); + CHECK_ARGUMENT(valid_mrle_dim, + fmt::format("mrle_dim({}) must be in range [0, {}]", mrle_dim, dim)); +} + void InnerIndexParameter::FromJson(const JsonType& json) { if (json.Contains(USE_REORDER_KEY)) { diff --git a/src/algorithm/inner_index_parameter.h b/src/algorithm/inner_index_parameter.h index b01895d8a2..32dd197b79 100644 --- a/src/algorithm/inner_index_parameter.h +++ b/src/algorithm/inner_index_parameter.h @@ -26,6 +26,12 @@ DEFINE_POINTER2(ExtraInfoDataCellParam, ExtraInfoDataCellParameter); DEFINE_POINTER2(FlattenInterfaceParam, FlattenInterfaceParameter); DEFINE_POINTER2(AttributeInvertedInterfaceParam, AttributeInvertedInterfaceParameter); +void +MapRaBitQSplitParam(const JsonType& external_json, JsonType& inner_json); + +void +ValidateMRLEDim(const JsonType& external_json, uint64_t dim); + class InnerIndexParameter : public Parameter { public: explicit InnerIndexParameter() = default; diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index eed4bc6d69..add150f930 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -225,10 +225,13 @@ Pyramid::build_by_odescent(const DatasetPtr& base) { std::memcpy(label_table_->label_table_.data(), data_ids, sizeof(LabelType) * data_num); base_codes_->BatchInsertVector(data_vectors, data_num); - if (use_reorder_) { + if (has_precise_codes()) { precise_codes_->BatchInsertVector(data_vectors, data_num); } - auto codes = use_reorder_ ? precise_codes_ : base_codes_; + if (raw_vector_ != nullptr) { + raw_vector_->BatchInsertVector(data_vectors, data_num); + } + auto codes = graph_codes(); if (thread_pool_ != nullptr && hierarchies_.size() > 1) { Vector> futures(allocator_); @@ -435,9 +438,12 @@ void Pyramid::Serialize(StreamWriter& writer) const { label_table_->Serialize(writer); base_codes_->Serialize(writer); - if (use_reorder_) { + if (has_precise_codes()) { precise_codes_->Serialize(writer); } + if (raw_vector_ != nullptr) { + raw_vector_->Serialize(writer); + } auto pyramid_param = std::dynamic_pointer_cast(create_param_ptr_); if (pyramid_param && pyramid_param->has_hierarchies) { @@ -485,13 +491,20 @@ Pyramid::collect_streaming_header() const { base_tag, StreamSerializationBlockCurrentVersion(base_tag), StreamSerializationTagCritical(base_tag)); - if (this->use_reorder_) { + if (this->has_precise_codes()) { auto tag = static_cast(StreamSerializationTag::HIGH_PRECISION_CODES); AppendStreamingManifestBlock(manifest, tag, StreamSerializationBlockCurrentVersion(tag), StreamSerializationTagCritical(tag)); } + if (this->raw_vector_ != nullptr) { + auto tag = static_cast(StreamSerializationTag::RAW_VECTOR); + AppendStreamingManifestBlock(manifest, + tag, + StreamSerializationBlockCurrentVersion(tag), + StreamSerializationTagCritical(tag)); + } AppendStreamingManifestBlock(manifest, hierarchy_tag, StreamSerializationBlockCurrentVersion(hierarchy_tag), @@ -530,13 +543,20 @@ Pyramid::serialize_streaming_body(StreamWriter& writer) const { writer, base_tag, StreamSerializationTagCritical(base_tag), [this](StreamWriter& w) { this->base_codes_->Serialize(w); }); - if (this->use_reorder_) { + if (this->has_precise_codes()) { auto tag = static_cast(StreamSerializationTag::HIGH_PRECISION_CODES); WriteStreamingBlock( writer, tag, StreamSerializationTagCritical(tag), [this](StreamWriter& w) { this->precise_codes_->Serialize(w); }); } + if (this->raw_vector_ != nullptr) { + 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 +621,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) { @@ -644,7 +665,7 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) loaded_base_codes = true; break; case StreamSerializationTag::HIGH_PRECISION_CODES: - if (this->use_reorder_) { + if (this->has_precise_codes()) { ReadSeekableBlockPayload( block_reader, block_header, [this](StreamReader& block) { this->precise_codes_->Deserialize(block); @@ -652,6 +673,15 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) loaded_precise_codes = true; } break; + case StreamSerializationTag::RAW_VECTOR: + if (this->raw_vector_ != nullptr) { + 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) { @@ -681,10 +711,14 @@ Pyramid::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) throw VsagException(ErrorType::READ_ERROR, "Pyramid streaming serialization required block is missing"); } - if (this->use_reorder_ && !loaded_precise_codes) { + if (this->has_precise_codes() && !loaded_precise_codes) { throw VsagException(ErrorType::READ_ERROR, "Pyramid streaming serialization precise codes block is missing"); } + if (this->raw_vector_ != nullptr && !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()); @@ -706,9 +740,12 @@ Pyramid::Deserialize(StreamReader& reader) { delete_count_.store(static_cast(label_table_->GetAllDeletedIds().size()), std::memory_order_relaxed); base_codes_->Deserialize(buffer_reader); - if (use_reorder_) { + if (has_precise_codes()) { precise_codes_->Deserialize(buffer_reader); } + if (raw_vector_ != nullptr) { + raw_vector_->Deserialize(buffer_reader); + } cur_element_count_ = base_codes_->TotalCount(); auto param_json = JsonType::Parse(basic_info[INDEX_PARAM].GetString()); @@ -746,13 +783,20 @@ Pyramid::ExportModel(const IndexCommonParam& param) const { "Export model's pyramid reorder config mismatched"); } this->base_codes_->ExportModel(index->base_codes_); - if (use_reorder_) { + if (has_precise_codes()) { if (index->precise_codes_ == nullptr) { throw VsagException(ErrorType::INTERNAL_ERROR, "Export model's pyramid precise codes is empty"); } this->precise_codes_->ExportModel(index->precise_codes_); } + if (raw_vector_ != nullptr) { + if (index->raw_vector_ == nullptr) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "Export model's pyramid raw vector is empty"); + } + this->raw_vector_->ExportModel(index->raw_vector_); + } index->current_memory_usage_ = index->CalSerializeSize(); return index; } @@ -783,10 +827,14 @@ Pyramid::Add(const DatasetPtr& base) { 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); - if (use_reorder_) { + if (has_precise_codes()) { precise_codes_->InsertVector(data_vectors + dim_ * i, valid_id_count + local_cur_element_count); } + if (raw_vector_ != nullptr) { + raw_vector_->InsertVector(data_vectors + dim_ * i, + valid_id_count + local_cur_element_count); + } valid_id_count++; data_biases.push_back(i); } else { @@ -816,9 +864,12 @@ Pyramid::resize(int64_t new_max_capacity) { pool_ = std::make_unique(1, allocator_, new_max_capacity, allocator_); label_table_->Resize(new_max_capacity); base_codes_->Resize(new_max_capacity); - if (use_reorder_) { + if (has_precise_codes()) { precise_codes_->Resize(new_max_capacity); } + if (raw_vector_ != nullptr) { + raw_vector_->Resize(new_max_capacity); + } points_mutex_->Resize(new_max_capacity); max_capacity_ = new_max_capacity; } @@ -876,6 +927,7 @@ static const std::string HGRAPH_PARAMS_TEMPLATE = { "{TYPE_KEY}": "{INDEX_TYPE_PYRAMID}", "{USE_REORDER_KEY}": false, + "{REORDER_SOURCE_KEY}": "{HGRAPH_REORDER_SOURCE_PRECISE}", "{GRAPH_KEY}": { "{IO_PARAMS_KEY}": { "{TYPE_KEY}": "{IO_TYPE_VALUE_BLOCK_MEMORY_IO}", @@ -904,7 +956,11 @@ static const std::string HGRAPH_PARAMS_TEMPLATE = "{TYPE_KEY}": "{QUANTIZATION_TYPE_VALUE_FP32}", "{SQ4_UNIFORM_QUANTIZATION_TRUNC_RATE_KEY}": 0.05, "{PCA_DIM_KEY}": 0, + "{MRLE_DIM_KEY}": 0, + "{RABITQ_QUANTIZATION_VERSION_KEY}": "standard", "{RABITQ_QUANTIZATION_BITS_PER_DIM_QUERY_KEY}": 32, + "{RABITQ_QUANTIZATION_BITS_PER_DIM_BASE_KEY}": 1, + "{RABITQ_QUANTIZATION_BITS_PER_DIM_FILTER_KEY}": 1, "{FAST_ENCODE_RABITQ_KEY}": true, "{FAST_ENCODE_RABITQ_ROUNDS_KEY}": 6, "{TQ_CHAIN_KEY}": "", @@ -929,6 +985,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}":[], @@ -942,9 +1010,14 @@ Pyramid::CheckAndMappingExternalParam(const JsonType& external_param, const ConstParamMap external_mapping = { {PYRAMID_EF_CONSTRUCTION, {EF_CONSTRUCTION_KEY}}, {PYRAMID_USE_REORDER, {USE_REORDER_KEY}}, + {HGRAPH_REORDER_SOURCE, {REORDER_SOURCE_KEY}}, {PYRAMID_BASE_QUANTIZATION_TYPE, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, TYPE_KEY}}, + {INDEX_TQ_CHAIN, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, TQ_CHAIN_KEY}}, + {INDEX_MRLE_DIM, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, MRLE_DIM_KEY}}, {PYRAMID_RABITQ_BITS_PER_DIM_BASE, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, RABITQ_QUANTIZATION_BITS_PER_DIM_BASE_KEY}}, + {RABITQ_BITS_PER_DIM_PRECISE, + {PRECISE_CODES_KEY, QUANTIZATION_PARAMS_KEY, RABITQ_QUANTIZATION_BITS_PER_DIM_BASE_KEY}}, {PYRAMID_RABITQ_BITS_PER_DIM_QUERY, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, RABITQ_QUANTIZATION_BITS_PER_DIM_QUERY_KEY}}, {PYRAMID_RABITQ_PCA_DIM, {BASE_CODES_KEY, QUANTIZATION_PARAMS_KEY, PCA_DIM_KEY}}, @@ -960,6 +1033,9 @@ Pyramid::CheckAndMappingExternalParam(const JsonType& external_param, {PYRAMID_PRECISE_QUANTIZATION_TYPE, {PRECISE_CODES_KEY, QUANTIZATION_PARAMS_KEY, TYPE_KEY}}, {PYRAMID_GRAPH_MAX_DEGREE, {GRAPH_KEY, GRAPH_PARAM_MAX_DEGREE_KEY}}, {PYRAMID_BASE_IO_TYPE, {BASE_CODES_KEY, IO_PARAMS_KEY, TYPE_KEY}}, + {HGRAPH_BASE_SUPPLEMENT_IO_TYPE, {BASE_CODES_KEY, SUPPLEMENT_IO_PARAMS_KEY, TYPE_KEY}}, + {HGRAPH_BASE_SUPPLEMENT_FILE_PATH, + {BASE_CODES_KEY, SUPPLEMENT_IO_PARAMS_KEY, IO_FILE_PATH_KEY}}, {PYRAMID_BUILD_ALPHA, {GRAPH_KEY, ODESCENT_PARAMETER_ALPHA}}, {PYRAMID_GRAPH_TYPE, {GRAPH_KEY, GRAPH_TYPE_KEY}}, {PYRAMID_GRAPH_STORAGE_TYPE, {GRAPH_KEY, GRAPH_STORAGE_TYPE_KEY}}, @@ -983,6 +1059,13 @@ Pyramid::CheckAndMappingExternalParam(const JsonType& external_param, std::string str = format_map(HGRAPH_PARAMS_TEMPLATE, DEFAULT_MAP); auto inner_json = JsonType::Parse(str); mapping_external_param_to_inner(external_param, external_mapping, inner_json); + MapRaBitQSplitParam(external_param, inner_json); + ValidateMRLEDim(external_param, common_param.dim_); + const bool requires_raw_vector = + inner_json[BASE_CODES_KEY][CODES_TYPE_KEY].GetString() == RABITQ_SPLIT_CODES and + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][TYPE_KEY].GetString() == + QUANTIZATION_TYPE_VALUE_TQ; + inner_json[STORE_RAW_VECTOR_KEY].SetBool(requires_raw_vector); auto pyramid_params = std::make_shared(); pyramid_params->FromJson(inner_json); return pyramid_params; @@ -991,9 +1074,12 @@ Pyramid::CheckAndMappingExternalParam(const JsonType& external_param, void Pyramid::Train(const DatasetPtr& base) { this->base_codes_->Train(base->GetFloat32Vectors(), base->GetNumElements()); - if (use_reorder_) { + if (has_precise_codes()) { this->precise_codes_->Train(base->GetFloat32Vectors(), base->GetNumElements()); } + if (raw_vector_ != nullptr) { + this->raw_vector_->Train(base->GetFloat32Vectors(), base->GetNumElements()); + } } std::vector Pyramid::Build(const DatasetPtr& base) { @@ -1058,15 +1144,19 @@ Pyramid::add_one_point(const Hierarchy& h, graph_node.ids_ = node->ids_; graph_node.Init(); - auto codes = use_reorder_ ? precise_codes_ : base_codes_; + auto codes = decodable_codes(); Vector decoded_vector(dim_, allocator_); for (const auto id : node->ids_) { bool need_release = false; const auto* buffer = codes->GetCodesById(id, need_release); - codes->Decode(buffer, decoded_vector.data()); + const bool decoded = codes->Decode(buffer, decoded_vector.data()); if (need_release) { codes->Release(buffer); } + if (not decoded) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "Pyramid graph promotion requires decodable vectors"); + } add_one_point(h, &graph_node, id, decoded_vector.data()); } @@ -1091,7 +1181,7 @@ Pyramid::add_one_point(const Hierarchy& h, search_param.find_duplicate = true; search_param.duplicate_query_id = inner_id; } - auto codes = use_reorder_ ? precise_codes_ : base_codes_; + auto codes = graph_codes(); bool update_entry_point; { std::scoped_lock entry_point_lock(entry_point_mutex_); @@ -1326,7 +1416,7 @@ Pyramid::CalcDistanceById(const float* query, int64_t id, bool calculate_precise std::shared_lock lock(resize_mutex_); auto flat = this->base_codes_; if (use_reorder_ && calculate_precise_distance) { - flat = this->precise_codes_; + flat = this->graph_codes(); } return InnerIndexInterface::calc_distance_by_id(query, id, flat); } @@ -1348,7 +1438,7 @@ Pyramid::CalDistanceById(const float* query, std::shared_lock lock(resize_mutex_); auto flat = this->base_codes_; if (use_reorder_ && calculate_precise_distance) { - flat = this->precise_codes_; + flat = this->graph_codes(); } std::vector validity; auto result = InnerIndexInterface::cal_distance_by_id(query, ids, count, flat, &validity); @@ -1361,13 +1451,17 @@ Pyramid::CalDistanceById(const float* query, void Pyramid::GetVectorByInnerId(InnerIdType inner_id, float* data) const { std::shared_lock lock(resize_mutex_); - auto codes = (use_reorder_) ? precise_codes_ : base_codes_; + auto codes = decodable_codes(); bool release = false; const auto* buffer = codes->GetCodesById(inner_id, release); - codes->Decode(buffer, data); + const bool decoded = codes->Decode(buffer, data); if (release) { codes->Release(buffer); } + if (not decoded) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "Pyramid vector source does not support decode"); + } } std::string diff --git a/src/algorithm/pyramid/pyramid.h b/src/algorithm/pyramid/pyramid.h index 8fd2c66b18..6df46db5a3 100644 --- a/src/algorithm/pyramid/pyramid.h +++ b/src/algorithm/pyramid/pyramid.h @@ -131,7 +131,8 @@ class Pyramid : public InnerIndexInterface { odescent_param_(pyramid_param->odescent_param), index_min_size_(pyramid_param->index_min_size), graph_type_(pyramid_param->graph_type), - support_duplicate_(pyramid_param->support_duplicate) { + support_duplicate_(pyramid_param->support_duplicate), + reorder_by_base_(pyramid_param->reorder_source == HGRAPH_REORDER_SOURCE_BASE) { base_codes_ = FlattenInterface::MakeInstance(pyramid_param->base_codes_param, common_param); if (pyramid_param->has_hierarchies) { for (const auto& h_param : pyramid_param->hierarchies) { @@ -164,9 +165,17 @@ class Pyramid : public InnerIndexInterface { points_mutex_ = std::make_shared(max_capacity_, allocator_); searcher_ = std::make_unique(common_param, points_mutex_); if (use_reorder_) { - precise_codes_ = - FlattenInterface::MakeInstance(pyramid_param->precise_codes_param, common_param); - reorder_ = std::make_shared(precise_codes_, allocator_); + if (reorder_by_base_) { + reorder_ = std::make_shared(base_codes_, allocator_); + } else { + precise_codes_ = FlattenInterface::MakeInstance(pyramid_param->precise_codes_param, + common_param); + reorder_ = std::make_shared(precise_codes_, allocator_); + } + } + if (pyramid_param->store_raw_vector) { + raw_vector_ = + FlattenInterface::MakeInstance(pyramid_param->raw_vector_param, common_param); } } @@ -362,11 +371,27 @@ class Pyramid : public InnerIndexInterface { QueryContext& ctx, uint64_t subindex_ef_search) const; + bool + has_precise_codes() const { + return use_reorder_ and not reorder_by_base_; + } + + FlattenInterfacePtr + graph_codes() const { + return has_precise_codes() ? precise_codes_ : base_codes_; + } + + FlattenInterfacePtr + decodable_codes() const { + return raw_vector_ != nullptr ? raw_vector_ : graph_codes(); + } + private: ODescentParameterPtr odescent_param_{nullptr}; // ODescent build parameters 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}; // original vectors for decode-only paths std::unique_ptr pool_ = nullptr; // pool of visited-lists for search MutexArrayPtr points_mutex_{nullptr}; // per-point locks for concurrent access @@ -375,6 +400,7 @@ class Pyramid : public InnerIndexInterface { int64_t cur_element_count_{0}; // number of vectors currently stored std::atomic delete_count_{0}; // number of deleted vectors bool support_duplicate_{false}; // whether to allow duplicate ids + bool reorder_by_base_{false}; // reorder directly from base codes mutable std::shared_mutex resize_mutex_; // guards resize operations std::mutex cur_element_count_mutex_; // guards cur_element_count_ updates diff --git a/src/algorithm/pyramid/pyramid_test.cpp b/src/algorithm/pyramid/pyramid_test.cpp index 3e7ba18c9e..9e7f097de9 100644 --- a/src/algorithm/pyramid/pyramid_test.cpp +++ b/src/algorithm/pyramid/pyramid_test.cpp @@ -15,6 +15,8 @@ #include "pyramid.h" +#include +#include #include #include "impl/allocator/safe_allocator.h" @@ -31,7 +33,7 @@ struct PyramidTestIndex { }; PyramidTestIndex -MakePyramidIndex(uint32_t index_min_size) { +MakePyramidIndex(uint32_t index_min_size, bool use_mrle_split = false) { PyramidTestIndex result; vsag::IndexCommonParam common_param; common_param.dim_ = PYRAMID_TEST_DIM; @@ -51,6 +53,17 @@ MakePyramidIndex(uint32_t index_min_size) { "index_min_size": 3 })"); external_param[vsag::PYRAMID_INDEX_MIN_SIZE].SetInt(index_min_size); + if (use_mrle_split) { + external_param[vsag::PYRAMID_BASE_QUANTIZATION_TYPE].SetString( + vsag::QUANTIZATION_TYPE_VALUE_TQ); + external_param[vsag::PYRAMID_PRECISE_QUANTIZATION_TYPE].SetString( + vsag::QUANTIZATION_TYPE_VALUE_RABITQ); + external_param[vsag::PYRAMID_USE_REORDER].SetBool(true); + external_param[vsag::INDEX_TQ_CHAIN].SetString("mrle, rabitq"); + external_param[vsag::INDEX_MRLE_DIM].SetInt(2); + external_param[vsag::PYRAMID_RABITQ_BITS_PER_DIM_BASE].SetInt(3); + external_param[vsag::RABITQ_BITS_PER_DIM_PRECISE].SetInt(5); + } auto param = vsag::Pyramid::CheckAndMappingExternalParam(external_param, common_param); result.index = std::make_shared(param, common_param); return result; @@ -157,3 +170,37 @@ TEST_CASE("Pyramid promotes flat node at index minimum size", "[ut][pyramid]") { REQUIRE(result->GetIds()[0] == ids[i]); } } + +TEST_CASE("Pyramid MRLE split retains vectors for flat node promotion", "[ut][pyramid][MRLE]") { + auto test_index = MakePyramidIndex(3, true); + const auto& index = test_index.index; + std::vector vectors = { + 0.0F, + 0.0F, + 0.0F, + 0.0F, + 1.0F, + 0.0F, + 0.0F, + 0.0F, + 0.0F, + 1.0F, + 0.0F, + 0.0F, + }; + std::vector ids = {100, 101, 102}; + std::vector paths(3, "tenant"); + + REQUIRE(index->Add(MakePyramidDataset(vectors.data(), ids.data(), paths.data(), 2)).empty()); + REQUIRE(index + ->Add(MakePyramidDataset( + vectors.data() + 2 * PYRAMID_TEST_DIM, ids.data() + 2, paths.data() + 2, 1)) + .empty()); + + REQUIRE(GetPyramidSubindexCount(index, "graph_subindexes") == 1); + for (int64_t i = 0; i < 3; ++i) { + std::array decoded{}; + index->GetVectorByInnerId(i, decoded.data()); + REQUIRE(std::equal(decoded.begin(), decoded.end(), vectors.begin() + i * PYRAMID_TEST_DIM)); + } +} diff --git a/src/algorithm/pyramid/pyramid_zparameters.cpp b/src/algorithm/pyramid/pyramid_zparameters.cpp index 5bd0b3b27a..d0486a4ef8 100644 --- a/src/algorithm/pyramid/pyramid_zparameters.cpp +++ b/src/algorithm/pyramid/pyramid_zparameters.cpp @@ -176,7 +176,7 @@ PyramidParameters::FromJson(const JsonType& json) { } this->use_reorder = json[USE_REORDER_KEY].GetBool(); - if (this->use_reorder) { + if (this->use_reorder and this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) { this->precise_codes_param = CreateFlattenParam(json[PRECISE_CODES_KEY]); } @@ -233,7 +233,7 @@ PyramidParameters::ToJson() const { json[USE_REORDER_KEY].SetBool(this->use_reorder); json[INDEX_MIN_SIZE].SetInt(index_min_size); json[SUPPORT_DUPLICATE].SetBool(support_duplicate); - if (this->use_reorder) { + if (this->use_reorder and this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) { json[PRECISE_CODES_KEY].SetJson(precise_codes_param->ToJson()); } if (this->has_hierarchies) { @@ -292,9 +292,13 @@ PyramidParameters::CheckCompatibility(const ParamPtr& other) const { return false; } CHECK_FIELD_EQ(*this, *p, use_reorder); - if (this->use_reorder) { + if (this->use_reorder and this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) { 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/algorithm/pyramid/pyramid_zparameters_test.cpp b/src/algorithm/pyramid/pyramid_zparameters_test.cpp index 35dc659b16..5a4accc7ec 100644 --- a/src/algorithm/pyramid/pyramid_zparameters_test.cpp +++ b/src/algorithm/pyramid/pyramid_zparameters_test.cpp @@ -451,3 +451,37 @@ TEST_CASE("Pyramid maps fast RaBitQ to base and precise quantizers", "[ut][Pyram REQUIRE_FALSE(precise_json["quantization_params"]["fast_encode_rabitq"].GetBool()); REQUIRE(precise_json["quantization_params"]["fast_encode_rabitq_rounds"].GetInt() == 11); } + +TEST_CASE("Pyramid maps MRLE RaBitQ split to base reorder", "[ut][PyramidParameters][MRLE]") { + auto param = vsag::JsonType::Parse(R"({ + "base_quantization_type": "tq", + "tq_chain": "mrle, rabitq", + "mrle_dim": 64, + "precise_quantization_type": "rabitq", + "rabitq_bits_per_dim_base": 3, + "rabitq_bits_per_dim_precise": 5, + "hierarchies": ["site"], + "use_reorder": true + })"); + + vsag::IndexCommonParam common_param; + common_param.dim_ = 128; + common_param.data_type_ = vsag::DataTypes::DATA_TYPE_FLOAT; + auto mapped = vsag::Pyramid::CheckAndMappingExternalParam(param, common_param); + auto typed_param = std::dynamic_pointer_cast(mapped); + + REQUIRE(typed_param != nullptr); + REQUIRE(typed_param->reorder_source == std::string("base")); + REQUIRE(typed_param->precise_codes_param == nullptr); + REQUIRE(typed_param->store_raw_vector); + REQUIRE(typed_param->raw_vector_param != nullptr); + const auto base_json = typed_param->base_codes_param->ToJson(); + const auto raw_json = typed_param->raw_vector_param->ToJson(); + REQUIRE(base_json["codes_type"].GetString() == std::string("rabitq_split")); + REQUIRE(base_json["quantization_params"]["type"].GetString() == std::string("tq")); + REQUIRE(base_json["quantization_params"]["tq_chain"].GetString() == std::string("mrle,rabitq")); + REQUIRE(base_json["quantization_params"]["mrle_dim"].GetInt() == 64); + REQUIRE(base_json["quantization_params"]["rabitq_bits_per_dim_filter"].GetInt() == 3); + REQUIRE(base_json["quantization_params"]["rabitq_bits_per_dim_base"].GetInt() == 8); + REQUIRE(raw_json["quantization_params"]["type"].GetString() == std::string("fp32")); +} diff --git a/src/analyzer/pyramid_analyzer.cpp b/src/analyzer/pyramid_analyzer.cpp index 69354d581c..2992b10e83 100644 --- a/src/analyzer/pyramid_analyzer.cpp +++ b/src/analyzer/pyramid_analyzer.cpp @@ -508,7 +508,7 @@ PyramidAnalyzer::calculate_groundtruth(const Vector& sample_datas, Vector ids_array(this->total_count_, allocator_); std::iota(ids_array.begin(), ids_array.end(), 0); - auto codes = pyramid_->use_reorder_ ? pyramid_->precise_codes_ : pyramid_->base_codes_; + auto codes = pyramid_->graph_codes(); for (uint32_t i = 0; i < sample_size; ++i) { if (i % 10 == 0) { @@ -793,7 +793,7 @@ PyramidAnalyzer::calculate_node_groundtruth(const IndexNode* node, return gt; } - auto codes = pyramid_->use_reorder_ ? pyramid_->precise_codes_ : pyramid_->base_codes_; + auto codes = pyramid_->graph_codes(); if (codes == nullptr) { return gt; } @@ -830,7 +830,7 @@ PyramidAnalyzer::search_single_node(const IndexNode* node, return result; } - auto codes = pyramid_->use_reorder_ ? pyramid_->precise_codes_ : pyramid_->base_codes_; + auto codes = pyramid_->graph_codes(); Vector distances(node_ids.size(), allocator_); auto computer = codes->FactoryComputer(query); codes->Query(distances.data(), computer, node_ids.data(), node_ids.size()); @@ -1021,7 +1021,7 @@ PyramidAnalyzer::get_node_neighbor_recall(const IndexNode* node, } auto graph = node->graph_; - auto codes = pyramid_->use_reorder_ ? pyramid_->precise_codes_ : pyramid_->base_codes_; + auto codes = pyramid_->graph_codes(); if (codes == nullptr) { return 0.0F; } diff --git a/src/datacell/flatten_datacell_test.cpp b/src/datacell/flatten_datacell_test.cpp index dbbd9fd81f..29af97a456 100644 --- a/src/datacell/flatten_datacell_test.cpp +++ b/src/datacell/flatten_datacell_test.cpp @@ -37,6 +37,7 @@ #include "impl/thread_pool/safe_thread_pool.h" #include "index_common_param.h" #include "quantization/rabitq_quantization/rabitq_quantizer.h" +#include "quantization/transform_quantization/transform_quantizer.h" #include "unittest.h" using namespace vsag; @@ -214,6 +215,81 @@ TEST_CASE("RaBitQSplitDataCell direct split compute", "[ut][RaBitQSplitDataCell] } } } +TEST_CASE("RaBitQSplitDataCell supports MRLE transform quantizer", + "[ut][RaBitQSplitDataCell][MRLE]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr uint64_t dim = 64; + constexpr uint64_t mrle_dim = 32; + constexpr InnerIdType count = 24; + auto vectors = fixtures::generate_vectors(count, dim); + auto query = fixtures::generate_vectors(1, dim, 71); + + auto param = std::make_shared(); + param->FromJson(JsonType::Parse(R"({ + "codes_type": "rabitq_split", + "io_params": { "type": "memory_io" }, + "quantization_params": { + "type": "tq", + "tq_chain": "mrle, rabitq", + "mrle_dim": 32, + "rabitq_version": "split", + "rabitq_bits_per_dim_query": 32, + "rabitq_bits_per_dim_base": 8, + "rabitq_bits_per_dim_filter": 3, + "fast_encode_rabitq": true + } + })")); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + + auto flatten = FlattenInterface::MakeInstance(param, common_param); + REQUIRE(flatten->GetQuantizerName() == std::string("tq")); + flatten->Train(vectors.data(), count); + flatten->Resize(count); + + auto optimized_build = std::dynamic_pointer_cast(flatten); + REQUIRE(optimized_build != nullptr); + auto finalize_pool = SafeThreadPool::FactoryDefaultThreadPool(); + finalize_pool->SetPoolSize(2); + FlattenOptimizedBuildContext build_context{finalize_pool, 2}; + REQUIRE(optimized_build->BeginOptimizedBuild(build_context)); + flatten->BatchInsertVector(vectors.data(), count); + + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + std::vector build_dists(count); + auto computer = flatten->FactoryComputer(query.data()); + flatten->Query(build_dists.data(), computer, ids.data(), count); + optimized_build->FinalizeOptimizedBuild(); + + std::vector split_dists(count); + flatten->Query(split_dists.data(), computer, ids.data(), count); + using QuantizerT = TransformQuantizer, + MetricType::METRIC_TYPE_L2SQR>; + auto* transform_computer = static_cast*>(computer.get()); + for (InnerIdType id = 0; id < count; ++id) { + bool need_release = false; + const auto* full_code = flatten->GetCodesById(id, need_release); + float merged_dist = 0.0F; + transform_computer->ComputeDist(full_code, &merged_dist); + if (need_release) { + flatten->Release(full_code); + } + REQUIRE(std::abs(build_dists[id] - split_dists[id]) <= 1e-4F); + REQUIRE(std::abs(split_dists[id] - merged_dist) <= 1e-5F); + } + + auto invalid_json = param->ToJson(); + invalid_json["quantization_params"]["tq_chain"].SetString("pca, rabitq"); + invalid_json["quantization_params"]["pca_dim"].SetInt(mrle_dim); + auto invalid_param = std::make_shared(); + invalid_param->FromJson(invalid_json); + REQUIRE_THROWS(FlattenInterface::MakeInstance(invalid_param, common_param)); +} + TEST_CASE("RaBitQSplitDataCell serialize and methods", "[ut][RaBitQSplitDataCell]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); constexpr uint64_t dim = 64; diff --git a/src/datacell/flatten_interface.cpp b/src/datacell/flatten_interface.cpp index f26b265c8b..d4b13a7b58 100644 --- a/src/datacell/flatten_interface.cpp +++ b/src/datacell/flatten_interface.cpp @@ -183,11 +183,7 @@ make_instance(const FlattenInterfaceParamPtr& param, const IndexCommonParam& com } if (actual_quant_type == QUANTIZATION_TYPE_VALUE_RABITQ) { if (param->name == RABITQ_SPLIT_DATA_CELL) { - if (is_transform_quantizer) { - throw VsagException(ErrorType::INVALID_ARGUMENT, - "rabitq split data cell does not support transform quantizer"); - } - return MakeRaBitQSplitDataCell(param, common_param); + return MakeRaBitQSplitDataCell(param, common_param, is_transform_quantizer); } return make_instance_with_tq, IOTemp, metric>( param, common_param, is_transform_quantizer); diff --git a/src/datacell/rabitq_split_datacell.h b/src/datacell/rabitq_split_datacell.h index f6cddb1fd2..aa6225bff5 100644 --- a/src/datacell/rabitq_split_datacell.h +++ b/src/datacell/rabitq_split_datacell.h @@ -38,6 +38,7 @@ #include "io/memory_io/memory_io.h" #include "io/memory_io/memory_io_parameter.h" #include "io/mmap_io/mmap_io_parameter.h" +#include "quantization/bottom_quantizer_accessor.h" #include "quantization/rabitq_quantization/rabitq_quantizer.h" #include "query_context.h" #include "storage/stream_reader.h" @@ -140,9 +141,19 @@ class RaBitQSplitCodeStorage { uint64_t code_size_{0}; }; -template +template > class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuildInterface { public: + using Accessor = BottomQuantizerAccessor; + using BottomQuantizer = typename Accessor::BottomQuantizerType; + using BottomComputer = typename Accessor::BottomComputerType; + + static_assert(std::is_same_v>, + "RaBitQSplitDataCell requires RaBitQuantizer as bottom quantizer"); + class OptimizedBuildComputer final : public ComputerInterface { public: OptimizedBuildComputer(uint64_t record_size, Allocator* allocator) @@ -166,9 +177,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil const IOParamPtr& supplement_io_param, const IndexCommonParam& common_param) : common_param_(common_param), allocator_(common_param.allocator_.get()) { - this->quantizer_ = - std::make_shared>(quantization_param, common_param); - if (not this->quantizer_->SupportSplitCodeStorage()) { + this->quantizer_ = std::make_shared(quantization_param, common_param); + if (not this->bottom_quantizer().SupportSplitCodeStorage()) { throw VsagException(ErrorType::INVALID_ARGUMENT, "rabitq split data cell requires rabitq_version=split, " "rabitq_bits_per_dim_query=32, and " @@ -203,7 +213,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil this->query_optimized_build_codes(result_dists, computer, idx, id_count); return; } - auto* comp = static_cast>*>(computer.get()); + auto* comp = this->get_bottom_computer(computer); if constexpr (not OneBitIOTmpl::InMemory or not SupplementIOTmpl::InMemory) { if (id_count > 1) { if constexpr (OneBitIOTmpl::InMemory and not SupplementIOTmpl::InMemory) { @@ -239,7 +249,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil this->query_optimized_build_codes(result_dists, computer, idx, id_count); return; } - auto* comp = static_cast>*>(computer.get()); + auto* comp = this->get_bottom_computer(computer); if constexpr (not OneBitIOTmpl::InMemory or not SupplementIOTmpl::InMemory) { if (id_count > 1) { if constexpr (OneBitIOTmpl::InMemory and not SupplementIOTmpl::InMemory) { @@ -278,7 +288,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil this->query_optimized_build_codes(result_dists, computer, idx, id_count); return; } - auto* comp = static_cast>*>(computer.get()); + auto* comp = this->get_bottom_computer(computer); for (uint32_t i = 0; i < this->prefetch_stride_code_ and i < id_count; ++i) { this->prefetch_full_code(idx[i]); } @@ -294,7 +304,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil float lower_bound = std::numeric_limits::max(); bool computed = false; try { - computed = this->quantizer_->ComputeDistWithOneBitLowerBound( + computed = this->bottom_quantizer().ComputeDistWithOneBitLowerBound( *comp, one_bit_code, &one_bit_dist, @@ -340,7 +350,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil } return; } - auto* comp = static_cast>*>(computer.get()); + auto* comp = this->get_bottom_computer(computer); this->add_filter_count(ctx, id_count); if constexpr (not OneBitIOTmpl::InMemory) { if (id_count > 1) { @@ -384,7 +394,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil auto* lower_bound2 = lower_bounds == nullptr ? nullptr : lower_bounds + i + 1; auto* lower_bound3 = lower_bounds == nullptr ? nullptr : lower_bounds + i + 2; auto* lower_bound4 = lower_bounds == nullptr ? nullptr : lower_bounds + i + 3; - this->quantizer_->ComputeDistsWithOneBitLowerBoundBatch4( + this->bottom_quantizer().ComputeDistsWithOneBitLowerBoundBatch4( *comp, code1, code2, @@ -440,7 +450,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil auto* lower_bound = lower_bounds == nullptr ? nullptr : lower_bounds + i; bool computed = false; try { - computed = this->quantizer_->ComputeDistWithOneBitLowerBound( + computed = this->bottom_quantizer().ComputeDistWithOneBitLowerBound( *comp, one_bit_code, result_dists + i, @@ -490,13 +500,14 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil bool BeginOptimizedBuild(const FlattenOptimizedBuildContext& context) override { - if (this->optimized_build_active_ or not this->quantizer_->SupportScalarCodeBuild()) { + if (this->optimized_build_active_ or + not this->bottom_quantizer().SupportScalarCodeBuild()) { return false; } auto io_param = std::make_shared(); auto build_codes = std::make_shared>(io_param, this->common_param_); - build_codes->SetCodeSize(this->quantizer_->GetScalarCodeSize()); + build_codes->SetCodeSize(this->bottom_quantizer().GetScalarCodeSize()); auto code_sums = std::make_unique>(this->allocator_); if (this->max_capacity_ > 0) { build_codes->Resize(this->max_capacity_); @@ -504,7 +515,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil } this->optimized_build_scalar_codes_ = build_codes; this->optimized_build_code_sums_ = std::move(code_sums); - this->optimized_build_record_size_ = this->quantizer_->GetScalarCodeSize(); + this->optimized_build_record_size_ = this->bottom_quantizer().GetScalarCodeSize(); this->optimized_build_context_ = context; this->optimized_build_active_ = true; return true; @@ -535,7 +546,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil "failed to read temporary scalar RaBitQ build code"); } try { - this->quantizer_->PackScalarCodeToSplitCode( + this->bottom_quantizer().PackScalarCodeToSplitCode( scalar_code, one_bit_code.data, supplement_code.data); this->x_bit_cell_->Write(one_bit_code.data, id); this->supplement_cell_->Write(supplement_code.data, id); @@ -681,7 +692,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil } float distance = 0.0F; try { - distance = this->quantizer_->ComputeScalarCodesDistance( + distance = this->bottom_quantizer().ComputeScalarCodesDistance( codes1, (*optimized_build_code_sums_)[id1], codes2, @@ -740,9 +751,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil this->quantizer_->Serialize(writer); ss.seekg(0, std::ios::beg); IOStreamReader reader(ss); - auto ptr = - std::dynamic_pointer_cast>( - other); + auto ptr = std::dynamic_pointer_cast< + RaBitQSplitDataCell>(other); if (ptr == nullptr) { throw VsagException(ErrorType::INTERNAL_ERROR, "Export model's rabitq split datacell failed"); @@ -832,7 +842,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil if (scalar_code == nullptr) { return false; } - this->quantizer_->PackScalarCode(scalar_code, codes); + memset(codes, 0, this->code_size_); + this->bottom_quantizer().PackScalarCode(scalar_code, codes); if (need_release) { this->optimized_build_scalar_codes_->Release(scalar_code); } @@ -845,7 +856,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil if (not one_bit_ok or not supplement_ok) { return false; } - this->quantizer_->MergeSplitCode(one_bit.data, supplement.data, codes); + memset(codes, 0, this->code_size_); + this->bottom_quantizer().MergeSplitCode(one_bit.data, supplement.data, codes); return true; } @@ -882,9 +894,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void MergeOther(const FlattenInterfacePtr& other, InnerIdType bias) override { - auto ptr = - std::dynamic_pointer_cast>( - other); + auto ptr = std::dynamic_pointer_cast< + RaBitQSplitDataCell>(other); if (ptr == nullptr) { throw VsagException(ErrorType::INTERNAL_ERROR, "Merge rabitq split datacell failed: not match type"); @@ -933,7 +944,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil uint64_t GetMemoryUsage() const override { - uint64_t memory = sizeof(RaBitQSplitDataCell); + uint64_t memory = + sizeof(RaBitQSplitDataCell); memory += this->x_bit_cell_->GetMemoryUsage(); memory += this->supplement_cell_->GetMemoryUsage(); if (this->optimized_build_scalar_codes_ != nullptr) { @@ -942,13 +954,13 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil if (this->optimized_build_code_sums_ != nullptr) { memory += this->optimized_build_code_sums_->capacity() * sizeof(uint64_t); } - memory += sizeof(RaBitQuantizer); + memory += sizeof(QuantizerT); return memory; } public: IndexCommonParam common_param_; - std::shared_ptr> quantizer_{nullptr}; + std::shared_ptr quantizer_{nullptr}; std::shared_ptr> x_bit_cell_{nullptr}; std::shared_ptr> supplement_cell_{nullptr}; std::shared_ptr> optimized_build_scalar_codes_{nullptr}; @@ -971,6 +983,22 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil uint64_t optimized_build_record_size_{0}; private: + BottomQuantizer& + bottom_quantizer() { + return Accessor::GetQuantizer(*this->quantizer_); + } + + const BottomQuantizer& + bottom_quantizer() const { + return Accessor::GetQuantizer(*this->quantizer_); + } + + BottomComputer* + get_bottom_computer(const ComputerInterfacePtr& computer) const { + auto* outer_computer = static_cast*>(computer.get()); + return &Accessor::GetComputer(*outer_computer); + } + static IOParamPtr SuffixIOParam(const IOParamPtr& io_param, const std::string& suffix) { if (io_param == nullptr) { @@ -1047,8 +1075,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void refresh_code_sizes() { this->code_size_ = static_cast(quantizer_->GetCodeSize()); - this->one_bit_code_size_ = quantizer_->GetOneBitCodeSize(); - this->supplement_code_size_ = quantizer_->GetSupplementCodeSize(); + this->one_bit_code_size_ = this->bottom_quantizer().GetOneBitCodeSize(); + this->supplement_code_size_ = this->bottom_quantizer().GetSupplementCodeSize(); this->x_bit_cell_->SetCodeSize(one_bit_code_size_); this->supplement_cell_->SetCodeSize(supplement_code_size_); } @@ -1057,8 +1085,12 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil write_encoded_vector(const float* vector, InnerIdType idx) { if (this->optimized_build_active_) { ByteBuffer scalar_code(this->optimized_build_record_size_, allocator_); + Vector transformed_input(this->allocator_); + const float* bottom_input = + Accessor::PrepareBottomInput(*this->quantizer_, vector, transformed_input); uint64_t code_sum = 0; - if (not this->quantizer_->EncodeOneToScalarCode(vector, scalar_code.data, code_sum)) { + if (not this->bottom_quantizer().EncodeOneToScalarCode( + bottom_input, scalar_code.data, code_sum)) { throw VsagException(ErrorType::INTERNAL_ERROR, "failed to encode temporary scalar RaBitQ build code"); } @@ -1070,7 +1102,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil this->quantizer_->EncodeOne(vector, full_code.data); ByteBuffer one_bit_code(one_bit_code_size_, allocator_); ByteBuffer supplement_code(supplement_code_size_, allocator_); - this->quantizer_->SplitCode(full_code.data, one_bit_code.data, supplement_code.data); + this->bottom_quantizer().SplitCode(full_code.data, one_bit_code.data, supplement_code.data); this->x_bit_cell_->Write(one_bit_code.data, idx); this->supplement_cell_->Write(supplement_code.data, idx); } @@ -1090,7 +1122,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil id_count); return; } - auto* comp = static_cast>*>(computer.get()); + auto* comp = this->get_bottom_computer(computer); for (InnerIdType i = 0; i < id_count; ++i) { bool need_release = false; const auto* scalar_code = @@ -1100,7 +1132,8 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil "failed to read temporary scalar RaBitQ build code"); } try { - this->quantizer_->ComputeDistWithScalarCode(*comp, scalar_code, result_dists + i); + this->bottom_quantizer().ComputeDistWithScalarCode( + *comp, scalar_code, result_dists + i); } catch (...) { if (need_release) { this->optimized_build_scalar_codes_->Release(scalar_code); @@ -1134,7 +1167,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil "failed to read temporary scalar RaBitQ build code"); } try { - result_dists[i] = this->quantizer_->ComputeScalarCodesDistance( + result_dists[i] = this->bottom_quantizer().ComputeScalarCodesDistance( query_code, query_sum, base_code, (*this->optimized_build_code_sums_)[idx[i]]); } catch (...) { if (need_release) { @@ -1236,7 +1269,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void query_one_bit_lower_bound_by_multiread(float* result_dists, float* lower_bounds, - Computer>* computer, + BottomComputer* computer, const InnerIdType* idx, InnerIdType id_count, QueryContext* ctx) const { @@ -1271,7 +1304,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil auto* lower_bound2 = lower_bounds == nullptr ? nullptr : lower_bounds + i + 1; auto* lower_bound3 = lower_bounds == nullptr ? nullptr : lower_bounds + i + 2; auto* lower_bound4 = lower_bounds == nullptr ? nullptr : lower_bounds + i + 3; - this->quantizer_->ComputeDistsWithOneBitLowerBoundBatch4( + this->bottom_quantizer().ComputeDistsWithOneBitLowerBoundBatch4( *computer, code1, code2, @@ -1315,7 +1348,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil for (; i < id_count; ++i) { auto* lower_bound = lower_bounds == nullptr ? nullptr : lower_bounds + i; const auto* one_bit_code = one_bit_codes.data + i * one_bit_code_size_; - bool computed = this->quantizer_->ComputeDistWithOneBitLowerBound( + bool computed = this->bottom_quantizer().ComputeDistWithOneBitLowerBound( *computer, one_bit_code, result_dists + i, @@ -1331,7 +1364,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void query_full_dist_by_multiread(float* result_dists, - Computer>* computer, + BottomComputer* computer, const InnerIdType* idx, InnerIdType id_count, QueryContext* ctx, @@ -1374,7 +1407,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void query_full_dist_by_supplement_multiread(float* result_dists, - Computer>* computer, + BottomComputer* computer, const InnerIdType* idx, InnerIdType id_count, QueryContext* ctx, @@ -1419,7 +1452,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void compute_full_dist_after_one_bit_failure(InnerIdType id, const uint8_t* one_bit_code, - Computer>* computer, + BottomComputer* computer, float* result_dist, float* lower_bound, QueryContext* ctx) const { @@ -1441,7 +1474,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil void compute_full_dist(const uint8_t* one_bit_code, const uint8_t* supplement_code, - Computer>* computer, + BottomComputer* computer, float* result_dist, QueryContext* ctx = nullptr, float hint_dist = std::numeric_limits::max()) const { @@ -1450,7 +1483,7 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil const bool has_hint = std::isfinite(hint_dist) and hint_dist < std::numeric_limits::max(); if (has_hint) { - computed = this->quantizer_->ComputeDistWithSplitCodeAndFilterDist( + computed = this->bottom_quantizer().ComputeDistWithSplitCodeAndFilterDist( *computer, one_bit_code, supplement_code, hint_dist, result_dist); } if (computed) { @@ -1458,17 +1491,17 @@ class RaBitQSplitDataCell : public FlattenInterface, public FlattenOptimizedBuil } else if (has_hint) { this->add_reorder_fallback_full_count(ctx, 1); } - if (not computed and not this->quantizer_->ComputeDistWithSplitCode( + if (not computed and not this->bottom_quantizer().ComputeDistWithSplitCode( *computer, one_bit_code, supplement_code, result_dist)) { ByteBuffer full_code(this->code_size_, allocator_); - this->quantizer_->MergeSplitCode(one_bit_code, supplement_code, full_code.data); + this->bottom_quantizer().MergeSplitCode(one_bit_code, supplement_code, full_code.data); computer->ComputeDist(full_code.data, result_dist); } } void compute_full_dist(InnerIdType id, - Computer>* computer, + BottomComputer* computer, float* result_dist, QueryContext* ctx = nullptr, float hint_dist = std::numeric_limits::max()) const { diff --git a/src/datacell/rabitq_split_datacell_factory.cpp b/src/datacell/rabitq_split_datacell_factory.cpp index 41f5972b37..33f6e269a4 100644 --- a/src/datacell/rabitq_split_datacell_factory.cpp +++ b/src/datacell/rabitq_split_datacell_factory.cpp @@ -20,27 +20,31 @@ namespace vsag { FlattenInterfacePtr make_ra_bit_q_split_data_cell_l2(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer); FlattenInterfacePtr make_ra_bit_q_split_data_cell_ip(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer); FlattenInterfacePtr make_ra_bit_q_split_data_cell_cosine(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer); FlattenInterfacePtr MakeRaBitQSplitDataCell(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param) { + const IndexCommonParam& common_param, + bool is_transform_quantizer) { if (common_param.metric_ == MetricType::METRIC_TYPE_L2SQR) { - return make_ra_bit_q_split_data_cell_l2(param, common_param); + return make_ra_bit_q_split_data_cell_l2(param, common_param, is_transform_quantizer); } if (common_param.metric_ == MetricType::METRIC_TYPE_IP) { - return make_ra_bit_q_split_data_cell_ip(param, common_param); + return make_ra_bit_q_split_data_cell_ip(param, common_param, is_transform_quantizer); } if (common_param.metric_ == MetricType::METRIC_TYPE_COSINE) { - return make_ra_bit_q_split_data_cell_cosine(param, common_param); + return make_ra_bit_q_split_data_cell_cosine(param, common_param, is_transform_quantizer); } return nullptr; } diff --git a/src/datacell/rabitq_split_datacell_factory.h b/src/datacell/rabitq_split_datacell_factory.h index 686e7f7e2a..1e7cbd1a1c 100644 --- a/src/datacell/rabitq_split_datacell_factory.h +++ b/src/datacell/rabitq_split_datacell_factory.h @@ -20,6 +20,7 @@ namespace vsag { FlattenInterfacePtr MakeRaBitQSplitDataCell(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer); } // namespace vsag diff --git a/src/datacell/rabitq_split_datacell_factory_cosine.cpp b/src/datacell/rabitq_split_datacell_factory_cosine.cpp index 765a3d739b..5f6577e4bf 100644 --- a/src/datacell/rabitq_split_datacell_factory_cosine.cpp +++ b/src/datacell/rabitq_split_datacell_factory_cosine.cpp @@ -18,8 +18,10 @@ namespace vsag { FlattenInterfacePtr make_ra_bit_q_split_data_cell_cosine(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param) { - return MakeRaBitQSplitDataCellForMetric(param, common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer) { + return MakeRaBitQSplitDataCellForMetric( + param, common_param, is_transform_quantizer); } } // namespace vsag diff --git a/src/datacell/rabitq_split_datacell_factory_impl.h b/src/datacell/rabitq_split_datacell_factory_impl.h index 49f1d22f9a..4a9cc15b72 100644 --- a/src/datacell/rabitq_split_datacell_factory_impl.h +++ b/src/datacell/rabitq_split_datacell_factory_impl.h @@ -30,28 +30,29 @@ ConvertRaBitQSplitIOParamType(const IOParamPtr& io_param, const std::string& typ return IOParameter::GetIOParameterByJson(json); } -template +template FlattenInterfacePtr MakeHomogeneousRaBitQSplitDataCell(const FlattenInterfaceParamPtr& param, const IndexCommonParam& common_param) { - return std::make_shared>( + return std::make_shared>( param->quantizer_parameter, param->io_parameter, param->supplement_io_parameter, common_param); } -template +template FlattenInterfacePtr -MakeRaBitQSplitDataCellForMetric(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param) { +MakeRaBitQSplitDataCellForMetricImpl(const FlattenInterfaceParamPtr& param, + const IndexCommonParam& common_param) { if (param->supplement_io_parameter != nullptr) { const auto& supplement_type = param->supplement_io_parameter->GetTypeName(); const auto& base_type = param->io_parameter->GetTypeName(); if (base_type == IO_TYPE_VALUE_BLOCK_MEMORY_IO and supplement_type == IO_TYPE_VALUE_ASYNC_IO) { #if HAVE_LIBAIO - return std::make_shared>( + return std::make_shared< + RaBitQSplitDataCell>( param->quantizer_parameter, param->io_parameter, param->supplement_io_parameter, @@ -59,7 +60,8 @@ MakeRaBitQSplitDataCellForMetric(const FlattenInterfaceParamPtr& param, #else auto buffer_supplement_io_param = ConvertRaBitQSplitIOParamType( param->supplement_io_parameter, IO_TYPE_VALUE_BUFFER_IO); - return std::make_shared>( + return std::make_shared< + RaBitQSplitDataCell>( param->quantizer_parameter, param->io_parameter, buffer_supplement_io_param, @@ -69,7 +71,8 @@ MakeRaBitQSplitDataCellForMetric(const FlattenInterfaceParamPtr& param, #if !HAVE_LIBAIO if (base_type == IO_TYPE_VALUE_BLOCK_MEMORY_IO and supplement_type == IO_TYPE_VALUE_BUFFER_IO) { - return std::make_shared>( + return std::make_shared< + RaBitQSplitDataCell>( param->quantizer_parameter, param->io_parameter, param->supplement_io_parameter, @@ -88,24 +91,46 @@ MakeRaBitQSplitDataCellForMetric(const FlattenInterfaceParamPtr& param, const auto& io_type = param->io_parameter->GetTypeName(); if (io_type == IO_TYPE_VALUE_BLOCK_MEMORY_IO) { - return MakeHomogeneousRaBitQSplitDataCell(param, common_param); + return MakeHomogeneousRaBitQSplitDataCell(param, + common_param); } if (io_type == IO_TYPE_VALUE_MEMORY_IO) { - return MakeHomogeneousRaBitQSplitDataCell(param, common_param); + return MakeHomogeneousRaBitQSplitDataCell(param, + common_param); } if (io_type == IO_TYPE_VALUE_BUFFER_IO) { - return MakeHomogeneousRaBitQSplitDataCell(param, common_param); + return MakeHomogeneousRaBitQSplitDataCell(param, + common_param); } if (io_type == IO_TYPE_VALUE_ASYNC_IO) { - return MakeHomogeneousRaBitQSplitDataCell(param, common_param); + return MakeHomogeneousRaBitQSplitDataCell(param, common_param); } if (io_type == IO_TYPE_VALUE_MMAP_IO) { - return MakeHomogeneousRaBitQSplitDataCell(param, common_param); + return MakeHomogeneousRaBitQSplitDataCell(param, common_param); } if (io_type == IO_TYPE_VALUE_READER_IO) { - return MakeHomogeneousRaBitQSplitDataCell(param, common_param); + return MakeHomogeneousRaBitQSplitDataCell(param, + common_param); } return nullptr; } +template +FlattenInterfacePtr +MakeRaBitQSplitDataCellForMetric(const FlattenInterfaceParamPtr& param, + const IndexCommonParam& common_param, + bool is_transform_quantizer) { + if (is_transform_quantizer) { + auto tq_param = + std::dynamic_pointer_cast(param->quantizer_parameter); + CHECK_ARGUMENT(tq_param != nullptr and tq_param->tq_chain_.size() == 1 and + tq_param->tq_chain_.front() == TRANSFORMER_TYPE_VALUE_MRLE, + "rabitq split transform quantizer requires tq_chain=\"mrle, rabitq\""); + using QuantizerT = TransformQuantizer, metric>; + return MakeRaBitQSplitDataCellForMetricImpl(param, common_param); + } + return MakeRaBitQSplitDataCellForMetricImpl>(param, + common_param); +} + } // namespace vsag diff --git a/src/datacell/rabitq_split_datacell_factory_ip.cpp b/src/datacell/rabitq_split_datacell_factory_ip.cpp index 5c604d5d4c..2c9dbed89b 100644 --- a/src/datacell/rabitq_split_datacell_factory_ip.cpp +++ b/src/datacell/rabitq_split_datacell_factory_ip.cpp @@ -18,8 +18,10 @@ namespace vsag { FlattenInterfacePtr make_ra_bit_q_split_data_cell_ip(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param) { - return MakeRaBitQSplitDataCellForMetric(param, common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer) { + return MakeRaBitQSplitDataCellForMetric( + param, common_param, is_transform_quantizer); } } // namespace vsag diff --git a/src/datacell/rabitq_split_datacell_factory_l2.cpp b/src/datacell/rabitq_split_datacell_factory_l2.cpp index 3a0ed71987..e9c1e6877b 100644 --- a/src/datacell/rabitq_split_datacell_factory_l2.cpp +++ b/src/datacell/rabitq_split_datacell_factory_l2.cpp @@ -18,8 +18,10 @@ namespace vsag { FlattenInterfacePtr make_ra_bit_q_split_data_cell_l2(const FlattenInterfaceParamPtr& param, - const IndexCommonParam& common_param) { - return MakeRaBitQSplitDataCellForMetric(param, common_param); + const IndexCommonParam& common_param, + bool is_transform_quantizer) { + return MakeRaBitQSplitDataCellForMetric( + param, common_param, is_transform_quantizer); } } // namespace vsag diff --git a/src/quantization/bottom_quantizer_accessor.h b/src/quantization/bottom_quantizer_accessor.h new file mode 100644 index 0000000000..2518ad2307 --- /dev/null +++ b/src/quantization/bottom_quantizer_accessor.h @@ -0,0 +1,78 @@ +// 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 "quantization/transform_quantization/transform_quantizer.h" + +namespace vsag { + +template +struct BottomQuantizerAccessor { + using BottomQuantizerType = QuantizerT; + using BottomComputerType = Computer; + + static QuantizerT& + GetQuantizer(QuantizerT& quantizer) { + return quantizer; + } + + static const QuantizerT& + GetQuantizer(const QuantizerT& quantizer) { + return quantizer; + } + + static BottomComputerType& + GetComputer(Computer& computer) { + return computer; + } + + static const float* + PrepareBottomInput(const QuantizerT&, const float* input, Vector&) { + return input; + } +}; + +template +struct BottomQuantizerAccessor> { + using TransformQuantizerType = TransformQuantizer; + using BottomQuantizerType = InnerQuantizerT; + using BottomComputerType = Computer; + + static InnerQuantizerT& + GetQuantizer(TransformQuantizerType& quantizer) { + return *quantizer.quantizer_; + } + + static const InnerQuantizerT& + GetQuantizer(const TransformQuantizerType& quantizer) { + return *quantizer.quantizer_; + } + + static BottomComputerType& + GetComputer(Computer& computer) { + return *computer.inner_computer_; + } + + static const float* + PrepareBottomInput(const TransformQuantizerType& quantizer, + const float* input, + Vector& scratch) { + scratch.resize(quantizer.GetTransformedDim()); + quantizer.TransformBaseVector(input, scratch.data()); + return scratch.data(); + } +}; + +} // namespace vsag diff --git a/src/quantization/transform_quantization/transform_quantizer.h b/src/quantization/transform_quantization/transform_quantizer.h index 3ee545279f..b6fe74ef20 100644 --- a/src/quantization/transform_quantization/transform_quantizer.h +++ b/src/quantization/transform_quantization/transform_quantizer.h @@ -106,6 +106,14 @@ class TransformQuantizer : public QuantizerGetDim(); + } + + void + TransformBaseVector(const float* input, float* output) const; + public: VectorTransformerPtr MakeTransformerInstance(std::string transform_str, @@ -235,12 +243,11 @@ TransformQuantizer::TrainImpl(const float* data, uint64_t cou } // 2. execute transform on original data - Vector transformed_data(this->dim_ * count, 0, this->allocator_); - Vector tmp_codes(this->code_size_, 0, this->allocator_); - transformed_data.assign(data, data + count * this->dim_); - for (auto i = 0; i < count; i++) { - ExecuteChainTransform( - transformed_data.data() + i * this->dim_, base_meta_offsets_.data(), tmp_codes.data()); + const uint64_t transformed_dim = this->GetTransformedDim(); + Vector transformed_data(transformed_dim * count, 0, this->allocator_); + for (uint64_t i = 0; i < count; ++i) { + this->TransformBaseVector(data + i * this->dim_, + transformed_data.data() + i * transformed_dim); } // 3. train quantizer based on transformed data @@ -257,15 +264,25 @@ TransformQuantizer::ExecuteChainTransform(float* prev_data, for (uint32_t i = 0; i < this->transform_chain_.size(); i++) { auto vector_transformer = this->transform_chain_[i]; - auto meta_offset = meta_offsets[i]; - auto meta = vector_transformer->Transform(prev_data, next_data.data()); - meta->EncodeMeta(codes + meta_offset); + if (codes != nullptr) { + meta->EncodeMeta(codes + meta_offsets[i]); + } - memcpy(prev_data, next_data.data(), this->dim_ * sizeof(float)); + memcpy(prev_data, next_data.data(), vector_transformer->GetOutputDim() * sizeof(float)); } } +template +void +TransformQuantizer::TransformBaseVector(const float* input, + float* output) const { + Vector data_buffer(this->dim_, 0, this->allocator_); + data_buffer.assign(input, input + this->dim_); + ExecuteChainTransform(data_buffer.data(), nullptr, nullptr); + memcpy(output, data_buffer.data(), this->GetTransformedDim() * sizeof(float)); +} + template bool TransformQuantizer::EncodeOneImpl(const float* data, uint8_t* codes) const { diff --git a/src/quantization/transform_quantization/transform_quantizer_parameter.cpp b/src/quantization/transform_quantization/transform_quantizer_parameter.cpp index e9d1b533dc..a221705365 100644 --- a/src/quantization/transform_quantization/transform_quantizer_parameter.cpp +++ b/src/quantization/transform_quantization/transform_quantizer_parameter.cpp @@ -15,6 +15,7 @@ #include "transform_quantizer_parameter.h" +#include "impl/transform/vector_transformer_parameter.h" #include "utils/param_compat_macros.h" namespace vsag { @@ -99,7 +100,18 @@ bool TransformQuantizerParameter::CheckCompatibility(const ParamPtr& other) const { PARAM_CAST_OR_RETURN(TransformQuantizerParameter, p, other); CHECK_FIELD_EQ(*this, *p, tq_chain_); - return this->base_quantizer_json_[TYPE_KEY].GetString() == - p->base_quantizer_json_[TYPE_KEY].GetString(); + + auto transformer_param = std::make_shared(); + transformer_param->FromJson(this->base_quantizer_json_); + auto other_transformer_param = std::make_shared(); + other_transformer_param->FromJson(p->base_quantizer_json_); + if (not transformer_param->CheckCompatibility(other_transformer_param)) { + return false; + } + + auto bottom_param = QuantizerParameter::GetQuantizerParameterByJson(this->base_quantizer_json_); + auto other_bottom_param = + QuantizerParameter::GetQuantizerParameterByJson(p->base_quantizer_json_); + return bottom_param->CheckCompatibility(other_bottom_param); } } // namespace vsag diff --git a/src/quantization/transform_quantization/transform_quantizer_parameter_test.cpp b/src/quantization/transform_quantization/transform_quantizer_parameter_test.cpp index 1222c99c9d..0f91fd435a 100644 --- a/src/quantization/transform_quantization/transform_quantizer_parameter_test.cpp +++ b/src/quantization/transform_quantization/transform_quantizer_parameter_test.cpp @@ -62,6 +62,15 @@ TEST_CASE("Transform Quantizer Parameter CheckCompatibility", "[ut][TransformQua TEST_COMPATIBILITY_CASE("different length", param_pca_fp32, param_pca_fht_fp32, false); TEST_COMPATIBILITY_CASE( "different space", param_pca_fht_fp32_no_space, param_pca_fht_fp32, true); + TEST_COMPATIBILITY_CASE("different transform parameter", + R"({"tq_chain":"mrle,rabitq","mrle_dim":64})", + R"({"tq_chain":"mrle,rabitq","mrle_dim":32})", + false); + TEST_COMPATIBILITY_CASE( + "different bottom parameter", + R"({"tq_chain":"mrle,rabitq","mrle_dim":64,"rabitq_bits_per_dim_base":3})", + R"({"tq_chain":"mrle,rabitq","mrle_dim":64,"rabitq_bits_per_dim_base":5})", + false); } TEST_CASE("TQ Parameter ToJson Test", "[ut][TransformQuantizerParameter]") { diff --git a/tests/test_hgraph_rabitq_split.cpp b/tests/test_hgraph_rabitq_split.cpp index f35e14cf58..66aed7708f 100644 --- a/tests/test_hgraph_rabitq_split.cpp +++ b/tests/test_hgraph_rabitq_split.cpp @@ -221,6 +221,24 @@ TEST_CASE("HGraph RaBitQ Split Homogeneous IO", "[ft][rabitq_split][hgraph]") { TestIndex::TestKnnSearch(index, dataset, kSplitSearchParam, 0.1F, true); } +TEST_CASE("HGraph MRLE RaBitQ Split", "[ft][rabitq_split][hgraph][MRLE]") { + using namespace fixtures; + constexpr int64_t dim = 128; + constexpr uint64_t base_count = 600; + + auto param = + HGraphRaBitQSplitTestIndex::GenerateBuildParam("l2", dim, "memory_io", "", 3, 5, true); + auto param_json = vsag::JsonType::Parse(param); + param_json["index_param"]["base_quantization_type"].SetString("tq"); + param_json["index_param"]["tq_chain"].SetString("mrle, rabitq"); + param_json["index_param"]["mrle_dim"].SetInt(64); + + auto index = TestIndex::TestFactory(HGraphRaBitQSplitTestIndex::name, param_json.Dump(), true); + auto dataset = HGraphRaBitQSplitTestIndex::pool.GetDatasetAndCreate(dim, base_count, "l2"); + TestIndex::TestBuildIndex(index, dataset, true); + TestIndex::TestKnnSearch(index, dataset, kSplitSearchParam, 0.05F, true); +} + TEST_CASE("HGraph RaBitQ Split ODescent optimized build", "[ft][rabitq_split][hgraph][odescent]") { using namespace fixtures; constexpr int64_t dim = 128; diff --git a/tests/test_pyramid.cpp b/tests/test_pyramid.cpp index 9896f636b0..8ca9282aaa 100644 --- a/tests/test_pyramid.cpp +++ b/tests/test_pyramid.cpp @@ -274,6 +274,46 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::PyramidTestIndex, REQUIRE(result.value()->GetDim() <= 5); } +TEST_CASE_PERSISTENT_FIXTURE(fixtures::PyramidTestIndex, + "Pyramid MRLE RaBitQ Split", + "[ft][pyramid][rabitq_split][MRLE]") { + constexpr int64_t dim = 64; + constexpr uint64_t count = 256; + + PyramidParam pyramid_param; + pyramid_param.no_build_levels = {0, 1, 2}; + pyramid_param.base_quantization_type = "tq"; + pyramid_param.precise_quantization_type = "rabitq"; + pyramid_param.use_reorder = true; + pyramid_param.rabitq_bits_per_dim_base = 3; + + auto param_json = + vsag::JsonType::Parse(GeneratePyramidBuildParametersString("l2", dim, pyramid_param)); + param_json["index_param"]["tq_chain"].SetString("mrle, rabitq"); + param_json["index_param"]["mrle_dim"].SetInt(32); + param_json["index_param"]["rabitq_bits_per_dim_precise"].SetInt(5); + + auto index = TestFactory("pyramid", param_json.Dump(), true); + auto dataset = pool.GetDatasetAndCreate(dim, count, "l2", /*with_path=*/true); + TestBuildIndex(index, dataset, true); + auto query = fixtures::get_one_query(dataset->query_, 0); + auto result = index->KnnSearch(query, 5, GeneratePyramidSearchParametersString(100)); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() > 0); + REQUIRE(result.value()->GetDim() <= 5); + REQUIRE_NOTHROW(index->GetStats()); + + auto restored = TestFactory("pyramid", param_json.Dump(), true); + TestSerializeBinarySet( + index, restored, dataset, GeneratePyramidSearchParametersString(100), true); + + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + auto streamed = TestFactory("pyramid", param_json.Dump(), true); + REQUIRE(streamed->DeserializeStreaming(stream).has_value()); + REQUIRE_NOTHROW(streamed->GetStats()); +} + TEST_CASE_PERSISTENT_FIXTURE(fixtures::PyramidTestIndex, "Pyramid Duplicate Path Semantics Same Path", "[ft][build][pyramid]") {