feat(search): add multi-query batch search support - #2607
Conversation
|
/label status/waiting-for-review |
|
Automated pull request review completed. Review effort: Submitted 4 inline comments. |
Merge Protections🟢 All 3 merge protections satisfied — ready to merge. Show 3 satisfied protections🟢 Require kind label
🟢 Require version label
🟢 Require linked issue for feature/bug PRs
|
LHT129
left a comment
There was a problem hiding this comment.
Overall this is a well-structured PR that adds multi-query batch KNN search support for HGraph and IVF. The code quality is high with thorough overflow guards, sentinel-based padding, and clear documentation updates.
Summary of findings:
-
[suggestion] Range search statistics regression: The extracted
search_range_with_requestmethod usesctx.stats->Dump()instead ofmci_result.MakeStatistics(stats).Dump(), losing the route field (brute_force/mci/graph) from range search statistics output. -
[suggestion] IVF batch path
Dataset::Make()per iteration: The temporary dataset allocation inside the per-query loop could be hoisted out for a minor performance improvement. -
[note]
HasActiveLabelnaming: The function is hardcoded for label-1but has a general-purpose signature. Consider renaming toHasActivePaddingLabel(). -
[note]
last_result_inner_idsnaming: The variable name is slightly misleading since it only captures reasoning-related inner IDs (single-query only). Consider renaming toreasoning_inner_ids.
Positive observations:
- Comprehensive overflow guards for
query_count * kand byte-level allocations - Sentinel pre-fill with
ids = -1anddists = +infis well-designed - Clean extraction of range search into
search_range_with_request - Good test coverage including empty index, batch KNN, batch range rejection, and IVF bucket routing
- Proper rejection of reasoning with batch queries
label_tabletracking of-1labels is correctly maintained across all mutation paths (Insert, Remove, Merge, UpdateLabel, ShrinkToFit, Deserialize)
There was a problem hiding this comment.
Pull request overview
Adds multi-query (batched) KNN search support to VSAG’s core indexes (notably HGraph and IVF) by allowing DatasetPtr queries with NumElements > 1, defining/clarifying result layout semantics, and extending regression coverage and API documentation accordingly.
Changes:
- Implement batched KNN execution paths for HGraph and IVF (with explicit rejection of batched range search).
- Standardize batch result layout to row-major
query_count x dimwith sentinel padding (id = -1), and enforce the “no external label-1” constraint for unambiguous padding. - Add/extend functional tests and update public API docs to describe single-query vs batched semantics.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_ivf.cpp | Adds multi-query KNN tests (including empty-index behavior) and batch routing behavior checks for IVF. |
| tests/test_hgraph.cpp | Adds/extends multi-query KNN tests and asserts multi-query range search is rejected for HGraph. |
| src/utils/timer.h | Adds Timer::Reset() declaration (used for per-query timeout tracking in batch search). |
| src/utils/timer.cpp | Implements Timer::Reset(). |
| src/index/index_impl.h | Adjusts empty-index short-circuit behavior to better support batch-query semantics. |
| src/impl/label_table/label_table.h | Tracks active external label -1 usage to gate batch KNN padding semantics. |
| src/impl/label_table/label_table.cpp | Wires allocator/maintenance of the active-padding-label tracker and rebuilds it on deserialize/merge. |
| src/algorithm/ivf/ivf.cpp | Implements IVF batch KNN behavior inside SearchWithRequest, including padding and overflow guards. |
| src/algorithm/hgraph/hgraph.h | Extends get_data to support indexed query access and adds offset overflow guards. |
| src/algorithm/hgraph/hgraph_serialize.cpp | Rebuilds active padding label tracking after legacy label-table deserialization paths. |
| src/algorithm/hgraph/hgraph_search.cpp | Implements HGraph batch KNN in SearchWithRequest, factors out range-search path, and adds padding/overflow handling. |
| include/vsag/search_request.h | Updates SearchRequest::query_ docs to describe single vs batched semantics and constraints. |
| include/vsag/index.h | Updates Index::SearchWithRequest result-shape documentation for single vs batched behaviors. |
| docs/docs/zh/src/api/search.md | Documents batched KNN availability/limitations in the Chinese API docs. |
| docs/docs/zh/src/api/index_class.md | Updates Chinese index API docs for batched KNN result reading and constraints. |
| docs/docs/zh/src/api/dataset.md | Updates Chinese dataset docs to explain batched result matrix layout and padding. |
| docs/docs/en/src/api/search.md | Documents batched KNN availability/limitations and clarifies IVF routing-only mode wording. |
| docs/docs/en/src/api/index_class.md | Updates English index API docs for batched KNN result reading and constraints. |
| docs/docs/en/src/api/dataset.md | Updates English dataset docs to explain batched result matrix layout and padding. |
Suppressed comments (1)
include/vsag/search_request.h:51
- Same indentation issue continues in the remainder of this bullet; keeping the alignment consistent avoids broken formatting in generated API docs.
* 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.
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1150 changed lines across 19 files).
Submitted 2 inline comments.
Reviewed commit 500b5ca.
132bc47 to
6744bab
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.
Suppressed comments (6)
include/vsag/index.h:343
- The API docs say batch-KNN padding uses an undefined distance value, but the implementations and tests consistently use +infinity (and rely on rejecting external label -1 so padding stays unambiguous). The comment should match the actual contract to avoid misleading callers.
* Queries that yielded fewer than dim neighbors are padded with
* sentinel entries (id = -1, distance value undefined). 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.
include/vsag/search_request.h:53
- The Doxygen comment has mis-indented
*markers in the Batched KNN bullet list (these lines start with an extra space), which can break formatting in generated docs.
* 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
include/vsag/search_request.h:54
- This line is also mis-indented (extra space before
*), which can break Doxygen formatting for the Batched KNN bullet list.
* label -1 to keep this padding unambiguous.
include/vsag/search_request.h:232
- The comment claims
bucket_ids_bypass currently supports only single-query, but IVF::SearchWithRequest enforcesbucket_ids_.size() == query->GetNumElements()and supports one bucket list per query (including batched queries). The documentation should reflect the actual supported shape.
* @details Currently only single-query is supported; outer vector must contain exactly one entry.
docs/docs/en/src/api/search.md:106
- This doc says IVF bucket-IDs bypass is single-query only, but IVF::SearchWithRequest requires the outer vector size to match the number of query vectors and supports one bucket list per query (including batches).
- Currently only single-query is supported; the outer vector must contain exactly one entry.
docs/docs/zh/src/api/search.md:100
- 这里写成了“仅支持单查询”,但 IVF::SearchWithRequest 要求 bucket_ids_ 外层长度与查询数一致,并支持为每个查询提供一个桶列表(包括批量查询)。文档应与实际行为一致。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。
6744bab to
c420ed1
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/index/index_impl.h:68
CHECK_EMPTY_INDEX_RETURN_EMPTY_DATASET_IF_SINGLE_QUERYtreats a null query as eligible for the empty-index shortcut, which makesSearchWithRequestreturn an empty dataset (instead of INVALID_ARGUMENT) whenrequest.query_ == nullptrand the index is empty. That creates inconsistent behavior depending on index size and can mask caller bugs.
#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(); \
}
include/vsag/search_request.h:54
- The
SearchRequest::query_doc comment has inconsistent leading whitespace on several lines (the*is misaligned). This looks like an accidental indentation change and makes generated docs harder to read.
* - 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
include/vsag/index.h:343
- The SearchWithRequest return contract says batched KNN padding has an undefined distance, but the implementation/tests/docs use
+infinity. This mismatch can mislead API consumers; the header should document the same sentinel semantics (id == -1,distance == +infinity).
* Queries that yielded fewer than dim neighbors are padded with
* sentinel entries (id = -1, distance value undefined). 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.
| @@ -2141,12 +2141,22 @@ IVF::SearchWithRequest(const SearchRequest& request) const { | |||
| CHECK_ARGUMENT(request.expected_labels_.empty(), | |||
LHT129
left a comment
There was a problem hiding this comment.
Review of PR #2607 (multi-query batch search for HGraph and IVF):
Previously reported issues — all confirmed fixed:
- MCI route info is now preserved via
mci_result.MakeStatistics(*ctx.stats).Dump()(hgraph_search.cpp:562) enable_reorderandrabitq_one_bit_searchare now properly guarded withuse_custom_distance ? false : ...(hgraph_search.cpp:494-496)Dataset::Make()is no longer called inside the per-query loop; result buffers are pre-allocated once
New observations:
- The
k > 0check is duplicated whenuse_custom_distanceis true (once in the else branch, once unconditionally after). Harmless but slightly redundant. - The IVF batch path reuses the existing
search_funclambda which does JSON Parse + Dump per query iteration — this is pre-existing code not touched by this PR, but worth noting for a future optimization pass.
Overall: The changes are well-structured, defensive (overflow checks, padding label validation, empty-index handling), and well-tested (HGraph multi-query KNN/range rejection, IVF multi-query KNN with sentinel padding, empty index batch). The refactoring of range search into search_range_with_request and extraction of apply_hops_limit improves readability.
| MCIHybridSearchResult mci_result(params, filter); | ||
| if (!use_custom_distance && params.brute_force_threshold > 0.0F && | ||
| mci_result.valid_ratio <= params.brute_force_threshold) { | ||
| search_result = this->brute_force_search<InnerSearchMode::RANGE_SEARCH>( |
There was a problem hiding this comment.
[suggestion] The k > 0 check is duplicated when use_custom_distance is true: the else branch checks it, and then the unconditional check after the if/else block also checks it. The else branch check is redundant and can be removed.
Current code:
} else {
CHECK_ARGUMENT(k > 0, "topk must be greater than 0");
}
CHECK_ARGUMENT(k > 0, fmt::format("k({}) must be greater than 0", k));Suggested:
}
CHECK_ARGUMENT(k > 0, fmt::format("k({}) must be greater than 0", k));| 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; |
There was a problem hiding this comment.
[suggestion] search_param.enable_rabitq_one_bit_search is set twice with the same expression in search_range_with_request: once before the routing loop (line ~450) and again after routing before the main search (line 496). The second assignment is redundant since the value does not change between the two assignments.
| auto new_label = std::get<1>(id_map(other->label_table_[i])); | ||
| auto new_inner_id = static_cast<InnerIdType>(i + current_total_count_u); | ||
| this->label_table_[i + current_total_count_u] = new_label; | ||
| if (new_label == -1) { |
There was a problem hiding this comment.
[note] MergeOther iterates over all entries from the source table (including logically deleted ones, since GetTotalCount() returns total_count_ which is not decremented on MarkRemove). If the source table has a deleted entry with label -1, that entry will be added to active_padding_label_ids_ in the merged table, incorrectly marking it as an active padding label. The source's deleted_ids_ are not transferred during merge, so the merged table has no way to know this entry was deleted.
Consider checking other->IsRemoved(i) for each entry before adding to active_padding_label_ids_, or skipping deleted entries from the source during merge.
| auto new_label = std::get<1>(id_map(other->label_table_[i])); | ||
| auto new_inner_id = static_cast<InnerIdType>(i + current_total_count_u); | ||
| this->label_table_[i + current_total_count_u] = new_label; | ||
| if (new_label == -1) { |
There was a problem hiding this comment.
[note] MergeOther iterates over all entries from the source table and adds entries with new_label == -1 to active_padding_label_ids_ without first checking other->IsRemoved(i). If the source table has deleted entries whose label is -1, those deleted entries will be incorrectly tracked as active padding labels, causing HasActivePaddingLabel() to return true when there are no actual active padding labels. This could incorrectly block batch queries on the merged index.
Suggested fix: add if (!other->IsRemoved(i)) guard before inserting into active_padding_label_ids_.
|
|
||
| static DatasetPtr | ||
| make_empty_dataset_with_stats(const SearchStatistics& stats) { | ||
| make_empty_dataset_with_stats() { |
There was a problem hiding this comment.
[critical] Removing the two-argument overload make_empty_dataset_with_stats(const SearchStatistics&) breaks the call site at line 253 in KnnSearch (iterator-based overload), which still calls make_empty_dataset_with_stats(stats). The no-argument overload creates a fresh default SearchStatistics and discards the caller's accumulated timing/state. Either restore the two-argument overload or update the call site to use the no-argument version (if discarding stats is intentional).
| 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 || |
There was a problem hiding this comment.
[suggestion] The overflow guard dim_ > 0 && index <= INT64_MAX / dim_ will reject any call when dim_ == 0 and data_type_ != DATA_TYPE_SPARSE, because the left side of || is false and dim_ > 0 short-circuits to false. For a non-sparse empty index (dim_ == 0), this CHECK_ARGUMENT would fire incorrectly, preventing legitimate queries that should return empty results.
Consider restructuring the condition so that dim_ == 0 is handled gracefully (e.g. return nullptr early, or use dim_ > 0 || data_type_ == DATA_TYPE_SPARSE and only apply the multiplication check when dim_ > 0).
| if (use_reverse_map_) { | ||
| label_remap_.InsertOrAssign(label, id); | ||
| } | ||
| const bool was_padding_label = id < label_table_.size() && label_table_[id] == -1; |
There was a problem hiding this comment.
[note] was_padding_label reads label_table_[id] outside the delete_ids_mutex_ critical section (acquired a few lines later). The mutex comment at line 458 documents it as protecting deleted_ids_, but label_table_ is also read here to decide whether to insert/erase from active_padding_label_ids_.
If InsertOrAssign can be called concurrently with UpdateLabel or Move (which modify label_table_ and active_padding_label_ids_ under the same mutex), this is a data race on label_table_[id]. In practice InsertOrAssign is likely single-threaded during build, so the risk is low, but the lock ordering is worth a comment to prevent future misuse.
| use_custom_distance ? false : params.rabitq_one_bit_search; | ||
| search_param.skip_ratio = params.skip_ratio; | ||
| search_param.skip_strategy_type = params.skip_strategy_type; | ||
|
|
There was a problem hiding this comment.
[P1] search_range_with_request missing mci_result.route != "mci" guard before reorder
The extracted search_range_with_request method applies reorder without checking whether MCI already produced the final result. The original inline code (before this refactor) guarded both reorder branches with mci_result.route != "mci":
// OLD (correct):
if (mci_result.route != "mci" and not brute_force_used and use_reorder_ and search_param.enable_reorder) { ... }
else if (mci_result.route != "mci" and not brute_force_used and search_param.enable_reorder and params.rabitq_one_bit_search) { ... }The new code at lines 499 and 507 drops this guard:
// NEW (bug):
if (not brute_force_used and use_reorder_ and search_param.enable_reorder) { ... }
else if (not brute_force_used and search_param.enable_reorder and params.rabitq_one_bit_search) { ... }When MCI produces the result (mci_result.route == "mci"), search_result already contains the MCI output, and reorder should be skipped — the KNN path (line 898) still correctly guards with mci_result.route != "mci". The range search path should do the same.
LHT129
left a comment
There was a problem hiding this comment.
Code Review Summary for PR #2607
Issue Found
[critical] IVF batch search: unchecked tl::expected dereference (src/algorithm/ivf/ivf.cpp:2007)
The batch search lambda calls SearchWithRequest(one_request) and immediately dereferences the result with ->GetDim() without checking .has_value(). If any single-query SearchWithRequest fails, this dereferences the error variant — undefined behavior. Check the result before accessing it.
Positive Observations
- HGraph batch search has proper overflow guards for both int64_t and size_t byte-level allocations
- RAII visited_list_guard properly used in both KNN and range search paths
- enable_reorder and enable_rabitq_one_bit_search correctly guarded with use_custom_distance
- HasActivePaddingLabel mechanism is a clean approach for disambiguating padding
- Test coverage includes multi-query KNN correctness, range search rejection, IVF bucket routing, and parallel search determinism
- Documentation updates consistently describe the new batch semantics
Several issues have already been flagged by other reviewers (SINDI regression, Doxygen alignment, missing limits include, etc). Please address the critical issue before merging.
| 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<size_t>::max() / sizeof(int64_t), |
There was a problem hiding this comment.
[critical] The batch search lambda at line 2189 dereferences one_result (a tl::expected<DatasetPtr, Error>) with ->GetDim(), ->GetIds(), and ->GetDistances() without first checking .has_value(). If SearchWithRequest returns an error for any single query (e.g. due to invalid parameters propagated from the batch request), this dereferences the error variant, which is undefined behavior.
The fix should check the result before accessing it:
auto one_result = this->SearchWithRequest(one_request);
if (not one_result.has_value()) {
// propagate error: either throw, or set an error flag and break
return;
}
const auto count = std::min(request.topk_, one_result.value()->GetDim());Note that in the parallel path (std::future<void>), errors cannot propagate from worker threads, so an atomic error flag or a std::future carrying the error status would be needed.
Signed-off-by: LHT129 <tianlan.lht@antgroup.com> Co-authored-by: opencode <opencode@anthropic.com> Signed-off-by: LHT129 <tianlan.lht@antgroup.com>
c420ed1 to
5357338
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.
Suppressed comments (9)
include/vsag/index.h:340
- The new batch contract is only documented for
SearchWithRequest, but HGraph and IVF's publicKnnSearchoverloads delegate here and now accept the same multi-query inputs. The overload documentation above still promisesnum_elements = 1and anum_elements * klayout, so callers of the primary KNN APIs receive an incorrect contract; update those overloads too.
* - 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
src/algorithm/hgraph/hgraph_search.cpp:116
- This iterator-only single-query check is after the empty-index early return at line 92.
IndexImpldeliberately forwards multi-query requests for the new batch handling, so a multi-query iterator call on an empty HGraph returns a single empty dataset instead of rejecting the unsupported iterator shape. Move the check before the early return.
CHECK_ARGUMENT(query->GetNumElements() == 1,
"iterator-based KnnSearch only supports single query (NumElements=1)");
src/algorithm/hgraph/hgraph_search.cpp:936
- This RaBitQ rerank call also omits
request.threshold_. As in the regular rerank branch, the post-rerank filter cannot recover eligible candidates discarded when the unfiltered top-k heap was formed, so thresholded batch KNN can return incomplete results. Pass the request threshold toreorder.
} 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);
src/algorithm/hgraph/hgraph_search.cpp:1026
- Unlike the single-query branch, which serializes
mci_result.MakeStatistics(stats), the batch path always serializes onlystats.Dump(). MCI-enabled batch searches therefore omit themci_hybrid_*, seed-count, and raw-CSR diagnostics from the returned dataset, making the result statistics incomplete. Preserve or explicitly aggregate the per-query MCI metadata for batch results.
if (query_count > 1) {
dataset_results->Statistics(stats.Dump());
}
src/algorithm/hgraph/hgraph_serialize.cpp:320
- The new active-padding rebuild is present for the modern label-info paths, but
deserialize_basic_info_v0_14still readslabel_table_directly at line 170 without rebuilding it. A legacy index containing an active external label-1will therefore bypass the new batch-safety check and can return an ambiguous-1result; rebuild the tracking set in that legacy path too.
this->label_table_->RebuildActivePaddingLabelIds();
docs/docs/en/src/api/search.md:106
- This constraint is now inconsistent with the IVF implementation: batch KNN accepts one bucket list per query and validates the outer vector against the query count. Document the batch form here; otherwise callers will be told to use a shape that the new implementation deliberately supports.
- Currently only single-query is supported; the outer vector must contain exactly one entry.
docs/docs/zh/src/api/search.md:100
- 此约束已与 IVF 实现不一致:批量 KNN 支持每个查询一个桶列表,并会校验外层向量与查询数一致。这里应记录批量形式,否则文档会要求调用方使用新实现不支持的形状。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。
src/algorithm/hgraph/hgraph_search.cpp:857
- Routing uses this separate
ep_search_param, but it never receivesbase_search_param.time_cost; the searcher only checksInnerSearchParam::time_costwhen enforcingtimeout_ms. Consequently hierarchical routing in this batch path is not timeout-bounded. Propagate the request timer to the routing parameter (and reset it at the start of each query if the timeout is intended to be per-query).
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;
src/algorithm/ivf/ivf.cpp:2145
- This batch branch invokes
this->SearchWithRequestonce per query, and each invocation creates a separateSearchStatistics; the outerstatsobject used for the final dataset is never updated. The returned batch result therefore reports zero distance evaluations and misses subquery timeouts even though the searches ran. Aggregate the subrequest statistics or refactor the batch path to share the query context.
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");
| make_empty_dataset_with_stats() { | ||
| SearchStatistics stats; | ||
| auto dataset_result = DatasetImpl::MakeEmptyDataset(); | ||
| dataset_result->Statistics(stats.Dump()); |
| this->reorder(raw_query, | ||
| this->get_reorder_codes(), | ||
| search_result, | ||
| k, | ||
| nullptr, | ||
| ctx, | ||
| rabitq_lower_bound_candidates_ptr); |
| if (visited_list != nullptr) { | ||
| pool->ReturnOne(visited_list); | ||
| } | ||
| } |
There was a problem hiding this comment.
[note] When element_count == 0 and query_count > 1 (line 835-841), the returned dataset has Dim(0). This is inconsistent with the normal batch layout where Dim is k. A caller that unconditionally reads dim = result->GetDim() and indexes with q_idx * dim + i would get 0 here, which differs from the documented rectangular query_count x k layout. Consider setting Dim(k) here for consistency with the non-empty batch path, or explicitly documenting this edge case in the API contract.
The same issue applies to the k == 0 early return at line 856-862.
| visited_list.reset(); | ||
| } | ||
| } | ||
| FilterPtr ft = this->create_search_filter(request.filter_, params.use_extra_info_filter); |
There was a problem hiding this comment.
[suggestion] In search_range_with_request, when brute_force_threshold triggers the brute-force path (line 727-732), the mci_result statistics object is left with its default-constructed state (route = "", valid_ratio = 0). The final statistics at line 779 use mci_result.MakeStatistics(*ctx.stats).Dump(), which will report an empty route string for this path. The old inline code set mci_result.route = "brute_force" before calling brute_force_search so that statistics correctly reflected the search path taken.
The KNN batch path correctly sets mci_result.route = "brute_force" at line 1161, but search_range_with_request at line 732 omits this assignment.
| check_bucket_result(batch_result.value(), 3, scan_buckets_count, buckets_count); | ||
| } | ||
|
|
||
| SECTION("batch routing ignores search-only options") { |
There was a problem hiding this comment.
[note] The test "batch routing ignores search-only options" at line 2377 uses RANGE_SEARCH mode with 2 queries and disable_bucket_scan params. This test passes because the bucket routing path (which handles disable_bucket_scan) returns early before reaching the range single-query validation. While this is correct behavior (bucket routing is a special mode that bypasses normal search), it may be worth adding a comment or making the test intent clearer — a reader might wonder why a 2-query RANGE_SEARCH succeeds when the API documents that range search only supports single queries.
LHT129
left a comment
There was a problem hiding this comment.
[suggestion] Overall, this PR is well-structured with solid engineering: the overflow guards, sentinel pre-fill, padding label tracking, and per-query entry point search are all correct and thorough. The existing 30+ comments already cover the critical issues (RAII visited list guards, IVF batch search implementation gaps, missing use_custom_distance guards). I added 3 additional notes:
hgraph_search.cpp:835—Dim(0)is returned whenelement_count == 0 && query_count > 1, while the single-query path returnsDim(element_count). Consider returningDim(element_count)consistently in both paths.hgraph_search.cpp:732—search_range_with_requestdoes not setmci_result.route = "brute_force"when falling through to brute force, unlike the KNN path.test_ivf.cpp:2377— The "batch routing ignores search-only options" test uses RANGE_SEARCH with 2 queries; consider adding a comment clarifying that bucket routing bypasses the range single-query validation.
The core batch KNN implementation is solid. The main areas to address are the existing critical comments (visited list RAII, IVF batch search, custom distance guards).
| /** | ||
| * @brief Pre-selected bucket IDs for bypassing IVF bucket routing (ClassifyDatasForSearch) | ||
| * @details The outer vector contains one entry per query vector. | ||
| * @details Currently only single-query is supported; outer vector must contain exactly one entry. |
There was a problem hiding this comment.
[note] The bucket_ids_ documentation says "Currently only single-query is supported; outer vector must contain exactly one entry", but the validation in index_impl.h (line 508) now allows bucket_ids_.size() != 1 for IVF indexes. The IVF batch path in ivf.cpp (lines 1987-2000) also handles per-query bucket_ids correctly for multi-query.
The docstring should be updated to reflect that IVF now supports multi-query bucket_ids_ in batch KNN mode.
Replacement for #1685, opened from the same current commit to start a clean review. This change adds multi-query batch KNN support for HGraph and IVF, updates API semantics and regression coverage, while retaining single-query behavior where batch result shapes are not representable. Closes #1684.