From 21e5552477de9d4323d3768ad0bfa91a40a975c5 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Thu, 23 Jul 2026 16:33:04 +0800 Subject: [PATCH 01/12] feat(ivf): support bucket-aligned precise codes Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- .../docs/en/src/advanced/new_serialization.md | 5 +- docs/docs/en/src/indexes/ivf.md | 9 + .../docs/zh/src/advanced/new_serialization.md | 6 +- docs/docs/zh/src/indexes/ivf.md | 7 + include/vsag/constants.h | 3 + src/algorithm/ivf/ivf.cpp | 232 +++++++++++++++--- src/algorithm/ivf/ivf.h | 12 +- src/algorithm/ivf/ivf_parameter.cpp | 26 ++ src/algorithm/ivf/ivf_parameter.h | 1 + src/algorithm/ivf/ivf_parameter_test.cpp | 97 ++++++++ src/constants.cpp | 3 + src/inner_string_params.h | 6 + src/storage/serialization_tags.h | 5 + tests/test_ivf.cpp | 129 +++++++++- 14 files changed, 506 insertions(+), 35 deletions(-) diff --git a/docs/docs/en/src/advanced/new_serialization.md b/docs/docs/en/src/advanced/new_serialization.md index f4389d24cd..3a8681d227 100644 --- a/docs/docs/en/src/advanced/new_serialization.md +++ b/docs/docs/en/src/advanced/new_serialization.md @@ -177,11 +177,14 @@ IVF writes these streaming blocks in order: | `ivf_bucket` | bucket datacell payloads for inverted lists | yes | | `ivf_partition_strategy` | partition strategy state, such as trained centroids | yes | | `label_table` | external labels and label remap | yes | -| `high_precision_codes` | reorder codes when IVF reorder is enabled | conditional | +| `high_precision_codes` | flat reorder codes when IVF reorder uses the `flat` layout | conditional | +| `ivf_precise_bucket` | bucket-aligned reorder codes when IVF reorder uses the `bucket` layout | conditional | | `attribute_filter` | optional attribute filter index | conditional | `DeserializeStreaming` restores the full in-memory IVF index. `Index::Load` can create the IVF index directly from streaming metadata and currently loads all emitted IVF blocks into memory. +The two precise-code blocks are mutually exclusive and are selected by +`precise_codes_layout`. ## SINDI Blocks diff --git a/docs/docs/en/src/indexes/ivf.md b/docs/docs/en/src/indexes/ivf.md index f5b96d99f3..eea5110daf 100644 --- a/docs/docs/en/src/indexes/ivf.md +++ b/docs/docs/en/src/indexes/ivf.md @@ -88,10 +88,19 @@ Build-time parameters live under `index_param`. See | `fast_encode_rabitq_rounds` | int | `6` | CAQ adjustment rounds; allowed range is `[1, 32]` | | `use_reorder` | bool | `false` | Keep a high-precision copy and re-rank after the coarse scan | | `precise_quantization_type` | string | `"fp32"` | Quantizer used for reordering (with `use_reorder: true`) | +| `precise_codes_layout` | string | `"flat"` | Storage layout for precise codes: `"flat"` keeps the legacy one-code-per-vector layout; `"bucket"` mirrors every basic posting in the same bucket and offset | | `base_io_type` | string | `"memory_io"` | Storage backend for coarse codes; supports `uring_io` when built with liburing | | `precise_io_type` | string | `"block_memory_io"` | Storage backend for precise codes (`memory_io`, `block_memory_io`, `mmap_io`, `buffer_io`, `async_io`, `uring_io`, `reader_io`) | | `precise_file_path` | string | `""` | File path when the precise IO type is disk-backed | +`precise_codes_layout: "bucket"` requires `use_reorder: true`. It supports +`memory_io`, `block_memory_io`, `mmap_io`, `buffer_io`, `async_io`, and `uring_io` +(when io_uring is available); +`reader_io` and `pqfs` precise quantization are not supported. When +`buckets_per_data` is greater than one, the precise vector is duplicated for every +basic posting, preserving exact bucket-offset alignment at the corresponding storage +cost. + A rule of thumb for `buckets_count` is `sqrt(N)` to `4 * sqrt(N)` where `N` is the corpus size. diff --git a/docs/docs/zh/src/advanced/new_serialization.md b/docs/docs/zh/src/advanced/new_serialization.md index 1eba4cc7eb..c272785c36 100644 --- a/docs/docs/zh/src/advanced/new_serialization.md +++ b/docs/docs/zh/src/advanced/new_serialization.md @@ -157,11 +157,13 @@ IVF 按顺序写入以下 streaming blocks： | `ivf_bucket` | 倒排列表使用的 bucket datacell 数据 | 是 | | `ivf_partition_strategy` | partition strategy 状态，例如已训练的中心点 | 是 | | `label_table` | 外部 label 和 label remap | 是 | -| `high_precision_codes` | IVF reorder 开启时的 reorder codes | 条件必需 | +| `high_precision_codes` | IVF reorder 使用 `flat` 布局时的精排 codes | 条件必需 | +| `ivf_precise_bucket` | IVF reorder 使用 `bucket` 布局时按桶对齐的精排 codes | 条件必需 | | `attribute_filter` | 开启属性过滤时写入的可选属性过滤索引 | 条件必需 | `DeserializeStreaming` 会恢复完整的内存 IVF 索引。`Index::Load` 可以直接从 streaming metadata -创建 IVF 索引对象，当前会把写出的 IVF blocks 都加载到内存中。 +创建 IVF 索引对象，当前会把写出的 IVF blocks 都加载到内存中。两种精排 codes block 互斥， +由 `precise_codes_layout` 选择。 ## SINDI Blocks diff --git a/docs/docs/zh/src/indexes/ivf.md b/docs/docs/zh/src/indexes/ivf.md index 9ec3a3f2d1..cb4ace5745 100644 --- a/docs/docs/zh/src/indexes/ivf.md +++ b/docs/docs/zh/src/indexes/ivf.md @@ -82,10 +82,17 @@ auto result = index->KnnSearch( | `fast_encode_rabitq_rounds` | int | `6` | CAQ 微调轮数，允许范围 `[1, 32]` | | `use_reorder` | bool | `false` | 是否保留高精度副本用于精排 | | `precise_quantization_type` | string | `"fp32"` | 精排量化类型（`use_reorder: true` 时使用） | +| `precise_codes_layout` | string | `"flat"` | 精排 codes 的存储布局：`"flat"` 保持旧的一向量一码布局；`"bucket"` 为 basic 的每个 posting 在相同 bucket 和 offset 保存一份高精度 code | | `base_io_type` | string | `"memory_io"` | 粗排向量的存储后端；以 liburing 构建时支持 `uring_io` | | `precise_io_type` | string | `"block_memory_io"` | 精排向量的存储后端（`memory_io`、`block_memory_io`、`mmap_io`、`buffer_io`、`async_io`、`uring_io`、`reader_io`） | | `precise_file_path` | string | `""` | 当精排 IO 为磁盘后端时的文件路径 | +`precise_codes_layout: "bucket"` 要求 `use_reorder: true`，支持 `memory_io`、 +`block_memory_io`、`mmap_io`、`buffer_io`、`async_io` 和 `uring_io` +（需要构建环境支持 io_uring），不支持 +`reader_io` 和 `pqfs` 精排量化。当 `buckets_per_data` 大于 1 时，每个 basic posting +都会保存一份高精度向量，从而保持完全相同的 bucket-offset 对齐，同时占用相应倍数的存储空间。 + `buckets_count` 的经验值一般为 `sqrt(N)` ~ `4 * sqrt(N)`，其中 `N` 是语料规模。 ## 检索参数 diff --git a/include/vsag/constants.h b/include/vsag/constants.h index 126dc9a643..40f9a07fce 100644 --- a/include/vsag/constants.h +++ b/include/vsag/constants.h @@ -280,6 +280,9 @@ extern const char* const IVF_PRECISE_CACHE_TOTAL_SIZE; extern const char* const IVF_PRECISE_QUANTIZATION_TYPE; extern const char* const IVF_PRECISE_IO_TYPE; extern const char* const IVF_PRECISE_FILE_PATH; +extern const char* const IVF_PRECISE_CODES_LAYOUT; +extern const char* const IVF_PRECISE_CODES_LAYOUT_FLAT; +extern const char* const IVF_PRECISE_CODES_LAYOUT_BUCKET; extern const char* const USE_ATTRIBUTE_FILTER; extern const char* const IVF_THREAD_COUNT; diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index d5479b4433..f26a963b06 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -58,6 +58,20 @@ namespace vsag { static constexpr BucketIdType INVALID_BUCKET_ID = static_cast(-1); +namespace { + +BucketDataCellParamPtr +make_precise_bucket_param(const IVFParameterPtr& param) { + auto precise_bucket_param = std::make_shared(); + precise_bucket_param->io_parameter = param->precise_codes_param->io_parameter; + precise_bucket_param->quantizer_parameter = param->precise_codes_param->quantizer_parameter; + precise_bucket_param->buckets_count = param->bucket_param->buckets_count; + precise_bucket_param->use_residual_ = false; + return precise_bucket_param; +} + +} // namespace + static constexpr const char* IVF_PARAMS_TEMPLATE = R"( { @@ -97,6 +111,7 @@ static constexpr const char* IVF_PARAMS_TEMPLATE = }, "{BUCKET_PER_DATA_KEY}": 1, "{USE_REORDER_KEY}": false, + "{PRECISE_CODES_LAYOUT_KEY}": "{PRECISE_CODES_LAYOUT_VALUE_FLAT}", "{PRECISE_CODES_KEY}": { "{IO_PARAMS_KEY}": { "{TYPE_KEY}": "{IO_TYPE_VALUE_BLOCK_MEMORY_IO}", @@ -233,6 +248,12 @@ IVF::CheckAndMappingExternalParam(const JsonType& external_param, USE_REORDER_KEY, }, }, + { + IVF_PRECISE_CODES_LAYOUT, + { + PRECISE_CODES_LAYOUT_KEY, + }, + }, { IVF_USE_RESIDUAL, { @@ -406,9 +427,16 @@ IVF::IVF(const IVFParameterPtr& param, const IndexCommonParam& common_param) common_param, param->ivf_partition_strategy_parameter); } if (this->use_reorder_) { - this->reorder_codes_ = - FlattenInterface::MakeInstance(param->precise_codes_param, common_param); - reorder_ = std::make_shared(this->reorder_codes_, allocator_); + if (param->precise_codes_layout == PRECISE_CODES_LAYOUT_VALUE_BUCKET) { + this->precise_bucket_ = + BucketInterface::MakeInstance(make_precise_bucket_param(param), common_param); + CHECK_ARGUMENT(this->precise_bucket_ != nullptr, + "unsupported IO or quantizer for IVF precise bucket"); + } else { + this->reorder_codes_ = + FlattenInterface::MakeInstance(param->precise_codes_param, common_param); + reorder_ = std::make_shared(this->reorder_codes_, allocator_); + } } if (param->bucket_param->use_residual_) { this->bucket_->SetStrategy(partition_strategy_); @@ -477,8 +505,11 @@ IVF::InitFeatures() { } bool has_fp32 = false; - if (use_reorder_ && reorder_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { - has_fp32 = true; + if (use_reorder_) { + const auto precise_quantizer_name = precise_bucket_ != nullptr + ? precise_bucket_->GetQuantizerName() + : reorder_codes_->GetQuantizerName(); + has_fp32 = precise_quantizer_name == QUANTIZATION_TYPE_VALUE_FP32; } if (name == QUANTIZATION_TYPE_VALUE_FP32 or has_fp32) { this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_CAL_DISTANCE_BY_ID); @@ -534,7 +565,11 @@ IVF::Train(const DatasetPtr& data) { const auto* data_ptr = train_data->GetFloat32Vectors(); this->bucket_->Train(data_ptr, sample_count); if (use_reorder_) { - this->reorder_codes_->Train(data->GetFloat32Vectors(), data->GetNumElements()); + if (precise_bucket_ != nullptr) { + this->precise_bucket_->Train(data->GetFloat32Vectors(), data->GetNumElements()); + } else { + this->reorder_codes_->Train(data->GetFloat32Vectors(), data->GetNumElements()); + } } this->is_trained_ = true; } @@ -546,6 +581,9 @@ IVF::Add(const DatasetPtr& base) { throw VsagException(ErrorType::INTERNAL_ERROR, "ivf index add without train error"); } this->bucket_->Unpack(); + if (precise_bucket_ != nullptr) { + this->precise_bucket_->Unpack(); + } auto num_element = base->GetNumElements(); const auto* ids = base->GetIds(); const auto* vectors = base->GetFloat32Vectors(); @@ -559,7 +597,7 @@ IVF::Add(const DatasetPtr& base) { bool need_cal_memory_usage = false; { std::lock_guard lock(label_lookup_mutex_); - if (use_reorder_) { + if (use_reorder_ and precise_bucket_ == nullptr) { this->reorder_codes_->BatchInsertVector(base->GetFloat32Vectors(), base->GetNumElements()); } @@ -579,8 +617,15 @@ IVF::Add(const DatasetPtr& base) { for (int64_t j = 0; j < buckets_per_data_; ++j) { const auto* data_ptr = vectors + i * dim_; auto idx = i * buckets_per_data_ + j; - InnerIdType offset_id = bucket_->InsertVector( - data_ptr, buckets[idx], idx + current_num * buckets_per_data_); + auto posting_id = static_cast(idx + current_num * buckets_per_data_); + InnerIdType offset_id; + if (precise_bucket_ != nullptr) { + // Publish the basic posting only after its precise mirror is ready. + offset_id = precise_bucket_->InsertVector(data_ptr, buckets[idx], posting_id); + bucket_->InsertVectorWithOffset(data_ptr, buckets[idx], posting_id, offset_id); + } else { + offset_id = bucket_->InsertVector(data_ptr, buckets[idx], posting_id); + } if (j == 0) { std::lock_guard lock(label_lookup_mutex_); location_map_[i + current_num] = @@ -614,6 +659,9 @@ IVF::Add(const DatasetPtr& base) { } } this->bucket_->Package(); + if (precise_bucket_ != nullptr) { + this->precise_bucket_->Package(); + } if (need_cal_memory_usage) { this->cal_memory_usage(); } @@ -807,11 +855,17 @@ IVF::GetNumElements() const { void IVF::Merge(const std::vector& merge_units) { this->bucket_->Unpack(); + if (precise_bucket_ != nullptr) { + this->precise_bucket_->Unpack(); + } for (const auto& unit : merge_units) { this->merge_one_unit(unit); } this->fill_location_map(); this->bucket_->Package(); + if (precise_bucket_ != nullptr) { + this->precise_bucket_->Package(); + } } std::pair @@ -867,7 +921,11 @@ IVF::Serialize(StreamWriter& writer) const { WRITE_DATACELL_WITH_NAME(writer, "label_table", label_table_); if (use_reorder_) { - WRITE_DATACELL_WITH_NAME(writer, "reorder_codes", reorder_codes_); + if (precise_bucket_ != nullptr) { + WRITE_DATACELL_WITH_NAME(writer, "precise_bucket", precise_bucket_); + } else { + WRITE_DATACELL_WITH_NAME(writer, "reorder_codes", reorder_codes_); + } } if (use_attribute_filter_) { @@ -954,7 +1012,9 @@ IVF::collect_streaming_header() const { StreamSerializationBlockCurrentVersion(label_tag), StreamSerializationTagCritical(label_tag)); if (this->use_reorder_) { - auto tag = static_cast(StreamSerializationTag::HIGH_PRECISION_CODES); + auto tag = static_cast(precise_bucket_ != nullptr + ? StreamSerializationTag::IVF_PRECISE_BUCKET + : StreamSerializationTag::HIGH_PRECISION_CODES); AppendStreamingManifestBlock(manifest, tag, StreamSerializationBlockCurrentVersion(tag), @@ -999,10 +1059,16 @@ IVF::serialize_streaming_body(StreamWriter& writer) const { this->label_table_->Serialize(w); }); if (this->use_reorder_) { - auto tag = static_cast(StreamSerializationTag::HIGH_PRECISION_CODES); + auto tag = static_cast(precise_bucket_ != nullptr + ? StreamSerializationTag::IVF_PRECISE_BUCKET + : StreamSerializationTag::HIGH_PRECISION_CODES); WriteStreamingBlock( writer, tag, StreamSerializationTagCritical(tag), [this](StreamWriter& w) { - this->reorder_codes_->Serialize(w); + if (this->precise_bucket_ != nullptr) { + this->precise_bucket_->Serialize(w); + } else { + this->reorder_codes_->Serialize(w); + } }); } if (this->use_attribute_filter_) { @@ -1054,6 +1120,10 @@ IVF::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) { this->total_elements_ = basic_info["total_elements"].GetInt(); this->use_reorder_ = basic_info["use_reorder"].GetBool(); this->is_trained_ = basic_info["is_trained"].GetBool(); + if (precise_bucket_ != nullptr and not basic_info.Contains(INDEX_PARAM)) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "IVF precise bucket requires persisted index parameters"); + } if (basic_info.Contains(INDEX_PARAM)) { auto index_param = std::make_shared(); index_param->FromString(basic_info[INDEX_PARAM].GetString()); @@ -1069,7 +1139,7 @@ IVF::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) { bool loaded_bucket = false; bool loaded_partition = false; bool loaded_label_table = false; - bool loaded_reorder_codes = false; + bool loaded_precise_codes = false; bool loaded_attribute_filter = false; while (true) { @@ -1114,12 +1184,21 @@ IVF::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) { loaded_label_table = true; break; case StreamSerializationTag::HIGH_PRECISION_CODES: - if (this->use_reorder_) { + if (this->use_reorder_ and this->reorder_codes_ != nullptr) { ReadSeekableBlockPayload( block_reader, block_header, [this](StreamReader& block) { this->reorder_codes_->Deserialize(block); }); - loaded_reorder_codes = true; + loaded_precise_codes = true; + } + break; + case StreamSerializationTag::IVF_PRECISE_BUCKET: + if (this->use_reorder_ and this->precise_bucket_ != nullptr) { + ReadSeekableBlockPayload( + block_reader, block_header, [this](StreamReader& block) { + this->precise_bucket_->Deserialize(block); + }); + loaded_precise_codes = true; } break; case StreamSerializationTag::ATTRIBUTE_FILTER: @@ -1213,7 +1292,7 @@ IVF::read_streaming_body(StreamReader& reader, const MetadataPtr& metadata) { throw VsagException(ErrorType::READ_ERROR, "IVF streaming serialization required block is missing"); } - if (this->use_reorder_ && !loaded_reorder_codes) { + if (this->use_reorder_ && !loaded_precise_codes) { throw VsagException(ErrorType::READ_ERROR, "IVF streaming serialization reorder block is missing"); } @@ -1242,6 +1321,10 @@ IVF::Deserialize(StreamReader& reader) { &reader, std::numeric_limits::max(), this->allocator_); if (footer == nullptr) { // old format, DON'T EDIT, remove in the future + if (precise_bucket_ != nullptr) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "legacy IVF serialization does not support precise bucket"); + } logger::debug("parse with v0.14 version format"); StreamReader::ReadObj(buffer_reader, this->total_elements_); @@ -1271,6 +1354,10 @@ IVF::Deserialize(StreamReader& reader) { this->total_elements_ = basic_info["total_elements"].GetInt(); this->use_reorder_ = basic_info["use_reorder"].GetBool(); this->is_trained_ = basic_info["is_trained"].GetBool(); + if (precise_bucket_ != nullptr and not basic_info.Contains(INDEX_PARAM)) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "IVF precise bucket requires persisted index parameters"); + } if (basic_info.Contains(INDEX_PARAM)) { auto param_str = basic_info[INDEX_PARAM].GetString(); auto index_param = std::make_shared(); @@ -1293,7 +1380,11 @@ IVF::Deserialize(StreamReader& reader) { READ_DATACELL_WITH_NAME(buffer_reader, "partition_strategy", this->partition_strategy_); READ_DATACELL_WITH_NAME(buffer_reader, "label_table", this->label_table_); if (use_reorder_) { - READ_DATACELL_WITH_NAME(buffer_reader, "reorder_codes", this->reorder_codes_); + if (precise_bucket_ != nullptr) { + READ_DATACELL_WITH_NAME(buffer_reader, "precise_bucket", this->precise_bucket_); + } else { + READ_DATACELL_WITH_NAME(buffer_reader, "reorder_codes", this->reorder_codes_); + } } if (use_attribute_filter_) { READ_DATACELL_WITH_NAME(buffer_reader, "attr_filter_index", this->attr_filter_index_); @@ -1459,20 +1550,73 @@ IVF::reorder(int64_t topk, QueryContext& ctx, ReasoningContext* reasoning_ctx, const std::optional& distance_threshold) const { - auto reorder_heap = - reorder_->Reorder(input, query, topk, ctx, nullptr, nullptr, distance_threshold); + auto reorder_heap = precise_bucket_ != nullptr + ? this->reorder_with_precise_bucket( + input, query, topk, ctx, distance_threshold) + : reorder_->Reorder( + input, query, topk, ctx, nullptr, nullptr, distance_threshold); auto dataset_results = this->pack_knn_result(reorder_heap, ctx.alloc); return dataset_results; } +DistHeapPtr +IVF::reorder_with_precise_bucket(const DistHeapPtr& input, + const float* query, + int64_t topk, + QueryContext& ctx, + const std::optional& distance_threshold) const { + Allocator* query_allocator = select_query_allocator(ctx.alloc, allocator_); + const uint64_t candidate_count = input == nullptr ? 0 : input->Size(); + topk = std::min(topk, static_cast(candidate_count)); + auto reorder_heap = std::make_shared>(query_allocator, topk); + if (candidate_count == 0 or topk == 0) { + return reorder_heap; + } + + if (ctx.stats != nullptr) { + ctx.stats->reorder_distance_count.fetch_add(static_cast(candidate_count), + std::memory_order_relaxed); + } + + auto computer = precise_bucket_->FactoryComputer(query); + const auto* candidates = input->GetData(); + for (uint64_t i = 0; i < candidate_count; ++i) { + const auto [coarse_distance, inner_id] = candidates[i]; + const auto [bucket_id, offset_id] = this->get_location(inner_id); + auto precise_distance = precise_bucket_->QueryOneById(computer, bucket_id, offset_id); + if (ctx.reasoning_ctx != nullptr) { + ctx.reasoning_ctx->RecordReorder(inner_id, coarse_distance, precise_distance); + } + if (distance_threshold.has_value() and + (not std::isfinite(precise_distance) or + precise_distance > distance_threshold.value())) { + continue; + } + if (reorder_heap->Size() < topk or precise_distance < reorder_heap->Top().first) { + reorder_heap->Push(precise_distance, inner_id); + if (reorder_heap->Size() > topk) { + if (ctx.reasoning_ctx != nullptr) { + ctx.reasoning_ctx->RecordReorderEviction(reorder_heap->Top().second, 0); + } + reorder_heap->Pop(); + } + } + } + return reorder_heap; +} + InnerIndexPtr IVF::ExportModel(const IndexCommonParam& param) const { auto index = std::make_shared(this->create_param_ptr_, param); IVFPartitionStrategy::Clone(this->partition_strategy_, index->partition_strategy_); this->bucket_->ExportModel(index->bucket_); if (use_reorder_) { - this->reorder_codes_->ExportModel(index->reorder_codes_); + if (precise_bucket_ != nullptr) { + this->precise_bucket_->ExportModel(index->precise_bucket_); + } else { + this->reorder_codes_->ExportModel(index->reorder_codes_); + } } index->is_trained_ = this->is_trained_; return index; @@ -1765,7 +1909,13 @@ IVF::merge_one_unit(const MergeUnit& unit) { other_index->bucket_->Package(); if (this->use_reorder_) { - this->reorder_codes_->MergeOther(other_index->reorder_codes_, this->total_elements_); + if (precise_bucket_ != nullptr) { + other_index->precise_bucket_->Unpack(); + this->precise_bucket_->MergeOther(other_index->precise_bucket_, bucket_bias); + other_index->precise_bucket_->Package(); + } else { + this->reorder_codes_->MergeOther(other_index->reorder_codes_, this->total_elements_); + } } this->total_elements_ += other_index->total_elements_; } @@ -1787,6 +1937,10 @@ IVF::check_merge_illegal(const vsag::MergeUnit& unit) const { this->use_reorder_, other_ivf_index->use_reorder_)); } + if ((other_ivf_index->precise_bucket_ == nullptr) != (this->precise_bucket_ == nullptr)) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "Merge Failed: IVF precise codes layout does not match"); + } auto cur_model = this->ExportModel(index->GetCommonParam()); std::stringstream ss1; std::stringstream ss2; @@ -2014,10 +2168,27 @@ void IVF::fill_location_map() { this->location_map_.resize(this->total_elements_ * buckets_per_data_); auto bucket_count = this->bucket_->bucket_count_; + if (precise_bucket_ != nullptr and precise_bucket_->GetBucketCount() != bucket_count) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "basic and precise bucket counts do not match"); + } for (BucketIdType i = 0; i < bucket_count; ++i) { auto* ids = this->bucket_->GetInnerIds(i); auto bucket_size = this->bucket_->GetBucketSize(i); + InnerIdType* precise_ids = nullptr; + if (precise_bucket_ != nullptr) { + auto precise_bucket_size = precise_bucket_->GetBucketSize(i); + if (precise_bucket_size != bucket_size) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "basic and precise bucket sizes do not match"); + } + precise_ids = precise_bucket_->GetInnerIds(i); + } for (uint64_t j = 0; j < bucket_size; ++j) { + if (precise_ids != nullptr and precise_ids[j] != ids[j]) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "basic and precise bucket inner ids do not match"); + } if (ids[j] == std::numeric_limits::max()) { continue; } @@ -2077,15 +2248,16 @@ IVF::CalDistanceById(const float* query, } } } - if (this->use_reorder_ && calculate_precise_distance) { + if (this->use_reorder_ && calculate_precise_distance && reorder_codes_ != nullptr) { auto computer = this->reorder_codes_->FactoryComputer(query); this->reorder_codes_->Query(distances, computer, inner_ids.data(), count); } else { - auto computer = this->bucket_->FactoryComputer(query); + auto codes = this->use_reorder_ && calculate_precise_distance ? precise_bucket_ : bucket_; + auto computer = codes->FactoryComputer(query); for (int64_t i = 0; i < count; ++i) { if (validity[i]) { auto [bucket_id, offset_id] = this->get_location(inner_ids[i]); - distances[i] = this->bucket_->QueryOneById(computer, bucket_id, offset_id); + distances[i] = codes->QueryOneById(computer, bucket_id, offset_id); } } } @@ -2107,15 +2279,16 @@ IVF::CalcDistanceById(const float* query, int64_t id, bool calculate_precise_dis if (not success) { return -1.0F; } - if (this->use_reorder_ && calculate_precise_distance) { + if (this->use_reorder_ && calculate_precise_distance && reorder_codes_ != nullptr) { float dist = 0.0F; auto computer = this->reorder_codes_->FactoryComputer(query); this->reorder_codes_->Query(&dist, computer, &inner_id, 1); return dist; } - auto computer = this->bucket_->FactoryComputer(query); + auto codes = this->use_reorder_ && calculate_precise_distance ? precise_bucket_ : bucket_; + auto computer = codes->FactoryComputer(query); auto [bucket_id, offset_id] = this->get_location(inner_id); - return this->bucket_->QueryOneById(computer, bucket_id, offset_id); + return codes->QueryOneById(computer, bucket_id, offset_id); } void @@ -2224,7 +2397,8 @@ IVF::cal_memory_usage() { auto memory = sizeof(IVF); memory += this->bucket_->GetMemoryUsage(); if (use_reorder_) { - memory += this->reorder_codes_->GetMemoryUsage(); + memory += precise_bucket_ != nullptr ? precise_bucket_->GetMemoryUsage() + : reorder_codes_->GetMemoryUsage(); } if (this->extra_info_size_ > 0 and this->extra_infos_ != nullptr) { memory += this->extra_infos_->GetMemoryUsage(); diff --git a/src/algorithm/ivf/ivf.h b/src/algorithm/ivf/ivf.h index 987b571da3..fefe3f54ac 100644 --- a/src/algorithm/ivf/ivf.h +++ b/src/algorithm/ivf/ivf.h @@ -227,6 +227,13 @@ class IVF : public InnerIndexInterface { ReasoningContext* reasoning_ctx = nullptr, const std::optional& distance_threshold = std::nullopt) const; + DistHeapPtr + reorder_with_precise_bucket(const DistHeapPtr& input, + const float* query, + int64_t topk, + QueryContext& ctx, + const std::optional& distance_threshold) const; + void AttachReasoningReport(const DatasetPtr& dataset_results, ReasoningContext* reasoning_ctx) const; @@ -287,8 +294,9 @@ class IVF : public InnerIndexInterface { int64_t total_elements_{0}; // total inserted (incl. deleted) bool is_trained_{false}; // true after Train() succeeds - FlattenInterfacePtr reorder_codes_{nullptr}; // high-precision codes for reranking - ReorderInterfacePtr reorder_{nullptr}; // reordering engine + FlattenInterfacePtr reorder_codes_{nullptr}; // legacy high-precision flat codes + BucketInterfacePtr precise_bucket_{nullptr}; // high-precision codes mirroring basic buckets + ReorderInterfacePtr reorder_{nullptr}; // flat-code reordering engine std::shared_ptr thread_pool_{nullptr}; // for parallel bucket scans diff --git a/src/algorithm/ivf/ivf_parameter.cpp b/src/algorithm/ivf/ivf_parameter.cpp index cbc4eca8b1..4b0c6eac6c 100644 --- a/src/algorithm/ivf/ivf_parameter.cpp +++ b/src/algorithm/ivf/ivf_parameter.cpp @@ -26,6 +26,30 @@ void IVFParameter::FromJson(const JsonType& json) { InnerIndexParameter::FromJson(json); + this->precise_codes_layout = PRECISE_CODES_LAYOUT_VALUE_FLAT; + if (json.Contains(PRECISE_CODES_LAYOUT_KEY)) { + this->precise_codes_layout = json[PRECISE_CODES_LAYOUT_KEY].GetString(); + } + CHECK_ARGUMENT( + this->precise_codes_layout == PRECISE_CODES_LAYOUT_VALUE_FLAT || + this->precise_codes_layout == PRECISE_CODES_LAYOUT_VALUE_BUCKET, + fmt::format("invalid precise_codes_layout: {}, supported values are \"{}\" and \"{}\"", + this->precise_codes_layout, + PRECISE_CODES_LAYOUT_VALUE_FLAT, + PRECISE_CODES_LAYOUT_VALUE_BUCKET)); + + if (this->precise_codes_layout == PRECISE_CODES_LAYOUT_VALUE_BUCKET) { + CHECK_ARGUMENT(this->use_reorder, "precise_codes_layout=bucket requires use_reorder=true"); + CHECK_ARGUMENT(this->reorder_source == HGRAPH_REORDER_SOURCE_PRECISE, + "precise_codes_layout=bucket requires reorder_source=precise"); + CHECK_ARGUMENT(this->precise_codes_param != nullptr && + this->precise_codes_param->name == FLATTEN_DATA_CELL, + "precise_codes_layout=bucket requires ordinary flatten precise_codes"); + CHECK_ARGUMENT(this->precise_codes_param->quantizer_parameter->GetTypeName() != + QUANTIZATION_TYPE_VALUE_PQFS, + "precise_codes_layout=bucket does not support pqfs precise quantization"); + } + if (json.Contains(BUCKET_PER_DATA_KEY)) { this->buckets_per_data = static_cast(json[BUCKET_PER_DATA_KEY].GetInt()); } @@ -77,6 +101,7 @@ IVFParameter::ToJson() const { json[BUCKET_PARAMS_KEY].SetJson(this->bucket_param->ToJson()); json[IVF_PARTITION_STRATEGY_PARAMS_KEY].SetJson( this->ivf_partition_strategy_parameter->ToJson()); + json[PRECISE_CODES_LAYOUT_KEY].SetString(this->precise_codes_layout); json[BUCKET_PER_DATA_KEY].SetInt(this->buckets_per_data); json[GRAPH_BUILD_THRESHOLD_KEY].SetInt(this->graph_build_threshold); return json; @@ -87,6 +112,7 @@ IVFParameter::CheckCompatibility(const ParamPtr& other) const { return false; } PARAM_CAST_OR_RETURN(IVFParameter, p, other); + CHECK_FIELD_EQ(*this, *p, precise_codes_layout); CHECK_FIELD_EQ(*this, *p, buckets_per_data); CHECK_FIELD_EQ(*this, *p, graph_build_threshold); CHECK_SUB_PARAM(*this, *p, bucket_param); diff --git a/src/algorithm/ivf/ivf_parameter.h b/src/algorithm/ivf/ivf_parameter.h index e8b5424123..1d627803e9 100644 --- a/src/algorithm/ivf/ivf_parameter.h +++ b/src/algorithm/ivf/ivf_parameter.h @@ -46,6 +46,7 @@ class IVFParameter : public InnerIndexParameter { public: BucketDataCellParamPtr bucket_param{nullptr}; IVFPartitionStrategyParametersPtr ivf_partition_strategy_parameter{nullptr}; + std::string precise_codes_layout{PRECISE_CODES_LAYOUT_VALUE_FLAT}; BucketIdType buckets_per_data{1}; int64_t train_sample_count{65536L}; GraphInterfaceParamPtr graph_param{nullptr}; diff --git a/src/algorithm/ivf/ivf_parameter_test.cpp b/src/algorithm/ivf/ivf_parameter_test.cpp index 143d9fa83d..05b1650816 100644 --- a/src/algorithm/ivf/ivf_parameter_test.cpp +++ b/src/algorithm/ivf/ivf_parameter_test.cpp @@ -32,6 +32,7 @@ struct IVFDefaultParam { bool use_reorder = true; std::string precise_codes_io_type = "block_memory_io"; std::string precise_codes_quantization_type = "fp32"; + std::string precise_codes_layout = "flat"; std::string partition_strategy_type = "ivf"; std::string ivf_train_type = "kmeans"; int buckets_per_data = 1; @@ -54,6 +55,7 @@ generate_ivf_param(const IVFDefaultParam& param) { "use_residual": {} }}, "use_reorder": {}, + "precise_codes_layout": "{}", "partition_strategy": {{ "partition_strategy_type": "{}", "ivf_train_type": "{}", @@ -79,6 +81,7 @@ generate_ivf_param(const IVFDefaultParam& param) { param.buckets_count, param.use_residual, param.use_reorder, + param.precise_codes_layout, param.partition_strategy_type, param.ivf_train_type, param.precise_codes_io_type, @@ -103,6 +106,7 @@ TEST_CASE("IVF Parameters Test", "[ut][IVFParameter]") { REQUIRE(param->use_reorder == true); REQUIRE(param->build_thread_count == 3); REQUIRE(param->precise_codes_param->quantizer_parameter->GetTypeName() == "fp32"); + REQUIRE(param->precise_codes_layout == "flat"); REQUIRE(param->train_sample_count == 65536L); index_param.ivf_train_type = "random"; @@ -147,6 +151,83 @@ TEST_CASE("IVF Parameters Test", "[ut][IVFParameter]") { REQUIRE(search_param.first_order_scan_ratio == 0.1f); } +TEST_CASE("IVF precise codes layout parameter", "[ut][IVFParameter]") { + SECTION("missing layout defaults to flat") { + IVFDefaultParam index_param; + auto param_json = vsag::JsonType::Parse(generate_ivf_param(index_param)); + param_json.Erase(vsag::PRECISE_CODES_LAYOUT_KEY); + + auto param = std::make_shared(); + param->FromJson(param_json); + + REQUIRE(param->precise_codes_layout == vsag::PRECISE_CODES_LAYOUT_VALUE_FLAT); + REQUIRE(param->ToJson()[vsag::PRECISE_CODES_LAYOUT_KEY].GetString() == + vsag::PRECISE_CODES_LAYOUT_VALUE_FLAT); + } + + SECTION("bucket layout supports multiple postings per data") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + index_param.buckets_per_data = 3; + + auto param = std::make_shared(); + param->FromString(generate_ivf_param(index_param)); + + REQUIRE(param->precise_codes_layout == vsag::PRECISE_CODES_LAYOUT_VALUE_BUCKET); + REQUIRE(param->buckets_per_data == 3); + vsag::ParameterTest::TestToJson(param); + } + + SECTION("reject invalid layout") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "invalid"; + auto param = std::make_shared(); + REQUIRE_THROWS(param->FromString(generate_ivf_param(index_param))); + } + + SECTION("bucket layout requires reorder") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + index_param.use_reorder = false; + auto param = std::make_shared(); + REQUIRE_THROWS(param->FromString(generate_ivf_param(index_param))); + } + + SECTION("bucket layout requires precise reorder source") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + auto param_json = vsag::JsonType::Parse(generate_ivf_param(index_param)); + param_json[vsag::REORDER_SOURCE_KEY].SetString(vsag::HGRAPH_REORDER_SOURCE_BASE); + + auto param = std::make_shared(); + REQUIRE_THROWS(param->FromJson(param_json)); + } + + SECTION("bucket layout requires ordinary flatten precise codes") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + auto param_json = vsag::JsonType::Parse(generate_ivf_param(index_param)); + param_json[vsag::PRECISE_CODES_KEY][vsag::CODES_TYPE_KEY].SetString( + vsag::RABITQ_SPLIT_CODES); + param_json[vsag::PRECISE_CODES_KEY][vsag::QUANTIZATION_PARAMS_KEY][vsag::TYPE_KEY] + .SetString(vsag::QUANTIZATION_TYPE_VALUE_RABITQ); + + auto param = std::make_shared(); + REQUIRE_THROWS(param->FromJson(param_json)); + } + + SECTION("bucket layout rejects pqfs") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + auto param_json = vsag::JsonType::Parse(generate_ivf_param(index_param)); + param_json[vsag::PRECISE_CODES_KEY][vsag::QUANTIZATION_PARAMS_KEY][vsag::TYPE_KEY] + .SetString(vsag::QUANTIZATION_TYPE_VALUE_PQFS); + + auto param = std::make_shared(); + REQUIRE_THROWS(param->FromJson(param_json)); + } +} + TEST_CASE("IVF maps RabitQ external parameters", "[ut][IVFParameter]") { auto external_param = vsag::JsonType::Parse(R"({ "base_quantization_type": "rabitq", @@ -207,6 +288,20 @@ TEST_CASE("IVF Parameters CheckCompatibility", "[ut][IVFParameter][CheckCompatib REQUIRE_FALSE(param->CheckCompatibility(std::make_shared())); } + SECTION("missing layout is compatible with explicit flat layout") { + IVFDefaultParam index_param; + auto legacy_json = vsag::JsonType::Parse(generate_ivf_param(index_param)); + legacy_json.Erase(vsag::PRECISE_CODES_LAYOUT_KEY); + + auto legacy_param = std::make_shared(); + legacy_param->FromJson(legacy_json); + auto explicit_flat_param = std::make_shared(); + explicit_flat_param->FromString(generate_ivf_param(index_param)); + + REQUIRE(legacy_param->CheckCompatibility(explicit_flat_param)); + REQUIRE(explicit_flat_param->CheckCompatibility(legacy_param)); + } + TEST_COMPATIBILITY_CASE("ivf buckets_count", buckets_count, 3, 4, false); TEST_COMPATIBILITY_CASE( "ivf bucket io type", buckect_io_type, "block_memory_io", "memory_io", true); @@ -224,6 +319,8 @@ TEST_CASE("IVF Parameters CheckCompatibility", "[ut][IVFParameter][CheckCompatib TEST_COMPATIBILITY_CASE( "ivf partition_strategy_type", partition_strategy_type, "ivf", "gno_imi", false); TEST_COMPATIBILITY_CASE("ivf ivf_train_type", ivf_train_type, "kmeans", "random", true); + TEST_COMPATIBILITY_CASE( + "ivf precise_codes_layout", precise_codes_layout, "flat", "bucket", false); TEST_COMPATIBILITY_CASE("ivf buckets_per_data", buckets_per_data, 3, 2, false); TEST_COMPATIBILITY_CASE("ivf use_attribute_filter", use_attribute_filter, true, false, false); } diff --git a/src/constants.cpp b/src/constants.cpp index 76b59f85da..3c7551c7e9 100644 --- a/src/constants.cpp +++ b/src/constants.cpp @@ -285,6 +285,9 @@ const char* const GNO_IMI_SECOND_ORDER_BUCKETS_COUNT = "second_order_buckets_cou const char* const IVF_PRECISE_QUANTIZATION_TYPE = "precise_quantization_type"; const char* const IVF_PRECISE_IO_TYPE = "precise_io_type"; const char* const IVF_PRECISE_FILE_PATH = "precise_file_path"; +const char* const IVF_PRECISE_CODES_LAYOUT = "precise_codes_layout"; +const char* const IVF_PRECISE_CODES_LAYOUT_FLAT = "flat"; +const char* const IVF_PRECISE_CODES_LAYOUT_BUCKET = "bucket"; const char* const USE_ATTRIBUTE_FILTER = "use_attribute_filter"; const char* const IVF_THREAD_COUNT = "thread_count"; diff --git a/src/inner_string_params.h b/src/inner_string_params.h index 6c2115a0f8..29c6992598 100644 --- a/src/inner_string_params.h +++ b/src/inner_string_params.h @@ -38,6 +38,9 @@ const char* const BUILD_THREAD_COUNT_KEY = "build_thread_count"; const char* const LABEL_REMAP_TYPE_KEY = "label_remap_type"; const char* const BASE_CODES_KEY = "base_codes"; const char* const PRECISE_CODES_KEY = "precise_codes"; +const char* const PRECISE_CODES_LAYOUT_KEY = "precise_codes_layout"; +const char* const PRECISE_CODES_LAYOUT_VALUE_FLAT = "flat"; +const char* const PRECISE_CODES_LAYOUT_VALUE_BUCKET = "bucket"; const char* const STORE_RAW_VECTOR_KEY = "store_raw_vector"; const char* const RAW_VECTOR_KEY = "raw_vector"; const char* const ATTR_HAS_BUCKETS_KEY = "has_buckets"; @@ -225,6 +228,9 @@ const std::unordered_map DEFAULT_MAP = { {"GRAPH_KEY", GRAPH_KEY}, {"BASE_CODES_KEY", BASE_CODES_KEY}, {"PRECISE_CODES_KEY", PRECISE_CODES_KEY}, + {"PRECISE_CODES_LAYOUT_KEY", PRECISE_CODES_LAYOUT_KEY}, + {"PRECISE_CODES_LAYOUT_VALUE_FLAT", PRECISE_CODES_LAYOUT_VALUE_FLAT}, + {"PRECISE_CODES_LAYOUT_VALUE_BUCKET", PRECISE_CODES_LAYOUT_VALUE_BUCKET}, {"HGRAPH_SUPPORT_DUPLICATE", HGRAPH_SUPPORT_DUPLICATE}, {"HGRAPH_DEDUPLICATE_STORAGE", HGRAPH_DEDUPLICATE_STORAGE}, {"HGRAPH_DUPLICATE_DISTANCE_THRESHOLD", HGRAPH_DUPLICATE_DISTANCE_THRESHOLD}, diff --git a/src/storage/serialization_tags.h b/src/storage/serialization_tags.h index e9b021f631..44c6534185 100644 --- a/src/storage/serialization_tags.h +++ b/src/storage/serialization_tags.h @@ -40,6 +40,7 @@ enum class StreamSerializationTag : uint32_t { PYRAMID_HIERARCHIES = 14, CODE_SLOT_MAP = 15, IVF_BUCKET_GRAPH = 16, + IVF_PRECISE_BUCKET = 17, }; inline const char* @@ -79,6 +80,8 @@ StreamSerializationTagName(uint32_t tag) { return "code_slot_map"; case StreamSerializationTag::IVF_BUCKET_GRAPH: return "ivf_bucket_graph"; + case StreamSerializationTag::IVF_PRECISE_BUCKET: + return "ivf_precise_bucket"; } return "unknown"; } @@ -100,6 +103,7 @@ StreamSerializationTagCritical(uint32_t tag) { case StreamSerializationTag::SINDI_TERM_ID_MAPPER: case StreamSerializationTag::PYRAMID_HIERARCHIES: case StreamSerializationTag::CODE_SLOT_MAP: + case StreamSerializationTag::IVF_PRECISE_BUCKET: return true; case StreamSerializationTag::ATTRIBUTE_FILTER: case StreamSerializationTag::EXTRA_INFO: @@ -131,6 +135,7 @@ StreamSerializationBlockCurrentVersion(uint32_t tag) { case StreamSerializationTag::SINDI_TERM_ID_MAPPER: case StreamSerializationTag::PYRAMID_HIERARCHIES: case StreamSerializationTag::CODE_SLOT_MAP: + case StreamSerializationTag::IVF_PRECISE_BUCKET: return kStreamSerializationBlockVersionV1; case StreamSerializationTag::IVF_BUCKET_GRAPH: return kStreamSerializationBlockVersionV1; diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 7c0125c987..3f79e93d5c 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -15,6 +15,7 @@ #include #include +#include #include #include "functest.h" @@ -192,6 +193,42 @@ namespace { using vsag::test::EraseStreamingBlock; using vsag::test::InsertUnknownStreamingBlock; +std::string +GenerateBucketPreciseParameters(int buckets_per_data, + const std::string& precise_io_type = "block_memory_io", + const std::string& precise_file_path = "", + int thread_count = 1) { + auto params = nlohmann::json::parse(IVFTestIndex::GenerateIVFBuildParametersString( + "l2", 16, "sq8,fp32", 16, "random", false, buckets_per_data, false, thread_count)); + params["index_param"]["precise_codes_layout"] = "bucket"; + params["index_param"]["precise_io_type"] = precise_io_type; + params["index_param"]["precise_file_path"] = precise_file_path; + return params.dump(); +} + +void +CheckBucketPreciseIndex(const TestIndex::IndexPtr& index, + const TestDatasetPtr& dataset, + const std::string& search_param) { + constexpr int64_t query_id = 3; + auto query = vsag::Dataset::Make(); + query->NumElements(1) + ->Dim(dataset->base_->GetDim()) + ->Float32Vectors(dataset->base_->GetFloat32Vectors() + query_id * dataset->base_->GetDim()) + ->Owner(false); + + auto result = index->KnnSearch(query, 1, search_param); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 1); + REQUIRE(result.value()->GetIds()[0] == dataset->base_->GetIds()[query_id]); + REQUIRE(std::abs(result.value()->GetDistances()[0]) < 2e-6F); + + auto distance = + index->CalcDistanceById(query->GetFloat32Vectors(), dataset->base_->GetIds()[query_id]); + REQUIRE(distance.has_value()); + REQUIRE(std::abs(distance.value()) < 2e-6F); +} + } // namespace TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, "IVF GetStatus", "[ft][ivf]") { @@ -311,7 +348,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, auto origin_size = vsag::Options::Instance().block_size_limit(); vsag::Options::Instance().set_block_size_limit(1024 * 1024 * 2); - auto param = IVFTestIndex::GenerateIVFBuildParametersString("l2", 16, "sq8", 32, "random"); + auto param = IVFTestIndex::GenerateIVFBuildParametersString("l2", 16, "sq8,fp32", 32, "random"); auto index = TestIndex::TestFactory(IVFTestIndex::name, param, true); auto dataset = IVFTestIndex::pool.GetDatasetAndCreate(16, 200, "l2"); TestIndex::TestBuildIndex(index, dataset, true); @@ -356,6 +393,96 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, vsag::Options::Instance().set_block_size_limit(origin_size); } + +TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, + "IVF bucket precise mirrors basic postings", + "[ft][ivf][reorder][serialize][pr]") { + constexpr int64_t dim = 16; + constexpr int64_t base_count = 128; + constexpr int64_t buckets_count = 16; + const auto buckets_per_data = GENERATE(1, 2); + const auto precise_io_type = GENERATE("block_memory_io", "buffer_io"); + INFO(fmt::format("buckets_per_data: {}", buckets_per_data)); + INFO(fmt::format("precise_io_type: {}", precise_io_type)); + + const auto precise_file_path = + precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); + const auto params = + GenerateBucketPreciseParameters(buckets_per_data, precise_io_type, precise_file_path); + const auto search_param = fmt::format(search_param_tmp, buckets_count); + auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); + auto index = TestFactory(name, params, true); + + TestContinueAdd(index, dataset, true); + CheckBucketPreciseIndex(index, dataset, search_param); + TestCalcDistanceById(index, dataset, 2e-6F, true); + TestBatchCalcDistanceById(index, dataset, 2e-6F, true); + + const auto restored_file_path = + precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); + const auto restored_params = + GenerateBucketPreciseParameters(buckets_per_data, precise_io_type, restored_file_path); + auto restored = TestFactory(name, restored_params, true); + TestSerializeBinarySet(index, restored, dataset, search_param, true); + CheckBucketPreciseIndex(restored, dataset, search_param); + TestCalcDistanceById(restored, dataset, 2e-6F, true); + TestBatchCalcDistanceById(restored, dataset, 2e-6F, true); +} + +TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, + "IVF bucket precise streaming", + "[ft][ivf][reorder][serialize][streaming][pr]") { + constexpr int64_t dim = 16; + constexpr int64_t base_count = 128; + constexpr int64_t buckets_count = 16; + const auto params = GenerateBucketPreciseParameters(2, "block_memory_io", "", 4); + const auto search_param = fmt::format(search_param_tmp, buckets_count); + auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); + auto index = TestFactory(name, params, true); + TestBuildIndex(index, dataset, true); + + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + const auto bytes = stream.str(); + + auto restored = TestFactory(name, params, true); + std::stringstream deserialize_stream(bytes); + REQUIRE(restored->DeserializeStreaming(deserialize_stream).has_value()); + CheckBucketPreciseIndex(restored, dataset, search_param); + TestBatchCalcDistanceById(restored, dataset, 2e-6F, true); + + std::stringstream load_stream(bytes); + auto loaded = vsag::Index::Load(load_stream, "{}"); + REQUIRE(loaded.has_value()); + CheckBucketPreciseIndex(loaded.value(), dataset, search_param); + TestBatchCalcDistanceById(loaded.value(), dataset, 2e-6F, true); + + auto missing_precise = + EraseStreamingBlock(bytes, vsag::StreamSerializationTag::IVF_PRECISE_BUCKET); + auto invalid_restored = TestFactory(name, params, true); + std::stringstream missing_deserialize_stream(missing_precise); + REQUIRE_FALSE(invalid_restored->DeserializeStreaming(missing_deserialize_stream).has_value()); + std::stringstream missing_load_stream(missing_precise); + REQUIRE_FALSE(vsag::Index::Load(missing_load_stream, "{}").has_value()); +} + +TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, + "IVF bucket precise merge", + "[ft][ivf][reorder][merge][pr]") { + constexpr int64_t dim = 16; + constexpr int64_t base_count = 128; + constexpr int64_t buckets_count = 16; + const auto params = GenerateBucketPreciseParameters(2); + const auto search_param = fmt::format(search_param_tmp, buckets_count); + auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); + auto model = TestFactory(name, params, true); + REQUIRE(model->Train(dataset->base_).has_value()); + + auto merged = TestMergeIndexWithSameModel(model, dataset, 3, true); + CheckBucketPreciseIndex(merged, dataset, search_param); + TestCalcDistanceById(merged, dataset, 2e-6F, true); + TestBatchCalcDistanceById(merged, dataset, 2e-6F, true); +} } // namespace fixtures static void From 8b6fa287eab975860650bd226314c7850842591d Mon Sep 17 00:00:00 2001 From: jac0626 Date: Thu, 23 Jul 2026 16:40:42 +0800 Subject: [PATCH 02/12] docs(ivf): add precise bucket configuration example Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- docs/docs/en/src/indexes/ivf.md | 31 +++++++++++++++++++++++++++++++ docs/docs/zh/src/indexes/ivf.md | 30 ++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/docs/docs/en/src/indexes/ivf.md b/docs/docs/en/src/indexes/ivf.md index eea5110daf..703e9cd5db 100644 --- a/docs/docs/en/src/indexes/ivf.md +++ b/docs/docs/en/src/indexes/ivf.md @@ -64,6 +64,37 @@ auto result = index->KnnSearch( R"({"ivf": {"scan_buckets_count": 16}})").value(); ``` +### Bucket-aligned precise codes + +For disk-backed precise reordering, enable the bucket layout while keeping the +existing precise quantization and IO parameters: + +```cpp +std::string params = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 128, + "index_param": { + "buckets_count": 256, + "buckets_per_data": 1, + "base_quantization_type": "sq8", + "partition_strategy_type": "ivf", + "ivf_train_type": "kmeans", + + "use_reorder": true, + "precise_quantization_type": "fp32", + "precise_codes_layout": "bucket", + "precise_io_type": "buffer_io", + "precise_file_path": "/data/ivf-precise.bin" + } +})"; + +auto index = vsag::Factory::CreateIndex("ivf", params).value(); +``` + +`precise_codes_layout` is the only new parameter in this example. Build and +search use the same APIs and search parameters as the regular IVF quick start. + ## Build parameters Build-time parameters live under `index_param`. See diff --git a/docs/docs/zh/src/indexes/ivf.md b/docs/docs/zh/src/indexes/ivf.md index cb4ace5745..76f2b6e9aa 100644 --- a/docs/docs/zh/src/indexes/ivf.md +++ b/docs/docs/zh/src/indexes/ivf.md @@ -59,6 +59,36 @@ auto result = index->KnnSearch( R"({"ivf": {"scan_buckets_count": 16}})").value(); ``` +### Bucket 对齐的精排 codes + +精排数据存储在磁盘上时，可以在复用现有精排量化与 IO 参数的基础上启用 bucket 布局： + +```cpp +std::string params = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 128, + "index_param": { + "buckets_count": 256, + "buckets_per_data": 1, + "base_quantization_type": "sq8", + "partition_strategy_type": "ivf", + "ivf_train_type": "kmeans", + + "use_reorder": true, + "precise_quantization_type": "fp32", + "precise_codes_layout": "bucket", + "precise_io_type": "buffer_io", + "precise_file_path": "/data/ivf-precise.bin" + } +})"; + +auto index = vsag::Factory::CreateIndex("ivf", params).value(); +``` + +这个示例中唯一新增的参数是 `precise_codes_layout`。构建、查询接口及查询参数与普通 IVF +快速开始示例完全相同。 + ## 构建参数 构建参数放在 `index_param` 下。完整列表请见 [索引参数](../resources/index_parameters.md)。 From d7d3a1c9e8c6891f94e76dade7155a924e61eef9 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Thu, 23 Jul 2026 16:44:37 +0800 Subject: [PATCH 03/12] docs(ivf): add runnable precise bucket example Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- docs/docs/en/src/indexes/ivf.md | 31 ---- docs/docs/zh/src/indexes/ivf.md | 30 ---- .../cpp/326_feature_ivf_precise_bucket.cpp | 134 ++++++++++++++++++ examples/cpp/CMakeLists.txt | 3 + examples/cpp/README.md | 1 + 5 files changed, 138 insertions(+), 61 deletions(-) create mode 100644 examples/cpp/326_feature_ivf_precise_bucket.cpp diff --git a/docs/docs/en/src/indexes/ivf.md b/docs/docs/en/src/indexes/ivf.md index 703e9cd5db..eea5110daf 100644 --- a/docs/docs/en/src/indexes/ivf.md +++ b/docs/docs/en/src/indexes/ivf.md @@ -64,37 +64,6 @@ auto result = index->KnnSearch( R"({"ivf": {"scan_buckets_count": 16}})").value(); ``` -### Bucket-aligned precise codes - -For disk-backed precise reordering, enable the bucket layout while keeping the -existing precise quantization and IO parameters: - -```cpp -std::string params = R"({ - "dtype": "float32", - "metric_type": "l2", - "dim": 128, - "index_param": { - "buckets_count": 256, - "buckets_per_data": 1, - "base_quantization_type": "sq8", - "partition_strategy_type": "ivf", - "ivf_train_type": "kmeans", - - "use_reorder": true, - "precise_quantization_type": "fp32", - "precise_codes_layout": "bucket", - "precise_io_type": "buffer_io", - "precise_file_path": "/data/ivf-precise.bin" - } -})"; - -auto index = vsag::Factory::CreateIndex("ivf", params).value(); -``` - -`precise_codes_layout` is the only new parameter in this example. Build and -search use the same APIs and search parameters as the regular IVF quick start. - ## Build parameters Build-time parameters live under `index_param`. See diff --git a/docs/docs/zh/src/indexes/ivf.md b/docs/docs/zh/src/indexes/ivf.md index 76f2b6e9aa..cb4ace5745 100644 --- a/docs/docs/zh/src/indexes/ivf.md +++ b/docs/docs/zh/src/indexes/ivf.md @@ -59,36 +59,6 @@ auto result = index->KnnSearch( R"({"ivf": {"scan_buckets_count": 16}})").value(); ``` -### Bucket 对齐的精排 codes - -精排数据存储在磁盘上时，可以在复用现有精排量化与 IO 参数的基础上启用 bucket 布局： - -```cpp -std::string params = R"({ - "dtype": "float32", - "metric_type": "l2", - "dim": 128, - "index_param": { - "buckets_count": 256, - "buckets_per_data": 1, - "base_quantization_type": "sq8", - "partition_strategy_type": "ivf", - "ivf_train_type": "kmeans", - - "use_reorder": true, - "precise_quantization_type": "fp32", - "precise_codes_layout": "bucket", - "precise_io_type": "buffer_io", - "precise_file_path": "/data/ivf-precise.bin" - } -})"; - -auto index = vsag::Factory::CreateIndex("ivf", params).value(); -``` - -这个示例中唯一新增的参数是 `precise_codes_layout`。构建、查询接口及查询参数与普通 IVF -快速开始示例完全相同。 - ## 构建参数 构建参数放在 `index_param` 下。完整列表请见 [索引参数](../resources/index_parameters.md)。 diff --git a/examples/cpp/326_feature_ivf_precise_bucket.cpp b/examples/cpp/326_feature_ivf_precise_bucket.cpp new file mode 100644 index 0000000000..554effa16c --- /dev/null +++ b/examples/cpp/326_feature_ivf_precise_bucket.cpp @@ -0,0 +1,134 @@ +// 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 + +namespace { + +constexpr int64_t NUM_VECTORS = 1000; +constexpr int64_t DIM = 32; + +} // namespace + +int +main() { + vsag::init(); + + /******************* Prepare Dataset *****************/ + std::vector ids(NUM_VECTORS); + std::vector vectors(NUM_VECTORS * DIM); + std::mt19937 rng(47); + std::uniform_real_distribution distribution; + for (int64_t i = 0; i < NUM_VECTORS; ++i) { + ids[i] = i; + } + for (auto& value : vectors) { + value = distribution(rng); + } + auto base = vsag::Dataset::Make() + ->NumElements(NUM_VECTORS) + ->Dim(DIM) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); + + /******************* Create Disk-Backed Precise Codes *****************/ + const auto unique_id = std::chrono::steady_clock::now().time_since_epoch().count(); + const auto temp_dir = std::filesystem::temp_directory_path() / + ("vsag-ivf-precise-bucket-" + std::to_string(unique_id)); + std::filesystem::create_directories(temp_dir); + const auto precise_file_path = (temp_dir / "precise.codes").generic_string(); + + // precise_codes_layout is the only parameter introduced by bucket-aligned precise storage. + // The existing precise quantizer and IO settings are reused. + const auto build_params = std::string(R"( + { + "dtype": "float32", + "metric_type": "l2", + "dim": 32, + "index_param": { + "buckets_count": 16, + "buckets_per_data": 1, + "base_quantization_type": "sq8", + "partition_strategy_type": "ivf", + "ivf_train_type": "kmeans", + "use_reorder": true, + "precise_quantization_type": "fp32", + "precise_codes_layout": "bucket", + "precise_io_type": "buffer_io", + "precise_file_path": ")") + + precise_file_path + R"(" + } + } + )"; + + auto create_result = vsag::Factory::CreateIndex("ivf", build_params); + if (not create_result.has_value()) { + std::cerr << "Create index failed: " << create_result.error().message << std::endl; + std::filesystem::remove_all(temp_dir); + return EXIT_FAILURE; + } + auto index = std::move(create_result.value()); + + auto build_result = index->Build(base); + if (not build_result.has_value()) { + std::cerr << "Build failed: " << build_result.error().message << std::endl; + index.reset(); + std::filesystem::remove_all(temp_dir); + return EXIT_FAILURE; + } + + /******************* Search With Precise Reordering *****************/ + auto query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(DIM) + ->Float32Vectors(vectors.data()) + ->Owner(false); + const auto search_params = R"( + { + "ivf": { + "scan_buckets_count": 16, + "factor": 4.0 + } + } + )"; + { + auto search_result = index->KnnSearch(query, 10, search_params); + if (not search_result.has_value()) { + std::cerr << "Search failed: " << search_result.error().message << std::endl; + index.reset(); + std::filesystem::remove_all(temp_dir); + return EXIT_FAILURE; + } + + std::cout << "Top-" << search_result.value()->GetDim() << " results:" << std::endl; + for (int64_t i = 0; i < search_result.value()->GetDim(); ++i) { + std::cout << " id=" << search_result.value()->GetIds()[i] + << " dist=" << search_result.value()->GetDistances()[i] << std::endl; + } + } + + index.reset(); + std::filesystem::remove_all(temp_dir); + return EXIT_SUCCESS; +} diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 899b1228e0..97dcbcd8f6 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -162,3 +162,6 @@ target_link_libraries(406_feature_read_cache vsag) add_executable(407_feature_ivf_read_cache 407_feature_ivf_read_cache.cpp) target_link_libraries(407_feature_ivf_read_cache vsag) + +add_executable(326_feature_ivf_precise_bucket 326_feature_ivf_precise_bucket.cpp) +target_link_libraries(326_feature_ivf_precise_bucket vsag) diff --git a/examples/cpp/README.md b/examples/cpp/README.md index 879bce936d..da9eeb161f 100644 --- a/examples/cpp/README.md +++ b/examples/cpp/README.md @@ -102,6 +102,7 @@ together when the directory is listed: | [`320_feature_extra_info.cpp`](320_feature_extra_info.cpp) | Attach per-vector extra info / payload. | | [`322_feature_hgraph_brute_force_threshold.cpp`](322_feature_hgraph_brute_force_threshold.cpp) | HGraph search-time `brute_force_threshold`: automatically switch to an exact scan under highly selective filters. | | [`324_feature_lazy_hgraph_extra_info.cpp`](324_feature_lazy_hgraph_extra_info.cpp) | LazyHGraph `extra_info` filtering across flat and graph phases. | +| [`326_feature_ivf_precise_bucket.cpp`](326_feature_ivf_precise_bucket.cpp) | Disk-backed IVF precise codes using the bucket-aligned layout. | ### Persistence (`4xx`) From e6ca3e9f8c8580f7febded575944966eb6d6f713 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Fri, 24 Jul 2026 11:53:51 +0800 Subject: [PATCH 04/12] fix(io): make non-continuous allocation thread-safe Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- .../noncontinuous_allocator.h | 7 +-- .../noncontinuous_io_test.cpp | 52 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/src/io/noncontinuous_io/noncontinuous_allocator.h b/src/io/noncontinuous_io/noncontinuous_allocator.h index ae7fbe6005..cf0cd79ab7 100644 --- a/src/io/noncontinuous_io/noncontinuous_allocator.h +++ b/src/io/noncontinuous_io/noncontinuous_allocator.h @@ -15,6 +15,7 @@ #pragma once +#include #include namespace vsag { @@ -83,8 +84,8 @@ class NonContinuousAllocator { Require(uint64_t size) { // 4k align size = (size + ALOGN_SIZE - 1) & ~(ALOGN_SIZE - 1); - NonContinuousArea area{last_offset_, size}; - last_offset_ += size; + auto offset = last_offset_.fetch_add(size, std::memory_order_relaxed); + NonContinuousArea area{offset, size}; return area; } @@ -93,7 +94,7 @@ class NonContinuousAllocator { Allocator* const allocator_{nullptr}; /// Last allocated offset for tracking non-overlapping areas. - uint64_t last_offset_{0}; + std::atomic last_offset_{0}; /// Alignment size for area allocation (4KB). static constexpr uint64_t ALOGN_SIZE = 4096; diff --git a/src/io/noncontinuous_io/noncontinuous_io_test.cpp b/src/io/noncontinuous_io/noncontinuous_io_test.cpp index e79917ab59..85ed444885 100644 --- a/src/io/noncontinuous_io/noncontinuous_io_test.cpp +++ b/src/io/noncontinuous_io/noncontinuous_io_test.cpp @@ -15,6 +15,11 @@ #include "io/noncontinuous_io/noncontinuous_io.h" +#include +#include +#include +#include + #include "impl/allocator/safe_allocator.h" #include "io/async_io/async_io.h" #include "io/buffer_io/buffer_io.h" @@ -89,3 +94,50 @@ TEST_CASE("NonContinuousIO Serialize Test", "[NonContinuousIO][ut]") { NonContinuousIOTestSerialize(); NonContinuousIOTestSerialize(); } + +TEST_CASE("NonContinuousAllocator allocates unique regions concurrently", + "[NonContinuousIO][ut][concurrent]") { + constexpr uint64_t thread_count = 16; + constexpr uint64_t allocations_per_thread = 4096; + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + NonContinuousAllocator non_continuous_allocator(allocator.get()); + std::atomic ready{0}; + std::atomic start{false}; + std::vector> thread_areas(thread_count); + std::vector threads; + threads.reserve(thread_count); + + for (uint64_t thread_id = 0; thread_id < thread_count; ++thread_id) { + threads.emplace_back([&, thread_id]() { + auto& areas = thread_areas[thread_id]; + areas.reserve(allocations_per_thread); + ready.fetch_add(1, std::memory_order_release); + while (not start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (uint64_t i = 0; i < allocations_per_thread; ++i) { + areas.emplace_back(non_continuous_allocator.Require(1)); + } + }); + } + + while (ready.load(std::memory_order_acquire) != thread_count) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } + + std::vector areas; + areas.reserve(thread_count * allocations_per_thread); + for (const auto& thread_area : thread_areas) { + areas.insert(areas.end(), thread_area.begin(), thread_area.end()); + } + std::sort(areas.begin(), areas.end(), [](const auto& lhs, const auto& rhs) { + return lhs.offset < rhs.offset; + }); + for (uint64_t i = 1; i < areas.size(); ++i) { + REQUIRE(areas[i - 1].offset + areas[i - 1].size <= areas[i].offset); + } +} From 667d73cb03103a28b260dd53ab8bfbbb19fa3a45 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Fri, 24 Jul 2026 14:43:15 +0800 Subject: [PATCH 05/12] fix(ivf): isolate unsafe disk-backed precise operations Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- .../docs/en/src/advanced/new_serialization.md | 4 +- docs/docs/en/src/indexes/ivf.md | 11 ++- .../docs/zh/src/advanced/new_serialization.md | 4 +- docs/docs/zh/src/indexes/ivf.md | 10 ++- src/algorithm/ivf/ivf.cpp | 14 ++++ src/algorithm/ivf/ivf.h | 14 ++++ src/algorithm/ivf/ivf_parameter.cpp | 15 ++++ src/algorithm/ivf/ivf_parameter.h | 3 + src/algorithm/ivf/ivf_parameter_test.cpp | 18 ++++ src/factory/factory.cpp | 10 ++- tests/test_ivf.cpp | 83 +++++++++++++++++++ 11 files changed, 177 insertions(+), 9 deletions(-) diff --git a/docs/docs/en/src/advanced/new_serialization.md b/docs/docs/en/src/advanced/new_serialization.md index 3a8681d227..d2973a5d24 100644 --- a/docs/docs/en/src/advanced/new_serialization.md +++ b/docs/docs/en/src/advanced/new_serialization.md @@ -184,7 +184,9 @@ IVF writes these streaming blocks in order: `DeserializeStreaming` restores the full in-memory IVF index. `Index::Load` can create the IVF index directly from streaming metadata and currently loads all emitted IVF blocks into memory. The two precise-code blocks are mutually exclusive and are selected by -`precise_codes_layout`. +`precise_codes_layout`. For file-backed bucket-aligned precise codes, create an IVF destination +with an independent `precise_file_path` and use `DeserializeStreaming`; static `Index::Load` is +rejected until it can accept an independent target path. ## SINDI Blocks diff --git a/docs/docs/en/src/indexes/ivf.md b/docs/docs/en/src/indexes/ivf.md index eea5110daf..5bb919a178 100644 --- a/docs/docs/en/src/indexes/ivf.md +++ b/docs/docs/en/src/indexes/ivf.md @@ -94,13 +94,18 @@ Build-time parameters live under `index_param`. See | `precise_file_path` | string | `""` | File path when the precise IO type is disk-backed | `precise_codes_layout: "bucket"` requires `use_reorder: true`. It supports -`memory_io`, `block_memory_io`, `mmap_io`, `buffer_io`, `async_io`, and `uring_io` -(when io_uring is available); -`reader_io` and `pqfs` precise quantization are not supported. When +`memory_io`, `block_memory_io`, `buffer_io`, `async_io`, and `uring_io` +(when io_uring is available); `mmap_io`, `reader_io`, and `pqfs` precise +quantization are not supported. When `buckets_per_data` is greater than one, the precise vector is duplicated for every basic posting, preserving exact bucket-offset alignment at the corresponding storage cost. +For file-backed bucket-aligned precise codes, `Clone`, `ExportModel`, `Merge`, and static +`Index::Load` are rejected because those operations cannot yet assign an independent target +file. To restore a streaming index on disk, create the destination IVF with a different +`precise_file_path` and call `DeserializeStreaming`. + A rule of thumb for `buckets_count` is `sqrt(N)` to `4 * sqrt(N)` where `N` is the corpus size. diff --git a/docs/docs/zh/src/advanced/new_serialization.md b/docs/docs/zh/src/advanced/new_serialization.md index c272785c36..53ab5c7cbd 100644 --- a/docs/docs/zh/src/advanced/new_serialization.md +++ b/docs/docs/zh/src/advanced/new_serialization.md @@ -163,7 +163,9 @@ IVF 按顺序写入以下 streaming blocks： `DeserializeStreaming` 会恢复完整的内存 IVF 索引。`Index::Load` 可以直接从 streaming metadata 创建 IVF 索引对象，当前会把写出的 IVF blocks 都加载到内存中。两种精排 codes block 互斥， -由 `precise_codes_layout` 选择。 +由 `precise_codes_layout` 选择。对于文件型 bucket 精排 codes，请使用独立的 +`precise_file_path` 创建 IVF 目标对象并调用 `DeserializeStreaming`；在静态 `Index::Load` +能够接收独立目标路径之前，该组合会被拒绝。 ## SINDI Blocks diff --git a/docs/docs/zh/src/indexes/ivf.md b/docs/docs/zh/src/indexes/ivf.md index cb4ace5745..136c66ca82 100644 --- a/docs/docs/zh/src/indexes/ivf.md +++ b/docs/docs/zh/src/indexes/ivf.md @@ -88,11 +88,15 @@ auto result = index->KnnSearch( | `precise_file_path` | string | `""` | 当精排 IO 为磁盘后端时的文件路径 | `precise_codes_layout: "bucket"` 要求 `use_reorder: true`，支持 `memory_io`、 -`block_memory_io`、`mmap_io`、`buffer_io`、`async_io` 和 `uring_io` -（需要构建环境支持 io_uring），不支持 -`reader_io` 和 `pqfs` 精排量化。当 `buckets_per_data` 大于 1 时，每个 basic posting +`block_memory_io`、`buffer_io`、`async_io` 和 `uring_io` +（需要构建环境支持 io_uring），不支持 `mmap_io`、`reader_io` 和 `pqfs` +精排量化。当 `buckets_per_data` 大于 1 时，每个 basic posting 都会保存一份高精度向量，从而保持完全相同的 bucket-offset 对齐，同时占用相应倍数的存储空间。 +对于文件型 bucket 精排 codes，暂不支持 `Clone`、`ExportModel`、`Merge` 和静态 +`Index::Load`，因为这些操作目前无法为目标索引指定独立文件。若需从 streaming 数据恢复磁盘索引， +请使用不同的 `precise_file_path` 创建目标 IVF，然后调用 `DeserializeStreaming`。 + `buckets_count` 的经验值一般为 `sqrt(N)` ~ `4 * sqrt(N)`，其中 `N` 是语料规模。 ## 检索参数 diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index f26a963b06..97ab122a42 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -413,6 +413,7 @@ IVF::IVF(const IVFParameterPtr& param, const IndexCommonParam& common_param) bucket_graphs_(common_param.allocator_.get()), common_param_(common_param), bucket_searcher_(std::make_shared()) { + this->disk_backed_precise_bucket_ = param->UsesDiskBackedPreciseBucket(); this->bucket_ = BucketInterface::MakeInstance(param->bucket_param, common_param); if (this->bucket_ == nullptr) { throw VsagException(ErrorType::INTERNAL_ERROR, "bucket init error"); @@ -526,6 +527,11 @@ IVF::InitFeatures() { IndexFeature::SUPPORT_EXPORT_MODEL, IndexFeature::SUPPORT_GET_MEMORY_USAGE, IndexFeature::SUPPORT_MERGE_INDEX}); + if (this->disk_backed_precise_bucket_) { + this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_CLONE, false); + this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_EXPORT_MODEL, false); + this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_MERGE_INDEX, false); + } if (this->bucket_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_PQFS) { this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_ADD_AFTER_BUILD, false); @@ -854,6 +860,10 @@ IVF::GetNumElements() const { void IVF::Merge(const std::vector& merge_units) { + if (this->disk_backed_precise_bucket_) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "Merge does not support disk-backed IVF precise buckets"); + } this->bucket_->Unpack(); if (precise_bucket_ != nullptr) { this->precise_bucket_->Unpack(); @@ -1608,6 +1618,10 @@ IVF::reorder_with_precise_bucket(const DistHeapPtr& input, InnerIndexPtr IVF::ExportModel(const IndexCommonParam& param) const { + if (this->disk_backed_precise_bucket_) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "ExportModel does not support disk-backed IVF precise buckets"); + } auto index = std::make_shared(this->create_param_ptr_, param); IVFPartitionStrategy::Clone(this->partition_strategy_, index->partition_strategy_); this->bucket_->ExportModel(index->bucket_); diff --git a/src/algorithm/ivf/ivf.h b/src/algorithm/ivf/ivf.h index fefe3f54ac..9d0c4ff5e8 100644 --- a/src/algorithm/ivf/ivf.h +++ b/src/algorithm/ivf/ivf.h @@ -105,8 +105,21 @@ class IVF : public InnerIndexInterface { [[nodiscard]] InnerIndexPtr ExportModel(const IndexCommonParam& param) const override; + [[nodiscard]] InnerIndexPtr + Clone(const IndexCommonParam& param) override { + if (this->disk_backed_precise_bucket_) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "Clone does not support disk-backed IVF precise buckets"); + } + return InnerIndexInterface::Clone(param); + } + [[nodiscard]] InnerIndexPtr Fork(const IndexCommonParam& param) override { + if (this->disk_backed_precise_bucket_) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "Clone does not support disk-backed IVF precise buckets"); + } return std::make_shared(this->create_param_ptr_, param); } @@ -297,6 +310,7 @@ class IVF : public InnerIndexInterface { FlattenInterfacePtr reorder_codes_{nullptr}; // legacy high-precision flat codes BucketInterfacePtr precise_bucket_{nullptr}; // high-precision codes mirroring basic buckets ReorderInterfacePtr reorder_{nullptr}; // flat-code reordering engine + bool disk_backed_precise_bucket_{false}; std::shared_ptr thread_pool_{nullptr}; // for parallel bucket scans diff --git a/src/algorithm/ivf/ivf_parameter.cpp b/src/algorithm/ivf/ivf_parameter.cpp index 4b0c6eac6c..c3b8f118ca 100644 --- a/src/algorithm/ivf/ivf_parameter.cpp +++ b/src/algorithm/ivf/ivf_parameter.cpp @@ -48,6 +48,10 @@ IVFParameter::FromJson(const JsonType& json) { CHECK_ARGUMENT(this->precise_codes_param->quantizer_parameter->GetTypeName() != QUANTIZATION_TYPE_VALUE_PQFS, "precise_codes_layout=bucket does not support pqfs precise quantization"); + CHECK_ARGUMENT( + this->precise_codes_param->io_parameter == nullptr || + this->precise_codes_param->io_parameter->GetTypeName() != IO_TYPE_VALUE_MMAP_IO, + "precise_codes_layout=bucket does not support mmap_io"); } if (json.Contains(BUCKET_PER_DATA_KEY)) { @@ -94,6 +98,17 @@ IVFParameter::FromJson(const JsonType& json) { } } +bool +IVFParameter::UsesDiskBackedPreciseBucket() const { + if (this->precise_codes_layout != PRECISE_CODES_LAYOUT_VALUE_BUCKET || + this->precise_codes_param == nullptr || + this->precise_codes_param->io_parameter == nullptr) { + return false; + } + const auto io_type = this->precise_codes_param->io_parameter->GetTypeName(); + return io_type != IO_TYPE_VALUE_MEMORY_IO && io_type != IO_TYPE_VALUE_BLOCK_MEMORY_IO; +} + JsonType IVFParameter::ToJson() const { JsonType json = InnerIndexParameter::ToJson(); diff --git a/src/algorithm/ivf/ivf_parameter.h b/src/algorithm/ivf/ivf_parameter.h index 1d627803e9..08d46f9019 100644 --- a/src/algorithm/ivf/ivf_parameter.h +++ b/src/algorithm/ivf/ivf_parameter.h @@ -43,6 +43,9 @@ class IVFParameter : public InnerIndexParameter { bool CheckCompatibility(const vsag::ParamPtr& other) const override; + [[nodiscard]] bool + UsesDiskBackedPreciseBucket() const; + public: BucketDataCellParamPtr bucket_param{nullptr}; IVFPartitionStrategyParametersPtr ivf_partition_strategy_parameter{nullptr}; diff --git a/src/algorithm/ivf/ivf_parameter_test.cpp b/src/algorithm/ivf/ivf_parameter_test.cpp index 05b1650816..0f4809f796 100644 --- a/src/algorithm/ivf/ivf_parameter_test.cpp +++ b/src/algorithm/ivf/ivf_parameter_test.cpp @@ -23,6 +23,7 @@ #include "quantization/rabitq_quantization/rabitq_quantizer_parameter.h" #include "unittest.h" #include "utils/util_functions.h" +#include "vsag_exception.h" struct IVFDefaultParam { std::string buckect_io_type = "block_memory_io"; @@ -226,6 +227,23 @@ TEST_CASE("IVF precise codes layout parameter", "[ut][IVFParameter]") { auto param = std::make_shared(); REQUIRE_THROWS(param->FromJson(param_json)); } + + SECTION("bucket layout rejects mmap io") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + index_param.precise_codes_io_type = "mmap_io"; + auto param_json = vsag::JsonType::Parse(generate_ivf_param(index_param)); + param_json[vsag::PRECISE_CODES_KEY][vsag::IO_PARAMS_KEY][vsag::IO_FILE_PATH_KEY].SetString( + "ivf_precise_mmap_test"); + + auto param = std::make_shared(); + try { + param->FromJson(param_json); + FAIL("mmap_io should be rejected for bucket-aligned precise codes"); + } catch (const vsag::VsagException& error) { + REQUIRE(error.error_.type == vsag::ErrorType::INVALID_ARGUMENT); + } + } } TEST_CASE("IVF maps RabitQ external parameters", "[ut][IVFParameter]") { diff --git a/src/factory/factory.cpp b/src/factory/factory.cpp index 0ffe1608ac..b254dc5a50 100644 --- a/src/factory/factory.cpp +++ b/src/factory/factory.cpp @@ -245,7 +245,15 @@ create_streaming_index_from_metadata(const MetadataPtr& metadata, return create_streaming_index(index_param, common_param); } if (index_name == INDEX_IVF) { - return create_streaming_index(index_param, common_param); + auto param = std::make_shared(); + param->FromJson(index_param); + if (param->UsesDiskBackedPreciseBucket()) { + throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, + "Index::Load does not support disk-backed IVF precise buckets"); + } + auto inner_index = std::make_shared(param, common_param); + return streaming_index_load_target{ + std::make_shared>(inner_index, common_param), inner_index}; } if (index_name == INDEX_PYRAMID) { return create_streaming_index(index_param, common_param); diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 3f79e93d5c..4b26ed8f8a 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -14,6 +14,7 @@ // limitations under the License. #include +#include #include #include #include @@ -440,6 +441,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); auto index = TestFactory(name, params, true); TestBuildIndex(index, dataset, true); + REQUIRE(index->ExportModel().has_value()); std::stringstream stream; REQUIRE(index->SerializeStreaming(stream).has_value()); @@ -466,6 +468,87 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, REQUIRE_FALSE(vsag::Index::Load(missing_load_stream, "{}").has_value()); } +TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, + "IVF disk bucket precise rejects aliased ownership operations", + "[ft][ivf][reorder][serialize][streaming][export][pr]") { + constexpr int64_t dim = 16; + constexpr int64_t base_count = 128; + const auto precise_file_path = dir.GenerateRandomFile(false); + const auto params = GenerateBucketPreciseParameters(1, "buffer_io", precise_file_path); + auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); + auto index = TestFactory(name, params, true); + TestBuildIndex(index, dataset, true); + REQUIRE_FALSE(index->CheckFeature(vsag::SUPPORT_CLONE)); + REQUIRE_FALSE(index->CheckFeature(vsag::SUPPORT_EXPORT_MODEL)); + REQUIRE_FALSE(index->CheckFeature(vsag::SUPPORT_MERGE_INDEX)); + + SECTION("rejects export model") { + auto result = index->ExportModel(); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::UNSUPPORTED_INDEX_OPERATION); + } + + SECTION("rejects clone") { + auto result = index->Clone(); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::UNSUPPORTED_INDEX_OPERATION); + } + + SECTION("rejects merge") { + auto result = index->Merge({}); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::UNSUPPORTED_INDEX_OPERATION); + } + + SECTION("rejects streaming load") { + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + std::stringstream load_stream(stream.str()); + auto result = vsag::Index::Load(load_stream, "{}"); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::UNSUPPORTED_INDEX_OPERATION); + } + + SECTION("rejects empty streaming load") { + const auto empty_file_path = dir.GenerateRandomFile(false); + const auto empty_params = GenerateBucketPreciseParameters(1, "buffer_io", empty_file_path); + auto empty_index = TestFactory(name, empty_params, true); + std::stringstream stream; + REQUIRE(empty_index->SerializeStreaming(stream).has_value()); + std::stringstream load_stream(stream.str()); + auto result = vsag::Index::Load(load_stream, "{}"); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::UNSUPPORTED_INDEX_OPERATION); + } + + SECTION("rejects mmap before creating the precise file") { + const auto mmap_file_path = dir.GenerateRandomFile(false); + const auto mmap_params = GenerateBucketPreciseParameters(1, "mmap_io", mmap_file_path); + auto result = vsag::Factory::CreateIndex(name, mmap_params); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + REQUIRE_FALSE(std::filesystem::exists(mmap_file_path)); + } + + SECTION("allows streaming deserialize into an independent precise file") { + std::stringstream stream; + REQUIRE(index->SerializeStreaming(stream).has_value()); + const auto restored_file_path = dir.GenerateRandomFile(false); + const auto restored_params = + GenerateBucketPreciseParameters(1, "buffer_io", restored_file_path); + auto restored = TestFactory(name, restored_params, true); + std::stringstream deserialize_stream(stream.str()); + REQUIRE(restored->DeserializeStreaming(deserialize_stream).has_value()); + const auto search_param = fmt::format(search_param_tmp, 16); + CheckBucketPreciseIndex(restored, dataset, search_param); + TestBatchCalcDistanceById(restored, dataset, 2e-6F, true); + } + + const auto search_param = fmt::format(search_param_tmp, 16); + CheckBucketPreciseIndex(index, dataset, search_param); + TestBatchCalcDistanceById(index, dataset, 2e-6F, true); +} + TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, "IVF bucket precise merge", "[ft][ivf][reorder][merge][pr]") { From b15adab095f7f0500a71ac7fd5f1cdf77d80b1f7 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Fri, 24 Jul 2026 17:24:50 +0800 Subject: [PATCH 06/12] test(ivf): bound precise bucket test memory Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- tests/test_ivf.cpp | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 4b26ed8f8a..5c7a0353df 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -194,6 +194,21 @@ namespace { using vsag::test::EraseStreamingBlock; using vsag::test::InsertUnknownStreamingBlock; +class BlockSizeLimitGuard { +public: + explicit BlockSizeLimitGuard(uint64_t block_size_limit) + : origin_size_(vsag::Options::Instance().block_size_limit()) { + vsag::Options::Instance().set_block_size_limit(block_size_limit); + } + + ~BlockSizeLimitGuard() { + vsag::Options::Instance().set_block_size_limit(origin_size_); + } + +private: + uint64_t origin_size_; +}; + std::string GenerateBucketPreciseParameters(int buckets_per_data, const std::string& precise_io_type = "block_memory_io", @@ -398,6 +413,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, "IVF bucket precise mirrors basic postings", "[ft][ivf][reorder][serialize][pr]") { + BlockSizeLimitGuard block_size_limit_guard(2ULL * 1024 * 1024); constexpr int64_t dim = 16; constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; @@ -433,6 +449,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, "IVF bucket precise streaming", "[ft][ivf][reorder][serialize][streaming][pr]") { + BlockSizeLimitGuard block_size_limit_guard(2ULL * 1024 * 1024); constexpr int64_t dim = 16; constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; @@ -471,6 +488,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, "IVF disk bucket precise rejects aliased ownership operations", "[ft][ivf][reorder][serialize][streaming][export][pr]") { + BlockSizeLimitGuard block_size_limit_guard(2ULL * 1024 * 1024); constexpr int64_t dim = 16; constexpr int64_t base_count = 128; const auto precise_file_path = dir.GenerateRandomFile(false); @@ -552,6 +570,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, "IVF bucket precise merge", "[ft][ivf][reorder][merge][pr]") { + BlockSizeLimitGuard block_size_limit_guard(2ULL * 1024 * 1024); constexpr int64_t dim = 16; constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; From 24301ec947f223dd53531716b7d4037cc21a95a4 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Mon, 27 Jul 2026 14:28:40 +0800 Subject: [PATCH 07/12] test(ivf): cover precise bucket batch distances Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- tests/test_ivf.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 5c7a0353df..cf9436a521 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -434,6 +434,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, CheckBucketPreciseIndex(index, dataset, search_param); TestCalcDistanceById(index, dataset, 2e-6F, true); TestBatchCalcDistanceById(index, dataset, 2e-6F, true); + TestMultiQueryBatchCalcDistanceById(index, dataset, 2e-6F, true); const auto restored_file_path = precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); From 4bcc5b1146cee0631e328372ec70db03c7576447 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Thu, 30 Jul 2026 11:18:08 +0800 Subject: [PATCH 08/12] fix(ivf): scope precise bucket to single posting Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- docs/docs/en/src/indexes/ivf.md | 8 ++---- docs/docs/zh/src/indexes/ivf.md | 6 ++-- src/algorithm/ivf/ivf.h | 2 +- src/algorithm/ivf/ivf_parameter.cpp | 10 ++++--- src/algorithm/ivf/ivf_parameter_test.cpp | 21 ++++++++++++-- tests/test_ivf.cpp | 36 +++++++++++++++--------- 6 files changed, 53 insertions(+), 30 deletions(-) diff --git a/docs/docs/en/src/indexes/ivf.md b/docs/docs/en/src/indexes/ivf.md index 5bb919a178..39436fe6e1 100644 --- a/docs/docs/en/src/indexes/ivf.md +++ b/docs/docs/en/src/indexes/ivf.md @@ -88,7 +88,7 @@ Build-time parameters live under `index_param`. See | `fast_encode_rabitq_rounds` | int | `6` | CAQ adjustment rounds; allowed range is `[1, 32]` | | `use_reorder` | bool | `false` | Keep a high-precision copy and re-rank after the coarse scan | | `precise_quantization_type` | string | `"fp32"` | Quantizer used for reordering (with `use_reorder: true`) | -| `precise_codes_layout` | string | `"flat"` | Storage layout for precise codes: `"flat"` keeps the legacy one-code-per-vector layout; `"bucket"` mirrors every basic posting in the same bucket and offset | +| `precise_codes_layout` | string | `"flat"` | Storage layout for precise codes: `"flat"` keeps the legacy one-code-per-vector layout; `"bucket"` stores the precise code in the same bucket and offset as its basic posting | | `base_io_type` | string | `"memory_io"` | Storage backend for coarse codes; supports `uring_io` when built with liburing | | `precise_io_type` | string | `"block_memory_io"` | Storage backend for precise codes (`memory_io`, `block_memory_io`, `mmap_io`, `buffer_io`, `async_io`, `uring_io`, `reader_io`) | | `precise_file_path` | string | `""` | File path when the precise IO type is disk-backed | @@ -96,10 +96,8 @@ Build-time parameters live under `index_param`. See `precise_codes_layout: "bucket"` requires `use_reorder: true`. It supports `memory_io`, `block_memory_io`, `buffer_io`, `async_io`, and `uring_io` (when io_uring is available); `mmap_io`, `reader_io`, and `pqfs` precise -quantization are not supported. When -`buckets_per_data` is greater than one, the precise vector is duplicated for every -basic posting, preserving exact bucket-offset alignment at the corresponding storage -cost. +quantization are not supported. The bucket layout currently requires +`buckets_per_data: 1`; configurations that assign one vector to multiple buckets are rejected. For file-backed bucket-aligned precise codes, `Clone`, `ExportModel`, `Merge`, and static `Index::Load` are rejected because those operations cannot yet assign an independent target diff --git a/docs/docs/zh/src/indexes/ivf.md b/docs/docs/zh/src/indexes/ivf.md index 136c66ca82..c0f037d34d 100644 --- a/docs/docs/zh/src/indexes/ivf.md +++ b/docs/docs/zh/src/indexes/ivf.md @@ -82,7 +82,7 @@ auto result = index->KnnSearch( | `fast_encode_rabitq_rounds` | int | `6` | CAQ 微调轮数，允许范围 `[1, 32]` | | `use_reorder` | bool | `false` | 是否保留高精度副本用于精排 | | `precise_quantization_type` | string | `"fp32"` | 精排量化类型（`use_reorder: true` 时使用） | -| `precise_codes_layout` | string | `"flat"` | 精排 codes 的存储布局：`"flat"` 保持旧的一向量一码布局；`"bucket"` 为 basic 的每个 posting 在相同 bucket 和 offset 保存一份高精度 code | +| `precise_codes_layout` | string | `"flat"` | 精排 codes 的存储布局：`"flat"` 保持旧的一向量一码布局；`"bucket"` 在 basic posting 的相同 bucket 和 offset 保存高精度 code | | `base_io_type` | string | `"memory_io"` | 粗排向量的存储后端；以 liburing 构建时支持 `uring_io` | | `precise_io_type` | string | `"block_memory_io"` | 精排向量的存储后端（`memory_io`、`block_memory_io`、`mmap_io`、`buffer_io`、`async_io`、`uring_io`、`reader_io`） | | `precise_file_path` | string | `""` | 当精排 IO 为磁盘后端时的文件路径 | @@ -90,8 +90,8 @@ auto result = index->KnnSearch( `precise_codes_layout: "bucket"` 要求 `use_reorder: true`，支持 `memory_io`、 `block_memory_io`、`buffer_io`、`async_io` 和 `uring_io` （需要构建环境支持 io_uring），不支持 `mmap_io`、`reader_io` 和 `pqfs` -精排量化。当 `buckets_per_data` 大于 1 时，每个 basic posting -都会保存一份高精度向量，从而保持完全相同的 bucket-offset 对齐，同时占用相应倍数的存储空间。 +精排量化。bucket 布局当前要求 `buckets_per_data: 1`；一个向量分配到多个 bucket +的配置会被拒绝。 对于文件型 bucket 精排 codes，暂不支持 `Clone`、`ExportModel`、`Merge` 和静态 `Index::Load`，因为这些操作目前无法为目标索引指定独立文件。若需从 streaming 数据恢复磁盘索引， diff --git a/src/algorithm/ivf/ivf.h b/src/algorithm/ivf/ivf.h index 9d0c4ff5e8..c3f56e1e8f 100644 --- a/src/algorithm/ivf/ivf.h +++ b/src/algorithm/ivf/ivf.h @@ -118,7 +118,7 @@ class IVF : public InnerIndexInterface { Fork(const IndexCommonParam& param) override { if (this->disk_backed_precise_bucket_) { throw VsagException(ErrorType::UNSUPPORTED_INDEX_OPERATION, - "Clone does not support disk-backed IVF precise buckets"); + "Fork does not support disk-backed IVF precise buckets"); } return std::make_shared(this->create_param_ptr_, param); } diff --git a/src/algorithm/ivf/ivf_parameter.cpp b/src/algorithm/ivf/ivf_parameter.cpp index c3b8f118ca..58fa053e1c 100644 --- a/src/algorithm/ivf/ivf_parameter.cpp +++ b/src/algorithm/ivf/ivf_parameter.cpp @@ -38,8 +38,14 @@ IVFParameter::FromJson(const JsonType& json) { PRECISE_CODES_LAYOUT_VALUE_FLAT, PRECISE_CODES_LAYOUT_VALUE_BUCKET)); + if (json.Contains(BUCKET_PER_DATA_KEY)) { + this->buckets_per_data = static_cast(json[BUCKET_PER_DATA_KEY].GetInt()); + } + if (this->precise_codes_layout == PRECISE_CODES_LAYOUT_VALUE_BUCKET) { CHECK_ARGUMENT(this->use_reorder, "precise_codes_layout=bucket requires use_reorder=true"); + CHECK_ARGUMENT(this->buckets_per_data == 1, + "precise_codes_layout=bucket requires buckets_per_data=1"); CHECK_ARGUMENT(this->reorder_source == HGRAPH_REORDER_SOURCE_PRECISE, "precise_codes_layout=bucket requires reorder_source=precise"); CHECK_ARGUMENT(this->precise_codes_param != nullptr && @@ -54,10 +60,6 @@ IVFParameter::FromJson(const JsonType& json) { "precise_codes_layout=bucket does not support mmap_io"); } - if (json.Contains(BUCKET_PER_DATA_KEY)) { - this->buckets_per_data = static_cast(json[BUCKET_PER_DATA_KEY].GetInt()); - } - this->bucket_param = std::make_shared(); CHECK_ARGUMENT(json.Contains(BUCKET_PARAMS_KEY), diff --git a/src/algorithm/ivf/ivf_parameter_test.cpp b/src/algorithm/ivf/ivf_parameter_test.cpp index 0f4809f796..7ef836259c 100644 --- a/src/algorithm/ivf/ivf_parameter_test.cpp +++ b/src/algorithm/ivf/ivf_parameter_test.cpp @@ -166,19 +166,34 @@ TEST_CASE("IVF precise codes layout parameter", "[ut][IVFParameter]") { vsag::PRECISE_CODES_LAYOUT_VALUE_FLAT); } - SECTION("bucket layout supports multiple postings per data") { + SECTION("bucket layout supports one posting per data") { IVFDefaultParam index_param; index_param.precise_codes_layout = "bucket"; - index_param.buckets_per_data = 3; auto param = std::make_shared(); param->FromString(generate_ivf_param(index_param)); REQUIRE(param->precise_codes_layout == vsag::PRECISE_CODES_LAYOUT_VALUE_BUCKET); - REQUIRE(param->buckets_per_data == 3); + REQUIRE(param->buckets_per_data == 1); vsag::ParameterTest::TestToJson(param); } + SECTION("bucket layout rejects multiple postings per data") { + IVFDefaultParam index_param; + index_param.precise_codes_layout = "bucket"; + index_param.buckets_per_data = 2; + + auto param = std::make_shared(); + try { + param->FromString(generate_ivf_param(index_param)); + FAIL("multiple postings should be rejected for bucket-aligned precise codes"); + } catch (const vsag::VsagException& error) { + REQUIRE(error.error_.type == vsag::ErrorType::INVALID_ARGUMENT); + REQUIRE(error.error_.message == + "precise_codes_layout=bucket requires buckets_per_data=1"); + } + } + SECTION("reject invalid layout") { IVFDefaultParam index_param; index_param.precise_codes_layout = "invalid"; diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index cf9436a521..c31ca62dd8 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -210,12 +210,11 @@ class BlockSizeLimitGuard { }; std::string -GenerateBucketPreciseParameters(int buckets_per_data, - const std::string& precise_io_type = "block_memory_io", +GenerateBucketPreciseParameters(const std::string& precise_io_type = "block_memory_io", const std::string& precise_file_path = "", int thread_count = 1) { auto params = nlohmann::json::parse(IVFTestIndex::GenerateIVFBuildParametersString( - "l2", 16, "sq8,fp32", 16, "random", false, buckets_per_data, false, thread_count)); + "l2", 16, "sq8,fp32", 16, "random", false, 1, false, thread_count)); params["index_param"]["precise_codes_layout"] = "bucket"; params["index_param"]["precise_io_type"] = precise_io_type; params["index_param"]["precise_file_path"] = precise_file_path; @@ -410,6 +409,18 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, vsag::Options::Instance().set_block_size_limit(origin_size); } +TEST_CASE("IVF bucket precise rejects multiple postings per data", "[ft][ivf][reorder][pr]") { + auto params = nlohmann::json::parse(GenerateBucketPreciseParameters()); + params["index_param"]["buckets_per_data"] = 2; + + auto result = vsag::Factory::CreateIndex(IVFTestIndex::name, params.dump()); + + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + REQUIRE(result.error().message.find( + "precise_codes_layout=bucket requires buckets_per_data=1") != std::string::npos); +} + TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, "IVF bucket precise mirrors basic postings", "[ft][ivf][reorder][serialize][pr]") { @@ -417,15 +428,12 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, constexpr int64_t dim = 16; constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; - const auto buckets_per_data = GENERATE(1, 2); const auto precise_io_type = GENERATE("block_memory_io", "buffer_io"); - INFO(fmt::format("buckets_per_data: {}", buckets_per_data)); INFO(fmt::format("precise_io_type: {}", precise_io_type)); const auto precise_file_path = precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); - const auto params = - GenerateBucketPreciseParameters(buckets_per_data, precise_io_type, precise_file_path); + const auto params = GenerateBucketPreciseParameters(precise_io_type, precise_file_path); const auto search_param = fmt::format(search_param_tmp, buckets_count); auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); auto index = TestFactory(name, params, true); @@ -439,7 +447,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, const auto restored_file_path = precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); const auto restored_params = - GenerateBucketPreciseParameters(buckets_per_data, precise_io_type, restored_file_path); + GenerateBucketPreciseParameters(precise_io_type, restored_file_path); auto restored = TestFactory(name, restored_params, true); TestSerializeBinarySet(index, restored, dataset, search_param, true); CheckBucketPreciseIndex(restored, dataset, search_param); @@ -454,7 +462,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, constexpr int64_t dim = 16; constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; - const auto params = GenerateBucketPreciseParameters(2, "block_memory_io", "", 4); + const auto params = GenerateBucketPreciseParameters("block_memory_io", "", 4); const auto search_param = fmt::format(search_param_tmp, buckets_count); auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); auto index = TestFactory(name, params, true); @@ -493,7 +501,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, constexpr int64_t dim = 16; constexpr int64_t base_count = 128; const auto precise_file_path = dir.GenerateRandomFile(false); - const auto params = GenerateBucketPreciseParameters(1, "buffer_io", precise_file_path); + const auto params = GenerateBucketPreciseParameters("buffer_io", precise_file_path); auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); auto index = TestFactory(name, params, true); TestBuildIndex(index, dataset, true); @@ -530,7 +538,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, SECTION("rejects empty streaming load") { const auto empty_file_path = dir.GenerateRandomFile(false); - const auto empty_params = GenerateBucketPreciseParameters(1, "buffer_io", empty_file_path); + const auto empty_params = GenerateBucketPreciseParameters("buffer_io", empty_file_path); auto empty_index = TestFactory(name, empty_params, true); std::stringstream stream; REQUIRE(empty_index->SerializeStreaming(stream).has_value()); @@ -542,7 +550,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, SECTION("rejects mmap before creating the precise file") { const auto mmap_file_path = dir.GenerateRandomFile(false); - const auto mmap_params = GenerateBucketPreciseParameters(1, "mmap_io", mmap_file_path); + const auto mmap_params = GenerateBucketPreciseParameters("mmap_io", mmap_file_path); auto result = vsag::Factory::CreateIndex(name, mmap_params); REQUIRE_FALSE(result.has_value()); REQUIRE(result.error().type == vsag::ErrorType::INVALID_ARGUMENT); @@ -554,7 +562,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, REQUIRE(index->SerializeStreaming(stream).has_value()); const auto restored_file_path = dir.GenerateRandomFile(false); const auto restored_params = - GenerateBucketPreciseParameters(1, "buffer_io", restored_file_path); + GenerateBucketPreciseParameters("buffer_io", restored_file_path); auto restored = TestFactory(name, restored_params, true); std::stringstream deserialize_stream(stream.str()); REQUIRE(restored->DeserializeStreaming(deserialize_stream).has_value()); @@ -575,7 +583,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, constexpr int64_t dim = 16; constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; - const auto params = GenerateBucketPreciseParameters(2); + const auto params = GenerateBucketPreciseParameters(); const auto search_param = fmt::format(search_param_tmp, buckets_count); auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); auto model = TestFactory(name, params, true); From b7e77a0cd85d77656fcf46c4a868cb3d404a784a Mon Sep 17 00:00:00 2001 From: jac0626 Date: Tue, 4 Aug 2026 11:56:12 +0800 Subject: [PATCH 09/12] fix(datacell): validate serialized bucket metadata Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- src/datacell/bucket_datacell.h | 19 +++++++ src/datacell/bucket_datacell_test.cpp | 81 +++++++++++++++++++++++++++ tests/test_ivf.cpp | 14 ++++- 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/src/datacell/bucket_datacell.h b/src/datacell/bucket_datacell.h index dcdc0b5b80..7404a29518 100644 --- a/src/datacell/bucket_datacell.h +++ b/src/datacell/bucket_datacell.h @@ -554,6 +554,25 @@ BucketDataCell::Deserialize(lvalue_or_rvalue re } } StreamReader::ReadVector(reader, this->bucket_sizes_); + if (this->bucket_sizes_.size() != static_cast(this->bucket_count_)) { + throw VsagException(ErrorType::INVALID_BINARY, + "serialized bucket size vector does not match bucket count"); + } + for (BucketIdType i = 0; i < this->bucket_count_; ++i) { + const auto bucket_size = static_cast(this->bucket_sizes_[i]); + if (this->inner_ids_[i].size() < bucket_size) { + throw VsagException( + ErrorType::INVALID_BINARY, + fmt::format("serialized bucket {} inner id count is smaller than bucket size", i)); + } + if (this->use_residual_ and this->metric_ == MetricType::METRIC_TYPE_L2SQR and + this->residual_bias_[i].size() < bucket_size) { + throw VsagException( + ErrorType::INVALID_BINARY, + fmt::format("serialized bucket {} residual bias count is smaller than bucket size", + i)); + } + } } template diff --git a/src/datacell/bucket_datacell_test.cpp b/src/datacell/bucket_datacell_test.cpp index a4d737cae4..f78fd90617 100644 --- a/src/datacell/bucket_datacell_test.cpp +++ b/src/datacell/bucket_datacell_test.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -200,6 +201,86 @@ TEST_CASE("BucketDataCell rejects invalid parameters", "[ut][BucketDataCell]") { REQUIRE(BucketInterface::MakeInstance(nullptr, common_param) == nullptr); } +TEST_CASE("BucketDataCell rejects inconsistent serialized metadata", "[ut][BucketDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr int64_t dim = 4; + constexpr const char* param_str = R"( + { + "io_params": { + "type": "memory_io" + }, + "quantization_params": { + "type": "fp32" + }, + "buckets_count": 1 + } + )"; + + auto make_bucket = [&]() { + auto param_json = JsonType::Parse(param_str); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + return BucketInterface::MakeInstance(param, common_param); + }; + + auto bucket = make_bucket(); + auto vectors = fixtures::generate_vectors(1, dim); + bucket->Train(vectors.data(), 1); + bucket->InsertVector(vectors.data(), 0, 0); + + std::stringstream stream; + IOStreamWriter writer(stream); + bucket->Serialize(writer); + const auto serialized = stream.str(); + + SECTION("inner id vector is shorter than bucket size") { + auto malformed = serialized; + constexpr InnerIdType invalid_bucket_size = 2; + std::memcpy(malformed.data() + malformed.size() - sizeof(invalid_bucket_size), + &invalid_bucket_size, + sizeof(invalid_bucket_size)); + + std::stringstream malformed_stream(malformed); + IOStreamReader reader(malformed_stream); + auto restored = make_bucket(); + try { + restored->Deserialize(reader); + FAIL("inconsistent inner id metadata should be rejected"); + } catch (const VsagException& error) { + REQUIRE(error.error_.type == ErrorType::INVALID_BINARY); + REQUIRE(error.error_.message == + "serialized bucket 0 inner id count is smaller than bucket size"); + } + } + + SECTION("bucket size vector does not cover every bucket") { + auto malformed = serialized; + constexpr uint64_t invalid_bucket_size_count = 0; + const uint64_t count_offset = + static_cast(malformed.size()) - sizeof(InnerIdType) - sizeof(uint64_t); + std::memcpy(malformed.data() + count_offset, + &invalid_bucket_size_count, + sizeof(invalid_bucket_size_count)); + + std::stringstream malformed_stream(malformed); + IOStreamReader reader(malformed_stream); + auto restored = make_bucket(); + try { + restored->Deserialize(reader); + FAIL("inconsistent bucket size metadata should be rejected"); + } catch (const VsagException& error) { + REQUIRE(error.error_.type == ErrorType::INVALID_BINARY); + REQUIRE(error.error_.message == + "serialized bucket size vector does not match bucket count"); + } + } +} + TEST_CASE("BucketDataCell supports RabitQ", "[ut][BucketDataCell]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); constexpr uint64_t dim = 64; diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index c31ca62dd8..95b0ea878f 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -212,12 +212,17 @@ class BlockSizeLimitGuard { std::string GenerateBucketPreciseParameters(const std::string& precise_io_type = "block_memory_io", const std::string& precise_file_path = "", - int thread_count = 1) { + int thread_count = 1, + bool enable_read_cache = false) { auto params = nlohmann::json::parse(IVFTestIndex::GenerateIVFBuildParametersString( "l2", 16, "sq8,fp32", 16, "random", false, 1, false, thread_count)); params["index_param"]["precise_codes_layout"] = "bucket"; params["index_param"]["precise_io_type"] = precise_io_type; params["index_param"]["precise_file_path"] = precise_file_path; + if (enable_read_cache) { + params["index_param"]["precise_enable_read_cache"] = true; + params["index_param"]["precise_cache_total_size"] = 128 * 1024; + } return params.dump(); } @@ -429,11 +434,14 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, constexpr int64_t base_count = 128; constexpr int64_t buckets_count = 16; const auto precise_io_type = GENERATE("block_memory_io", "buffer_io"); + const bool enable_read_cache = precise_io_type == std::string("buffer_io"); INFO(fmt::format("precise_io_type: {}", precise_io_type)); + INFO(fmt::format("enable_read_cache: {}", enable_read_cache)); const auto precise_file_path = precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); - const auto params = GenerateBucketPreciseParameters(precise_io_type, precise_file_path); + const auto params = + GenerateBucketPreciseParameters(precise_io_type, precise_file_path, 1, enable_read_cache); const auto search_param = fmt::format(search_param_tmp, buckets_count); auto dataset = pool.GetDatasetAndCreate(dim, base_count, "l2"); auto index = TestFactory(name, params, true); @@ -447,7 +455,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, const auto restored_file_path = precise_io_type == std::string("buffer_io") ? dir.GenerateRandomFile(false) : std::string(); const auto restored_params = - GenerateBucketPreciseParameters(precise_io_type, restored_file_path); + GenerateBucketPreciseParameters(precise_io_type, restored_file_path, 1, enable_read_cache); auto restored = TestFactory(name, restored_params, true); TestSerializeBinarySet(index, restored, dataset, search_param, true); CheckBucketPreciseIndex(restored, dataset, search_param); From c5f4a5961ab38385b441679ad33cc9c661544517 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Tue, 4 Aug 2026 14:49:45 +0800 Subject: [PATCH 10/12] perf(ivf): batch precise bucket disk reads Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- src/algorithm/ivf/ivf.cpp | 47 ++- src/datacell/bucket_datacell.h | 167 ++++++++++ src/datacell/bucket_datacell_test.cpp | 284 ++++++++++++++++++ src/datacell/bucket_interface.h | 14 + src/io/noncontinuous_io/noncontinuous_io.h | 59 +++- .../noncontinuous_io_test.cpp | 134 +++++++++ tests/test_ivf.cpp | 4 +- 7 files changed, 698 insertions(+), 11 deletions(-) diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index 97ab122a42..962ba77ef1 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -1591,10 +1591,23 @@ IVF::reorder_with_precise_bucket(const DistHeapPtr& input, auto computer = precise_bucket_->FactoryComputer(query); const auto* candidates = input->GetData(); + Vector bucket_ids(candidate_count, query_allocator); + Vector offset_ids(candidate_count, query_allocator); + Vector precise_distances(candidate_count, query_allocator); + for (uint64_t i = 0; i < candidate_count; ++i) { + const auto [bucket_id, offset_id] = this->get_location(candidates[i].second); + bucket_ids[i] = bucket_id; + offset_ids[i] = offset_id; + } + precise_bucket_->Query(precise_distances.data(), + computer, + bucket_ids.data(), + offset_ids.data(), + static_cast(candidate_count), + &ctx); for (uint64_t i = 0; i < candidate_count; ++i) { const auto [coarse_distance, inner_id] = candidates[i]; - const auto [bucket_id, offset_id] = this->get_location(inner_id); - auto precise_distance = precise_bucket_->QueryOneById(computer, bucket_id, offset_id); + const auto precise_distance = precise_distances[i]; if (ctx.reasoning_ctx != nullptr) { ctx.reasoning_ctx->RecordReorder(inner_id, coarse_distance, precise_distance); } @@ -2265,13 +2278,37 @@ IVF::CalDistanceById(const float* query, if (this->use_reorder_ && calculate_precise_distance && reorder_codes_ != nullptr) { auto computer = this->reorder_codes_->FactoryComputer(query); this->reorder_codes_->Query(distances, computer, inner_ids.data(), count); + } else if (this->use_reorder_ && calculate_precise_distance && precise_bucket_ != nullptr) { + auto computer = this->precise_bucket_->FactoryComputer(query); + Vector bucket_ids(allocator_); + Vector offset_ids(allocator_); + Vector result_indices(allocator_); + bucket_ids.reserve(count); + offset_ids.reserve(count); + result_indices.reserve(count); + for (int64_t i = 0; i < count; ++i) { + if (validity[i]) { + auto [bucket_id, offset_id] = this->get_location(inner_ids[i]); + bucket_ids.emplace_back(bucket_id); + offset_ids.emplace_back(offset_id); + result_indices.emplace_back(i); + } + } + Vector valid_distances(result_indices.size(), allocator_); + this->precise_bucket_->Query(valid_distances.data(), + computer, + bucket_ids.data(), + offset_ids.data(), + static_cast(result_indices.size())); + for (uint64_t i = 0; i < result_indices.size(); ++i) { + distances[result_indices[i]] = valid_distances[i]; + } } else { - auto codes = this->use_reorder_ && calculate_precise_distance ? precise_bucket_ : bucket_; - auto computer = codes->FactoryComputer(query); + auto computer = this->bucket_->FactoryComputer(query); for (int64_t i = 0; i < count; ++i) { if (validity[i]) { auto [bucket_id, offset_id] = this->get_location(inner_ids[i]); - distances[i] = codes->QueryOneById(computer, bucket_id, offset_id); + distances[i] = this->bucket_->QueryOneById(computer, bucket_id, offset_id); } } } diff --git a/src/datacell/bucket_datacell.h b/src/datacell/bucket_datacell.h index 7404a29518..27d464c46f 100644 --- a/src/datacell/bucket_datacell.h +++ b/src/datacell/bucket_datacell.h @@ -27,6 +27,7 @@ #include "quantization/product_quantization/pq_fastscan_quantizer.h" #include "simd/fp32_simd.h" #include "utils/byte_buffer.h" +#include "utils/timer.h" namespace vsag { @@ -67,6 +68,22 @@ class BucketDataCell : public BucketInterface { return this->query_one_by_id(comp, bucket_id, offset_id); } + void + Query(float* result_dists, + const ComputerInterfacePtr& computer, + const BucketIdType* bucket_ids, + const InnerIdType* offset_ids, + InnerIdType id_count, + QueryContext* ctx = nullptr) override { + if (id_count > 0 and GetQuantizerName() == QUANTIZATION_TYPE_VALUE_PQFS) { + throw VsagException(ErrorType::INTERNAL_ERROR, + "PQFastScan doesn't support ComputeDist, only support " + "ComputeBatchDist"); + } + auto comp = static_cast*>(computer.get()); + this->query(result_dists, comp, bucket_ids, offset_ids, id_count, ctx); + } + ComputerInterfacePtr FactoryComputer(const void* query) override; @@ -169,6 +186,14 @@ class BucketDataCell : public BucketInterface { const BucketIdType& bucket_id, const InnerIdType& offset_id); + inline void + query(float* result_dists, + Computer* computer, + const BucketIdType* bucket_ids, + const InnerIdType* offset_ids, + InnerIdType id_count, + QueryContext* ctx); + inline void encode_vector(const void* vector, BucketIdType bucket_id, ByteBuffer& codes, float& res_score); @@ -279,6 +304,148 @@ BucketDataCell::query_one_by_id( return ret; } +template +void +BucketDataCell::query(float* result_dists, + Computer* computer, + const BucketIdType* bucket_ids, + const InnerIdType* offset_ids, + InnerIdType id_count, + QueryContext* ctx) { + if (id_count == 0) { + return; + } + + struct QueryRequest { + BucketIdType bucket_id; + InnerIdType offset_id; + InnerIdType result_index; + }; + + Allocator* search_alloc = select_query_allocator(ctx, allocator_); + Vector requests(search_alloc); + requests.reserve(id_count); + for (InnerIdType i = 0; i < id_count; ++i) { + check_valid_bucket_id(bucket_ids[i]); + requests.emplace_back(QueryRequest{bucket_ids[i], offset_ids[i], i}); + } + std::sort(requests.begin(), requests.end(), [](const auto& lhs, const auto& rhs) { + if (lhs.bucket_id != rhs.bucket_id) { + return lhs.bucket_id < rhs.bucket_id; + } + if (lhs.offset_id != rhs.offset_id) { + return lhs.offset_id < rhs.offset_id; + } + return lhs.result_index < rhs.result_index; + }); + + uint64_t io_count = 0; + double io_cost_ms = 0.0F; + Vector unique_offsets(search_alloc); + Vector read_sizes(search_alloc); + Vector read_offsets(search_alloc); + unique_offsets.reserve(id_count); + read_sizes.reserve(id_count); + read_offsets.reserve(id_count); + ByteBuffer codes(static_cast(id_count) * code_size_, search_alloc); + Vector unique_dists(id_count, 0.0F, search_alloc); + uint64_t group_begin = 0; + while (group_begin < requests.size()) { + const auto bucket_id = requests[group_begin].bucket_id; + uint64_t group_end = group_begin + 1; + while (group_end < requests.size() and requests[group_end].bucket_id == bucket_id) { + ++group_end; + } + + std::shared_lock lock(this->bucket_mutexes_[bucket_id]); + for (uint64_t i = group_begin; i < group_end; ++i) { + const auto offset_id = requests[i].offset_id; + if (offset_id >= this->bucket_sizes_[bucket_id]) { + throw VsagException(ErrorType::INVALID_ARGUMENT, "invalid offset id for bucket"); + } + if (this->inner_ids_[bucket_id][offset_id] == EMPTY_INNER_ID) { + throw VsagException( + ErrorType::INVALID_ARGUMENT, + fmt::format("visited empty offset in bucket: bucket_id={}, offset_id={}", + bucket_id, + offset_id)); + } + } + + unique_offsets.clear(); + for (uint64_t i = group_begin; i < group_end; ++i) { + const auto offset_id = requests[i].offset_id; + if (unique_offsets.empty() or unique_offsets.back() != offset_id) { + unique_offsets.emplace_back(offset_id); + } + } + + read_sizes.clear(); + read_offsets.clear(); + for (const auto offset_id : unique_offsets) { + if (not read_offsets.empty() and + read_offsets.back() + read_sizes.back() == + static_cast(offset_id) * code_size_ and + code_size_ <= static_cast(std::numeric_limits::max()) and + read_sizes.back() <= + static_cast(std::numeric_limits::max()) - code_size_) { + read_sizes.back() += code_size_; + } else { + read_sizes.emplace_back(code_size_); + read_offsets.emplace_back(static_cast(offset_id) * code_size_); + } + } + + bool read_success = false; + double group_io_cost_ms = 0.0F; + { + Timer timer(group_io_cost_ms); + read_success = this->datas_[bucket_id].MultiRead( + codes.data, read_sizes.data(), read_offsets.data(), read_sizes.size()); + } + if (not read_success) { + throw VsagException(ErrorType::READ_ERROR, "failed to batch read bucket data"); + } + io_count += read_sizes.size(); + io_cost_ms += group_io_cost_ms; + + computer->ScanBatchDists(unique_offsets.size(), codes.data, unique_dists.data()); + if (use_residual_) { + Vector centroid(this->quantizer_->GetDim(), search_alloc); + strategy_->GetCentroid(bucket_id, centroid); + auto ip_distance = FP32ComputeIP( + computer->raw_query_.data(), centroid.data(), this->quantizer_->GetDim()); + if (metric_ == MetricType::METRIC_TYPE_L2SQR) { + ip_distance *= 2; + for (uint64_t i = 0; i < unique_offsets.size(); ++i) { + unique_dists[i] -= residual_bias_[bucket_id][unique_offsets[i]]; + } + } + for (uint64_t i = 0; i < unique_offsets.size(); ++i) { + unique_dists[i] -= ip_distance; + } + } + + uint64_t unique_index = 0; + for (uint64_t i = group_begin; i < group_end; ++i) { + if (i > group_begin and requests[i].offset_id != requests[i - 1].offset_id) { + ++unique_index; + } + result_dists[requests[i].result_index] = unique_dists[unique_index]; + } + group_begin = group_end; + } + + if constexpr (not IOTmpl::InMemory) { + if (ctx != nullptr and ctx->stats != nullptr) { + ctx->stats->io_cnt.fetch_add(static_cast(io_count), + std::memory_order_relaxed); + ctx->stats->io_time_ms.fetch_add(static_cast(io_cost_ms), + std::memory_order_relaxed); + } + } +} + template void BucketDataCell::scan_bucket_by_id(float* result_dists, diff --git a/src/datacell/bucket_datacell_test.cpp b/src/datacell/bucket_datacell_test.cpp index f78fd90617..f8f282b77d 100644 --- a/src/datacell/bucket_datacell_test.cpp +++ b/src/datacell/bucket_datacell_test.cpp @@ -33,6 +33,34 @@ using namespace vsag; +namespace { + +class FixedCentroidPartitionStrategy : public IVFPartitionStrategy { +public: + FixedCentroidPartitionStrategy(const IndexCommonParam& common_param, BucketIdType bucket_count) + : IVFPartitionStrategy(common_param, bucket_count) { + } + + void + Train(const DatasetPtr) override { + is_trained_ = true; + } + + Vector + ClassifyDatas(const void*, int64_t count, BucketIdType, QueryContext*) const override { + return Vector(count, 0, allocator_); + } + + void + GetCentroid(BucketIdType bucket_id, Vector& centroid) override { + for (uint64_t i = 0; i < centroid.size(); ++i) { + centroid[i] = static_cast((bucket_id + 1) * (i + 1)) * 0.01F; + } + } +}; + +} // namespace + namespace vsag { class BucketInterfaceTest { public: @@ -281,6 +309,262 @@ TEST_CASE("BucketDataCell rejects inconsistent serialized metadata", "[ut][Bucke } } +TEST_CASE("BucketDataCell batch query", "[ut][BucketDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto query_allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr int64_t dim = 8; + constexpr uint64_t bucket_count = 3; + constexpr uint64_t vectors_per_bucket = 4; + constexpr uint64_t base_count = bucket_count * vectors_per_bucket; + const auto io_type = GENERATE(std::string("memory_io"), std::string("buffer_io")); + auto vectors = fixtures::generate_vectors(base_count, dim); + auto queries = fixtures::generate_vectors(1, dim, 41); + fixtures::TempDir temp_dir("vsag_bucket_batch_query_test"); + + auto make_bucket = [&]() { + auto file_path = temp_dir.GenerateRandomFile(false); + auto param_json = JsonType::Parse(fmt::format( + R"({{ + "io_params": {{ + "type": "{}", + "file_path": "{}" + }}, + "quantization_params": {{ + "type": "fp32" + }}, + "buckets_count": {} + }})", + io_type, + file_path, + bucket_count)); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + bucket->Train(vectors.data(), base_count); + return bucket; + }; + + auto bucket = make_bucket(); + for (uint64_t bucket_id = 0; bucket_id < bucket_count; ++bucket_id) { + for (uint64_t offset_id = 0; offset_id < vectors_per_bucket; ++offset_id) { + auto inner_id = bucket_id * vectors_per_bucket + offset_id; + bucket->InsertVector(vectors.data() + inner_id * dim, + static_cast(bucket_id), + static_cast(inner_id)); + } + } + + std::vector bucket_ids{2, 0, 1, 0, 2, 1, 0}; + std::vector offset_ids{3, 1, 2, 2, 3, 0, 1}; + std::vector expected(bucket_ids.size()); + std::vector actual(bucket_ids.size()); + auto computer = bucket->FactoryComputer(queries.data()); + for (uint64_t i = 0; i < bucket_ids.size(); ++i) { + expected[i] = bucket->QueryOneById(computer, bucket_ids[i], offset_ids[i]); + } + + SearchStatistics stats; + QueryContext ctx{query_allocator.get(), &stats}; + bucket->Query(actual.data(), + computer, + bucket_ids.data(), + offset_ids.data(), + static_cast(bucket_ids.size()), + &ctx); + REQUIRE(actual == expected); + if (io_type == "buffer_io") { + // Four contiguous read ranges remain after sorting and de-duplicating the locations. + REQUIRE(stats.io_cnt.load(std::memory_order_relaxed) == 4); + } + + REQUIRE_NOTHROW(bucket->Query(nullptr, ComputerInterfacePtr{}, nullptr, nullptr, 0, &ctx)); + + SECTION("invalid bucket is rejected") { + BucketIdType invalid_bucket_id = -1; + InnerIdType offset_id = 0; + float dist = 0.0F; + REQUIRE_THROWS(bucket->Query(&dist, computer, &invalid_bucket_id, &offset_id, 1)); + } + + SECTION("invalid offset is rejected") { + BucketIdType bucket_id = 0; + InnerIdType invalid_offset_id = 100; + float dist = 0.0F; + REQUIRE_THROWS(bucket->Query(&dist, computer, &bucket_id, &invalid_offset_id, 1)); + } + + SECTION("hole is rejected") { + auto sparse_bucket = make_bucket(); + sparse_bucket->InsertVectorWithOffset(vectors.data() + 2 * dim, 0, 2, 2); + auto sparse_computer = sparse_bucket->FactoryComputer(queries.data()); + BucketIdType bucket_id = 0; + InnerIdType hole_offset_id = 0; + float dist = 0.0F; + REQUIRE_THROWS( + sparse_bucket->Query(&dist, sparse_computer, &bucket_id, &hole_offset_id, 1)); + } +} + +TEST_CASE("BucketDataCell batch query preserves residual correction", "[ut][BucketDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr int64_t dim = 8; + constexpr uint64_t bucket_count = 2; + constexpr uint64_t base_count = 8; + auto vectors = fixtures::generate_vectors(base_count, dim); + auto queries = fixtures::generate_vectors(1, dim, 43); + MetricType metrics[] = { + MetricType::METRIC_TYPE_L2SQR, MetricType::METRIC_TYPE_IP, MetricType::METRIC_TYPE_COSINE}; + + for (auto metric : metrics) { + auto param_json = JsonType::Parse(R"({ + "io_params": { + "type": "memory_io" + }, + "quantization_params": { + "type": "fp32" + }, + "buckets_count": 2, + "use_residual": true + })"); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = metric; + auto bucket = BucketInterface::MakeInstance(param, common_param); + auto strategy = + std::make_shared(common_param, bucket_count); + bucket->SetStrategy(strategy); + bucket->Train(vectors.data(), base_count); + for (uint64_t i = 0; i < base_count; ++i) { + bucket->InsertVector(vectors.data() + i * dim, + static_cast(i % bucket_count), + static_cast(i)); + } + + std::vector bucket_ids{1, 0, 1, 0, 1}; + std::vector offset_ids{2, 3, 0, 1, 2}; + std::vector expected(bucket_ids.size()); + std::vector actual(bucket_ids.size()); + auto computer = bucket->FactoryComputer(queries.data()); + for (uint64_t i = 0; i < bucket_ids.size(); ++i) { + expected[i] = bucket->QueryOneById(computer, bucket_ids[i], offset_ids[i]); + } + bucket->Query(actual.data(), + computer, + bucket_ids.data(), + offset_ids.data(), + static_cast(bucket_ids.size())); + for (uint64_t i = 0; i < actual.size(); ++i) { + REQUIRE(std::abs(actual[i] - expected[i]) < 1e-5F); + } + } +} + +TEST_CASE("BucketDataCell batch query handles all bucket quantizers", "[ut][BucketDataCell]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr int64_t dim = 64; + constexpr uint64_t train_count = 300; + constexpr uint64_t bucket_count = 2; + constexpr uint64_t vectors_per_bucket = 64; + auto vectors = fixtures::generate_vectors(train_count, dim); + auto queries = fixtures::generate_vectors(1, dim, 47); + fixtures::TempDir temp_dir("vsag_bucket_batch_quantizers_test"); + + const std::vector> quantizers = { + {"fp32", R"({"type": "fp32"})"}, + {"sq8", R"({"type": "sq8"})"}, + {"sq4", R"({"type": "sq4"})"}, + {"sq4_uniform", R"({"type": "sq4_uniform"})"}, + {"sq8_uniform", R"({"type": "sq8_uniform"})"}, + {"bf16", R"({"type": "bf16"})"}, + {"fp16", R"({"type": "fp16"})"}, + {"pq", R"({"type": "pq", "pq_dim": 8, "pq_bits": 8})"}, + {"pqfs", R"({"type": "pqfs", "pq_dim": 8})"}, + {"rabitq", + R"({"type": "rabitq", "rabitq_bits_per_dim_query": 32, "rabitq_bits_per_dim_base": 1})"}, + }; + + for (const auto& io_type : {std::string("memory_io"), std::string("buffer_io")}) { + for (const auto& [quantizer_name, quantizer_json] : quantizers) { + CAPTURE(io_type, quantizer_name); + JsonType param_json; + JsonType io_json; + io_json["type"].SetString(io_type); + io_json["file_path"].SetString(temp_dir.GenerateRandomFile(false)); + param_json["io_params"].SetJson(io_json); + param_json["quantization_params"].SetJson(JsonType::Parse(quantizer_json)); + param_json["buckets_count"].SetInt(bucket_count); + auto param = std::make_shared(); + param->FromJson(param_json); + + IndexCommonParam common_param; + common_param.allocator_ = allocator; + common_param.dim_ = dim; + common_param.metric_ = MetricType::METRIC_TYPE_L2SQR; + auto bucket = BucketInterface::MakeInstance(param, common_param); + bucket->Train(vectors.data(), train_count); + for (uint64_t bucket_id = 0; bucket_id < bucket_count; ++bucket_id) { + for (uint64_t offset_id = 0; offset_id < vectors_per_bucket; ++offset_id) { + auto inner_id = bucket_id * vectors_per_bucket + offset_id; + bucket->InsertVector(vectors.data() + inner_id * dim, + static_cast(bucket_id), + static_cast(inner_id)); + } + } + // PQFS queries require its normal IVF package state; Package is a no-op for others. + bucket->Package(); + + auto computer = bucket->FactoryComputer(queries.data()); + std::vector bucket_ids{1, 0, 1, 0, 1, 0}; + std::vector offset_ids{33, 1, 2, 34, 33, 2}; + std::vector actual(bucket_ids.size()); + if (quantizer_name == "pqfs") { + try { + bucket->Query(actual.data(), + computer, + bucket_ids.data(), + offset_ids.data(), + static_cast(bucket_ids.size())); + FAIL("PQFS batch point query should preserve QueryOneById rejection"); + } catch (const VsagException& error) { + REQUIRE(error.error_.type == ErrorType::INTERNAL_ERROR); + REQUIRE(error.error_.message == + "PQFastScan doesn't support ComputeDist, only support " + "ComputeBatchDist"); + } + continue; + } + + std::vector expected(bucket_ids.size()); + for (uint64_t i = 0; i < expected.size(); ++i) { + expected[i] = bucket->QueryOneById(computer, bucket_ids[i], offset_ids[i]); + } + SearchStatistics stats; + QueryContext ctx{nullptr, &stats}; + bucket->Query(actual.data(), + computer, + bucket_ids.data(), + offset_ids.data(), + static_cast(bucket_ids.size()), + &ctx); + for (uint64_t i = 0; i < actual.size(); ++i) { + REQUIRE(std::abs(actual[i] - expected[i]) < 1e-5F); + } + if (io_type == "buffer_io") { + REQUIRE(stats.io_cnt.load(std::memory_order_relaxed) == 4); + } + } + } +} + TEST_CASE("BucketDataCell supports RabitQ", "[ut][BucketDataCell]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); constexpr uint64_t dim = 64; diff --git a/src/datacell/bucket_interface.h b/src/datacell/bucket_interface.h index a40a87213a..9568c89647 100644 --- a/src/datacell/bucket_interface.h +++ b/src/datacell/bucket_interface.h @@ -21,6 +21,7 @@ #include "bucket_datacell_parameter.h" #include "index_common_param.h" #include "quantization/computer.h" +#include "query_context.h" #include "storage/stream_reader.h" #include "storage/stream_writer.h" #include "typing.h" @@ -47,6 +48,19 @@ class BucketInterface { const BucketIdType& bucket_id, const InnerIdType& offset_id) = 0; + virtual void + Query(float* result_dists, + const ComputerInterfacePtr& computer, + const BucketIdType* bucket_ids, + const InnerIdType* offset_ids, + InnerIdType id_count, + QueryContext* ctx = nullptr) { + (void)ctx; + for (InnerIdType i = 0; i < id_count; ++i) { + result_dists[i] = QueryOneById(computer, bucket_ids[i], offset_ids[i]); + } + } + virtual ComputerInterfacePtr FactoryComputer(const void* query) = 0; diff --git a/src/io/noncontinuous_io/noncontinuous_io.h b/src/io/noncontinuous_io/noncontinuous_io.h index f0d1855479..8d9632de38 100644 --- a/src/io/noncontinuous_io/noncontinuous_io.h +++ b/src/io/noncontinuous_io/noncontinuous_io.h @@ -173,12 +173,61 @@ class NonContinuousIO : public BasicIO> { */ bool MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const { - bool ret = true; - for (uint64_t i = 0; i < count; i++) { - ret &= this->ReadImpl(sizes[i], offsets[i], datas); - datas += sizes[i]; + if (count == 0) { + return true; } - return ret; + if (sizes == nullptr or offsets == nullptr) { + return false; + } + + bool has_data = false; + for (uint64_t i = 0; i < count; ++i) { + if (offsets[i] > this->size_ or sizes[i] > this->size_ - offsets[i]) { + return false; + } + has_data |= sizes[i] > 0; + } + if (not has_data) { + return true; + } + if (datas == nullptr) { + return false; + } + + std::vector physical_sizes; + std::vector physical_offsets; + physical_sizes.reserve(count); + physical_offsets.reserve(count); + for (uint64_t i = 0; i < count; ++i) { + uint64_t remaining_size = sizes[i]; + uint64_t logical_offset = offsets[i]; + if (remaining_size == 0) { + continue; + } + + auto area_it = this->get_area(logical_offset); + while (remaining_size > 0) { + if (area_it == areas_.end()) { + return false; + } + const auto& area = area_it->first; + uint64_t logical_area_start = area_it->second - area.size; + if (logical_offset < logical_area_start or logical_offset >= area_it->second) { + return false; + } + uint64_t fragment_size = std::min(remaining_size, area_it->second - logical_offset); + physical_sizes.emplace_back(fragment_size); + physical_offsets.emplace_back(area.offset + logical_offset - logical_area_start); + logical_offset += fragment_size; + remaining_size -= fragment_size; + ++area_it; + } + } + + return inner_io_->MultiRead(datas, + physical_sizes.data(), + physical_offsets.data(), + static_cast(physical_sizes.size())); } /** diff --git a/src/io/noncontinuous_io/noncontinuous_io_test.cpp b/src/io/noncontinuous_io/noncontinuous_io_test.cpp index 85ed444885..374216fb7c 100644 --- a/src/io/noncontinuous_io/noncontinuous_io_test.cpp +++ b/src/io/noncontinuous_io/noncontinuous_io_test.cpp @@ -17,6 +17,9 @@ #include #include +#include +#include +#include #include #include @@ -27,6 +30,48 @@ #include "io/mmap_io/mmap_io.h" #include "unittest.h" namespace vsag { +struct TrackingIOState { + std::vector data_; + uint64_t multi_read_calls_{0}; + std::vector last_sizes_; + std::vector last_offsets_; +}; + +class TrackingIO : public BasicIO { +public: + static constexpr bool InMemory = true; + static constexpr bool SkipDeserialize = false; + + TrackingIO(std::shared_ptr state, Allocator* allocator) + : BasicIO(allocator), state_(std::move(state)) { + } + + void + WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) { + state_->data_.resize(std::max(state_->data_.size(), offset + size)); + std::memcpy(state_->data_.data() + offset, data, size); + this->size_ = std::max(this->size_, offset + size); + } + + bool + MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const { + ++state_->multi_read_calls_; + state_->last_sizes_.assign(sizes, sizes + count); + state_->last_offsets_.assign(offsets, offsets + count); + for (uint64_t i = 0; i < count; ++i) { + if (offsets[i] > state_->data_.size() or sizes[i] > state_->data_.size() - offsets[i]) { + return false; + } + std::memcpy(datas, state_->data_.data() + offsets[i], sizes[i]); + datas += sizes[i]; + } + return true; + } + +private: + std::shared_ptr state_; +}; + template class NonContinuousIOTest { public: @@ -95,6 +140,95 @@ TEST_CASE("NonContinuousIO Serialize Test", "[NonContinuousIO][ut]") { NonContinuousIOTestSerialize(); } +TEST_CASE("NonContinuousIO batches physical fragments", "[NonContinuousIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto non_continuous_allocator = std::make_unique(allocator.get()); + auto state = std::make_shared(); + NonContinuousIOTest test; + std::unique_ptr> io(test.CreateNonContinuousIO( + non_continuous_allocator.get(), allocator.get(), state, allocator.get())); + std::unique_ptr> spacer(test.CreateNonContinuousIO( + non_continuous_allocator.get(), allocator.get(), state, allocator.get())); + + constexpr uint64_t page_size = 4096; + std::vector logical_data(page_size * 3); + for (uint64_t i = 0; i < logical_data.size(); ++i) { + logical_data[i] = static_cast(i % 251); + } + std::vector spacer_data(page_size, 0xFF); + io->Write(logical_data.data(), page_size, 0); + spacer->Write(spacer_data.data(), page_size, 0); + io->Write(logical_data.data() + page_size, page_size, page_size); + spacer->Write(spacer_data.data(), page_size, page_size); + io->Write(logical_data.data() + page_size * 2, page_size, page_size * 2); + + std::vector sizes{12, 0, 5, 12, 8}; + std::vector offsets{ + page_size * 2 - 4, page_size * 3, 7, page_size * 2 - 4, page_size - 2}; + std::vector output(37); + REQUIRE(io->MultiRead(output.data(), sizes.data(), offsets.data(), sizes.size())); + + std::vector expected; + expected.reserve(output.size()); + for (uint64_t i = 0; i < sizes.size(); ++i) { + expected.insert(expected.end(), + logical_data.begin() + offsets[i], + logical_data.begin() + offsets[i] + sizes[i]); + } + REQUIRE(output == expected); + REQUIRE(state->multi_read_calls_ == 1); + REQUIRE(state->last_sizes_ == std::vector{4, 8, 5, 4, 8, 2, 6}); + REQUIRE(state->last_offsets_ == std::vector{page_size * 3 - 4, + page_size * 4, + 7, + page_size * 3 - 4, + page_size * 4, + page_size - 2, + page_size * 2}); +} + +TEST_CASE("NonContinuousIO validates batch ranges and empty reads", "[NonContinuousIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto non_continuous_allocator = std::make_unique(allocator.get()); + auto state = std::make_shared(); + NonContinuousIOTest test; + std::unique_ptr> io(test.CreateNonContinuousIO( + non_continuous_allocator.get(), allocator.get(), state, allocator.get())); + + constexpr uint64_t page_size = 4096; + std::vector data(page_size, 0x5A); + io->Write(data.data(), data.size(), 0); + + REQUIRE(io->MultiRead(nullptr, nullptr, nullptr, 0)); + REQUIRE(state->multi_read_calls_ == 0); + + std::vector empty_sizes{0, 0, 0}; + std::vector empty_offsets{0, page_size / 2, page_size}; + REQUIRE(io->MultiRead(nullptr, empty_sizes.data(), empty_offsets.data(), empty_sizes.size())); + REQUIRE(state->multi_read_calls_ == 0); + + std::vector invalid_sizes{4, 1}; + std::vector invalid_offsets{0, page_size}; + std::vector output(5, 0xA5); + REQUIRE_FALSE(io->MultiRead( + output.data(), invalid_sizes.data(), invalid_offsets.data(), invalid_sizes.size())); + REQUIRE(output == std::vector(5, 0xA5)); + REQUIRE(state->multi_read_calls_ == 0); + + uint64_t overflow_size = 16; + uint64_t overflow_offset = std::numeric_limits::max() - 7; + REQUIRE_FALSE(io->MultiRead(&output[0], &overflow_size, &overflow_offset, 1)); + REQUIRE(state->multi_read_calls_ == 0); + + uint64_t zero_size = 0; + uint64_t past_end_offset = page_size + 1; + REQUIRE_FALSE(io->MultiRead(nullptr, &zero_size, &past_end_offset, 1)); + REQUIRE_FALSE(io->MultiRead(nullptr, &overflow_size, &zero_size, 1)); + REQUIRE_FALSE(io->MultiRead(output.data(), nullptr, &zero_size, 1)); + REQUIRE_FALSE(io->MultiRead(output.data(), &zero_size, nullptr, 1)); + REQUIRE(state->multi_read_calls_ == 0); +} + TEST_CASE("NonContinuousAllocator allocates unique regions concurrently", "[NonContinuousIO][ut][concurrent]") { constexpr uint64_t thread_count = 16; diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 95b0ea878f..1700d3c058 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -221,7 +221,9 @@ GenerateBucketPreciseParameters(const std::string& precise_io_type = "block_memo params["index_param"]["precise_file_path"] = precise_file_path; if (enable_read_cache) { params["index_param"]["precise_enable_read_cache"] = true; - params["index_param"]["precise_cache_total_size"] = 128 * 1024; + // BucketDataCell divides the cache budget across the 16 buckets. Allocate one + // 128-KiB cache page per bucket so this case exercises the read-cache path. + params["index_param"]["precise_cache_total_size"] = 16 * 128 * 1024; } return params.dump(); } From 0eacadaded249aa7e97dccb4fa3d9fe67167625b Mon Sep 17 00:00:00 2001 From: jac0626 Date: Wed, 5 Aug 2026 14:49:27 +0800 Subject: [PATCH 11/12] test(ivf): cover precise bucket thresholds Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- tests/test_ivf.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 1700d3c058..4970b04214 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -245,6 +245,12 @@ CheckBucketPreciseIndex(const TestIndex::IndexPtr& index, REQUIRE(result.value()->GetIds()[0] == dataset->base_->GetIds()[query_id]); REQUIRE(std::abs(result.value()->GetDistances()[0]) < 2e-6F); + auto threshold_search_param = nlohmann::json::parse(search_param); + threshold_search_param["threshold"] = -1.0F; + auto threshold_result = index->KnnSearch(query, 1, threshold_search_param.dump()); + REQUIRE(threshold_result.has_value()); + REQUIRE(threshold_result.value()->GetDim() == 0); + auto distance = index->CalcDistanceById(query->GetFloat32Vectors(), dataset->base_->GetIds()[query_id]); REQUIRE(distance.has_value()); From f79de195222d2d2aad2df6fe271f0b99ca78b1c0 Mon Sep 17 00:00:00 2001 From: jac0626 Date: Wed, 5 Aug 2026 16:04:48 +0800 Subject: [PATCH 12/12] style(ivf): format threshold reorder path Signed-off-by: jac0626 Assisted-by: Codex:gpt-5 --- src/algorithm/ivf/ivf.cpp | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index 962ba77ef1..a315e98391 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -1560,11 +1560,10 @@ IVF::reorder(int64_t topk, QueryContext& ctx, ReasoningContext* reasoning_ctx, const std::optional& distance_threshold) const { - auto reorder_heap = precise_bucket_ != nullptr - ? this->reorder_with_precise_bucket( - input, query, topk, ctx, distance_threshold) - : reorder_->Reorder( - input, query, topk, ctx, nullptr, nullptr, distance_threshold); + auto reorder_heap = + precise_bucket_ != nullptr + ? this->reorder_with_precise_bucket(input, query, topk, ctx, distance_threshold) + : reorder_->Reorder(input, query, topk, ctx, nullptr, nullptr, distance_threshold); auto dataset_results = this->pack_knn_result(reorder_heap, ctx.alloc); return dataset_results; @@ -1611,9 +1610,8 @@ IVF::reorder_with_precise_bucket(const DistHeapPtr& input, if (ctx.reasoning_ctx != nullptr) { ctx.reasoning_ctx->RecordReorder(inner_id, coarse_distance, precise_distance); } - if (distance_threshold.has_value() and - (not std::isfinite(precise_distance) or - precise_distance > distance_threshold.value())) { + if (distance_threshold.has_value() and (not std::isfinite(precise_distance) or + precise_distance > distance_threshold.value())) { continue; } if (reorder_heap->Size() < topk or precise_distance < reorder_heap->Top().first) {