diff --git a/docs/docs/en/src/api/dataset.md b/docs/docs/en/src/api/dataset.md index f47f95d00f..154ab3e2bc 100644 --- a/docs/docs/en/src/api/dataset.md +++ b/docs/docs/en/src/api/dataset.md @@ -116,9 +116,12 @@ if (result.has_value()) { } ``` -For KNN, `GetNumElements()` is `1` and the ids/distances arrays have length `k`. For range search, -the number of matches is reported through the result's dimension. See -[k-Nearest Neighbor Search](../guide/knn_search.md). +For single-query KNN, `GetNumElements()` is `1` and the ids/distances arrays have length `k`. HGraph +and IVF batched KNN returns a row-major `GetNumElements() x GetDim()` matrix; slot `i` of row `q` is +at `q * GetDim() + i`. Short rows are padded with `id == -1` and infinite distance. Batched KNN +rejects an index containing external label `-1` to keep this padding unambiguous. Range search +supports one query only. +See [k-Nearest Neighbor Search](../guide/knn_search.md). ## `SparseVector` diff --git a/docs/docs/en/src/api/index_class.md b/docs/docs/en/src/api/index_class.md index eb24fd7938..75966f6ecc 100644 --- a/docs/docs/en/src/api/index_class.md +++ b/docs/docs/en/src/api/index_class.md @@ -129,13 +129,15 @@ See `examples/cpp/303_feature_remove.cpp`. ## Search -The recommended entry point is [`SearchWithRequest`](#searchwithrequest), which takes a single +The recommended entry point is [`SearchWithRequest`](#searchwithrequest), which takes a [`SearchRequest`](search.md#searchrequest) carrying the query, mode, top-k / radius, and any filters. The older per-argument `KnnSearch` / `RangeSearch` overloads remain for compatibility. -Every search returns a `DatasetPtr`: for KNN, `num_elements == 1` and `ids` / `distances` have -length `k`; for range search, the result length is the number of matches. See [Dataset](dataset.md) -for how to read results. +Every search returns a `DatasetPtr`: single-query KNN has `num_elements == 1`; HGraph and IVF +also support batched KNN, returning a row-major `num_elements x dim` matrix. `dim` is the returned +row width and can be clamped below the requested top-k. Missing entries are padded with `id == -1`; +batched KNN rejects an index containing external label `-1` to keep this padding unambiguous. Range +search accepts a single query only. See [Dataset](dataset.md). ### `SearchWithRequest` diff --git a/docs/docs/en/src/api/search.md b/docs/docs/en/src/api/search.md index 5019a05438..96a6e794a1 100644 --- a/docs/docs/en/src/api/search.md +++ b/docs/docs/en/src/api/search.md @@ -34,7 +34,7 @@ enum class SearchMode { | Field | Type | Default | Meaning | |-------|------|---------|---------| -| `query_` | `DatasetPtr` | `nullptr` | The query. IVF KNN requests support multiple query vectors; other requests allow one. | +| `query_` | `DatasetPtr` | `nullptr` | The query. HGraph and IVF support contiguous multi-query KNN batches; range search supports one query only. | | `mode_` | `SearchMode` | `KNN_SEARCH` | KNN vs. range search. | | `topk_` | `int64_t` | `10` | Neighbors to return (KNN mode). Must be positive. | | `radius_` | `float` | `0.5` | Distance threshold (range mode). Non-negative. | @@ -81,12 +81,12 @@ Reordering is automatically disabled in callback mode. Other indexes do not supp ### IVF bucket routing -IVF accepts `{"ivf":{"scan_buckets_count":N,"disable_bucket_scan":true}}` through -`params_str_`. This routing-only mode returns the `N` selected bucket IDs per query in the -result `Dataset` instead of vector labels. `NumElements()` equals the number of queries, -`Dim()` equals `scan_buckets_count`, `GetIds()` contains bucket IDs (with `-1` for empty -slots), and `GetDistances()` has distances to bucket centroids. No vector scan is performed, -so filters, `topk`, range limits, reordering, and reasoning options are ignored. +IVF accepts `{"ivf":{"scan_buckets_count":N,"disable_bucket_scan":true}}` in `params_str_`. +This routing-only mode returns `N` bucket IDs per query instead of vector labels. +`NumElements()` is the query count, `Dim()` is `scan_buckets_count`, `GetIds()` contains bucket IDs +(`-1` for empty slots), and `GetDistances()` contains distances to the corresponding bucket +centroids. It does not scan bucket vectors, so filters, `topk`, range limits, reordering, and +reasoning options are ignored. ### IVF bucket IDs bypass @@ -101,13 +101,9 @@ selection and scans only the specified buckets. **Semantics:** - **Empty** (default): Use normal bucket routing via `ClassifyDatasForSearch`. - **Non-empty**: Skip routing and scan only the provided bucket IDs. -- **Batch IVF KNN**: One outer `bucket_ids_` entry per query vector. - Result is rectangular: `NumElements()` is query count, `Dim()` is `topk_`. - Missing neighbors are `-1` with infinite distance. **Constraints:** -- Batch IVF search supports KNN only; custom query distance and reasoning labels are unsupported. -- A non-empty outer vector must contain exactly one non-empty entry per query vector. +- Currently only single-query is supported; the outer vector must contain exactly one entry. - Each bucket ID must be in range `[0, bucket_count)`. - Duplicate bucket IDs are rejected. - Incompatible with `disable_bucket_scan` mode. diff --git a/docs/docs/zh/src/api/dataset.md b/docs/docs/zh/src/api/dataset.md index 0f07e77dab..b7ddb46351 100644 --- a/docs/docs/zh/src/api/dataset.md +++ b/docs/docs/zh/src/api/dataset.md @@ -115,8 +115,10 @@ if (result.has_value()) { } ``` -对 KNN，`GetNumElements()` 为 `1`，ids/distances 数组长度为 `k`。对范围搜索，命中数通过结果的维度报告。 -见 [k-近邻搜索](../guide/knn_search.md)。 +对单查询 KNN，`GetNumElements()` 为 `1`，ids/distances 数组长度为 `k`。HGraph 和 IVF 的批量 KNN 返回 +按行主序排列的 `GetNumElements() x GetDim()` 矩阵，行 `q` 的第 `i` 个槽位位于 +`q * GetDim() + i`。不足一行的结果以 `id == -1` 和无穷距离填充。为保持填充语义明确，批量 KNN 会拒绝 +包含外部 label `-1` 的索引。范围搜索仅支持单个查询。见 [k-近邻搜索](../guide/knn_search.md)。 ## `SparseVector` diff --git a/docs/docs/zh/src/api/index_class.md b/docs/docs/zh/src/api/index_class.md index 3f9a2b3900..daa8bd9937 100644 --- a/docs/docs/zh/src/api/index_class.md +++ b/docs/docs/zh/src/api/index_class.md @@ -127,12 +127,14 @@ using WriteFuncType = std::function; ## 搜索 -推荐的入口是 [`SearchWithRequest`](#searchwithrequest)，它接收单个 +推荐的入口是 [`SearchWithRequest`](#searchwithrequest)，它接收 [`SearchRequest`](search.md#searchrequest)，其中携带查询、模式、top-k / 半径以及各类过滤器。较旧的 逐参数 `KnnSearch` / `RangeSearch` 重载为兼容性保留。 -每次搜索都返回一个 `DatasetPtr`：对 KNN，`num_elements == 1`，`ids` / `distances` 长度为 `k`；对范围 -搜索，结果长度即命中数。如何读取结果见 [Dataset](dataset.md)。 +每次搜索都返回一个 `DatasetPtr`：单查询 KNN 的 `num_elements == 1`；HGraph 和 IVF 还支持批量 KNN， +返回按行主序排列的 `num_elements x dim` 矩阵。`dim` 是实际返回的行宽，可能小于请求的 top-k；缺失条目 +以 `id == -1` 填充；为保持填充语义明确，批量 KNN 会拒绝包含外部 label `-1` 的索引。范围搜索只接受单个查询。如何读取结果见 +[Dataset](dataset.md)。 ### `SearchWithRequest` diff --git a/docs/docs/zh/src/api/search.md b/docs/docs/zh/src/api/search.md index 410010e194..ba9403c873 100644 --- a/docs/docs/zh/src/api/search.md +++ b/docs/docs/zh/src/api/search.md @@ -13,7 +13,7 @@ ```cpp vsag::SearchRequest request; -request.query_ = query; // 含单个查询向量的 DatasetPtr +request.query_ = query; // 单个查询，或 HGraph/IVF 的连续 KNN 查询批次 request.mode_ = vsag::SearchMode::KNN_SEARCH; request.topk_ = 10; request.params_str_ = R"({"hgraph": {"ef_search": 100}})"; @@ -34,7 +34,7 @@ enum class SearchMode { | 字段 | 类型 | 默认值 | 含义 | |------|------|--------|------| -| `query_` | `DatasetPtr` | `nullptr` | 查询。IVF KNN 请求支持多个查询向量，其他请求只允许一个。 | +| `query_` | `DatasetPtr` | `nullptr` | 查询。HGraph 和 IVF 的 KNN 支持连续的多查询批次；范围搜索只支持单个查询。 | | `mode_` | `SearchMode` | `KNN_SEARCH` | KNN 还是范围搜索。 | | `topk_` | `int64_t` | `10` | 要返回的邻居数（KNN 模式）。必须为正。 | | `radius_` | `float` | `0.5` | 距离阈值（范围模式）。非负。 | @@ -95,13 +95,9 @@ reasoning 选项均会被忽略。 **语义：** - **空**（默认）：使用正常的桶路由（`ClassifyDatasForSearch`）。 - **非空**：跳过路由，仅扫描提供的桶 ID。 -- **批量 IVF KNN**：`bucket_ids_` 外层每项对应一个查询向量。 - 结果为矩形：`NumElements()` 是查询数，`Dim()` 是 `topk_`。 - 缺失邻居用 `-1` 和无穷距离填充。 **约束：** -- 批量 IVF 搜索仅支持 KNN；不支持自定义查询距离和 reasoning labels。 -- 非空外层向量必须为每个查询向量提供一个非空条目。 +- 当前仅支持单查询；外层向量必须恰好包含一个条目。 - 每个桶 ID 必须在 `[0, bucket_count)` 范围内。 - 重复的桶 ID 会被拒绝。 - 与 `disable_bucket_scan` 模式不兼容。 diff --git a/include/vsag/index.h b/include/vsag/index.h index 608e5bc8b7..46206791d0 100644 --- a/include/vsag/index.h +++ b/include/vsag/index.h @@ -323,11 +323,24 @@ class Index { /** * @brief Performing search with request on index - * + * * @param request @see SearchRequest - * @return result contains - * - num_elements: 1 - * - ids, distances: length is (num_elements * k) + * @return result contains + * - single-query requests: num_elements = 1, dim = actual returned + * result count (<= request.topk_ for KNN, and <= request.limited_size_ + * for RANGE_SEARCH when limited_size_ > 0). May be < topk_ when filters + * reject candidates. + * - batched KNN requests, when supported by the implementation: + * num_elements = query->GetNumElements(), + * dim = implementation-defined returned row width. HGraph clamps it to + * min(request.topk_, GetNumElements()), while IVF preserves + * request.topk_. Callers MUST read `dim` from the returned dataset. + * ids/distances are stored row-major with length (num_elements * dim). + * Queries that yielded fewer than dim neighbors are padded with + * sentinel entries (id = -1, distance = +infinity). Batch KNN rejects + * an index containing external label -1 to keep this padding unambiguous. + * Callers MUST check `ids[i] == -1` + * to detect padding and MUST read `dim` from the returned dataset. */ [[nodiscard]] virtual tl::expected SearchWithRequest(const SearchRequest& request) const { diff --git a/include/vsag/search_request.h b/include/vsag/search_request.h index 1f42987251..2891caa2c0 100644 --- a/include/vsag/search_request.h +++ b/include/vsag/search_request.h @@ -39,11 +39,21 @@ enum class SearchMode { class SearchRequest { public: // basic params - /** + /** * @brief Query dataset containing the vector or vectors to search for - * @details This DatasetPtr holds the query vector used for similarity search. - * IVF KNN requests and supported AnalyzeIndexBySearch implementations accept - * multiple query vectors; other requests allow one. + * @details This DatasetPtr holds the query data used for similarity search. + * - Single query: Set NumElements to 1 with one vector. Supported by all + * search modes (KNN_SEARCH, RANGE_SEARCH). + * - Batched KNN: Set NumElements to the number of queries, with vectors + * stored contiguously. Supported by HGraph::SearchWithRequest and + * IVF::SearchWithRequest; results are returned with NumElements = + * query_count and a row-major Dim determined by the implementation + * (which can be less than topk when the index is smaller). Queries that yield + * fewer neighbors than the returned Dim are padded with sentinel entries + * (id = -1, distance = +infinity). Batch KNN rejects an index containing external + * label -1 to keep this padding unambiguous. + * - Batched RANGE_SEARCH is not supported; implementations MUST reject + * NumElements > 1 for range mode. */ DatasetPtr query_{nullptr}; @@ -219,8 +229,9 @@ class SearchRequest { /** * @brief Pre-selected bucket IDs for bypassing IVF bucket routing (ClassifyDatasForSearch) - * @details The outer vector contains one entry per query vector. - * Inner vector contains ordered bucket IDs (caller is responsible for ordering). + * @details Supports one ordered bucket-ID list per query vector. For a single query, + * the outer vector must contain exactly one entry. + * Inner vectors contain ordered bucket IDs (caller is responsible for ordering). * When non-empty with at least one ID, skips ClassifyDatasForSearch and searches only * the specified buckets. Empty means "use default bucket routing". */ diff --git a/src/algorithm/hgraph/hgraph.h b/src/algorithm/hgraph/hgraph.h index f83ee4fb4f..32471b7e5c 100644 --- a/src/algorithm/hgraph/hgraph.h +++ b/src/algorithm/hgraph/hgraph.h @@ -16,6 +16,7 @@ #pragma once #include +#include #include #include #include @@ -267,17 +268,21 @@ class HGraph : public InnerIndexInterface { * pointer is absent. */ const void* - get_data(const DatasetPtr& dataset, uint32_t index = 0) const { + get_data(const DatasetPtr& dataset, int64_t index = 0) const { + CHECK_ARGUMENT(index >= 0, "query index must be non-negative"); + CHECK_ARGUMENT(data_type_ == DataTypes::DATA_TYPE_SPARSE || + (dim_ > 0 && index <= std::numeric_limits::max() / dim_), + "query offset exceeds int64_t range"); if (data_type_ == DataTypes::DATA_TYPE_FLOAT) { auto* ptr = dataset->GetFloat32Vectors(); - return ptr ? ptr + static_cast(index) * dim_ : nullptr; + return ptr ? ptr + index * dim_ : nullptr; } else if (data_type_ == DataTypes::DATA_TYPE_INT8) { auto* ptr = dataset->GetInt8Vectors(); - return ptr ? ptr + static_cast(index) * dim_ : nullptr; + return ptr ? ptr + index * dim_ : nullptr; } else if (data_type_ == DataTypes::DATA_TYPE_FP16 || data_type_ == DataTypes::DATA_TYPE_BF16) { auto* ptr = dataset->GetFloat16Vectors(); - return ptr ? ptr + static_cast(index) * dim_ : nullptr; + return ptr ? ptr + index * dim_ : nullptr; } else if (data_type_ == DataTypes::DATA_TYPE_SPARSE) { auto* ptr = dataset->GetSparseVectors(); return ptr ? ptr + index : nullptr; @@ -626,6 +631,12 @@ class HGraph : public InnerIndexInterface { QueryContext* ctx, const std::optional& threshold = std::nullopt) const; + DatasetPtr + search_range_with_request(const SearchRequest& request, + const HGraphSearchParameters& params, + const FilterPtr& filter, + QueryContext& ctx) const; + private: /// Reorder the candidate heap using precise codes, updating in-place. void diff --git a/src/algorithm/hgraph/hgraph_search.cpp b/src/algorithm/hgraph/hgraph_search.cpp index 6a9eeeb6cc..743eaf6307 100644 --- a/src/algorithm/hgraph/hgraph_search.cpp +++ b/src/algorithm/hgraph/hgraph_search.cpp @@ -29,16 +29,33 @@ namespace vsag { static DatasetPtr -make_empty_dataset_with_stats(const SearchStatistics& stats) { +make_empty_dataset_with_stats() { + SearchStatistics stats; auto dataset_result = DatasetImpl::MakeEmptyDataset(); dataset_result->Statistics(stats.Dump()); return dataset_result; } static DatasetPtr -make_empty_dataset_with_stats() { - SearchStatistics stats; - return make_empty_dataset_with_stats(stats); +make_empty_dataset_with_stats(const SearchStatistics& stats) { + auto dataset_result = DatasetImpl::MakeEmptyDataset(); + dataset_result->Statistics(stats.Dump()); + return dataset_result; +} + +static void +apply_hops_limit(InnerSearchParam& search_param, const HGraphSearchParameters& params) { + if (static_cast(params.hops_limit) <= static_cast(params.ef_search)) { + search_param.hops_limit = std::numeric_limits::max(); + if (params.hops_limit != std::numeric_limits::max()) { + logger::warn( + fmt::format("hops_limit({}) is not greater than ef_search({}), ignoring hops_limit", + params.hops_limit, + params.ef_search)); + } + return; + } + search_param.hops_limit = params.hops_limit; } DatasetPtr @@ -101,6 +118,10 @@ HGraph::KnnSearch(const DatasetPtr& query, } k = std::min(k, GetNumElements()); + // Iterator state is maintained per query, so this overload remains single-query only. + CHECK_ARGUMENT(query->GetNumElements() == 1, + "iterator-based KnnSearch only supports single query (NumElements=1)"); + FilterPtr ft = this->create_search_filter(filter, params.use_extra_info_filter); if (iter_ctx == nullptr) { @@ -415,6 +436,132 @@ HGraph::RangeSearch(const DatasetPtr& query, return this->SearchWithRequest(req); } +DatasetPtr +HGraph::search_range_with_request(const SearchRequest& request, + const HGraphSearchParameters& params, + const FilterPtr& filter, + QueryContext& ctx) const { + InnerSearchParam search_param; + search_param.ep = this->entry_point_id_; + search_param.topk = 1; + search_param.ef = 1; + search_param.is_inner_id_allowed = nullptr; + search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search; + if (params.enable_time_record) { + search_param.time_cost = std::make_shared(); + search_param.time_cost->SetThreshold(params.timeout_ms); + ctx.stats->is_timeout.store(false, std::memory_order_relaxed); + } + + struct visited_list_guard { + std::shared_ptr pool; + VisitedListPtr visited_list; + ~visited_list_guard() { + if (visited_list != nullptr) { + pool->ReturnOne(visited_list); + } + } + }; + visited_list_guard vt_guard{this->pool_, this->pool_->TakeOne()}; + auto& vt = vt_guard.visited_list; + const auto* raw_query = get_data(request.query_); + ctx.distance_phase = DistanceEvaluationPhase::ROUTING; + for (auto i = static_cast(this->route_graphs_.size() - 1); i >= 0; --i) { + auto result = this->search_one_graph( + raw_query, this->route_graphs_[i], this->basic_flatten_codes_, search_param, vt, &ctx); + if (!result->Empty()) { + search_param.ep = result->Top().second; + } + } + ctx.distance_phase = DistanceEvaluationPhase::APPROXIMATE; + + if (request.enable_attribute_filter_ and this->attr_filter_index_ != nullptr) { + auto& schema = this->attr_filter_index_->field_type_map_; + auto expr = AstParse(request.attribute_filter_str_, &schema); + auto executor = Executor::MakeInstance(this->allocator_, expr, this->attr_filter_index_); + executor->Init(); + search_param.executors.emplace_back(executor); + } + + search_param.ef = std::max(params.ef_search, request.limited_size_); + search_param.is_inner_id_allowed = filter; + search_param.radius = request.radius_; + search_param.search_mode = RANGE_SEARCH; + search_param.consider_duplicate = true; + search_param.range_search_limit_size = static_cast(request.limited_size_); + search_param.parallel_search_thread_count = params.parallel_search_thread_count; + search_param.enable_reorder = params.enable_reorder; + search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search; + search_param.skip_ratio = params.skip_ratio; + search_param.skip_strategy_type = params.skip_strategy_type; + + apply_hops_limit(search_param, params); + DistanceRecordVector rabitq_lower_bound_candidates(ctx.alloc); + auto* rabitq_lower_bound_candidates_ptr = + search_param.enable_rabitq_one_bit_search and use_reorder_ and + search_param.enable_reorder and reorder_by_base_ + ? &rabitq_lower_bound_candidates + : nullptr; + + DistHeapPtr search_result; + bool brute_force_used = false; + MCIHybridSearchResult mci_result(params, filter); + if (params.brute_force_threshold > 0.0F && + mci_result.valid_ratio <= params.brute_force_threshold) { + search_result = this->brute_force_search( + raw_query, filter, request.limited_size_, request.radius_, &ctx); + brute_force_used = true; + mci_result.route = "brute_force"; + } else { + mci_result = this->try_mci_search(request, params, filter, raw_query, search_param, &ctx); + if (mci_result.route == "mci") { + search_result = std::move(mci_result.result); + } + } + if (search_result == nullptr) { + search_result = this->search_one_graph(raw_query, + this->bottom_graph_, + this->basic_flatten_codes_, + search_param, + vt, + &ctx, + rabitq_lower_bound_candidates_ptr); + } + + if (mci_result.route != "mci" && not brute_force_used && use_reorder_ && + search_param.enable_reorder) { + this->reorder(raw_query, + this->get_reorder_codes(), + search_result, + request.limited_size_, + nullptr, + ctx, + rabitq_lower_bound_candidates_ptr); + } else if (mci_result.route != "mci" && not brute_force_used && search_param.enable_reorder && + params.rabitq_one_bit_search) { + this->reorder(raw_query, + this->basic_flatten_codes_, + search_result, + request.limited_size_, + nullptr, + ctx); + } + + while (not search_result->Empty() and + search_result->Top().first > request.radius_ + THRESHOLD_ERROR) { + search_result->Pop(); + } + if (request.limited_size_ > 0) { + while (search_result->Size() > static_cast(request.limited_size_)) { + search_result->Pop(); + } + } + + auto result = this->pack_knn_result_with_extra_info(search_result, ctx.alloc); + result->Statistics(mci_result.MakeStatistics(*ctx.stats).Dump()); + return result; +} + [[nodiscard]] DatasetPtr HGraph::SearchWithRequest(const SearchRequest& request) const { ValidateSearchThreshold(request.threshold_); @@ -432,19 +579,33 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { if (use_custom_distance) { CHECK_ARGUMENT(request.distance_batch_size_ > 0, "distance_batch_size must be greater than 0"); - CHECK_ARGUMENT(not is_range, "HGraph custom distance only supports KNN search"); + CHECK_ARGUMENT(not is_range, "HGraph custom query distance only supports KNN search"); } if (is_range) { + // Range search remains single-query only (validate_range_args enforces NumElements==1). if (not use_custom_distance) { this->validate_range_args(query, request.radius_, request.limited_size_); + } else { + CHECK_ARGUMENT(query != nullptr, "query dataset cannot be null"); + CHECK_ARGUMENT(query->GetNumElements() == 1, + "HGraph range search only supports a single query"); } } else { + // KNN search supports multi-query batch: validate_knn_args enforces NumElements==1, + // so use inline checks that allow NumElements >= 1. if (not use_custom_distance) { - this->validate_knn_args(query, k); - } else { - CHECK_ARGUMENT(k > 0, "topk must be greater than 0"); + CHECK_ARGUMENT(query != nullptr, "query dataset cannot be null"); + if (data_type_ != DataTypes::DATA_TYPE_SPARSE) { + CHECK_ARGUMENT( + query->GetDim() == dim_, + fmt::format( + "query.dim({}) must be equal to index.dim({})", query->GetDim(), dim_)); + } + CHECK_ARGUMENT(get_data(query) != nullptr, + "query vector storage must match index data type"); } + CHECK_ARGUMENT(k > 0, fmt::format("k({}) must be greater than 0", k)); } auto params = HGraphSearchParameters::FromJson(request.params_str_); @@ -470,11 +631,42 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { shared_lock = this->acquire_global_read_lock(); } const auto element_count = GetNumElements(); + if (use_custom_distance && query != nullptr) { + CHECK_ARGUMENT(query->GetNumElements() == 1, + "HGraph custom query distance only supports a single query"); + } + int64_t query_count = use_custom_distance ? 1 : query->GetNumElements(); + CHECK_ARGUMENT(query_count >= 1, + fmt::format("query count({}) must be at least 1", query_count)); + if (is_range) { + CHECK_ARGUMENT(query_count == 1, "range search only supports single query (NumElements=1)"); + } + if (query_count > 1) { + CHECK_ARGUMENT(request.expected_labels_.empty(), + "reasoning (expected_labels_) is only supported for single-query search"); + CHECK_ARGUMENT(!this->label_table_->HasActivePaddingLabel(), + "batch KNN does not support an index containing external label -1"); + } if (element_count == 0) { + if (!is_range && query_count > 1) { + auto result = create_fast_dataset(0, ctx.alloc); + std::get<0>(result)->NumElements(query_count); + std::get<0>(result)->Dim(0); + std::get<0>(result)->Statistics(stats.Dump()); + return std::get<0>(result); + } return make_empty_dataset_with_stats(); } k = std::min(k, element_count); + if (!is_range && query_count > 1 && k == 0) { + auto result = create_fast_dataset(0, ctx.alloc); + std::get<0>(result)->NumElements(query_count); + std::get<0>(result)->Dim(0); + std::get<0>(result)->Statistics(stats.Dump()); + return std::get<0>(result); + } + // Setup reasoning context (KNN only) std::shared_ptr reasoning_ctx; if (not is_range and not request.expected_labels_.empty()) { @@ -520,112 +712,142 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { } reasoning_ctx->SetTrueDistance(inner_id, dist); } + ctx.reasoning_ctx = reasoning_ctx.get(); } - InnerSearchParam search_param; - search_param.ep = this->entry_point_id_; - search_param.topk = 1; - search_param.ef = 1; - search_param.is_inner_id_allowed = nullptr; - search_param.enable_rabitq_one_bit_search = - use_custom_distance ? false : params.rabitq_one_bit_search; - search_param.distance_batch_func = request.distance_batch_func_; - search_param.distance_batch_size = request.distance_batch_size_; - - if (search_param.ep == INVALID_ENTRY_POINT) { + if (this->entry_point_id_ == INVALID_ENTRY_POINT) { + if (query_count > 1) { + // Return batch-shaped empty result preserving documented layout. + CHECK_ARGUMENT( + query_count <= std::numeric_limits::max() / std::max(k, (int64_t)1), + fmt::format("query_count({}) * k({}) would overflow", query_count, k)); + int64_t batch_count = query_count * k; + auto [empty_ds, empty_dists, empty_ids] = create_fast_dataset(batch_count, ctx.alloc); + std::fill_n(empty_dists, batch_count, std::numeric_limits::infinity()); + std::fill_n(empty_ids, batch_count, -1); + empty_ds->NumElements(query_count); + empty_ds->Dim(k); + empty_ds->Statistics(stats.Dump()); + return empty_ds; + } return make_empty_dataset_with_stats(); } - struct visited_list_guard { - std::shared_ptr pool; - VisitedListPtr visited_list; - - void - Release() { - if (visited_list != nullptr) { - pool->ReturnOne(visited_list); - visited_list.reset(); - } - } + FilterPtr ft = this->create_search_filter(request.filter_, params.use_extra_info_filter); - ~visited_list_guard() { - Release(); - } - }; - visited_list_guard vt_guard{this->pool_, this->pool_->TakeOne()}; - auto& vt = vt_guard.visited_list; + if (is_range) { + return this->search_range_with_request(request, params, ft, ctx); + } - const auto* raw_query = use_custom_distance ? nullptr : get_data(query); - ctx.distance_phase = DistanceEvaluationPhase::ROUTING; - for (auto i = static_cast(this->route_graphs_.size() - 1); i >= 0; --i) { - auto result = this->search_one_graph( - raw_query, this->route_graphs_[i], this->basic_flatten_codes_, search_param, vt, &ctx); - // An unrankable route seed can still bridge to finite bottom-layer results. - if (not result->Empty()) { - search_param.ep = result->Top().second; - } + // ---- KNN search: multi-query batch path (PR #1685) ---- + + // Build a shared base search_param; per-query fields (ep) are set inside the loop. + InnerSearchParam base_search_param; + base_search_param.is_inner_id_allowed = ft; + base_search_param.distance_threshold = request.threshold_; + base_search_param.ef = std::max(params.ef_search, k); + base_search_param.topk = static_cast(base_search_param.ef); + if (params.topk_factor > 1.0F) { + base_search_param.topk = + std::min(base_search_param.topk, + static_cast(static_cast(k) * params.topk_factor)); } - ctx.distance_phase = DistanceEvaluationPhase::APPROXIMATE; + base_search_param.consider_duplicate = true; + base_search_param.enable_reorder = use_custom_distance ? false : params.enable_reorder; + base_search_param.enable_rabitq_one_bit_search = + use_custom_distance ? false : params.rabitq_one_bit_search; + base_search_param.skip_ratio = params.skip_ratio; + base_search_param.skip_strategy_type = params.skip_strategy_type; + base_search_param.distance_batch_func = request.distance_batch_func_; + base_search_param.distance_batch_size = request.distance_batch_size_; + if (params.enable_time_record) { + base_search_param.time_cost = std::make_shared(); + base_search_param.time_cost->SetThreshold(params.timeout_ms); + stats.is_timeout.store(false, std::memory_order_relaxed); + } + base_search_param.parallel_search_thread_count = params.parallel_search_thread_count; - FilterPtr ft = this->create_search_filter(request.filter_, params.use_extra_info_filter); + // hops_limit only takes effect when it's greater than ef_search + apply_hops_limit(base_search_param, params); if (request.enable_attribute_filter_ and this->attr_filter_index_ != nullptr) { auto& schema = this->attr_filter_index_->field_type_map_; auto expr = AstParse(request.attribute_filter_str_, &schema); auto executor = Executor::MakeInstance(this->allocator_, expr, this->attr_filter_index_); executor->Init(); - search_param.executors.emplace_back(executor); + base_search_param.executors.emplace_back(executor); } - if (is_range) { - search_param.ef = std::max(params.ef_search, request.limited_size_); - search_param.is_inner_id_allowed = ft; - search_param.radius = request.radius_; - search_param.search_mode = RANGE_SEARCH; - search_param.consider_duplicate = true; - search_param.range_search_limit_size = static_cast(request.limited_size_); - search_param.parallel_search_thread_count = params.parallel_search_thread_count; - search_param.enable_reorder = use_custom_distance ? false : params.enable_reorder; - search_param.enable_rabitq_one_bit_search = - use_custom_distance ? false : params.rabitq_one_bit_search; - } else { - search_param.ef = std::max(params.ef_search, k); - search_param.is_inner_id_allowed = ft; - search_param.distance_threshold = request.threshold_; - search_param.topk = static_cast(search_param.ef); - if (params.topk_factor > 1.0F) { - search_param.topk = - std::min(search_param.topk, - static_cast(static_cast(k) * params.topk_factor)); - } - search_param.enable_reorder = use_custom_distance ? false : params.enable_reorder; - search_param.consider_duplicate = true; - search_param.enable_rabitq_one_bit_search = - use_custom_distance ? false : params.rabitq_one_bit_search; - if (params.enable_time_record) { - search_param.time_cost = std::make_shared(); - search_param.time_cost->SetThreshold(params.timeout_ms); - stats.is_timeout.store(false, std::memory_order_relaxed); - } - search_param.parallel_search_thread_count = params.parallel_search_thread_count; - - if (static_cast(params.hops_limit) <= static_cast(params.ef_search)) { - search_param.hops_limit = std::numeric_limits::max(); - if (params.hops_limit != std::numeric_limits::max()) { - logger::warn(fmt::format( - "hops_limit({}) is not greater than ef_search({}), ignoring hops_limit", - params.hops_limit, - params.ef_search)); - } - } else { - search_param.hops_limit = params.hops_limit; + // Single-query preserves the original "dim = actual result count" contract; multi-query + // uses a fixed query_count x k rectangular layout. Guard the multiplication against overflow. + int64_t total_result_count = 0; + if (query_count > 1) { + CHECK_ARGUMENT( + query_count <= std::numeric_limits::max() / k, + fmt::format("query_count({}) * k({}) would overflow int64_t", query_count, k)); + total_result_count = query_count * k; + } + // Validate that byte-level allocations do not overflow size_t. + if (total_result_count > 0) { + constexpr auto k_id_size = sizeof(int64_t); + constexpr auto k_dist_size = sizeof(float); + CHECK_ARGUMENT(total_result_count <= std::numeric_limits::max() / k_id_size, + fmt::format("total_result_count({}) * sizeof(int64_t) would overflow size_t", + total_result_count)); + CHECK_ARGUMENT(total_result_count <= std::numeric_limits::max() / k_dist_size, + fmt::format("total_result_count({}) * sizeof(float) would overflow size_t", + total_result_count)); + if (extra_info_size_ > 0) { + constexpr auto k_extra_size = sizeof(char); + CHECK_ARGUMENT( + total_result_count <= std::numeric_limits::max() / + (static_cast(extra_info_size_) * k_extra_size), + fmt::format("total_result_count({}) * extra_info_size({}) would overflow size_t", + total_result_count, + extra_info_size_)); } } + auto [dataset_results, dists, ids] = create_fast_dataset(total_result_count, ctx.alloc); + char* extra_infos = nullptr; + if (query_count > 1 && extra_info_size_ > 0 && this->extra_infos_ != nullptr) { + extra_infos = static_cast(ctx.alloc->Allocate( + static_cast(extra_info_size_) * static_cast(total_result_count))); + std::memset(extra_infos, + 0, + static_cast(static_cast(extra_info_size_) * + static_cast(total_result_count))); + dataset_results->ExtraInfos(extra_infos); + dataset_results->ExtraInfoSize(static_cast(extra_info_size_)); + } - search_param.skip_ratio = params.skip_ratio; - search_param.skip_strategy_type = params.skip_strategy_type; + // Pre-fill sentinels: ids = -1 (authoritative signal for "no result") and + // dists = +infinity (unambiguous for inner-product / cosine metrics that may produce + // negative distances). Callers MUST detect padding via ids[i] == -1 rather than by + // distance comparison. + std::fill_n(dists, total_result_count, std::numeric_limits::infinity()); + std::fill_n(ids, total_result_count, -1); + + Vector reasoning_result_inner_ids(this->allocator_); + struct visited_list_guard { + std::shared_ptr pool; + VisitedListPtr visited_list; + ~visited_list_guard() { + if (visited_list != nullptr) { + pool->ReturnOne(visited_list); + } + } + }; + visited_list_guard vt_guard{this->pool_, this->pool_->TakeOne()}; + auto& vt = vt_guard.visited_list; + + // Hoist per-query search_param and rabitq candidate buffer out of the loop: + // the searcher only mutates `duplicate_id` (declared `mutable` on the const + // InnerSearchParam&) and callers only tweak `ep` per query, so a single instance + // reused across queries avoids copying the base_search_param (including its + // `executors` vector) on every iteration. + InnerSearchParam search_param = base_search_param; DistanceRecordVector rabitq_lower_bound_candidates(ctx.alloc); auto* rabitq_lower_bound_candidates_ptr = search_param.enable_rabitq_one_bit_search and use_reorder_ and @@ -633,146 +855,192 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { ? &rabitq_lower_bound_candidates : nullptr; - DistHeapPtr search_result; - bool brute_force_used = false; - MCIHybridSearchResult mci_result(params, ft); - if (not use_custom_distance) { - if (params.brute_force_threshold > 0.0F and - mci_result.valid_ratio <= params.brute_force_threshold) { - if (is_range) { - search_result = this->brute_force_search( - raw_query, ft, request.limited_size_, request.radius_, &ctx); - } else { + InnerSearchParam ep_search_param; + ep_search_param.ep = this->entry_point_id_; + ep_search_param.topk = 1; + ep_search_param.ef = 1; + ep_search_param.is_inner_id_allowed = nullptr; + ep_search_param.enable_rabitq_one_bit_search = + use_custom_distance ? false : params.rabitq_one_bit_search; + ep_search_param.distance_batch_func = request.distance_batch_func_; + ep_search_param.distance_batch_size = request.distance_batch_size_; + + for (int64_t q_idx = 0; q_idx < query_count; ++q_idx) { + const auto* raw_query = use_custom_distance ? nullptr : get_data(query, q_idx); + + // Reset per-query mutable state before each query. + search_param.duplicate_id = -1; + // Per-query entry point search through hierarchical graphs. + ctx.distance_phase = DistanceEvaluationPhase::ROUTING; + ep_search_param.ep = this->entry_point_id_; + rabitq_lower_bound_candidates.clear(); + + for (auto i = static_cast(this->route_graphs_.size() - 1); i >= 0; --i) { + auto result = this->search_one_graph(raw_query, + this->route_graphs_[i], + this->basic_flatten_codes_, + ep_search_param, + vt, + &ctx); + if (not result->Empty()) { + ep_search_param.ep = result->Top().second; + } + } + search_param.ep = ep_search_param.ep; + ctx.distance_phase = DistanceEvaluationPhase::APPROXIMATE; + if (search_param.time_cost != nullptr) { + search_param.time_cost->Reset(); + } + + DistHeapPtr search_result; + bool brute_force_used = false; + MCIHybridSearchResult mci_result(params, ft); + if (not use_custom_distance) { + if (params.brute_force_threshold > 0.0F && + mci_result.valid_ratio <= params.brute_force_threshold) { search_result = this->brute_force_search( raw_query, ft, k, 0.0F, &ctx, request.threshold_); - } - brute_force_used = true; - mci_result.route = "brute_force"; - } else { - mci_result = this->try_mci_search(request, params, ft, raw_query, search_param, &ctx); - if (mci_result.route == "mci") { - search_result = std::move(mci_result.result); + brute_force_used = true; + mci_result.route = "brute_force"; } else { - search_result = this->search_one_graph(raw_query, - this->bottom_graph_, - this->basic_flatten_codes_, - search_param, - vt, - &ctx, - rabitq_lower_bound_candidates_ptr); + mci_result = + this->try_mci_search(request, params, ft, raw_query, search_param, &ctx); + if (mci_result.route == "mci") { + search_result = std::move(mci_result.result); + } else { + search_result = this->search_one_graph(raw_query, + this->bottom_graph_, + this->basic_flatten_codes_, + search_param, + vt, + &ctx, + rabitq_lower_bound_candidates_ptr); + } } + } else { + search_result = this->search_one_graph(raw_query, + this->bottom_graph_, + this->basic_flatten_codes_, + search_param, + vt, + &ctx, + rabitq_lower_bound_candidates_ptr); } - } else { - search_result = this->search_one_graph(raw_query, - this->bottom_graph_, - this->basic_flatten_codes_, - search_param, - vt, - &ctx, - rabitq_lower_bound_candidates_ptr); - } - vt_guard.Release(); - // Reorder - if (mci_result.route != "mci" and not brute_force_used and use_reorder_ and - search_param.enable_reorder) { - auto limit = is_range ? request.limited_size_ : k; - auto reorder_threshold = is_range ? std::nullopt : request.threshold_; - this->reorder(raw_query, - this->get_reorder_codes(), - search_result, - limit, - nullptr, - ctx, - rabitq_lower_bound_candidates_ptr, - reorder_threshold); - } else if (mci_result.route != "mci" and not brute_force_used and - search_param.enable_reorder and params.rabitq_one_bit_search) { - auto limit = is_range ? request.limited_size_ : k; - auto reorder_threshold = is_range ? std::nullopt : request.threshold_; - this->reorder(raw_query, - this->basic_flatten_codes_, - search_result, - limit, - nullptr, - ctx, - nullptr, - reorder_threshold); - } + if (mci_result.route != "mci" && !brute_force_used && use_reorder_ && + search_param.enable_reorder) { + this->reorder(raw_query, + this->get_reorder_codes(), + search_result, + k, + nullptr, + ctx, + rabitq_lower_bound_candidates_ptr, + request.threshold_); + } else if (mci_result.route != "mci" && !brute_force_used && search_param.enable_reorder && + params.rabitq_one_bit_search) { + this->reorder(raw_query, + this->basic_flatten_codes_, + search_result, + k, + nullptr, + ctx, + nullptr, + request.threshold_); + } - // Trim and pack results - if (is_range) { - while (not search_result->Empty() and - search_result->Top().first > request.radius_ + THRESHOLD_ERROR) { + DistanceRecordVector finite_records(ctx.alloc); + finite_records.reserve(search_result->Size()); + while (not search_result->Empty()) { + const auto record = search_result->Top(); search_result->Pop(); - } - if (request.limited_size_ > 0) { - while (search_result->Size() > static_cast(request.limited_size_)) { - search_result->Pop(); + if (not std::isnan(record.first) and + (not request.threshold_.has_value() or std::isfinite(record.first))) { + finite_records.push_back(record); } } - auto result = this->pack_knn_result_with_extra_info(search_result, ctx.alloc); - result->Statistics(mci_result.MakeStatistics(stats).Dump()); - return result; - } + for (const auto& record : finite_records) { + search_result->Push(record); + } + filter_search_result_by_threshold(search_result, request.threshold_, ctx.alloc); + while (search_result->Size() > k) { + search_result->Pop(); + } - // NaN is unordered and cannot be returned. Infinity remains a valid legacy result only when - // threshold filtering is absent; the searcher has already kept it out of threshold heaps. - DistanceRecordVector finite_records(ctx.alloc); - finite_records.reserve(search_result->Size()); - while (not search_result->Empty()) { - const auto record = search_result->Top(); - search_result->Pop(); - if (not std::isnan(record.first) and - (not request.threshold_.has_value() or std::isfinite(record.first))) { - finite_records.push_back(record); + // Single-query preserves the original contract: an empty result returns an empty dataset. + if (query_count == 1 && search_result->Empty()) { + auto dataset_result = DatasetImpl::MakeEmptyDataset(); + dataset_result->Statistics(mci_result.MakeStatistics(stats).Dump()); + if (reasoning_ctx) { + reasoning_ctx->DiagnoseExpectedTargets(); + dataset_result->Reasoning(reasoning_ctx->GenerateReport()); + } + return dataset_result; } - } - for (const auto& record : finite_records) { - search_result->Push(record); - } - filter_search_result_by_threshold(search_result, request.threshold_, ctx.alloc); - while (search_result->Size() > static_cast(k)) { - search_result->Pop(); - } - // return an empty dataset directly if searcher returns nothing - if (search_result->Empty()) { - auto dataset_result = DatasetImpl::MakeEmptyDataset(); - dataset_result->Statistics(mci_result.MakeStatistics(stats).Dump()); + auto count = static_cast(search_result->Size()); if (reasoning_ctx) { - reasoning_ctx->DiagnoseExpectedTargets(); - dataset_result->Reasoning(reasoning_ctx->GenerateReport()); + reasoning_result_inner_ids.resize(static_cast(count)); } - return dataset_result; - } - auto count = static_cast(search_result->Size()); - Vector result_inner_ids(static_cast(count), this->allocator_); - - auto [dataset_results, dists, ids] = create_fast_dataset(count, ctx.alloc); - char* extra_infos = nullptr; - if (extra_info_size_ > 0 && this->extra_infos_ != nullptr) { - extra_infos = - static_cast(ctx.alloc->Allocate(extra_info_size_ * search_result->Size())); - dataset_results->ExtraInfos(extra_infos) - ->ExtraInfoSize(static_cast(extra_info_size_)); - } - for (int64_t j = count - 1; j >= 0; --j) { - const auto& top = search_result->Top(); - dists[j] = top.first; - ids[j] = this->label_table_->GetLabelById(top.second); - result_inner_ids[j] = top.second; - if (extra_infos != nullptr) { - this->extra_infos_->GetExtraInfoById(top.second, extra_infos + extra_info_size_ * j); + if (query_count == 1) { + // Single-query path may shrink the dataset to the actual neighbor count. + if (dataset_results->GetDim() != count) { + auto [single_results, single_dists, single_ids] = + create_fast_dataset(count, ctx.alloc); + dataset_results = single_results; + dists = single_dists; + ids = single_ids; + } + if (extra_info_size_ > 0 && this->extra_infos_ != nullptr && count > 0) { + extra_infos = static_cast(ctx.alloc->Allocate( + static_cast(extra_info_size_) * static_cast(count))); + dataset_results->ExtraInfos(extra_infos); + dataset_results->ExtraInfoSize(static_cast(extra_info_size_)); + } + for (int64_t j = count - 1; j >= 0; --j) { + const auto& top = search_result->Top(); + dists[j] = top.first; + ids[j] = this->label_table_->GetLabelById(top.second); + if (reasoning_ctx) { + reasoning_result_inner_ids[static_cast(j)] = top.second; + } + if (extra_infos != nullptr) { + this->extra_infos_->GetExtraInfoById(top.second, + extra_infos + extra_info_size_ * j); + } + search_result->Pop(); + } + dataset_results->Statistics(mci_result.MakeStatistics(stats).Dump()); + } else { + int64_t offset = q_idx * k; + for (int64_t j = count - 1; j >= 0; --j) { + const auto& top = search_result->Top(); + dists[offset + j] = top.first; + ids[offset + j] = this->label_table_->GetLabelById(top.second); + if (reasoning_ctx) { + reasoning_result_inner_ids[static_cast(j)] = top.second; + } + if (extra_infos != nullptr) { + this->extra_infos_->GetExtraInfoById( + top.second, extra_infos + extra_info_size_ * (offset + j)); + } + search_result->Pop(); + } } - search_result->Pop(); } - dataset_results->Statistics(mci_result.MakeStatistics(stats).Dump()); - // Generate reasoning report if reasoning context was created + dataset_results->NumElements(query_count); + if (query_count > 1) { + dataset_results->Dim(k); + } + if (query_count > 1) { + dataset_results->Statistics(stats.Dump()); + } + + // Generate reasoning report if reasoning context was created. if (reasoning_ctx) { - reasoning_ctx->MarkResult(result_inner_ids); + reasoning_ctx->MarkResult(reasoning_result_inner_ids); reasoning_ctx->DiagnoseExpectedTargets(); dataset_results->Reasoning(reasoning_ctx->GenerateReport()); } diff --git a/src/algorithm/hgraph/hgraph_serialize.cpp b/src/algorithm/hgraph/hgraph_serialize.cpp index d8faf5feee..4c38acd1fe 100644 --- a/src/algorithm/hgraph/hgraph_serialize.cpp +++ b/src/algorithm/hgraph/hgraph_serialize.cpp @@ -317,6 +317,7 @@ HGraph::deserialize_label_info(StreamReader& reader) const { this->label_table_->Deserialize(reader); } else { StreamReader::ReadVector(reader, this->label_table_->label_table_); + this->label_table_->RebuildActivePaddingLabelIds(); uint64_t size; StreamReader::ReadObj(reader, size); this->label_table_->ResetRemap(size); @@ -566,6 +567,7 @@ HGraph::deserialize_label_info_streaming(StreamReader& reader) const { this->label_table_->Deserialize(reader); } else { StreamReader::ReadVector(reader, this->label_table_->label_table_); + this->label_table_->RebuildActivePaddingLabelIds(); uint64_t size; StreamReader::ReadObj(reader, size); this->label_table_->ResetRemap(size); diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index 557939f7a0..4d1cd14501 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -2141,12 +2141,22 @@ IVF::SearchWithRequest(const SearchRequest& request) const { CHECK_ARGUMENT(request.expected_labels_.empty(), "IVF batch search does not support expected labels"); CHECK_ARGUMENT(request.topk_ > 0, "topk must be greater than 0"); + CHECK_ARGUMENT(!this->label_table_->HasActivePaddingLabel(), + "batch KNN does not support an index containing external label -1"); CHECK_ARGUMENT(query->GetFloat32Vectors() != nullptr, "query float32 vectors cannot be null"); CHECK_ARGUMENT(query->GetDim() == this->dim_, "query dimension must match index dimension"); const auto num_queries = query->GetNumElements(); + CHECK_ARGUMENT( + num_queries <= std::numeric_limits::max() / request.topk_, + fmt::format( + "num_queries({}) * topk({}) would overflow int64_t", num_queries, request.topk_)); const auto total_slots = num_queries * request.topk_; + CHECK_ARGUMENT(total_slots <= std::numeric_limits::max() / sizeof(int64_t), + "batch result id allocation would overflow size_t"); + CHECK_ARGUMENT(total_slots <= std::numeric_limits::max() / sizeof(float), + "batch result distance allocation would overflow size_t"); auto* alloc = select_query_allocator(ctx.alloc, this->allocator_); auto* ids = static_cast(alloc->Allocate(sizeof(int64_t) * total_slots)); auto* distances = static_cast(alloc->Allocate(sizeof(float) * total_slots)); @@ -2177,6 +2187,7 @@ IVF::SearchWithRequest(const SearchRequest& request) const { } one_request.params_str_ = json.Dump(); auto one_result = this->SearchWithRequest(one_request); + CHECK_ARGUMENT(one_result != nullptr, "IVF batch search returned an empty result"); const auto count = std::min(request.topk_, one_result->GetDim()); if (count > 0) { std::copy_n(one_result->GetIds(), count, ids + query_idx * request.topk_); @@ -2219,6 +2230,14 @@ IVF::SearchWithRequest(const SearchRequest& request) const { param.executors.emplace_back(executor); } } + CHECK_ARGUMENT(query != nullptr, "query dataset cannot be null"); + CHECK_ARGUMENT(query->GetNumElements() > 0, "query count must be greater than 0"); + CHECK_ARGUMENT(query->GetDim() == this->dim_, "query dimension must match index dimension"); + if (is_range) { + CHECK_ARGUMENT(query->GetNumElements() == 1, + "IVF range search only supports a single query"); + } + std::shared_ptr reasoning_ctx; if (not request.expected_labels_.empty()) { reasoning_ctx = std::make_shared(this->allocator_); @@ -2322,6 +2341,7 @@ IVF::SearchWithRequest(const SearchRequest& request) const { // Reordered searches defer the finite bound to exact distances, but bucket selection still // needs threshold-mode state so non-finite approximations cannot consume the rerank pool. param.distance_threshold = request.threshold_; + auto search_result = this->search(query, param, ctx, reasoning_ctx.get()); if (reorder_enabled) { auto result = reorder(request.threshold_.has_value() ? param.topk : request.topk_, diff --git a/src/impl/label_table/label_table.cpp b/src/impl/label_table/label_table.cpp index 14fd7ef320..5708027132 100644 --- a/src/impl/label_table/label_table.cpp +++ b/src/impl/label_table/label_table.cpp @@ -44,6 +44,7 @@ LabelTable::LabelTable(Allocator* allocator, label_remap_(allocator, label_remap_type), allocator_(allocator), deleted_ids_(allocator), + active_padding_label_ids_(allocator), source_id_table_(0, allocator) { (void)compress_redundant_data; deleted_ids_filter_ = std::make_shared(deleted_ids_, delete_ids_mutex_); @@ -142,6 +143,7 @@ LabelTable::MarkRemove(const std::vector& labels) { std::scoped_lock wlock(delete_ids_mutex_); for (const auto& id : ids) { if (this->deleted_ids_.insert(id).second) { + active_padding_label_ids_.erase(id); ++removed_count; } } @@ -151,6 +153,7 @@ LabelTable::MarkRemove(const std::vector& labels) { void LabelTable::Deserialize(StreamReader& reader) { StreamReader::ReadVector(reader, label_table_); + RebuildActivePaddingLabelIds(); if (use_reverse_map_) { this->label_remap_.Clear(); this->label_remap_.Reserve(label_table_.size()); @@ -180,12 +183,21 @@ LabelTable::MergeOther(const LabelTablePtr& other, const IdMapFunction& id_map) auto new_label = std::get<1>(id_map(other->label_table_[i])); auto new_inner_id = static_cast(i + current_total_count_u); this->label_table_[i + current_total_count_u] = new_label; + if (new_label == -1 && !other->IsRemoved(static_cast(i))) { + std::scoped_lock wlock(delete_ids_mutex_); + active_padding_label_ids_.insert(new_inner_id); + } this->label_remap_.InsertOrAssign(new_label, new_inner_id); } } else { for (uint64_t i = 0; i < other_size_u; ++i) { auto new_label = std::get<1>(id_map(other->label_table_[i])); this->label_table_[i + current_total_count_u] = new_label; + if (new_label == -1 && !other->IsRemoved(static_cast(i))) { + std::scoped_lock wlock(delete_ids_mutex_); + active_padding_label_ids_.insert( + static_cast(i + current_total_count_u)); + } } } total_count_ += static_cast(other_size_u); diff --git a/src/impl/label_table/label_table.h b/src/impl/label_table/label_table.h index 7c328ff8fc..9de386f90d 100644 --- a/src/impl/label_table/label_table.h +++ b/src/impl/label_table/label_table.h @@ -53,10 +53,17 @@ class LabelTable { if (use_reverse_map_) { label_remap_.InsertOrAssign(label, id); } + std::scoped_lock wlock(delete_ids_mutex_); + const bool was_padding_label = id < label_table_.size() && label_table_[id] == -1; if (id + 1 > label_table_.size()) { label_table_.resize(id + 1); } label_table_[id] = label; + if (label == -1 && deleted_ids_.count(id) == 0) { + active_padding_label_ids_.insert(id); + } else if (was_padding_label) { + active_padding_label_ids_.erase(id); + } total_count_++; } @@ -98,7 +105,9 @@ class LabelTable { void EraseFromDeletedIds(InnerIdType id) { std::scoped_lock wlock(delete_ids_mutex_); - deleted_ids_.erase(id); + if (id < label_table_.size() && deleted_ids_.erase(id) > 0 && label_table_[id] == -1) { + active_padding_label_ids_.insert(id); + } } /** @@ -128,6 +137,23 @@ class LabelTable { bool CheckLabel(LabelType label) const; + bool + HasActivePaddingLabel() const { + std::shared_lock rlock(delete_ids_mutex_); + return !active_padding_label_ids_.empty(); + } + + void + RebuildActivePaddingLabelIds() { + std::scoped_lock wlock(delete_ids_mutex_); + active_padding_label_ids_.clear(); + for (InnerIdType id = 0; id < label_table_.size(); ++id) { + if (label_table_[id] == -1 && deleted_ids_.count(id) == 0) { + active_padding_label_ids_.insert(id); + } + } + } + void UpdateLabel(LabelType old_label, LabelType new_label) { // 1. check whether new_label is occupied @@ -142,6 +168,16 @@ class LabelTable { for (size_t i = 0; i < label_table_.size(); ++i) { if (label_table_[i] == old_label) { label_table_[i] = new_label; + if (old_label == -1 || new_label == -1) { + std::scoped_lock wlock(delete_ids_mutex_); + if (deleted_ids_.count(static_cast(i)) == 0) { + if (new_label == -1) { + active_padding_label_ids_.insert(static_cast(i)); + } else { + active_padding_label_ids_.erase(static_cast(i)); + } + } + } found = true; } } @@ -185,6 +221,7 @@ class LabelTable { void Deserialize(lvalue_or_rvalue reader) { StreamReader::ReadVector(reader, label_table_); + RebuildActivePaddingLabelIds(); if (use_reverse_map_) { this->label_remap_.Clear(); this->label_remap_.Reserve(label_table_.size()); @@ -349,8 +386,12 @@ class LabelTable { std::scoped_lock wlock(delete_ids_mutex_); from_removed = deleted_ids_.erase(from) > 0; deleted_ids_.erase(to); + active_padding_label_ids_.erase(from); + active_padding_label_ids_.erase(to); if (from_removed) { deleted_ids_.insert(to); + } else if (label_table_[from] == -1) { + active_padding_label_ids_.insert(to); } } @@ -365,6 +406,17 @@ class LabelTable { void ShrinkToFit(InnerIdType capacity) { + { + std::scoped_lock wlock(delete_ids_mutex_); + for (auto it = active_padding_label_ids_.begin(); + it != active_padding_label_ids_.end();) { + if (*it >= capacity) { + it = active_padding_label_ids_.erase(it); + } else { + ++it; + } + } + } // Avoid a full-table copy for small removals; vector storage is still compacted by BruteForce. if (capacity <= label_table_.capacity() / 2) { try { @@ -391,13 +443,15 @@ class LabelTable { } { std::scoped_lock wlock(delete_ids_mutex_); + active_padding_label_ids_.erase(inner_id); deleted_ids_.erase(inner_id); } total_count_.fetch_sub(1); } private: - UnorderedSet deleted_ids_; // Record deleted ids. + UnorderedSet deleted_ids_; // Record deleted ids. + UnorderedSet active_padding_label_ids_; FilterPtr deleted_ids_filter_{nullptr}; // Filter to filter out deleted ids. mutable std::shared_mutex delete_ids_mutex_; // Mutex to protect deleted_ids_. diff --git a/src/index/index_impl.h b/src/index/index_impl.h index 9941525d94..8e57b70e25 100644 --- a/src/index/index_impl.h +++ b/src/index/index_impl.h @@ -61,6 +61,11 @@ class IndexImpl : public Index { if ((query)->GetNumElements() == 0) { \ return make_empty_search_result(); \ } +#define CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters) \ + if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters) && \ + ((query) == nullptr || (query)->GetNumElements() <= 1)) { \ + return make_empty_search_result(); \ + } #define CHECK_IMMUTABLE_INDEX(operation_str) \ if (this->inner_index_->immutable_.load(std::memory_order_acquire)) { \ return tl::unexpected(Error(ErrorType::UNSUPPORTED_INDEX_OPERATION, \ @@ -315,9 +320,7 @@ class IndexImpl : public Index { return tl::unexpected(threshold_validation.error()); } CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->KnnSearch(query, k, parameters, invalid)); } @@ -331,9 +334,7 @@ class IndexImpl : public Index { return tl::unexpected(threshold_validation.error()); } CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->KnnSearch(query, k, parameters, filter)); } @@ -347,9 +348,7 @@ class IndexImpl : public Index { return tl::unexpected(threshold_validation.error()); } CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->KnnSearch(query, k, parameters, filter)); } @@ -360,9 +359,7 @@ class IndexImpl : public Index { return tl::unexpected(threshold_validation.error()); } CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(search_param.parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, search_param.parameters); if (search_param.is_iter_filter) { SAFE_CALL(return this->inner_index_->KnnSearch(query, k, @@ -389,9 +386,7 @@ class IndexImpl : public Index { return tl::unexpected(threshold_validation.error()); } CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->KnnSearch( query, k, parameters, filter, nullptr, iter_ctx, is_last_filter)); } @@ -428,9 +423,7 @@ class IndexImpl : public Index { const std::string& parameters, int64_t limited_size = -1) const override { CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->RangeSearch(query, radius, parameters, limited_size)); } @@ -441,9 +434,7 @@ class IndexImpl : public Index { BitsetPtr invalid, int64_t limited_size = -1) const override { CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->RangeSearch( query, radius, parameters, invalid, limited_size)); } @@ -455,9 +446,7 @@ class IndexImpl : public Index { const std::function& filter, int64_t limited_size = -1) const override { CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->RangeSearch( query, radius, parameters, filter, limited_size)); } @@ -469,9 +458,7 @@ class IndexImpl : public Index { const FilterPtr& filter, int64_t limited_size = -1) const override { CHECK_QUERY_RETURN_EMPTY_DATASET(query); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { - return make_empty_search_result(); - } + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(query, parameters); SAFE_CALL(return this->inner_index_->RangeSearch( query, radius, parameters, filter, limited_size)); } @@ -558,9 +545,9 @@ class IndexImpl : public Index { } } SAFE_CALL(ValidateSearchThreshold(request.threshold_); - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(request.params_str_)) { - return make_empty_search_result(); - } return this->inner_index_->SearchWithRequest(request)); + CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERY(request.query_, + request.params_str_); + return this->inner_index_->SearchWithRequest(request)); } tl::expected diff --git a/src/utils/timer.cpp b/src/utils/timer.cpp index 87245a0e03..ca37a12ad0 100644 --- a/src/utils/timer.cpp +++ b/src/utils/timer.cpp @@ -48,6 +48,11 @@ Timer::SetThreshold(double threshold) { threshold_ = threshold; } +void +Timer::Reset() { + start = std::chrono::steady_clock::now(); +} + Timer::~Timer() { auto finish = std::chrono::steady_clock::now(); std::chrono::duration duration = finish - start; diff --git a/src/utils/timer.h b/src/utils/timer.h index 13fab25084..e8cbbb3aad 100644 --- a/src/utils/timer.h +++ b/src/utils/timer.h @@ -14,7 +14,9 @@ // limitations under the License. #pragma once + #include +#include namespace vsag { class Timer { @@ -33,6 +35,9 @@ class Timer { void SetThreshold(double threshold); + void + Reset(); + bool CheckOvertime(); diff --git a/tests/test_hgraph.cpp b/tests/test_hgraph.cpp index e3ef05917b..e1e452239c 100644 --- a/tests/test_hgraph.cpp +++ b/tests/test_hgraph.cpp @@ -943,6 +943,38 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::HGraphTestIndex, auto param = GenerateHGraphBuildParametersString(build_param); auto index = TestFactory(name, param, true); auto dataset = pool.GetDatasetAndCreate(dim, base_count, metric_type); + auto batch_query = vsag::Dataset::Make(); + batch_query->NumElements(2) + ->Dim(dataset->query_->GetDim()) + ->Float32Vectors(dataset->query_->GetFloat32Vectors()) + ->Owner(false); + auto batch_result = index->KnnSearch(batch_query, 10, search_param); + REQUIRE(batch_result.has_value()); + REQUIRE(batch_result.value()->GetNumElements() == batch_query->GetNumElements()); + REQUIRE(batch_result.value()->GetDim() == 0); + + auto removed_index = TestFactory(name, param, true); + TestIndex::TestBuildIndex(removed_index, dataset, true); + std::vector removed_ids(dataset->base_->GetIds(), + dataset->base_->GetIds() + dataset->base_->GetNumElements()); + auto remove_result = removed_index->Remove(removed_ids, vsag::RemoveMode::MARK_REMOVE); + REQUIRE(remove_result.has_value()); + REQUIRE(remove_result.value() == removed_ids.size()); + auto removed_batch_result = removed_index->KnnSearch(batch_query, 10, search_param); + REQUIRE(removed_batch_result.has_value()); + REQUIRE(removed_batch_result.value()->GetNumElements() == batch_query->GetNumElements()); + REQUIRE(removed_batch_result.value()->GetDim() == 0); + + vsag::SearchRequest batch_range_request; + batch_range_request.mode_ = vsag::SearchMode::RANGE_SEARCH; + batch_range_request.query_ = batch_query; + batch_range_request.radius_ = 10.0F; + batch_range_request.limited_size_ = 1; + batch_range_request.params_str_ = search_param; + auto batch_range_result = index->SearchWithRequest(batch_range_request); + REQUIRE_FALSE(batch_range_result.has_value()); + auto direct_batch_range_result = index->RangeSearch(batch_query, 10.0F, search_param, 1); + REQUIRE_FALSE(direct_batch_range_result.has_value()); TestGetMinAndMaxId(index, dataset, false); TestKnnSearch(index, dataset, search_param, recall, false); TestKnnSearchIter(index, dataset, search_param, recall, false); @@ -3174,7 +3206,9 @@ TestHGraphReverseEdges(const fixtures::HGraphTestIndexPtr& test_index, for (auto metric_type : resource->metric_types) { for (auto dim : resource->dims) { - for (auto& [base_quantization_str, recall] : resource->test_cases) { + for (const auto& test_case : resource->test_cases) { + const auto& base_quantization_str = test_case.first; + const auto& recall = test_case.second; INFO(fmt::format("metric_type: {}, dim: {}, base_quantization_str: {}", metric_type, dim, @@ -4300,3 +4334,149 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::HGraphTestIndex, TestIndex::TestBuildIndex(cache_index, dataset, true); HGraphTestIndex::TestGeneral(cache_index, dataset, search_param, 0.98f); } + +static void +TestHGraphMultiQueryKnnSearch(const fixtures::HGraphTestIndexPtr& test_index, + const fixtures::HGraphResourcePtr& resource) { + using namespace fixtures; + auto search_param_str = fmt::format(search_param_tmp, 200, false); + + for (auto metric_type : resource->metric_types) { + for (auto dim : resource->dims) { + for (const auto& test_case : resource->test_cases) { + const auto& base_quantization_str = test_case.first; + INFO(fmt::format("metric_type: {}, dim: {}, base_quantization_str: {}", + metric_type, + dim, + base_quantization_str)); + + if (HGraphTestIndex::IsRaBitQ(base_quantization_str) && + dim < fixtures::RABITQ_MIN_RACALL_DIM) { + dim = fixtures::RABITQ_MIN_RACALL_DIM; + } + + HGraphTestIndex::HGraphBuildParam build_param( + metric_type, dim, base_quantization_str); + auto param = HGraphTestIndex::GenerateHGraphBuildParametersString(build_param); + + auto index = TestIndex::TestFactory(test_index->name, param, true); + auto dataset = HGraphTestIndex::pool.GetDatasetAndCreate( + dim, resource->base_count, metric_type); + + TestIndex::TestBuildIndex(index, dataset, true); + + const int64_t num_queries = 5; + const int64_t k = 10; + int64_t dim_val = dataset->query_->GetDim(); + int64_t available_queries = dataset->query_->GetNumElements(); + + std::vector multi_query_data(num_queries * dim_val); + const float* original_queries = dataset->query_->GetFloat32Vectors(); + for (int64_t i = 0; i < num_queries; ++i) { + int64_t src_idx = i % available_queries; + std::copy(original_queries + src_idx * dim_val, + original_queries + (src_idx + 1) * dim_val, + multi_query_data.data() + i * dim_val); + } + + auto multi_query = vsag::Dataset::Make(); + multi_query->NumElements(num_queries) + ->Dim(dim_val) + ->Float32Vectors(multi_query_data.data()) + ->Owner(false); + + auto multi_result = index->KnnSearch(multi_query, k, search_param_str); + REQUIRE(multi_result.has_value()); + REQUIRE(multi_result.value()->GetNumElements() == num_queries); + REQUIRE(multi_result.value()->GetDim() == k); + + const auto* multi_ids = multi_result.value()->GetIds(); + for (int64_t q_idx = 0; q_idx < num_queries; ++q_idx) { + auto single_query = vsag::Dataset::Make(); + single_query->NumElements(1) + ->Dim(dim_val) + ->Float32Vectors(multi_query_data.data() + q_idx * dim_val) + ->Owner(false); + auto single_result = index->KnnSearch(single_query, k, search_param_str); + REQUIRE(single_result.has_value()); + + int64_t single_count = single_result.value()->GetDim(); + const auto* single_ids = single_result.value()->GetIds(); + int64_t offset = q_idx * k; + // Without filter, all queries should return exactly k results. + REQUIRE(single_count == k); + for (int64_t i = 0; i < k; ++i) { + REQUIRE(multi_ids[offset + i] == single_ids[i]); + } + } + } + } + } +} + +HGRAPH_PR_DAILY_CASE("HGraph Multi-Query Knn Search", + "[ft][search][hgraph]", + TestHGraphMultiQueryKnnSearch) + +static void +TestHGraphMultiQueryRangeSearch(const fixtures::HGraphTestIndexPtr& test_index, + const fixtures::HGraphResourcePtr& resource) { + using namespace fixtures; + auto search_param_str = fmt::format(search_param_tmp, 200, false); + + for (auto metric_type : resource->metric_types) { + for (auto dim : resource->dims) { + for (const auto& test_case : resource->test_cases) { + const auto& base_quantization_str = test_case.first; + INFO(fmt::format("metric_type: {}, dim: {}, base_quantization_str: {}", + metric_type, + dim, + base_quantization_str)); + + if (HGraphTestIndex::IsRaBitQ(base_quantization_str) && + dim < fixtures::RABITQ_MIN_RACALL_DIM) { + dim = fixtures::RABITQ_MIN_RACALL_DIM; + } + + HGraphTestIndex::HGraphBuildParam build_param( + metric_type, dim, base_quantization_str); + auto param = HGraphTestIndex::GenerateHGraphBuildParametersString(build_param); + + auto index = TestIndex::TestFactory(test_index->name, param, true); + auto dataset = HGraphTestIndex::pool.GetDatasetAndCreate( + dim, resource->base_count, metric_type); + + TestIndex::TestBuildIndex(index, dataset, true); + + const int64_t num_queries = 3; + int64_t dim_val = dataset->query_->GetDim(); + int64_t available_queries = dataset->query_->GetNumElements(); + const float radius = 0.5F; + const int64_t limited_size = 10; + + std::vector multi_query_data(num_queries * dim_val); + const float* original_queries = dataset->query_->GetFloat32Vectors(); + for (int64_t i = 0; i < num_queries; ++i) { + int64_t src_idx = i % available_queries; + std::copy(original_queries + src_idx * dim_val, + original_queries + (src_idx + 1) * dim_val, + multi_query_data.data() + i * dim_val); + } + + auto multi_query = vsag::Dataset::Make(); + multi_query->NumElements(num_queries) + ->Dim(dim_val) + ->Float32Vectors(multi_query_data.data()) + ->Owner(false); + + auto multi_result = + index->RangeSearch(multi_query, radius, search_param_str, limited_size); + REQUIRE_FALSE(multi_result.has_value()); + } + } + } +} + +HGRAPH_PR_DAILY_CASE("HGraph Multi-Query Range Search", + "[ft][search][hgraph]", + TestHGraphMultiQueryRangeSearch) diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 8740fd257b..72a31b31bf 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -1523,6 +1523,96 @@ TestIVFSearchDisableReorder(const fixtures::IVFResourcePtr& resource) { IVF_PR_DAILY_CASE("IVF Search Disable Reorder", "[ft][search][ivf]", TestIVFSearchDisableReorder) +static void +TestIVFMultiQueryKnnSearch(const fixtures::IVFResourcePtr& resource) { + using namespace fixtures; + const std::vector> test_cases = {{"fp32", 0.0F}}; + ForEachIVFCase( + resource, + test_cases, + [resource](const auto& metric_type, + auto dim, + const auto& train_type, + const auto& base_quantization_str, + auto /* recall */) { + auto build_param = IVFTestIndex::GenerateIVFBuildParametersString( + metric_type, dim, base_quantization_str, 210, train_type); + auto index = TestIndex::TestFactory(IVFTestIndex::name, build_param, true); + auto dataset = + IVFTestIndex::pool.GetDatasetAndCreate(dim, resource->base_count, metric_type); + TestIndex::TestBuildIndex(index, dataset, true); + + const int64_t query_count = 3; + const int64_t query_dim = dataset->query_->GetDim(); + const int64_t k = dataset->base_->GetNumElements() + 1; + const float* source_queries = dataset->query_->GetFloat32Vectors(); + std::vector queries(query_count * query_dim); + for (int64_t q_idx = 0; q_idx < query_count; ++q_idx) { + std::copy_n(source_queries + q_idx * query_dim, + query_dim, + queries.data() + q_idx * query_dim); + } + + auto multi_query = vsag::Dataset::Make(); + multi_query->NumElements(query_count) + ->Dim(query_dim) + ->Float32Vectors(queries.data()) + ->Owner(false); + const auto search_param = fmt::format(search_param_tmp, 210); + auto multi_result = index->KnnSearch(multi_query, k, search_param); + REQUIRE(multi_result.has_value()); + REQUIRE(multi_result.value()->GetNumElements() == query_count); + REQUIRE(multi_result.value()->GetDim() == k); + + const auto* multi_ids = multi_result.value()->GetIds(); + const auto* multi_dists = multi_result.value()->GetDistances(); + for (int64_t q_idx = 0; q_idx < query_count; ++q_idx) { + auto single_query = vsag::Dataset::Make(); + single_query->NumElements(1) + ->Dim(query_dim) + ->Float32Vectors(queries.data() + q_idx * query_dim) + ->Owner(false); + auto single_result = index->KnnSearch(single_query, k, search_param); + REQUIRE(single_result.has_value()); + + const int64_t single_count = single_result.value()->GetDim(); + REQUIRE(single_count < k); + const int64_t offset = q_idx * k; + for (int64_t i = 0; i < single_count; ++i) { + REQUIRE(multi_ids[offset + i] == single_result.value()->GetIds()[i]); + REQUIRE(multi_dists[offset + i] == single_result.value()->GetDistances()[i]); + } + for (int64_t i = single_count; i < k; ++i) { + REQUIRE(multi_ids[offset + i] == -1); + REQUIRE(multi_dists[offset + i] == std::numeric_limits::infinity()); + } + } + }); +} + +IVF_PR_DAILY_CASE("IVF Multi-Query Knn Search", "[ft][search][ivf]", TestIVFMultiQueryKnnSearch) + +TEST_CASE("IVF Multi-Query Knn Search Empty Index", "[ft][search][ivf][pr]") { + constexpr int64_t dim = 16; + constexpr int64_t query_count = 2; + constexpr int64_t k = 3; + const auto build_param = + fixtures::IVFTestIndex::GenerateIVFBuildParametersString("l2", dim, "fp32", 8, "kmeans"); + auto index = fixtures::TestIndex::TestFactory(fixtures::IVFTestIndex::name, build_param, true); + std::vector queries(query_count * dim, 0.0F); + auto batch_query = vsag::Dataset::Make(); + batch_query->NumElements(query_count)->Dim(dim)->Float32Vectors(queries.data())->Owner(false); + + auto result = index->KnnSearch(batch_query, k, fmt::format(fixtures::search_param_tmp, 8)); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetNumElements() == query_count); + REQUIRE(result.value()->GetDim() == k); + for (int64_t i = 0; i < query_count * k; ++i) { + REQUIRE(result.value()->GetIds()[i] == -1); + REQUIRE(result.value()->GetDistances()[i] == std::numeric_limits::infinity()); + } +} + static void TestIVFBuildWithLargeK(const fixtures::IVFResourcePtr& resource) { using namespace fixtures; @@ -2284,6 +2374,22 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, check_bucket_result(batch_result.value(), 3, scan_buckets_count, buckets_count); } + SECTION("batch routing ignores search-only options") { + std::vector batch(dim * 2); + std::memcpy(batch.data(), query_vector.data(), dim * sizeof(float)); + std::memcpy(batch.data() + dim, query_vector.data(), dim * sizeof(float)); + auto batch_query = vsag::Dataset::Make(); + batch_query->NumElements(2)->Dim(dim)->Float32Vectors(batch.data())->Owner(false); + vsag::SearchRequest batch_req; + batch_req.query_ = batch_query; + batch_req.mode_ = vsag::SearchMode::RANGE_SEARCH; + batch_req.topk_ = 1; + batch_req.params_str_ = route_params; + auto batch_result = index->SearchWithRequest(batch_req); + REQUIRE(batch_result.has_value()); + check_bucket_result(batch_result.value(), 2, scan_buckets_count, buckets_count); + } + SECTION("reject zero queries") { auto zero_query = vsag::Dataset::Make(); zero_query->NumElements(0)->Dim(dim)->Owner(false);