Skip to content

feat(core, indexer): support group_by searching and fix hnsw sparse#527

Open
JalinWang wants to merge 38 commits into
alibaba:mainfrom
JalinWang:feat/group_by
Open

feat(core, indexer): support group_by searching and fix hnsw sparse#527
JalinWang wants to merge 38 commits into
alibaba:mainfrom
JalinWang:feat/group_by

Conversation

@JalinWang

@JalinWang JalinWang commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR implements group_by search (a.k.a. group-by deduplication) across the core interface, algorithm, and DB layers of zvec. Users can now retrieve top-K results per group instead of globally, with support for fetch_vector, is_linear, and bf_pks query modes on supported index types. group_by is mutually exclusive with refiner search and is rejected when refiner_param is set.

Motivation

Prior to this PR, group_by was either silently ignored (IVF, Vamana) or produced incorrect/empty results (DiskAnn, sparse fetch_vector). This PR brings consistent, well-tested group_by behavior to all index types.

Current Status

Unsupported index types now fail fast with IndexError_Unsupported and an error log, instead of silently returning empty or incorrect results.

Support Reject
flat (dense) diskann
flat (sparse) ivf
hnsw (dense) vamana
hnsw (sparse)
hnsw_rabitq

GroupBy search mode compatibility:

Search mode with group_by Status
fetch_vector is_linear bf_pks
refiner search (refiner_param) Rejected

We will gradually enable GroupBy searching for them.

Previous Status

Index Type group_by (graph search) group_by + is_linear group_by + bf_pks fetch_vector in group_by Notes
flat (dense) N/A (always linear) Yes Yes Yes Full support via FlatStreamerContext
flat_sparse N/A (always linear) Yes Yes No Known limitation; search_group/search_group_p_keys don't emit vectors
hnsw (dense) Yes Yes Yes Yes Graph search via HnswAlgorithm::search, bf via HnswStreamer::search_bf_impl and search_bf_by_p_keys_impl
hnsw_sparse Yes Yes Yes Partial (no sparse vector in group result) topk_to_group_result uses get_vector_meta(id) only, not full sparse data
hnsw_rabitq Yes Yes Yes Yes Both graph and bf paths handle group_by; HnswRabitqContext has full group result + fetch
diskann Yes (buggy) Yes (buggy) Yes (buggy) No Empty results — linear_search/keys_search never populated group_topk_heaps_; knn_search had typo at L985 writing to wrong heap;
and the fetch vector doesn’t work in normal search
ivf No No No No IVFSearcherContext has no group_by_search(), no group_topk_heaps. group_by is silently ignored
vamana No No No No VamanaContext has no group support. topk_to_result only builds flat results. group_by is silently ignored

Key Changes

Core Interface Layer

  • index_param.h: Added GroupByParam struct (group_topk, group_count, group_by callback) and group_by_param field on BaseIndexQueryParam.
  • index.h / index.cc: Added group_doc_list_ to SearchResult. Introduced for_each_doc helper to uniformly iterate over flat and grouped results for score normalization and vector reverting. Added _set_group_by_on_context static helper called by supported index types at the end of _prepare_for_search. Added a fast rejection for group_by combined with refiner search.
  • Per-index _prepare_for_search: FlatIndex, HNSWIndex, HNSWRabitqIndex call _set_group_by_on_context. IVFIndex, DiskAnnIndex, and VamanaIndex add early rejection checks for group_by.

Algorithm Layer

  • flat_sparse_search.h: Populates IndexSparseDocument in group_by results when fetch_vector is enabled (previously missing).
  • hnsw_sparse_context.h: Fetches full sparse vector data via get_sparse_data + SparseUtility::ReverseSparseFormat in group results (previously only stored get_vector_meta).

DB Layer

  • engine_helper.hpp: Translates DB-layer vector_column_params::GroupByParams to core-layer GroupByParam when building the engine query param.
  • vector_column_indexer.cc: Returns GroupVectorIndexResults when search_result.group_doc_list_ is populated.
  • combined_vector_column_indexer.cc: Merges group_by results across multiple index blocks — adjusts doc keys by block offset, merges groups by group_id, sorts within groups by score (respecting metric direction), and truncates to group_topk / group_count.

Tests

  • index_group_by_test.cc (new, 520 lines): Data-driven GroupByInterfaceTest fixture at the core interface layer with RunOk/RunRejected methods. Covers dense (flat, hnsw, hnsw_rabitq — graph, linear, bf_pks, fetch_vector), sparse (flat, hnsw — graph, linear, bf_pks, fetch_vector), and unsupported (vamana, ivf, diskann) index types. Also verifies that group_by and refiner search are mutually exclusive.
  • vector_column_indexer_test.cc (+257 lines): Data-driven GroupByIndexerTest fixture at the DB indexer layer, mirroring the core test pattern. Covers dense (flat, hnsw — graph, linear, bf_pks, fetch_vector), sparse (flat, hnsw), and unsupported (ivf, diskann with optional plugin skip).

JalinWang added 15 commits June 18, 2026 16:12
…ssing

Added for_each_doc helper to unify iteration over flat doc_list_ or grouped
group_doc_list_, eliminating 5 of 6 duplicated if(has_group_by) blocks in
_dense_search and _sparse_search. Reduces code duplication and makes the
group_by logic easier to maintain.

The one remaining dual-rail (reformer normalize) is preserved because the
batch API requires different handling for grouped vs flat results.
Move DiskAnn group_by support to feat/group_by_diskann branch.
On this branch, DiskAnn explicitly rejects group_by with
IndexError_Unsupported, matching Vamana and IVF behavior.
@JalinWang JalinWang changed the title Feat/group by feat(core, indexer): support group_by searching Jun 25, 2026
@JalinWang JalinWang marked this pull request as ready for review June 25, 2026 08:08
@JalinWang JalinWang removed the request for review from chinaux June 25, 2026 08:23
@JalinWang JalinWang changed the title feat(core, indexer): support group_by searching feat(core, indexer): support group_by searching and fix hsnw sparse Jun 25, 2026
@JalinWang JalinWang requested review from egolearner and feihongxu0824 and removed request for zhourrr June 25, 2026 08:24
uint32_t group_count =
query_params.group_by ? query_params.group_by->group_count : 0;
if (group_count > 0 && merged_group_docs.size() > group_count) {
merged_group_docs.resize(group_count);

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.

这里截断的依据是什么,好像没看到merged_group_docs有排过序?

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.

确实掉了,已加上;后续单独pr重构此函数

Comment thread src/core/interface/index.cc Outdated
(search_param->group_by_param && search_param->group_by_param->group_by);
if (has_group_by) {
result->group_doc_list_ = std::move(
const_cast<core::IndexGroupDocumentList &>(context->group_result()));

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.

minor:可以考虑给context弄一个mutable_group_result方法,这样就不需要const_cast

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.

done

// Return grouped results when group_by is active
if (!search_result.group_doc_list_.empty()) {
auto result = std::make_shared<GroupVectorIndexResults>(
std::move(search_result.group_doc_list_));

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.

reverted_vector_list_ 和 reverted_sparse_values_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.

done

std::min(256u, vamana_search_param->prefetch_lines);
params.set(core::PARAM_VAMANA_STREAMER_PL, real_search_pl);
context->update(params);
_set_group_by_on_context(search_param, context);

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.

vamana不支持groupby,这里还需要set context吗?

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.

rm

return core::IndexError_Runtime;
}

auto &base_result = context->result();

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.

groupby search的情况下,这个result应该是空的

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.

确实,refiner和group by我先互斥掉吧

@JalinWang JalinWang changed the title feat(core, indexer): support group_by searching and fix hsnw sparse feat(core, indexer): support group_by searching and fix hnsw sparse Jun 26, 2026
const auto &ivf_search_param =
std::dynamic_pointer_cast<IVFQueryParam>(search_param);

if (search_param->group_by_param && search_param->group_by_param->group_by) {

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.

ivf应该也可以支持group by? 对参与距离计算的doc获取group id,然后做处理。

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.

理论上是的,但我打算后面单独pr补这几个reject group by的索引,比如目前diskann的支持已经写好了,等这个合了我就提

Comment thread src/core/interface/index.cc Outdated
namespace {

template <typename Fn>
void for_each_doc(SearchResult *result, bool has_group_by, Fn &&fn) {

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.

把has_group_by放到模板参数?

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.

这样判断逻辑就得加到主代码里面,我可能还是倾向于保持现状?
image

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.

本质上来说,for_each_doc拆成两个函数比较合理,本身也没有共同逻辑

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.

添加group_reverted_vector_list_后这个helper引用点太少了,我直接inline掉了

doc.set_key(block_offsets_[i] + doc.key());
}
// Merge into existing group or create new one
auto it = group_merge_map.find(group.group_id());

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.

可以使用operator [],不存在自动创建

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.

重构后不需要区分no exist状态了,done

// Return merged group_by results if any
if (!merged_group_docs.empty()) {
// Sort docs within each group by score and truncate to group_topk
bool lower_is_better =

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.

抽个公共实现吧,然后原先向量结果合并的逻辑也可以复用

Image

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会让逻辑不太清晰暂时没放:
https://github.com/JalinWang/zvec/compare/feat/group_by...JalinWang:zvec:codex/combined-search-merger-refactor?expand=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.

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.

merge了

reverted_sparse_values.clear();
}

std::vector<size_t> indices(docs.size());

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.

为什么先排序indices再设置,可以直接排序docs然后做截断吗?

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.

reverted_vectors 和 reverted_sparse_values 要和docs保持顺序一致

@JalinWang JalinWang Jul 2, 2026

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.

重构后抽象了一个ResultDoc保存Reverted结果,不再维护三个分开的数组,直接用sort了,done


struct SearchResult {
core::IndexDocumentList doc_list_;
core::IndexGroupDocumentList group_doc_list_;

@egolearner egolearner Jul 1, 2026

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.

这样一一对应好点?另外之前为什么vector/sparse要分开,能合成一个vector<string>

  core::IndexGroupDocumentList group_doc_list_;
  std::vector<std::vector<std::string>> group_reverted_vector_list_;
  std::vector<std::vector<std::string>> group_reverted_sparse_values_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.

之前应该是因为hybrid搜索可能两个都要返回,感觉确实可以合并,我看看

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.

先添加了group_reverted_vector_list_/group_reverted_sparse_values_list_实现对应,后续单拆一个Pr合并Reverted vector/sparse vector list

if (has_group_by) {
for (auto &group : result->group_doc_list_) {
auto *docs = group.mutable_docs();
core::IndexDocumentList doc_list(docs->begin(), docs->end());

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.

这个doc_list需要吗,normalize直接传docs是不是就好了

merged_docs.reserve(merged_docs.size() + docs.size());

for (size_t doc_idx = 0; doc_idx < docs.size(); ++doc_idx) {
auto doc = docs[doc_idx];

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.

这里为何不和VectorResultAccumulator::AddBlock的实现一样,使用std::move

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants