feat(core, indexer): support group_by searching and fix hnsw sparse#527
feat(core, indexer): support group_by searching and fix hnsw sparse#527JalinWang wants to merge 38 commits into
Conversation
…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.
| 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); |
There was a problem hiding this comment.
这里截断的依据是什么,好像没看到merged_group_docs有排过序?
There was a problem hiding this comment.
确实掉了,已加上;后续单独pr重构此函数
| (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())); |
There was a problem hiding this comment.
minor:可以考虑给context弄一个mutable_group_result方法,这样就不需要const_cast
| // 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_)); |
There was a problem hiding this comment.
reverted_vector_list_ 和 reverted_sparse_values_list_ 应该还是需要的吧
| 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); |
There was a problem hiding this comment.
vamana不支持groupby,这里还需要set context吗?
| return core::IndexError_Runtime; | ||
| } | ||
|
|
||
| auto &base_result = context->result(); |
There was a problem hiding this comment.
groupby search的情况下,这个result应该是空的
There was a problem hiding this comment.
确实,refiner和group by我先互斥掉吧
| 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) { |
There was a problem hiding this comment.
ivf应该也可以支持group by? 对参与距离计算的doc获取group id,然后做处理。
There was a problem hiding this comment.
理论上是的,但我打算后面单独pr补这几个reject group by的索引,比如目前diskann的支持已经写好了,等这个合了我就提
| namespace { | ||
|
|
||
| template <typename Fn> | ||
| void for_each_doc(SearchResult *result, bool has_group_by, Fn &&fn) { |
There was a problem hiding this comment.
本质上来说,for_each_doc拆成两个函数比较合理,本身也没有共同逻辑
There was a problem hiding this comment.
添加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()); |
There was a problem hiding this comment.
可以使用operator [],不存在自动创建
There was a problem hiding this comment.
重构后不需要区分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 = |
There was a problem hiding this comment.
这个重构我也单独抽了一版,感觉review会让逻辑不太清晰暂时没放:
https://github.com/JalinWang/zvec/compare/feat/group_by...JalinWang:zvec:codex/combined-search-merger-refactor?expand=1
师兄看看如何觉得合适的话我把它仍过来吧
There was a problem hiding this comment.
| reverted_sparse_values.clear(); | ||
| } | ||
|
|
||
| std::vector<size_t> indices(docs.size()); |
There was a problem hiding this comment.
为什么先排序indices再设置,可以直接排序docs然后做截断吗?
There was a problem hiding this comment.
reverted_vectors 和 reverted_sparse_values 要和docs保持顺序一致
There was a problem hiding this comment.
重构后抽象了一个ResultDoc保存Reverted结果,不再维护三个分开的数组,直接用sort了,done
|
|
||
| struct SearchResult { | ||
| core::IndexDocumentList doc_list_; | ||
| core::IndexGroupDocumentList group_doc_list_; |
There was a problem hiding this comment.
这样一一对应好点?另外之前为什么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_;
There was a problem hiding this comment.
之前应该是因为hybrid搜索可能两个都要返回,感觉确实可以合并,我看看
There was a problem hiding this comment.
先添加了group_reverted_vector_list_/group_reverted_sparse_values_list_实现对应,后续单拆一个Pr合并Reverted vector/sparse vector list
…tor' into feat/group_by
| if (has_group_by) { | ||
| for (auto &group : result->group_doc_list_) { | ||
| auto *docs = group.mutable_docs(); | ||
| core::IndexDocumentList doc_list(docs->begin(), docs->end()); |
There was a problem hiding this comment.
这个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]; |
There was a problem hiding this comment.
这里为何不和VectorResultAccumulator::AddBlock的实现一样,使用std::move


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, andbf_pksquery modes on supported index types.group_byis mutually exclusive with refiner search and is rejected whenrefiner_paramis 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_Unsupportedand an error log, instead of silently returning empty or incorrect results.GroupBy search mode compatibility:
group_byfetch_vectoris_linearbf_pksrefiner_param)We will gradually enable GroupBy searching for them.
Previous Status
and the fetch vector doesn’t work in normal search
Key Changes
Core Interface Layer
index_param.h: AddedGroupByParamstruct (group_topk,group_count,group_bycallback) andgroup_by_paramfield onBaseIndexQueryParam.index.h/index.cc: Addedgroup_doc_list_toSearchResult. Introducedfor_each_dochelper to uniformly iterate over flat and grouped results for score normalization and vector reverting. Added_set_group_by_on_contextstatic helper called by supported index types at the end of_prepare_for_search. Added a fast rejection forgroup_bycombined with refiner search._prepare_for_search:FlatIndex,HNSWIndex,HNSWRabitqIndexcall_set_group_by_on_context.IVFIndex,DiskAnnIndex, andVamanaIndexadd early rejection checks for group_by.Algorithm Layer
flat_sparse_search.h: PopulatesIndexSparseDocumentin group_by results whenfetch_vectoris enabled (previously missing).hnsw_sparse_context.h: Fetches full sparse vector data viaget_sparse_data+SparseUtility::ReverseSparseFormatin group results (previously only storedget_vector_meta).DB Layer
engine_helper.hpp: Translates DB-layervector_column_params::GroupByParamsto core-layerGroupByParamwhen building the engine query param.vector_column_indexer.cc: ReturnsGroupVectorIndexResultswhensearch_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 bygroup_id, sorts within groups by score (respecting metric direction), and truncates togroup_topk/group_count.Tests
index_group_by_test.cc(new, 520 lines): Data-drivenGroupByInterfaceTestfixture at the core interface layer withRunOk/RunRejectedmethods. 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 thatgroup_byand refiner search are mutually exclusive.vector_column_indexer_test.cc(+257 lines): Data-drivenGroupByIndexerTestfixture 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).