Skip to content

feat(search): add multi-query batch search support - #2607

Open
LHT129 wants to merge 1 commit into
antgroup:mainfrom
LHT129:batch_query_review_reset_1685
Open

feat(search): add multi-query batch search support#2607
LHT129 wants to merge 1 commit into
antgroup:mainfrom
LHT129:batch_query_review_reset_1685

Conversation

@LHT129

@LHT129 LHT129 commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

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.

Copilot AI lite review requested due to automatic review settings August 3, 2026 09:46
@vsag-bot

vsag-bot commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

/label status/waiting-for-review
/waiting-on reviewer
/request-review @jiaweizone
/request-review @wxyucs
/request-review @inabao

@vsag-bot

vsag-bot commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Automated pull request review completed.

Review effort: high (1428 changed lines across 19 files).

Submitted 4 inline comments.
Review: #2607 (review)

@mergify

mergify Bot commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Merge Protections

🟢 All 3 merge protections satisfied — ready to merge.

Show 3 satisfied protections

🟢 Require kind label

  • label~=^kind/

🟢 Require version label

  • label~=^version/

🟢 Require linked issue for feature/bug PRs

  • body~=(?im)(?:^|[\s\-\*])(?:close[sd]?|fix(?:e[sd])?|resolve[sd]?)\s*:?\s+(?:#\d+|[\w.\-]+/[\w.\-]+#\d+|https?://github\.com/[\w.\-]+/[\w.\-]+/issues/\d+)

Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
Comment thread src/algorithm/ivf/ivf.cpp Outdated
Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/algorithm/hgraph/hgraph_search.cpp
Comment thread src/impl/label_table/label_table.h Outdated
Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated

@LHT129 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. [suggestion] Range search statistics regression: The extracted search_range_with_request method uses ctx.stats->Dump() instead of mci_result.MakeStatistics(stats).Dump(), losing the route field (brute_force/mci/graph) from range search statistics output.

  2. [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.

  3. [note] HasActiveLabel naming: The function is hardcoded for label -1 but has a general-purpose signature. Consider renaming to HasActivePaddingLabel().

  4. [note] last_result_inner_ids naming: The variable name is slightly misleading since it only captures reasoning-related inner IDs (single-query only). Consider renaming to reasoning_inner_ids.

Positive observations:

  • Comprehensive overflow guards for query_count * k and byte-level allocations
  • Sentinel pre-fill with ids = -1 and dists = +inf is 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_table tracking of -1 labels is correctly maintained across all mutation paths (Insert, Remove, Merge, UpdateLabel, ShrinkToFit, Deserialize)

@LHT129 LHT129 added kind/feature Brand-new functionality or capabilities 引入全新的功能、新特性或新能力 version/1.1 labels Aug 3, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 dim with 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.

Comment thread src/utils/timer.h
Comment thread include/vsag/index.h
Comment thread include/vsag/search_request.h Outdated

@vsag-bot vsag-bot left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated inline review completed.

Review effort: high (1150 changed lines across 19 files).
Submitted 2 inline comments.

Reviewed commit 500b5ca.

Comment thread include/vsag/search_request.h
Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
Comment thread src/algorithm/hgraph/hgraph.h
Comment thread src/algorithm/ivf/ivf.cpp Outdated
Comment thread src/algorithm/hgraph/hgraph_search.cpp

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 enforces bucket_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_ 外层长度与查询数一致,并支持为每个查询提供一个桶列表(包括批量查询)。文档应与实际行为一致。
- 当前仅支持单查询;外层向量必须恰好包含一个条目。

Copilot AI review requested due to automatic review settings August 20, 2026 07:53
@LHT129
LHT129 force-pushed the batch_query_review_reset_1685 branch from 6744bab to c420ed1 Compare August 20, 2026 07:53

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_QUERY treats a null query as eligible for the empty-index shortcut, which makes SearchWithRequest return an empty dataset (instead of INVALID_ARGUMENT) when request.query_ == nullptr and 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.

@LHT129 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review comments for PR #2607 (batch query support).

Comment thread src/algorithm/ivf/ivf.cpp
@@ -2141,12 +2141,22 @@ IVF::SearchWithRequest(const SearchRequest& request) const {
CHECK_ARGUMENT(request.expected_labels_.empty(),

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

test

@LHT129 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review of PR #2607 (multi-query batch search for HGraph and IVF):

Previously reported issues — all confirmed fixed:

  1. MCI route info is now preserved via mci_result.MakeStatistics(*ctx.stats).Dump() (hgraph_search.cpp:562)
  2. enable_reorder and rabitq_one_bit_search are now properly guarded with use_custom_distance ? false : ... (hgraph_search.cpp:494-496)
  3. Dataset::Make() is no longer called inside the per-query loop; result buffers are pre-allocated once

New observations:

  • The k > 0 check is duplicated when use_custom_distance is true (once in the else branch, once unconditionally after). Harmless but slightly redundant.
  • The IVF batch path reuses the existing search_func lambda 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>(

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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));

Comment thread src/algorithm/hgraph/hgraph_search.cpp Outdated
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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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) {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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) {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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() {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 ||

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/algorithm/ivf/ivf.cpp
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),

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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>
Copilot AI review requested due to automatic review settings August 21, 2026 07:29
@LHT129
LHT129 force-pushed the batch_query_review_reset_1685 branch from c420ed1 to 5357338 Compare August 21, 2026 07:29

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 public KnnSearch overloads delegate here and now accept the same multi-query inputs. The overload documentation above still promises num_elements = 1 and a num_elements * k layout, 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. IndexImpl deliberately 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 to reorder.
        } 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 only stats.Dump(). MCI-enabled batch searches therefore omit the mci_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_14 still reads label_table_ directly at line 170 without rebuilding it. A legacy index containing an active external label -1 will therefore bypass the new batch-safety check and can return an ambiguous -1 result; 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 receives base_search_param.time_cost; the searcher only checks InnerSearchParam::time_cost when enforcing timeout_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->SearchWithRequest once per query, and each invocation creates a separate SearchStatistics; the outer stats object 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");

Comment on lines +32 to 35
make_empty_dataset_with_stats() {
SearchStatistics stats;
auto dataset_result = DatasetImpl::MakeEmptyDataset();
dataset_result->Statistics(stats.Dump());
Comment on lines +927 to +933
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);
}
}

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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);

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Comment thread tests/test_ivf.cpp
check_bucket_result(batch_result.value(), 3, scan_buckets_count, buckets_count);
}

SECTION("batch routing ignores search-only options") {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 LHT129 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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:

  1. hgraph_search.cpp:835Dim(0) is returned when element_count == 0 && query_count > 1, while the single-query path returns Dim(element_count). Consider returning Dim(element_count) consistently in both paths.
  2. hgraph_search.cpp:732search_range_with_request does not set mci_result.route = "brute_force" when falling through to brute force, unlike the KNN path.
  3. 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).

Comment thread include/vsag/search_request.h Outdated
/**
* @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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/feature Brand-new functionality or capabilities 引入全新的功能、新特性或新能力 module/api module/docs module/index module/testing size/XXL version/1.1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] Add multi-query batch search support

3 participants