xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

This commit is contained in:
Jalin Wang 2026-07-07 17:15:12 +08:00 committed by GitHub
parent c35d24e215
commit 468c565f86
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
28 changed files with 1507 additions and 536 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/core/algorithm/diskann/diskann_context.h
View File

@ -94,6 +94,14 @@ class DiskAnnContext : public IndexContext,
return group_results_[idx];
}
virtual IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
virtual IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}
virtual uint32_t magic(void) const override {
return magic_;
}

405
src/core/algorithm/diskann/diskann_visit_filter.h
View File

@ -13,407 +13,4 @@
// limitations under the License.
#pragma once
#include <algorithm>
#include <chrono>
#include <cstdint>
#include <limits>
#include <random>
#include <tuple>
#include <ailego/container/bloom_filter.h>
#include <ailego/utility/bitset_helper.h>
#include <zvec/core/framework/index_error.h>
#include <zvec/core/framework/index_logger.h>
namespace zvec {
namespace core {
struct VisitFilterHeader {
VisitFilterHeader() : maxDocCnt(0), maxScanNum(0) {}
uint64_t maxDocCnt;
uint64_t maxScanNum;
};
constexpr int PROXIMA_DISKANN_VISITFILTER_CUSTOM_PARAMS_INDEX_NEGPROB = 0;
class VisitBloomFilter {
public:
static constexpr int mode = 1;
static constexpr int N = 5;
struct Context {
Context()
: mt(std::chrono::system_clock::now().time_since_epoch().count()) {};
VisitFilterHeader h;
std::mt19937 mt;
ailego::BloomFilter<N> *filter{nullptr};
int offset[N] = {0};
};
#define BLOOM_FILTER_HASH_BITS_OFFSETS(i) \
i + c->offset[0], i + c->offset[1], i + c->offset[2], i + c->offset[3], \
i + c->offset[4]
VisitBloomFilter() = delete;
inline static void set_visited(Context *c, id_t idx) {
c->filter->force_insert(BLOOM_FILTER_HASH_BITS_OFFSETS(idx));
return;
}
inline static void *get_visited(Context *, id_t) {
// TODO
return nullptr;
}
inline static bool visited(Context *c, id_t idx) {
return c->filter->has(BLOOM_FILTER_HASH_BITS_OFFSETS(idx));
}
inline static int set_max_scan_num(Context *c, uint64_t maxScanNum) {
if (maxScanNum == c->h.maxScanNum) {
return 0;
}
c->h.maxScanNum = maxScanNum;
if (c->filter->reset(maxScanNum, c->filter->probability()) != 0) {
LOG_ERROR("reset BloomFilter failed");
return IndexError_Runtime;
}
genRandomHashBits(c);
return 0;
}
inline static void clear(Context *c) {
c->filter->clear();
return;
}
inline static bool reset(Context *c, uint64_t maxDocCnt,
uint64_t max_scan_num) {
if (ailego_unlikely(maxDocCnt > c->h.maxDocCnt ||
max_scan_num > c->h.maxScanNum)) {
// Create a new one, if failed, we can reuse the old one
auto filter = new (std::nothrow) ailego::BloomFilter<VisitBloomFilter::N>(
max_scan_num, c->filter->probability());
if (ailego_unlikely(filter == nullptr)) {
LOG_ERROR("reset bloomfilter failed, maxScanNum %zu prob %f",
(size_t)max_scan_num, c->filter->probability());
c->filter->clear();
return false;
}
delete c->filter;
c->filter = filter;
c->h.maxScanNum = max_scan_num;
c->h.maxDocCnt = maxDocCnt;
genRandomHashBits(c);
}
return true;
}
inline static void genRandomHashBits(Context *c) {
std::uniform_int_distribution<int> dt(0, c->h.maxDocCnt);
for (size_t i = 0; i < sizeof(c->offset) / sizeof(c->offset[0]); ++i) {
int r = dt(c->mt);
size_t j = 0;
do { // gen distinct number
for (j = 0; j < i; ++j) {
if (c->offset[j] == r) {
r = dt(c->mt);
break;
}
}
} while (j < i);
c->offset[i] = r;
}
std::sort(c->offset, c->offset + N);
}
template <class... T>
static int init(Context *, void **ctx, uint64_t maxDocCnt,
uint64_t maxScanNum, std::tuple<T...> &&tpl) {
Context *c = new (std::nothrow) Context;
if (c == nullptr) {
LOG_ERROR("New memory in initVisitBitMap failed");
return IndexError_NoMemory;
}
c->h.maxDocCnt = maxDocCnt;
c->h.maxScanNum = maxScanNum;
float p =
std::get<PROXIMA_DISKANN_VISITFILTER_CUSTOM_PARAMS_INDEX_NEGPROB>(tpl);
c->filter = new (std::nothrow)
ailego::BloomFilter<VisitBloomFilter::N>(maxScanNum, p);
if (c->filter == nullptr) {
LOG_ERROR("New BloomFilter failed, reuse old one");
return IndexError_NoMemory;
}
genRandomHashBits(c);
*ctx = c;
return 0;
}
inline static void destroy(Context *c) {
delete c->filter;
delete c;
}
#undef BLOOM_FILTER_HASH_BITS_OFFSETS
}; // end of VisitBloomFilter
class VisitBitMap {
public:
static constexpr int mode = 2;
struct Context {
VisitFilterHeader h;
ailego::BitsetHelper bitset;
char *buf{nullptr};
};
VisitBitMap() = delete;
inline static void set_visited(Context *c, id_t idx) {
c->bitset.set(idx);
return;
}
inline static void *get_visited(Context *c, id_t idx) {
return &c->buf[idx >> 3];
}
inline static bool visited(Context *c, id_t idx) {
return c->bitset.test(idx);
}
inline static int set_max_scan_num(Context *c, uint64_t maxScanNum) {
c->h.maxScanNum = maxScanNum;
return 0;
}
inline static void clear(Context *c) {
c->bitset.clear();
return;
}
inline static bool reset(Context *c, uint64_t maxDocCnt,
uint64_t maxScanNum) {
if (ailego_unlikely(maxDocCnt > c->h.maxDocCnt ||
maxScanNum > c->h.maxScanNum)) {
uint64_t len = ((maxDocCnt + 31) >> 5) << 2; // round to uint32_t
auto buf = new (std::nothrow) char[len];
if (buf == nullptr) {
LOG_ERROR("New memory in initVisitBitMap failed");
c->bitset.clear();
return false;
}
c->h.maxDocCnt = maxDocCnt;
c->h.maxScanNum = maxScanNum;
delete[] c->buf;
c->buf = buf;
memset(c->buf, 0, len);
c->bitset.mount(c->buf, len);
}
return true;
}
template <class... T>
static int init(Context *, void **ctx, uint64_t maxDocCnt,
uint64_t maxScanNum, std::tuple<T...> &&tpl) {
(void)tpl; // unsed warning
Context *c = new (std::nothrow) Context;
if (c == nullptr) {
LOG_ERROR("New memory in initVisitBitMap failed");
return IndexError_NoMemory;
}
c->h.maxDocCnt = maxDocCnt;
c->h.maxScanNum = maxScanNum;
uint64_t len = ((maxDocCnt + 31) >> 5) << 2; // round to uint32_t
c->buf = new (std::nothrow) char[len];
if (c->buf == nullptr) {
LOG_ERROR("New memory in initVisitBitMap failed, reuse old one");
delete c;
return IndexError_NoMemory;
}
memset(c->buf, 0, len);
c->bitset.mount(c->buf, len);
*ctx = c;
return 0;
}
inline static void destroy(Context *c) {
delete[] c->buf;
delete c;
}
}; // end of VisitBitMap
class VisitByteMap {
public:
static constexpr int mode = 3;
struct Context {
VisitFilterHeader h;
uint8_t curNum{0};
uint8_t *arr{nullptr};
};
VisitByteMap() = delete;
inline static void set_visited(Context *c, id_t idx) {
c->arr[idx] = c->curNum;
return;
}
inline static void *get_visited(Context *c, id_t idx) {
return c->arr + idx;
}
inline static bool visited(Context *c, id_t idx) {
return c->arr[idx] == c->curNum;
}
inline static int set_max_scan_num(Context *c, uint64_t maxScanNum) {
c->h.maxScanNum = maxScanNum;
return 0;
}
inline static void clear(Context *c) {
c->curNum++;
if (c->curNum == 0) {
memset(c->arr, 0, c->h.maxDocCnt * sizeof(uint8_t));
c->curNum = 1;
}
return;
}
inline static bool reset(Context *c, uint64_t maxDocCnt,
uint64_t maxScanNum) {
if (ailego_unlikely(maxDocCnt > c->h.maxDocCnt ||
maxScanNum > c->h.maxScanNum)) {
auto arr = new (std::nothrow) uint8_t[maxDocCnt];
if (arr != nullptr) {
memset(arr, 0, maxDocCnt * sizeof(uint8_t));
c->curNum = 1;
c->h.maxDocCnt = maxDocCnt;
c->h.maxScanNum = maxScanNum;
delete[] c->arr;
c->arr = arr;
return true;
}
LOG_ERROR("New memory in initVisitByteMap failed, reuse old one");
}
return true;
}
template <class... T>
static int init(Context *, void **ctx, uint64_t maxDocCnt,
uint64_t maxScanNum, std::tuple<T...> &&tpl) {
(void)tpl; // unsed warning
Context *c = new (std::nothrow) Context;
if (c == nullptr) {
LOG_ERROR("New memory in initVisitByteMap failed");
return IndexError_NoMemory;
}
c->h.maxDocCnt = maxDocCnt;
c->h.maxScanNum = maxScanNum;
c->arr = new (std::nothrow) uint8_t[maxDocCnt];
if (c->arr == nullptr) {
LOG_ERROR("New memory in initVisitByteMap failed");
delete c;
return IndexError_NoMemory;
}
memset(c->arr, 0, maxDocCnt * sizeof(uint8_t));
c->curNum = 1;
*ctx = c;
return 0;
}
inline static void destroy(Context *c) {
delete[] c->arr;
delete c;
}
}; // end of VisitByteMap
#define PROXIMA_DISKANN_VISITFILTER_SWITCH_CASE(cls, impl, ctx, ...) \
case cls::mode: \
return cls::impl(static_cast<cls::Context *>(ctx), ##__VA_ARGS__);
#define PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(impl, ...) \
switch (mode_) { \
PROXIMA_DISKANN_VISITFILTER_SWITCH_CASE(VisitBloomFilter, impl, ctx_, \
##__VA_ARGS__) \
PROXIMA_DISKANN_VISITFILTER_SWITCH_CASE(VisitBitMap, impl, ctx_, \
##__VA_ARGS__) \
PROXIMA_DISKANN_VISITFILTER_SWITCH_CASE(VisitByteMap, impl, ctx_, \
##__VA_ARGS__) \
}
// visit list will be called with high frequency,
// so using switch instead of std::function or virtual class
// funtion point, lambda, virtual class all cannot be inlined
class VisitFilter {
public:
enum Mode {
Default = 0,
BloomFilter = VisitBloomFilter::mode,
BitMap = VisitBitMap::mode,
ByteMap = VisitByteMap::mode
};
VisitFilter() : mode_(0), ctx_(nullptr) {};
inline bool visited(id_t idx) {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(visited, idx);
return true; // place holder
}
inline void set_visited(id_t idx) {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(set_visited, idx);
}
inline void *get_visited(id_t idx) {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(get_visited, idx);
return nullptr; // place holder
}
inline int set_max_scan_num(id_t idx) {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(set_max_scan_num, idx);
return 0; // place holder
}
inline void clear() {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(clear);
}
inline bool reset(uint64_t maxDocCnt, uint64_t maxScanNum) {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(reset, maxDocCnt, maxScanNum);
return true;
}
inline void destroy() {
if (ctx_ != nullptr) {
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(destroy);
}
}
int init(int mode, uint64_t maxDocCnt, uint64_t maxScanNum,
float negativeProbility) {
mode_ = mode;
PROXIMA_DISKANN_VISITFILTER_CALL_IMPL(init, &ctx_, maxDocCnt, maxScanNum,
std::make_tuple(negativeProbility));
return 0; // place holder
}
int get_mode(void) const {
return mode_;
}
private:
VisitFilter(const VisitFilter &) = delete;
VisitFilter &operator=(const VisitFilter &) = delete;
int mode_{0U}; // custom data for each method
void *ctx_{nullptr};
};
} // namespace core
} // namespace zvec
#include "utility/visit_filter.h"

8
src/core/algorithm/flat/flat_searcher_context.h
View File

@ -67,6 +67,14 @@ class FlatSearcherContext : public IndexSearcher::Context {
return group_results_[idx];
}
IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}
//! Update the parameters of context
int update(const ailego::Params & /*params*/) override {
return 0;

8
src/core/algorithm/flat/flat_streamer_context.h
View File

@ -67,6 +67,14 @@ class FlatStreamerContext : public IndexStreamer::Context {
return group_results_[idx];
}
IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}
//! Update the parameters of context
int update(const ailego::Params & /*params*/) override {
return 0;

6
src/core/algorithm/flat_sparse/flat_sparse_context.h
View File

@ -101,7 +101,11 @@ class FlatSparseContext : public IndexContext {
return group_results_[idx];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) {
IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}

23
src/core/algorithm/flat_sparse/flat_sparse_search.h
View File

@ -113,6 +113,29 @@ static inline int FlatSearch(const uint32_t *sparse_count,
auto group_result =
ConvertGroupMapToResult(std::move(group_heap), ctx->group_num());
// Populate sparse vector data when fetch_vector is enabled
if (ctx->fetch_vector()) {
for (auto &group_doc : group_result) {
for (auto &doc : *group_doc.mutable_docs()) {
node_id_t id = entity->get_id(doc.key());
if (id != kInvalidNodeId) {
IndexSparseDocument sparse_doc;
IndexStorage::MemoryBlock vec_block;
entity->get_sparse_vector(id, vec_block);
const void *sparse_data = vec_block.data();
if (sparse_data != nullptr) {
SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc,
entity->sparse_unit_size());
}
// Reconstruct doc with sparse vector data
doc = IndexDocument(doc.key(), doc.score(), id, nullptr,
sparse_doc);
}
}
}
}
ctx->mutable_group_result(q)->swap(group_result);
}
} else {

10
src/core/algorithm/hnsw/hnsw_context.h
View File

@ -80,6 +80,14 @@ class HnswContext : public IndexContext {
return group_results_[idx];
}
IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}
uint32_t magic(void) const override {
return magic_;
}
@ -585,4 +593,4 @@ class HnswContext : public IndexContext {
};
} // namespace core
} // namespace zvec
} // namespace zvec

8
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h
View File

@ -79,6 +79,14 @@ class HnswRabitqContext : public IndexContext {
return group_results_[idx];
}
IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}
uint32_t magic(void) const override {
return magic_;
}

19
src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h
View File

@ -73,6 +73,14 @@ class HnswSparseContext : public IndexContext {
return group_results_[idx];
}
IndexGroupDocumentList *mutable_group_result(void) override {
return &group_results_[0];
}
IndexGroupDocumentList *mutable_group_result(size_t idx) override {
return &group_results_[idx];
}
uint32_t magic(void) const override {
return magic_;
}
@ -241,8 +249,17 @@ class HnswSparseContext : public IndexContext {
node_id_t id = group_topk_list[i].second[j].first;
if (fetch_vector_) {
IndexSparseDocument sparse_doc;
IndexStorage::MemoryBlock vec_block;
entity_->get_sparse_data(id, vec_block);
const void *sparse_data = vec_block.data();
if (sparse_data != nullptr) {
SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc,
entity_->sparse_unit_size());
}
group_results_[idx][i].mutable_docs()->emplace_back(
entity_->get_key(id), score, id, entity_->get_vector_meta(id));
entity_->get_key(id), score, id, entity_->get_vector_meta(id),
sparse_doc);
} else {
group_results_[idx][i].mutable_docs()->emplace_back(
entity_->get_key(id), score, id);

183
src/core/interface/index.cc
View File

@ -21,6 +21,14 @@
namespace zvec::core_interface {
namespace {
bool has_group_by_search(const BaseIndexQueryParam::Pointer &search_param) {
return search_param->group_by_param && search_param->group_by_param->group_by;
}
} // namespace
// eliminate the pre-alloc of the context pool
thread_local static std::array<core::IndexContext::Pointer,
(magic_enum::enum_count<IndexType>() - 1) * 2>
@ -496,6 +504,17 @@ int Index::Search(const VectorData &vector_data,
return core::IndexError_Runtime;
}
const bool has_group_by = has_group_by_search(search_param);
if (has_group_by && is_group_by_unsupported_index(param_.index_type)) {
LOG_ERROR("group_by search is not supported for this index type");
return core::IndexError_Unsupported;
}
if (search_param->refiner_param != nullptr && has_group_by) {
LOG_ERROR("group_by search is not supported with refiner");
return core::IndexError_Unsupported;
}
if (!is_trained_ && this->Train() != 0) {
LOG_ERROR("Failed to train index");
return core::IndexError_Runtime;
@ -519,7 +538,7 @@ int Index::Search(const VectorData &vector_data,
return ret;
}
// dense support refiner, but sparse doesn't
// dense supports refiner, but sparse doesn't
int ret = 0;
if (search_param->refiner_param == nullptr) {
ret = _dense_search(vector_data, search_param, result, context);
@ -561,8 +580,8 @@ int Index::Search(const VectorData &vector_data,
flat_search_param->bf_pks = std::make_shared<std::vector<uint64_t>>(keys);
ret = reference_index->Search(vector_data, flat_search_param, result);
context->reset();
}
context->reset();
return ret;
}
@ -720,7 +739,6 @@ int Index::_dense_search(const VectorData &vector_data,
}
vector = new_vector.data();
}
// TODO: group by
if (search_param->bf_pks != nullptr) {
// should we eliminate the copy of bf_pks?
if (streamer_->search_bf_by_p_keys_impl(
@ -740,32 +758,84 @@ int Index::_dense_search(const VectorData &vector_data,
return core::IndexError_Runtime;
}
}
result->doc_list_ = std::move(context->result());
// Retrieve group_by results if applicable
bool has_group_by =
(search_param->group_by_param && search_param->group_by_param->group_by);
if (has_group_by) {
auto *group_result = context->mutable_group_result();
if (group_result == nullptr) {
LOG_ERROR("Failed to retrieve group_by result");
return core::IndexError_Runtime;
}
result->group_doc_list_ = std::move(*group_result);
} else {
result->doc_list_ = std::move(context->result());
}
if (metric_->support_normalize()) {
for (uint32_t i = 0; i < result->doc_list_.size(); ++i) {
metric_->normalize(result->doc_list_[i].mutable_score());
if (has_group_by) {
for (auto &group : result->group_doc_list_) {
for (auto &doc : *group.mutable_docs()) {
metric_->normalize(doc.mutable_score());
}
}
} else {
for (auto &doc : result->doc_list_) {
metric_->normalize(doc.mutable_score());
}
}
}
if (reformer_) {
if (reformer_->normalize(dense_vector.data, input_vector_meta_,
result->doc_list_) != 0) {
LOG_ERROR("Failed to normalize vector");
return core::IndexError_Runtime;
}
if (context->fetch_vector() && reformer_->need_revert()) {
// TODO: use std::pmr to optimize memory allocation
result->reverted_vector_list_.resize(context->result().size());
for (uint32_t i = 0; i < context->result().size(); ++i) {
std::string &reverted_vector = result->reverted_vector_list_[i];
reverted_vector.resize(input_vector_meta_.dimension() *
input_vector_meta_.unit_size());
if (reformer_->revert(context->result()[i].vector(), new_meta,
&reverted_vector) != 0) {
LOG_ERROR("Failed to revert vector");
if (has_group_by) {
for (auto &group : result->group_doc_list_) {
auto *docs = group.mutable_docs();
if (reformer_->normalize(dense_vector.data, input_vector_meta_,
*docs) != 0) {
LOG_ERROR("Failed to normalize vector");
return core::IndexError_Runtime;
}
}
} else {
if (reformer_->normalize(dense_vector.data, input_vector_meta_,
result->doc_list_) != 0) {
LOG_ERROR("Failed to normalize vector");
return core::IndexError_Runtime;
}
}
if (context->fetch_vector() && reformer_->need_revert()) {
int revert_err = 0;
auto revert_one = [&](const void *vec, std::vector<std::string> *out) {
if (revert_err) return;
std::string reverted_vector;
reverted_vector.resize(input_vector_meta_.dimension() *
input_vector_meta_.unit_size());
if (reformer_->revert(vec, new_meta, &reverted_vector) != 0) {
LOG_ERROR("Failed to revert vector");
revert_err = core::IndexError_Runtime;
return;
}
out->push_back(std::move(reverted_vector));
};
auto revert_docs = [&](auto &docs, std::vector<std::string> &out) {
out.reserve(docs.size());
for (auto &doc : docs) {
revert_one(doc.vector(), &out);
}
};
if (has_group_by) {
result->group_reverted_vector_list_.reserve(
result->group_doc_list_.size());
for (auto &group : result->group_doc_list_) {
std::vector<std::string> group_vectors;
revert_docs(*group.mutable_docs(), group_vectors);
result->group_reverted_vector_list_.push_back(
std::move(group_vectors));
}
} else {
revert_docs(result->doc_list_, result->reverted_vector_list_);
}
if (revert_err) return revert_err;
}
}
@ -819,23 +889,41 @@ int Index::_sparse_search(const VectorData &vector_data,
return core::IndexError_Runtime;
}
}
result->doc_list_ = std::move(context->result());
// Retrieve group_by results if applicable
const bool has_group_by = has_group_by_search(search_param);
if (has_group_by) {
auto *group_result = context->mutable_group_result();
if (group_result == nullptr) {
LOG_ERROR("Failed to retrieve group_by result");
return core::IndexError_Runtime;
}
result->group_doc_list_ = std::move(*group_result);
} else {
result->doc_list_ = std::move(context->result());
}
if (metric_->support_normalize()) {
for (uint32_t i = 0; i < result->doc_list_.size(); ++i) {
metric_->normalize(result->doc_list_[i].mutable_score());
if (has_group_by) {
for (auto &group : result->group_doc_list_) {
for (auto &doc : *group.mutable_docs()) {
metric_->normalize(doc.mutable_score());
}
}
} else {
for (auto &doc : result->doc_list_) {
metric_->normalize(doc.mutable_score());
}
}
}
if (reformer_) {
// TODO: no need to call reformer_->normalize() when sparse?
if (context->fetch_vector() && reformer_->need_revert()) {
// TODO: use std::pmr to optimize memory allocation
auto &result_doc_list = context->result();
result->reverted_sparse_values_list_.resize(result_doc_list.size());
for (uint32_t i = 0; i < result_doc_list.size(); ++i) {
auto &result_doc = result_doc_list[i].sparse_doc();
std::string &reverted_sparse_values =
result->reverted_sparse_values_list_[i];
int revert_err = 0;
auto revert_one = [&](const core::IndexDocument &doc,
std::vector<std::string> *out) {
if (revert_err) return;
auto &result_doc = doc.sparse_doc();
std::string reverted_sparse_values;
reverted_sparse_values.resize(result_doc.sparse_count() *
input_vector_meta_.unit_size());
if (reformer_->revert(result_doc.sparse_count(),
@ -845,9 +933,30 @@ int Index::_sparse_search(const VectorData &vector_data,
result_doc.sparse_values().data()),
new_meta, &reverted_sparse_values) != 0) {
LOG_ERROR("Failed to revert sparse vector");
return core::IndexError_Runtime;
revert_err = core::IndexError_Runtime;
return;
}
out->push_back(std::move(reverted_sparse_values));
};
auto revert_docs = [&](auto &docs, std::vector<std::string> &out) {
out.reserve(docs.size());
for (auto &doc : docs) {
revert_one(doc, &out);
}
};
if (has_group_by) {
result->group_reverted_sparse_values_list_.reserve(
result->group_doc_list_.size());
for (auto &group : result->group_doc_list_) {
std::vector<std::string> group_sparse_values;
revert_docs(*group.mutable_docs(), group_sparse_values);
result->group_reverted_sparse_values_list_.push_back(
std::move(group_sparse_values));
}
} else {
revert_docs(result->doc_list_, result->reverted_sparse_values_list_);
}
if (revert_err) return revert_err;
}
}
return 0;
@ -921,6 +1030,16 @@ int Index::_get_coarse_search_topk(
return floor(search_param->topk * scale_factor);
}
void Index::_set_group_by_on_context(
const BaseIndexQueryParam::Pointer &search_param,
core::IndexContext::Pointer &context) {
if (search_param->group_by_param && search_param->group_by_param->group_by) {
context->set_group_by(search_param->group_by_param->group_by);
context->set_group_params(search_param->group_by_param->group_count,
search_param->group_by_param->group_topk);
}
}
std::string Index::get_metric_name(MetricType metric_type, bool is_sparse) {
if (is_sparse) {
switch (metric_type) {

5
src/core/interface/indexes/diskann_index.cc
View File

@ -269,6 +269,11 @@ int DiskAnnIndex::_prepare_for_search(
return core::IndexError_Runtime;
}
if (search_param->group_by_param && search_param->group_by_param->group_by) {
LOG_ERROR("group_by search is not supported for DiskAnn index");
return core::IndexError_Unsupported;
}
context->set_topk(diskann_search_param->topk);
// Propagate the query-time beam-search list size into the context. Must be

1
src/core/interface/indexes/flat_index.cc
View File

@ -63,6 +63,7 @@ int FlatIndex::_prepare_for_search(
if (flat_search_param->radius > 0.0f) {
context->set_threshold(flat_search_param->radius);
}
_set_group_by_on_context(search_param, context);
return 0;
}

1
src/core/interface/indexes/hnsw_index.cc
View File

@ -168,6 +168,7 @@ int HNSWIndex::_prepare_for_search(
std::min(256u, hnsw_search_param->prefetch_lines);
params.set(core::PARAM_HNSW_STREAMER_PL, real_search_pl);
context->update(params);
_set_group_by_on_context(search_param, context);
return 0;
}

1
src/core/interface/indexes/hnsw_rabitq_index.cc
View File

@ -121,6 +121,7 @@ int HNSWRabitqIndex::_prepare_for_search(
std::max(1u, std::min(2048u, hnsw_search_param->ef_search));
params.set(core::PARAM_HNSW_RABITQ_STREAMER_EF, real_search_ef);
context->update(params);
_set_group_by_on_context(search_param, context);
return 0;
#endif // RABITQ_SUPPORTED
}

5
src/core/interface/indexes/ivf_index.cc
View File

@ -214,6 +214,11 @@ int IVFIndex::_prepare_for_search(
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) {
LOG_ERROR("group_by search is not supported for IVF index");
return core::IndexError_Unsupported;
}
context->set_topk(ivf_search_param->topk);
context->set_fetch_vector(ivf_search_param->fetch_vector);
if (ivf_search_param->filter) {

5
src/core/interface/indexes/vamana_index.cc
View File

@ -72,6 +72,11 @@ int VamanaIndex::_prepare_for_search(
return core::IndexError_Runtime;
}
if (search_param->group_by_param && search_param->group_by_param->group_by) {
LOG_ERROR("group_by search is not supported for Vamana index");
return core::IndexError_Unsupported;
}
if (vamana_search_param->ef_search == 0 ||
vamana_search_param->ef_search > 2048) {
LOG_ERROR(

6
src/core/utility/mmap_file_read_storage.cc
View File

@ -237,7 +237,9 @@ class MMapFileReadStorage : public IndexStorage {
}
int close(void) override {
file_ptr_->close();
if (file_ptr_) {
file_ptr_->close();
}
file_ptr_ = nullptr;
segments_.clear();
return 0;
@ -294,4 +296,4 @@ class MMapFileReadStorage : public IndexStorage {
INDEX_FACTORY_REGISTER_STORAGE(MMapFileReadStorage);
} // namespace core
} // namespace zvec
} // namespace zvec

354
src/db/index/column/vector_column/combined_vector_column_indexer.cc
View File

@ -15,9 +15,244 @@
#include <algorithm>
#include <cstdint>
#include <numeric>
#include <unordered_map>
namespace zvec {
namespace {
bool IsBetterScore(MetricType metric_type, float lhs, float rhs) {
switch (metric_type) {
case MetricType::IP:
return lhs > rhs;
case MetricType::L2:
case MetricType::COSINE:
default:
return lhs < rhs;
}
}
bool HasRevertedValues(const std::vector<std::string> &values) {
return std::any_of(values.begin(), values.end(),
[](const auto &value) { return !value.empty(); });
}
bool HasRevertedValues(const std::vector<std::vector<std::string>> &values) {
return std::any_of(values.begin(), values.end(), [](const auto &group) {
return HasRevertedValues(group);
});
}
struct ResultDoc {
core::IndexDocument doc;
// Keep fetched/reverted payloads attached to the doc while sorting and
// truncating, so parallel result vectors cannot drift out of sync.
std::string reverted_vector;
std::string reverted_sparse_values;
};
class VectorResultAccumulator {
public:
// Collect plain topk results from each block after translating block-local
// doc IDs back to segment-level IDs.
void AddBlock(uint32_t block_offset, VectorIndexResults *results) {
auto &docs = results->docs();
auto &reverted_vectors = results->reverted_vector_list();
auto &reverted_sparse_values = results->reverted_sparse_values_list();
docs_.reserve(docs_.size() + docs.size());
for (size_t i = 0; i < docs.size(); ++i) {
auto doc = std::move(docs[i]);
doc.set_key(block_offset + doc.key());
ResultDoc result_doc{std::move(doc), {}, {}};
if (i < reverted_vectors.size()) {
result_doc.reverted_vector = std::move(reverted_vectors[i]);
}
if (i < reverted_sparse_values.size()) {
result_doc.reverted_sparse_values =
std::move(reverted_sparse_values[i]);
}
docs_.emplace_back(std::move(result_doc));
}
}
IndexResults::Ptr Finish(bool is_sparse, MetricType metric_type,
uint32_t topk) {
// Finish turns accumulated block docs into the public result format:
// rank all docs globally, keep topk, then split ResultDoc back into the
// doc list and optional reverted payload lists expected by
// VectorIndexResults.
std::sort(docs_.begin(), docs_.end(),
[metric_type](const ResultDoc &lhs, const ResultDoc &rhs) {
return IsBetterScore(metric_type, lhs.doc.score(),
rhs.doc.score());
});
if (docs_.size() > topk) {
docs_.resize(topk);
}
core::IndexDocumentList doc_list;
std::vector<std::string> reverted_vector_list;
std::vector<std::string> reverted_sparse_values_list;
doc_list.reserve(docs_.size());
reverted_vector_list.reserve(docs_.size());
reverted_sparse_values_list.reserve(docs_.size());
for (auto &doc : docs_) {
doc_list.emplace_back(std::move(doc.doc));
reverted_vector_list.emplace_back(std::move(doc.reverted_vector));
reverted_sparse_values_list.emplace_back(
std::move(doc.reverted_sparse_values));
}
if (!HasRevertedValues(reverted_vector_list)) {
reverted_vector_list.clear();
}
if (!HasRevertedValues(reverted_sparse_values_list)) {
reverted_sparse_values_list.clear();
}
return std::make_unique<VectorIndexResults>(
is_sparse, std::move(doc_list), std::move(reverted_vector_list),
std::move(reverted_sparse_values_list));
}
private:
std::vector<ResultDoc> docs_;
};
class GroupResultAccumulator {
private:
struct GroupResult {
std::string group_id;
std::vector<ResultDoc> docs;
};
public:
// Merge same-named groups across blocks. The per-doc payload stays inside
// ResultDoc until the final GroupVectorIndexResults is materialized.
void AddBlock(uint32_t block_offset, GroupVectorIndexResults *results) {
auto &groups = results->groups();
auto &reverted_vectors = results->reverted_vector_list();
auto &reverted_sparse_values = results->reverted_sparse_values_list();
for (size_t group_idx = 0; group_idx < groups.size(); ++group_idx) {
auto &group = groups[group_idx];
auto *docs = group.mutable_docs();
auto &merged_docs = docs_by_group_[group.group_id()];
merged_docs.reserve(merged_docs.size() + docs->size());
for (size_t doc_idx = 0; doc_idx < docs->size(); ++doc_idx) {
auto doc = std::move((*docs)[doc_idx]);
doc.set_key(block_offset + doc.key());
ResultDoc result_doc{std::move(doc), {}, {}};
if (group_idx < reverted_vectors.size() &&
doc_idx < reverted_vectors[group_idx].size()) {
result_doc.reverted_vector =
std::move(reverted_vectors[group_idx][doc_idx]);
}
if (group_idx < reverted_sparse_values.size() &&
doc_idx < reverted_sparse_values[group_idx].size()) {
result_doc.reverted_sparse_values =
std::move(reverted_sparse_values[group_idx][doc_idx]);
}
merged_docs.emplace_back(std::move(result_doc));
}
}
}
bool empty() const {
return docs_by_group_.empty();
}
IndexResults::Ptr Finish(MetricType metric_type, uint32_t group_topk,
uint32_t group_count) {
// Finish first ranks docs inside each merged group and trims group_topk.
// It then ranks groups by their best remaining doc, trims group_count, and
// finally expands ResultDoc back into GroupVectorIndexResults payloads.
std::vector<GroupResult> groups;
groups.reserve(docs_by_group_.size());
for (auto &[group_id, docs] : docs_by_group_) {
if (docs.empty()) {
continue;
}
std::sort(docs.begin(), docs.end(),
[metric_type](const ResultDoc &lhs, const ResultDoc &rhs) {
return IsBetterScore(metric_type, lhs.doc.score(),
rhs.doc.score());
});
if (group_topk > 0 && docs.size() > group_topk) {
docs.resize(group_topk);
}
groups.emplace_back(GroupResult{group_id, std::move(docs)});
}
std::sort(groups.begin(), groups.end(),
[metric_type](const GroupResult &lhs, const GroupResult &rhs) {
if (lhs.docs.empty() || rhs.docs.empty()) {
return !lhs.docs.empty() && rhs.docs.empty();
}
const float lhs_score = lhs.docs[0].doc.score();
const float rhs_score = rhs.docs[0].doc.score();
if (lhs_score == rhs_score) {
return lhs.group_id < rhs.group_id;
}
return IsBetterScore(metric_type, lhs_score, rhs_score);
});
if (group_count > 0 && groups.size() > group_count) {
groups.resize(group_count);
}
core::IndexGroupDocumentList group_list;
std::vector<std::vector<std::string>> reverted_vector_list;
std::vector<std::vector<std::string>> reverted_sparse_values_list;
group_list.reserve(groups.size());
reverted_vector_list.reserve(groups.size());
reverted_sparse_values_list.reserve(groups.size());
for (auto &group : groups) {
core::GroupIndexDocument group_doc;
group_doc.set_group_id(group.group_id);
auto *docs = group_doc.mutable_docs();
docs->reserve(group.docs.size());
std::vector<std::string> group_reverted_vectors;
std::vector<std::string> group_reverted_sparse_values;
group_reverted_vectors.reserve(group.docs.size());
group_reverted_sparse_values.reserve(group.docs.size());
for (auto &doc : group.docs) {
docs->emplace_back(std::move(doc.doc));
group_reverted_vectors.emplace_back(std::move(doc.reverted_vector));
group_reverted_sparse_values.emplace_back(
std::move(doc.reverted_sparse_values));
}
group_list.emplace_back(std::move(group_doc));
reverted_vector_list.emplace_back(std::move(group_reverted_vectors));
reverted_sparse_values_list.emplace_back(
std::move(group_reverted_sparse_values));
}
if (!HasRevertedValues(reverted_vector_list)) {
reverted_vector_list.clear();
}
if (!HasRevertedValues(reverted_sparse_values_list)) {
reverted_sparse_values_list.clear();
}
return std::make_unique<GroupVectorIndexResults>(
std::move(group_list), std::move(reverted_vector_list),
std::move(reverted_sparse_values_list));
}
private:
std::unordered_map<std::string, std::vector<ResultDoc>> docs_by_group_;
};
} // namespace
CombinedVectorColumnIndexer::CombinedVectorColumnIndexer(
const std::vector<VectorColumnIndexer::Ptr> &indexers,
const std::vector<VectorColumnIndexer::Ptr> &normal_indexers,
@ -54,9 +289,12 @@ CombinedVectorColumnIndexer::CombinedVectorColumnIndexer(
Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
const vector_column_params::VectorData &vector_data,
const vector_column_params::QueryParams &query_params) {
core::IndexDocumentList doc_list;
std::vector<std::string> reverted_vector_list;
std::vector<std::string> reverted_sparse_values_list;
// Search runs each block with block-local query params, then folds those
// partial results into one segment-level result. The accumulators keep doc
// IDs and fetched/reverted payloads aligned while final sorting and
// truncation are deferred until every block has been searched.
VectorResultAccumulator vector_results;
GroupResultAccumulator group_results;
// query_params.bf_pks is segment level, here we need to convert it to block
// level
@ -105,6 +343,9 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
need_refine = true;
}
// Rewrite segment-level query state to the current block: filters and
// group_by callbacks see segment IDs, while the underlying block indexer
// searches with block-local doc IDs.
const IndexFilter *filter{nullptr};
auto per_block_filter =
BlockOffsetFilter{query_params.filter, block_offsets_[i]};
@ -153,101 +394,30 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
}
auto index_results = result.value();
GroupVectorIndexResults *group_index_results =
dynamic_cast<GroupVectorIndexResults *>(index_results.get());
if (group_index_results != nullptr) {
group_results.AddBlock(block_offsets_[i], group_index_results);
continue;
}
VectorIndexResults *vector_index_results =
dynamic_cast<VectorIndexResults *>(index_results.get());
const auto &sub_docs = vector_index_results->docs();
for (size_t j = 0; j < sub_docs.size(); ++j) {
auto doc = sub_docs[j];
doc.set_key(block_offsets_[i] + sub_docs[j].key());
doc_list.emplace_back(std::move(doc));
if (vector_index_results != nullptr) {
vector_results.AddBlock(block_offsets_[i], vector_index_results);
}
auto &&temp_vector_list = vector_index_results->reverted_vector_list();
reverted_vector_list.insert(
reverted_vector_list.end(),
std::make_move_iterator(temp_vector_list.begin()),
std::make_move_iterator(temp_vector_list.end()));
auto &&temp_sparse_list =
vector_index_results->reverted_sparse_values_list();
reverted_sparse_values_list.insert(
reverted_sparse_values_list.end(),
std::make_move_iterator(temp_sparse_list.begin()),
std::make_move_iterator(temp_sparse_list.end()));
}
if (doc_list.empty()) {
// return empty result
return std::make_unique<VectorIndexResults>(
field_schema_.is_sparse_vector(), std::move(doc_list),
std::move(reverted_vector_list),
std::move(reverted_sparse_values_list));
if (!group_results.empty()) {
const uint32_t group_topk =
query_params.group_by ? query_params.group_by->group_topk : 0;
const uint32_t group_count =
query_params.group_by ? query_params.group_by->group_count : 0;
return group_results.Finish(metric_type_, group_topk, group_count);
}
std::vector<size_t> indices(doc_list.size());
std::iota(indices.begin(), indices.end(), 0);
std::sort(indices.begin(), indices.end(),
[this, &doc_list](size_t lhs, size_t rhs) {
const auto &lhs_doc = doc_list[lhs];
const auto &rhs_doc = doc_list[rhs];
if (this->metric_type_ == MetricType::L2) {
return lhs_doc.score() < rhs_doc.score();
} else if (this->metric_type_ == MetricType::IP) {
return lhs_doc.score() > rhs_doc.score();
} else if (this->metric_type_ == MetricType::COSINE) {
return lhs_doc.score() < rhs_doc.score();
} else {
// default
return lhs_doc.score() < rhs_doc.score();
}
});
// doc_list
std::vector<core::IndexDocument> sorted_doc_list(doc_list.size());
for (size_t i = 0; i < indices.size(); ++i) {
sorted_doc_list[i] = std::move(doc_list[indices[i]]);
}
doc_list = std::move(sorted_doc_list);
// reverted_vector_list
if (!reverted_vector_list.empty()) {
std::vector<std::string> sorted_reverted_vector_list(
reverted_vector_list.size());
for (size_t i = 0; i < indices.size(); ++i) {
if (indices[i] < reverted_vector_list.size()) {
sorted_reverted_vector_list[i] =
std::move(reverted_vector_list[indices[i]]);
}
}
reverted_vector_list = std::move(sorted_reverted_vector_list);
}
// reverted_sparse_values_list
if (!reverted_sparse_values_list.empty()) {
std::vector<std::string> sorted_reverted_sparse_vector_list(
reverted_sparse_values_list.size());
for (size_t i = 0; i < indices.size(); ++i) {
if (indices[i] < reverted_sparse_values_list.size()) {
sorted_reverted_sparse_vector_list[i] =
std::move(reverted_sparse_values_list[indices[i]]);
}
}
reverted_sparse_values_list = std::move(sorted_reverted_sparse_vector_list);
}
// truncate to topk
if (doc_list.size() > query_params.topk) doc_list.resize(query_params.topk);
if (reverted_vector_list.size() > query_params.topk)
reverted_vector_list.resize(query_params.topk);
if (reverted_sparse_values_list.size() > query_params.topk)
reverted_sparse_values_list.resize(query_params.topk);
return std::make_unique<VectorIndexResults>(
field_schema_.is_sparse_vector(), std::move(doc_list),
std::move(reverted_vector_list), std::move(reverted_sparse_values_list));
return vector_results.Finish(field_schema_.is_sparse_vector(), metric_type_,
query_params.topk);
}
Result<vector_column_params::VectorDataBuffer>

10
src/db/index/column/vector_column/engine_helper.hpp
View File

@ -120,6 +120,16 @@ class ProximaEngineHelper {
std::make_shared<core_interface::RefinerParam>(rp);
}
}
if (db_query_params.group_by) {
engine_query_param->group_by_param =
std::make_shared<core_interface::GroupByParam>();
engine_query_param->group_by_param->group_topk =
db_query_params.group_by->group_topk;
engine_query_param->group_by_param->group_count =
db_query_params.group_by->group_count;
engine_query_param->group_by_param->group_by =
db_query_params.group_by->group_by;
}
return engine_query_param;
}

9
src/db/index/column/vector_column/vector_column_indexer.cc
View File

@ -198,6 +198,15 @@ Result<IndexResults::Ptr> VectorColumnIndexer::Search(
Status::InternalError("Failed to search vector"));
}
// 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_),
std::move(search_result.group_reverted_vector_list_),
std::move(search_result.group_reverted_sparse_values_list_));
return result;
}
auto result = std::make_shared<VectorIndexResults>(
is_sparse_, std::move(search_result.doc_list_),
std::move(search_result.reverted_vector_list_),

12
src/db/index/column/vector_column/vector_index_results.h
View File

@ -241,7 +241,9 @@ class GroupVectorIndexResults : public IndexResults {
std::vector<std::vector<std::string>> &&reverted_sparse_values_list)
: groups_(std::move(group_list)),
reverted_vector_list_(std::move(reverted_vector_list)),
reverted_sparse_values_list_(std::move(reverted_sparse_values_list)) {}
reverted_sparse_values_list_(std::move(reverted_sparse_values_list)) {
init_count();
}
public:
IndexResults::IteratorUPtr create_iterator() override {
@ -258,6 +260,14 @@ class GroupVectorIndexResults : public IndexResults {
return groups_;
}
std::vector<std::vector<std::string>> &reverted_vector_list() {
return reverted_vector_list_;
}
std::vector<std::vector<std::string>> &reverted_sparse_values_list() {
return reverted_sparse_values_list_;
}
private:
const core::IndexDocument &document(size_t group_index,
size_t doc_index) const {

12
src/include/zvec/core/framework/index_context.h
View File

@ -161,6 +161,16 @@ class IndexContext {
return this->group_result();
}
//! Retrieve mutable search group result
virtual IndexGroupDocumentList *mutable_group_result(void) {
return nullptr;
}
//! Retrieve mutable search group result with index
virtual IndexGroupDocumentList *mutable_group_result(size_t /*idx*/) {
return this->mutable_group_result();
}
//! Update the parameters of context
virtual int update(const ailego::Params & /*params*/) {
return IndexError_NotImplemented;
@ -262,4 +272,4 @@ class IndexContext {
};
} // namespace core
} // namespace zvec
} // namespace zvec

15
src/include/zvec/core/interface/index.h
View File

@ -95,9 +95,13 @@ struct VectorDataBuffer {
struct SearchResult {
core::IndexDocumentList doc_list_;
core::IndexGroupDocumentList group_doc_list_;
// use string to manage memory
std::vector<std::string> reverted_vector_list_{};
std::vector<std::string> reverted_sparse_values_list_{};
// Grouped reverted values, aligned with group_doc_list_.
std::vector<std::vector<std::string>> group_reverted_vector_list_{};
std::vector<std::vector<std::string>> group_reverted_sparse_values_list_{};
};
class Index {
@ -172,6 +176,11 @@ class Index {
static std::string get_metric_name(MetricType metric_type, bool is_sparse);
static bool is_group_by_unsupported_index(IndexType index_type) {
return index_type == IndexType::kIVF || index_type == IndexType::kDiskAnn ||
index_type == IndexType::kVamana;
}
protected:
int _sparse_fetch(const uint32_t doc_id,
VectorDataBuffer *vector_data_buffer);
@ -195,6 +204,12 @@ class Index {
virtual int _get_coarse_search_topk(
const BaseIndexQueryParam::Pointer &search_param);
//! Helper: set group_by on context from the query param (common for all
//! index types). Call this at the end of _prepare_for_search.
static void _set_group_by_on_context(
const BaseIndexQueryParam::Pointer &search_param,
core::IndexContext::Pointer &context);
protected:
friend class IndexFactory;
Index() = default;

9
src/include/zvec/core/interface/index_param.h
View File

@ -15,6 +15,7 @@
#pragma once
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <vector>
@ -166,6 +167,13 @@ struct RefinerParam {
std::shared_ptr<Index> reference_index = nullptr;
};
// --- GroupBy Parameters ---
struct GroupByParam {
uint32_t group_topk{0};
uint32_t group_count{0};
std::function<std::string(uint64_t key)> group_by{};
};
// --- Query Parameters (can be passed to search methods) ---
class BaseIndexQueryParam {
public:
@ -180,6 +188,7 @@ class BaseIndexQueryParam {
float radius = 0.0f;
bool is_linear = false;
RefinerParam::Pointer refiner_param = nullptr;
std::shared_ptr<GroupByParam> group_by_param = nullptr;
virtual Pointer Clone() const = 0;
};

539
tests/core/interface/index_group_by_test.cc Normal file
View File

@ -0,0 +1,539 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <numeric>
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "tests/test_util.h"
#if RABITQ_SUPPORTED
#include "core/algorithm/hnsw_rabitq/rabitq_converter.h"
#include "zvec/core/framework/index_provider.h"
#endif
#include "zvec/core/interface/index.h"
#include "zvec/core/interface/index_factory.h"
#include "zvec/core/interface/index_param.h"
#include "zvec/core/interface/index_param_builders.h"
using namespace zvec::core_interface;
namespace {
constexpr uint32_t kDimension = 4;
constexpr uint32_t kNumDocs = 12;
constexpr uint32_t kNumGroups = 3;
constexpr uint32_t kGroupTopk = 2;
constexpr uint32_t kSearchTopk = 100;
struct GroupByCase {
std::string name;
BaseIndexParam::Pointer index_param;
BaseIndexQueryParam::Pointer query_param;
bool is_sparse = false;
uint32_t dimension = kDimension;
bool with_refiner = false;
};
std::shared_ptr<std::vector<uint64_t>> AllPks() {
auto pks = std::make_shared<std::vector<uint64_t>>();
pks->reserve(kNumDocs);
for (uint32_t i = 0; i < kNumDocs; ++i) {
pks->push_back(i);
}
return pks;
}
void AttachGroupBy(const BaseIndexQueryParam::Pointer &query_param) {
query_param->group_by_param = std::make_shared<GroupByParam>();
query_param->group_by_param->group_count = kNumGroups;
query_param->group_by_param->group_topk = kGroupTopk;
query_param->group_by_param->group_by = [](uint64_t key) {
return std::to_string(key % kNumGroups);
};
}
BaseIndexParam::Pointer DenseFlatParam(uint32_t dimension = kDimension) {
return FlatIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(dimension)
.WithIsSparse(false)
.Build();
}
BaseIndexParam::Pointer SparseFlatParam() {
return FlatIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithIsSparse(true)
.Build();
}
BaseIndexParam::Pointer DenseHnswParam(uint32_t dimension = kDimension) {
return HNSWIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(dimension)
.WithIsSparse(false)
.WithEFConstruction(100)
.Build();
}
BaseIndexParam::Pointer SparseHnswParam() {
return HNSWIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithIsSparse(true)
.WithEFConstruction(100)
.Build();
}
BaseIndexQueryParam::Pointer FlatQuery(bool fetch_vector = false) {
return FlatQueryParamBuilder()
.with_topk(kSearchTopk)
.with_fetch_vector(fetch_vector)
.build();
}
BaseIndexQueryParam::Pointer FlatQuery(bool fetch_vector, bool is_linear,
bool with_bf_pks) {
auto builder = FlatQueryParamBuilder()
.with_topk(kSearchTopk)
.with_fetch_vector(fetch_vector)
.with_is_linear(is_linear);
if (with_bf_pks) {
builder.with_bf_pks(AllPks());
}
return builder.build();
}
BaseIndexQueryParam::Pointer HnswQuery(bool fetch_vector = false,
bool is_linear = false,
bool with_bf_pks = false) {
auto builder = HNSWQueryParamBuilder()
.with_topk(kSearchTopk)
.with_ef_search(kSearchTopk)
.with_fetch_vector(fetch_vector)
.with_is_linear(is_linear);
if (with_bf_pks) {
builder.with_bf_pks(AllPks());
}
return builder.build();
}
#if RABITQ_SUPPORTED
BaseIndexParam::Pointer DenseHnswRabitqParam(uint32_t dimension) {
using namespace zvec::ailego;
using namespace zvec::core;
constexpr size_t kTrainCount = 500;
auto holder =
std::make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(
dimension);
for (size_t i = 0; i < kTrainCount; ++i) {
NumericalVector<float> vec(dimension, static_cast<float>(i));
EXPECT_TRUE(holder->emplace(i, vec));
}
auto index_meta =
std::make_shared<IndexMeta>(IndexMeta::DataType::DT_FP32, dimension);
index_meta->set_metric("InnerProduct", 0, Params());
RabitqConverter converter;
EXPECT_EQ(0, converter.init(*index_meta, Params()));
EXPECT_EQ(0, converter.train(holder));
std::shared_ptr<IndexReformer> reformer;
EXPECT_EQ(0, converter.to_reformer(&reformer));
return HNSWRabitqIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(dimension)
.WithIsSparse(false)
.WithEFConstruction(100)
.WithProvider(holder)
.WithReformer(reformer)
.Build();
}
BaseIndexQueryParam::Pointer HnswRabitqQuery(bool fetch_vector = false,
bool is_linear = false,
bool with_bf_pks = false) {
auto builder = HNSWRabitqQueryParamBuilder()
.with_topk(kSearchTopk)
.with_ef_search(kSearchTopk)
.with_fetch_vector(fetch_vector)
.with_is_linear(is_linear);
if (with_bf_pks) {
builder.with_bf_pks(AllPks());
}
return builder.build();
}
#endif
#if DISKANN_SUPPORTED
BaseIndexParam::Pointer DenseDiskAnnParam(uint32_t dimension = kDimension) {
return DiskAnnIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(dimension)
.WithIsSparse(false)
.WithMaxDegree(32)
.WithListSize(kSearchTopk)
.WithPqChunkNum(0)
.Build();
}
BaseIndexQueryParam::Pointer DiskAnnQuery(bool fetch_vector = false,
bool is_linear = false,
bool with_bf_pks = false) {
auto query = std::make_shared<DiskAnnQueryParam>();
query->topk = kSearchTopk;
query->list_size = kSearchTopk;
query->fetch_vector = fetch_vector;
query->is_linear = is_linear;
if (with_bf_pks) {
query->bf_pks = AllPks();
}
return query;
}
#endif
class GroupByInterfaceTest : public ::testing::Test {
protected:
void RunOk(const GroupByCase &test_case) {
Run(test_case, /*expect_error=*/false);
}
void RunRejected(const GroupByCase &test_case) {
Run(test_case, /*expect_error=*/true);
}
private:
struct QueryHolder {
std::vector<float> values;
std::vector<uint32_t> indices;
VectorData data;
};
void Run(const GroupByCase &test_case, bool expect_error) {
const std::string index_name = "test_groupby_" + test_case.name;
const std::string source_index_name = index_name + "_source";
zvec::test_util::RemoveTestFiles(index_name + "*");
zvec::test_util::RemoveTestFiles(source_index_name + "*");
auto source = IndexFactory::CreateAndInitIndex(*FlatSourceParam(test_case));
ASSERT_NE(nullptr, source) << test_case.name;
ASSERT_EQ(0, source->Open(source_index_name,
{StorageOptions::StorageType::kMMAP, true}))
<< test_case.name;
for (uint32_t i = 0; i < kNumDocs; ++i) {
AddDoc(source, i, test_case);
}
ASSERT_EQ(0, source->Train()) << test_case.name;
auto index = IndexFactory::CreateAndInitIndex(*test_case.index_param);
ASSERT_NE(nullptr, index) << test_case.name;
ASSERT_EQ(
0, index->Open(index_name, {StorageOptions::StorageType::kMMAP, true}))
<< test_case.name;
ASSERT_EQ(0, index->Merge({source}, IndexFilter())) << test_case.name;
auto query_param = test_case.query_param->Clone();
AttachGroupBy(query_param);
if (test_case.with_refiner) {
query_param->refiner_param = std::make_shared<RefinerParam>();
query_param->refiner_param->scale_factor_ = 1.0f;
query_param->refiner_param->reference_index = source;
}
auto query = MakeQuery(test_case);
SearchResult result;
const int ret = index->Search(query.data, query_param, &result);
if (expect_error) {
ASSERT_NE(0, ret) << test_case.name;
} else {
ASSERT_EQ(0, ret) << test_case.name;
AssertGroupedResult(result, query_param, test_case);
}
ASSERT_EQ(0, index->Close()) << test_case.name;
ASSERT_EQ(0, source->Close()) << test_case.name;
zvec::test_util::RemoveTestFiles(index_name + "*");
zvec::test_util::RemoveTestFiles(source_index_name + "*");
}
BaseIndexParam::Pointer FlatSourceParam(const GroupByCase &test_case) {
if (test_case.is_sparse) {
return SparseFlatParam();
}
return DenseFlatParam(test_case.dimension);
}
void AddDoc(const Index::Pointer &index, uint32_t key,
const GroupByCase &test_case) {
std::vector<float> values(test_case.dimension, static_cast<float>(key));
if (test_case.is_sparse) {
std::vector<uint32_t> indices(test_case.dimension);
std::iota(indices.begin(), indices.end(), 0u);
VectorData data{
SparseVector{test_case.dimension, indices.data(), values.data()}};
ASSERT_EQ(0, index->Add(data, key)) << key;
return;
}
VectorData data{DenseVector{values.data()}};
ASSERT_EQ(0, index->Add(data, key)) << key;
}
QueryHolder MakeQuery(const GroupByCase &test_case) {
QueryHolder holder;
holder.values.assign(test_case.dimension, 1.0f);
if (test_case.is_sparse) {
holder.indices.resize(test_case.dimension);
std::iota(holder.indices.begin(), holder.indices.end(), 0u);
holder.data = VectorData{SparseVector{
test_case.dimension, holder.indices.data(), holder.values.data()}};
} else {
holder.data = VectorData{DenseVector{holder.values.data()}};
}
return holder;
}
void AssertGroupedResult(const SearchResult &result,
const BaseIndexQueryParam::Pointer &query_param,
const GroupByCase &test_case) {
ASSERT_TRUE(result.doc_list_.empty());
ASSERT_EQ(kNumGroups, result.group_doc_list_.size());
std::set<std::string> group_ids;
for (const auto &group : result.group_doc_list_) {
group_ids.insert(group.group_id());
ASSERT_LE(group.docs().size(), kGroupTopk);
ASSERT_GE(group.docs().size(), 1u);
const uint32_t expected_mod = std::stoul(group.group_id());
for (const auto &doc : group.docs()) {
ASSERT_EQ(expected_mod, doc.key() % kNumGroups);
}
for (size_t i = 1; i < group.docs().size(); ++i) {
ASSERT_GE(group.docs()[i - 1].score(), group.docs()[i].score());
}
}
for (uint32_t group = 0; group < kNumGroups; ++group) {
ASSERT_TRUE(group_ids.count(std::to_string(group)) > 0);
}
if (!query_param->fetch_vector) {
return;
}
if (test_case.is_sparse) {
AssertSparseVectorsFetched(result, test_case.dimension);
} else {
AssertDenseVectorsFetched(result, test_case.dimension, test_case.name);
}
}
void AssertDenseVectorsFetched(const SearchResult &result, uint32_t dimension,
const std::string &case_name = "") {
const bool has_reverted = !result.group_reverted_vector_list_.empty();
if (has_reverted) {
ASSERT_EQ(result.group_doc_list_.size(),
result.group_reverted_vector_list_.size());
}
for (size_t group_idx = 0; group_idx < result.group_doc_list_.size();
++group_idx) {
const auto &group = result.group_doc_list_[group_idx];
const std::vector<std::string> *group_vectors = nullptr;
if (has_reverted) {
group_vectors = &result.group_reverted_vector_list_[group_idx];
ASSERT_EQ(group.docs().size(), group_vectors->size());
}
for (size_t doc_idx = 0; doc_idx < group.docs().size(); ++doc_idx) {
const auto &doc = group.docs()[doc_idx];
const float expected = static_cast<float>(doc.key());
const float *vector = nullptr;
if (has_reverted) {
vector =
reinterpret_cast<const float *>((*group_vectors)[doc_idx].data());
} else if (doc.vector() != nullptr) {
vector = reinterpret_cast<const float *>(doc.vector());
} else {
// DiskAnn stores fetched vectors in vector_string_ rather than
// the raw pointer field.
ASSERT_FALSE(doc.vector_string().empty())
<< case_name << " key=" << doc.key();
vector = reinterpret_cast<const float *>(doc.vector_string().data());
}
for (uint32_t i = 0; i < dimension; ++i) {
ASSERT_FLOAT_EQ(expected, vector[i])
<< case_name << " key=" << doc.key() << " i=" << i;
}
}
}
}
void AssertSparseVectorsFetched(const SearchResult &result,
uint32_t dimension) {
const bool has_reverted =
!result.group_reverted_sparse_values_list_.empty();
if (has_reverted) {
ASSERT_EQ(result.group_doc_list_.size(),
result.group_reverted_sparse_values_list_.size());
}
for (size_t group_idx = 0; group_idx < result.group_doc_list_.size();
++group_idx) {
const auto &group = result.group_doc_list_[group_idx];
const std::vector<std::string> *group_sparse_values = nullptr;
if (has_reverted) {
group_sparse_values =
&result.group_reverted_sparse_values_list_[group_idx];
ASSERT_EQ(group.docs().size(), group_sparse_values->size());
}
for (size_t doc_idx = 0; doc_idx < group.docs().size(); ++doc_idx) {
const auto &doc = group.docs()[doc_idx];
const auto &sparse = doc.sparse_doc();
ASSERT_EQ(dimension, sparse.sparse_count());
const auto *indices =
reinterpret_cast<const uint32_t *>(sparse.sparse_indices().data());
const float *values = nullptr;
if (has_reverted) {
values = reinterpret_cast<const float *>(
(*group_sparse_values)[doc_idx].data());
} else {
values =
reinterpret_cast<const float *>(sparse.sparse_values().data());
}
const float expected = static_cast<float>(doc.key());
for (uint32_t i = 0; i < dimension; ++i) {
ASSERT_EQ(i, indices[i]);
ASSERT_FLOAT_EQ(expected, values[i]);
}
}
}
}
};
} // namespace
TEST_F(GroupByInterfaceTest, Dense) {
std::vector<GroupByCase> cases{
{"dense_flat_graph", DenseFlatParam(), FlatQuery()},
{"dense_flat_linear", DenseFlatParam(),
FlatQuery(/*fetch_vector=*/false, /*is_linear=*/true,
/*with_bf_pks=*/false)},
{"dense_flat_bf_pks", DenseFlatParam(),
FlatQuery(/*fetch_vector=*/false, /*is_linear=*/false,
/*with_bf_pks=*/true)},
{"dense_flat_fetch_vector", DenseFlatParam(),
FlatQuery(/*fetch_vector=*/true, /*is_linear=*/false,
/*with_bf_pks=*/false)},
{"dense_hnsw_graph", DenseHnswParam(), HnswQuery()},
{"dense_hnsw_linear", DenseHnswParam(),
HnswQuery(/*fetch_vector=*/false, /*is_linear=*/true)},
{"dense_hnsw_bf_pks", DenseHnswParam(),
HnswQuery(/*fetch_vector=*/false, /*is_linear=*/false,
/*with_bf_pks=*/true)},
{"dense_hnsw_fetch_vector", DenseHnswParam(),
HnswQuery(/*fetch_vector=*/true)},
#if RABITQ_SUPPORTED
{"dense_hnsw_rabitq_graph", DenseHnswRabitqParam(64), HnswRabitqQuery(),
/*is_sparse=*/false, /*dimension=*/64},
{"dense_hnsw_rabitq_linear", DenseHnswRabitqParam(64),
HnswRabitqQuery(/*fetch_vector=*/false, /*is_linear=*/true),
/*is_sparse=*/false, /*dimension=*/64},
{"dense_hnsw_rabitq_bf_pks", DenseHnswRabitqParam(64),
HnswRabitqQuery(/*fetch_vector=*/false, /*is_linear=*/false,
/*with_bf_pks=*/true),
/*is_sparse=*/false, /*dimension=*/64},
// Note: fetch_vector is not supported for RabitQ because the entity
// stores quantized binary data (not original float vectors), and
// RabitqReformer does not implement revert().
#endif
};
for (const auto &test_case : cases) {
RunOk(test_case);
}
}
TEST_F(GroupByInterfaceTest, Sparse) {
std::vector<GroupByCase> cases{
{"sparse_flat_graph", SparseFlatParam(), FlatQuery(),
/*is_sparse=*/true},
{"sparse_hnsw_graph", SparseHnswParam(), HnswQuery(),
/*is_sparse=*/true},
{"sparse_hnsw_linear", SparseHnswParam(),
HnswQuery(/*fetch_vector=*/false, /*is_linear=*/true),
/*is_sparse=*/true},
{"sparse_hnsw_bf_pks", SparseHnswParam(),
HnswQuery(/*fetch_vector=*/false, /*is_linear=*/false,
/*with_bf_pks=*/true),
/*is_sparse=*/true},
{"sparse_hnsw_fetch_vector", SparseHnswParam(),
HnswQuery(/*fetch_vector=*/true), /*is_sparse=*/true},
};
for (const auto &test_case : cases) {
RunOk(test_case);
}
}
TEST_F(GroupByInterfaceTest, UnsupportedIndexTypes) {
std::vector<GroupByCase> cases{
{"unsupported_vamana",
VamanaIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(kDimension)
.WithIsSparse(false)
.WithMaxDegree(32)
.WithSearchListSize(100)
.WithAlpha(1.2f)
.Build(),
VamanaQueryParamBuilder()
.with_topk(kSearchTopk)
.with_ef_search(kSearchTopk)
.build()},
{"unsupported_ivf",
IVFIndexParamBuilder()
.WithMetricType(MetricType::kInnerProduct)
.WithDataType(DataType::DT_FP32)
.WithDimension(kDimension)
.WithIsSparse(false)
.WithNList(4)
.Build(),
IVFQueryParamBuilder().with_topk(kSearchTopk).build()},
{"unsupported_refiner", DenseHnswParam(), HnswQuery(),
/*is_sparse=*/false,
/*dimension=*/kDimension,
/*with_refiner=*/true},
#if DISKANN_SUPPORTED
{"unsupported_diskann_graph", DenseDiskAnnParam(), DiskAnnQuery()},
{"unsupported_diskann_linear", DenseDiskAnnParam(),
DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/true)},
{"unsupported_diskann_bf_pks", DenseDiskAnnParam(),
DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/false,
/*with_bf_pks=*/true)},
{"unsupported_diskann_fetch_vector", DenseDiskAnnParam(),
DiskAnnQuery(/*fetch_vector=*/true)},
#endif
};
for (const auto &test_case : cases) {
RunRejected(test_case);
}
}

1
tests/db/index/CMakeLists.txt
View File

@ -39,6 +39,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS})
LIBS zvec
zvec_ailego
zvec_proto
core_knn_diskann
core_metric_static
core_utility_static
core_quantizer_static

378
tests/db/index/column/vector_column_group_by_test.cc Normal file
View File

@ -0,0 +1,378 @@
// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstdint>
#include <numeric>
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/vector_column/combined_vector_column_indexer.h"
#include "db/index/column/vector_column/vector_column_indexer.h"
#include "db/index/column/vector_column/vector_column_params.h"
#include "tests/test_util.h"
#include "zvec/db/index_params.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
namespace {
constexpr uint32_t kGbDimension = 4;
constexpr uint32_t kGbNumDocs = 12;
constexpr uint32_t kGbNumGroups = 3;
constexpr uint32_t kGbGroupTopk = 2;
constexpr uint32_t kGbSearchTopk = 100;
constexpr uint32_t kGbSparseCount = 5;
struct GroupByCase {
std::string name;
IndexParams::Ptr index_params;
QueryParams::Ptr query_params; // core-layer query params; nullptr for sparse
bool is_sparse = false;
uint32_t dimension = kGbDimension;
bool optional = false; // skip when plugin unavailable (e.g. DiskAnn)
bool with_bf_pks = false;
bool fetch_vector = false;
};
std::unique_ptr<vector_column_params::GroupByParams> MakeGroupByParams(
uint32_t group_count = kGbNumGroups) {
return std::make_unique<vector_column_params::GroupByParams>(
kGbGroupTopk, group_count, [](uint64_t key) -> std::string {
return std::to_string(key % kGbNumGroups);
});
}
std::unique_ptr<vector_column_params::GroupByParams> MakeSegmentGroupByParams(
uint32_t group_topk, uint32_t group_count) {
return std::make_unique<vector_column_params::GroupByParams>(
group_topk, group_count,
[](uint64_t key) -> std::string { return key < 2 ? "low" : "high"; });
}
std::vector<uint64_t> AllPks() {
std::vector<uint64_t> pks(kGbNumDocs);
std::iota(pks.begin(), pks.end(), 0ull);
return pks;
}
class GroupByIndexerTest : public ::testing::Test {
protected:
void RunOk(const GroupByCase &tc) {
Run(tc, /*expect_error=*/false);
}
void RunRejected(const GroupByCase &tc) {
Run(tc, /*expect_error=*/true);
}
private:
struct QueryHolder {
std::vector<float> dense;
std::vector<uint32_t> sparse_indices;
std::vector<float> sparse_values;
vector_column_params::VectorData data;
};
void Run(const GroupByCase &tc, bool expect_error) {
const std::string path = "test_groupby_" + tc.name + ".index";
zvec::test_util::RemoveTestFiles(path);
auto indexer = OpenIndexer(tc, path);
if (indexer == nullptr) {
zvec::test_util::RemoveTestFiles(path);
return; // optional plugin unavailable
}
InsertDocs(indexer, tc);
QueryHolder holder = MakeQuery(tc);
vector_column_params::QueryParams qp = MakeQueryParams(tc);
auto results = indexer->Search(holder.data, qp);
if (expect_error) {
ASSERT_FALSE(results.has_value())
<< "group_by should be rejected for " << tc.name;
} else {
ASSERT_TRUE(results.has_value()) << tc.name;
AssertGroupedResult(results.value().get(), tc);
}
indexer->Close();
zvec::test_util::RemoveTestFiles(path);
}
static FieldSchema MakeSchema(const GroupByCase &tc) {
if (tc.is_sparse) {
return FieldSchema("test", DataType::SPARSE_VECTOR_FP32, false,
tc.index_params);
}
return FieldSchema("test", DataType::VECTOR_FP32, tc.dimension, false,
tc.index_params);
}
static VectorColumnIndexer::Ptr OpenIndexer(const GroupByCase &tc,
const std::string &path) {
auto indexer = std::make_shared<VectorColumnIndexer>(path, MakeSchema(tc));
if (!indexer->Open(vector_column_params::ReadOptions{true, true}).ok()) {
return nullptr;
}
return indexer;
}
static void InsertDocs(const VectorColumnIndexer::Ptr &indexer,
const GroupByCase &tc) {
for (uint32_t i = 0; i < kGbNumDocs; ++i) {
if (tc.is_sparse) {
std::vector<uint32_t> indices(kGbSparseCount);
std::vector<float> values(kGbSparseCount);
for (uint32_t j = 0; j < kGbSparseCount; ++j) {
indices[j] = i * kGbSparseCount + j;
values[j] = static_cast<float>(i + 1);
}
vector_column_params::SparseVector sv{kGbSparseCount, indices.data(),
values.data()};
ASSERT_TRUE(
indexer->Insert(vector_column_params::VectorData{sv}, i).ok());
} else {
std::vector<float> vec(tc.dimension, static_cast<float>(i));
vector_column_params::DenseVector dv{vec.data()};
ASSERT_TRUE(
indexer->Insert(vector_column_params::VectorData{dv}, i).ok());
}
}
}
static QueryHolder MakeQuery(const GroupByCase &tc) {
QueryHolder h;
if (tc.is_sparse) {
h.sparse_indices.resize(kGbSparseCount);
h.sparse_values.assign(kGbSparseCount, 1.0f);
std::iota(h.sparse_indices.begin(), h.sparse_indices.end(), 0u);
h.data =
vector_column_params::VectorData{vector_column_params::SparseVector{
kGbSparseCount, h.sparse_indices.data(), h.sparse_values.data()}};
} else {
h.dense.assign(tc.dimension, 1.0f);
h.data = vector_column_params::VectorData{
vector_column_params::DenseVector{h.dense.data()}};
}
return h;
}
static vector_column_params::QueryParams MakeQueryParams(
const GroupByCase &tc) {
vector_column_params::QueryParams qp;
qp.topk = kGbSearchTopk;
qp.filter = nullptr;
qp.fetch_vector = tc.fetch_vector;
qp.query_params = tc.query_params;
if (tc.with_bf_pks) {
qp.bf_pks = {AllPks()};
}
qp.group_by = MakeGroupByParams();
return qp;
}
static void AssertGroupedResult(IndexResults *results,
const GroupByCase &tc) {
auto *group_results = dynamic_cast<GroupVectorIndexResults *>(results);
ASSERT_TRUE(group_results)
<< "Expected GroupVectorIndexResults for " << tc.name;
ASSERT_EQ(kGbNumGroups, group_results->groups().size()) << tc.name;
std::set<std::string> group_ids;
for (const auto &group : group_results->groups()) {
group_ids.insert(group.group_id());
ASSERT_LE(group.docs().size(), kGbGroupTopk) << tc.name;
ASSERT_GE(group.docs().size(), 1u) << tc.name;
const uint32_t expected_mod = std::stoul(group.group_id());
for (const auto &doc : group.docs()) {
ASSERT_EQ(expected_mod, doc.key() % kGbNumGroups)
<< tc.name << " doc " << doc.key();
}
for (size_t j = 1; j < group.docs().size(); ++j) {
ASSERT_GE(group.docs()[j - 1].score(), group.docs()[j].score())
<< tc.name << " group " << group.group_id();
}
}
for (uint32_t g = 0; g < kGbNumGroups; ++g) {
ASSERT_TRUE(group_ids.count(std::to_string(g)) > 0)
<< tc.name << " missing group " << g;
}
auto iter = group_results->create_iterator();
size_t total = 0;
while (iter->valid()) {
if (tc.fetch_vector && !tc.is_sparse) {
const auto vector_data = iter->vector();
const auto &dense_vector =
std::get<vector_column_params::DenseVector>(vector_data.vector);
const float *vector =
reinterpret_cast<const float *>(dense_vector.data);
const float expected = static_cast<float>(iter->doc_id());
for (uint32_t i = 0; i < tc.dimension; ++i) {
ASSERT_FLOAT_EQ(expected, vector[i])
<< tc.name << " doc " << iter->doc_id() << " i " << i;
}
}
total++;
iter->next();
}
ASSERT_EQ(group_results->count(), total) << tc.name;
}
};
} // namespace
TEST_F(GroupByIndexerTest, Dense) {
auto hnsw_linear_qp = std::make_shared<HnswQueryParams>(300);
hnsw_linear_qp->set_is_linear(true);
std::vector<GroupByCase> cases{
{"dense_flat_graph", std::make_shared<FlatIndexParams>(MetricType::IP),
std::make_shared<QueryParams>(IndexType::FLAT)},
{"dense_hnsw_graph",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
std::make_shared<HnswQueryParams>(300)},
{"dense_hnsw_linear",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
hnsw_linear_qp},
{"dense_hnsw_bf_pks",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
std::make_shared<HnswQueryParams>(300),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/false, /*with_bf_pks=*/true},
{"dense_hnsw_fetch_vector",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
std::make_shared<HnswQueryParams>(300),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/false, /*with_bf_pks=*/false, /*fetch_vector=*/true},
{"dense_hnsw_fp16_fetch_vector",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100,
QuantizeType::FP16),
std::make_shared<HnswQueryParams>(300),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/false, /*with_bf_pks=*/false, /*fetch_vector=*/true},
};
for (const auto &tc : cases) {
RunOk(tc);
}
}
TEST_F(GroupByIndexerTest, CombinedSortsGroupsBeforeTruncating) {
// Build two blocks where the best group is in the later block; group_count
// must be applied after cross-block group sorting, not merge order.
const std::string block0_path = "test_groupby_combined_block0.index";
const std::string block1_path = "test_groupby_combined_block1.index";
zvec::test_util::RemoveTestFiles(block0_path);
zvec::test_util::RemoveTestFiles(block1_path);
auto index_params = std::make_shared<FlatIndexParams>(MetricType::IP);
FieldSchema schema("test", DataType::VECTOR_FP32, kGbDimension, false,
index_params);
auto block0 = std::make_shared<VectorColumnIndexer>(block0_path, schema);
auto block1 = std::make_shared<VectorColumnIndexer>(block1_path, schema);
ASSERT_TRUE(block0->Open(vector_column_params::ReadOptions{true, true}).ok());
ASSERT_TRUE(block1->Open(vector_column_params::ReadOptions{true, true}).ok());
auto insert_dense = [](const VectorColumnIndexer::Ptr &indexer,
uint32_t doc_id, float value) {
std::vector<float> vec(kGbDimension, value);
vector_column_params::DenseVector dense{vec.data()};
ASSERT_TRUE(
indexer->Insert(vector_column_params::VectorData{dense}, doc_id).ok());
};
insert_dense(block0, 0, 0.0f);
insert_dense(block0, 1, 1.0f);
insert_dense(block1, 0, 10.0f);
insert_dense(block1, 1, 11.0f);
std::vector<BlockMeta> blocks{
BlockMeta(0, BlockType::VECTOR_INDEX, 0, 1, 2, {"test"}),
BlockMeta(1, BlockType::VECTOR_INDEX, 2, 3, 2, {"test"}),
};
SegmentMeta segment_meta;
CombinedVectorColumnIndexer combined({block0, block1}, {}, schema,
segment_meta, blocks, MetricType::IP);
std::vector<float> query(kGbDimension, 1.0f);
vector_column_params::DenseVector dense_query{query.data()};
vector_column_params::QueryParams query_params;
query_params.topk = kGbSearchTopk;
query_params.query_params = std::make_shared<QueryParams>(IndexType::FLAT);
query_params.group_by = MakeSegmentGroupByParams(/*group_topk=*/1,
/*group_count=*/1);
auto results = combined.Search(vector_column_params::VectorData{dense_query},
query_params);
ASSERT_TRUE(results.has_value());
auto *group_results =
dynamic_cast<GroupVectorIndexResults *>(results.value().get());
ASSERT_TRUE(group_results);
ASSERT_EQ(1u, group_results->groups().size());
ASSERT_EQ("high", group_results->groups()[0].group_id());
ASSERT_EQ(1u, group_results->groups()[0].docs().size());
ASSERT_EQ(3u, group_results->groups()[0].docs()[0].key());
ASSERT_FLOAT_EQ(44.0f, group_results->groups()[0].docs()[0].score());
ASSERT_TRUE(block0->Close().ok());
ASSERT_TRUE(block1->Close().ok());
zvec::test_util::RemoveTestFiles(block0_path);
zvec::test_util::RemoveTestFiles(block1_path);
}
TEST_F(GroupByIndexerTest, Sparse) {
std::vector<GroupByCase> cases{
{"sparse_flat_graph", std::make_shared<FlatIndexParams>(MetricType::IP),
/*query_params=*/nullptr,
/*is_sparse=*/true},
{"sparse_hnsw_graph",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
/*query_params=*/nullptr,
/*is_sparse=*/true},
};
for (const auto &tc : cases) {
RunOk(tc);
}
}
TEST_F(GroupByIndexerTest, UnsupportedIndexTypes) {
std::vector<GroupByCase> cases{
{"unsupported_ivf", std::make_shared<IVFIndexParams>(MetricType::IP, 4),
std::make_shared<IVFQueryParams>(4)},
{"unsupported_diskann",
std::make_shared<DiskAnnIndexParams>(MetricType::IP),
std::make_shared<DiskAnnQueryParams>(),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/true},
};
for (const auto &tc : cases) {
RunRejected(tc);
}
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif

2
tests/db/index/column/vector_column_indexer_test.cc
View File

@ -2684,4 +2684,4 @@ TEST(VectorColumnIndexerTest, Refiner) {
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
#endif