diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index f8a736c..dd824ff 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -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_; } diff --git a/src/core/algorithm/diskann/diskann_visit_filter.h b/src/core/algorithm/diskann/diskann_visit_filter.h index b0b8799..dd4aff7 100644 --- a/src/core/algorithm/diskann/diskann_visit_filter.h +++ b/src/core/algorithm/diskann/diskann_visit_filter.h @@ -13,407 +13,4 @@ // limitations under the License. #pragma once -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -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 *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( - 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 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 - static int init(Context *, void **ctx, uint64_t maxDocCnt, - uint64_t maxScanNum, std::tuple &&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(tpl); - c->filter = new (std::nothrow) - ailego::BloomFilter(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 - static int init(Context *, void **ctx, uint64_t maxDocCnt, - uint64_t maxScanNum, std::tuple &&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 - static int init(Context *, void **ctx, uint64_t maxDocCnt, - uint64_t maxScanNum, std::tuple &&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(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" diff --git a/src/core/algorithm/flat/flat_searcher_context.h b/src/core/algorithm/flat/flat_searcher_context.h index 01183e0..a04e53a 100644 --- a/src/core/algorithm/flat/flat_searcher_context.h +++ b/src/core/algorithm/flat/flat_searcher_context.h @@ -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; diff --git a/src/core/algorithm/flat/flat_streamer_context.h b/src/core/algorithm/flat/flat_streamer_context.h index 839409f..6561787 100644 --- a/src/core/algorithm/flat/flat_streamer_context.h +++ b/src/core/algorithm/flat/flat_streamer_context.h @@ -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; diff --git a/src/core/algorithm/flat_sparse/flat_sparse_context.h b/src/core/algorithm/flat_sparse/flat_sparse_context.h index 6f03f3b..fa62397 100644 --- a/src/core/algorithm/flat_sparse/flat_sparse_context.h +++ b/src/core/algorithm/flat_sparse/flat_sparse_context.h @@ -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]; } diff --git a/src/core/algorithm/flat_sparse/flat_sparse_search.h b/src/core/algorithm/flat_sparse/flat_sparse_search.h index f77c28d..11c5429 100644 --- a/src/core/algorithm/flat_sparse/flat_sparse_search.h +++ b/src/core/algorithm/flat_sparse/flat_sparse_search.h @@ -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 { diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h index cf18295..b7ac79c 100644 --- a/src/core/algorithm/hnsw/hnsw_context.h +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h index cf8a649..ea2c86b 100644 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h @@ -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_; } diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h index 83df7f2..56ae1cb 100644 --- a/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h @@ -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); diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 441df9d..6c69530 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -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() - 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>(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 *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 &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 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 *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 &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 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) { diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 7615ed3..bf1233a 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -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 diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc index baf8e22..032f366 100644 --- a/src/core/interface/indexes/flat_index.cc +++ b/src/core/interface/indexes/flat_index.cc @@ -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; } diff --git a/src/core/interface/indexes/hnsw_index.cc b/src/core/interface/indexes/hnsw_index.cc index 1c0dfca..9756289 100644 --- a/src/core/interface/indexes/hnsw_index.cc +++ b/src/core/interface/indexes/hnsw_index.cc @@ -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; } diff --git a/src/core/interface/indexes/hnsw_rabitq_index.cc b/src/core/interface/indexes/hnsw_rabitq_index.cc index 0c77b05..c555389 100644 --- a/src/core/interface/indexes/hnsw_rabitq_index.cc +++ b/src/core/interface/indexes/hnsw_rabitq_index.cc @@ -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 } diff --git a/src/core/interface/indexes/ivf_index.cc b/src/core/interface/indexes/ivf_index.cc index bb68828..9df540f 100644 --- a/src/core/interface/indexes/ivf_index.cc +++ b/src/core/interface/indexes/ivf_index.cc @@ -214,6 +214,11 @@ int IVFIndex::_prepare_for_search( const auto &ivf_search_param = std::dynamic_pointer_cast(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) { diff --git a/src/core/interface/indexes/vamana_index.cc b/src/core/interface/indexes/vamana_index.cc index 8591ff9..f40991b 100644 --- a/src/core/interface/indexes/vamana_index.cc +++ b/src/core/interface/indexes/vamana_index.cc @@ -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( diff --git a/src/core/utility/mmap_file_read_storage.cc b/src/core/utility/mmap_file_read_storage.cc index 04605bf..e909687 100644 --- a/src/core/utility/mmap_file_read_storage.cc +++ b/src/core/utility/mmap_file_read_storage.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.cc b/src/db/index/column/vector_column/combined_vector_column_indexer.cc index fd79be3..5e3fe80 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.cc +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.cc @@ -15,9 +15,244 @@ #include #include #include +#include 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 &values) { + return std::any_of(values.begin(), values.end(), + [](const auto &value) { return !value.empty(); }); +} + +bool HasRevertedValues(const std::vector> &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 reverted_vector_list; + std::vector 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( + is_sparse, std::move(doc_list), std::move(reverted_vector_list), + std::move(reverted_sparse_values_list)); + } + + private: + std::vector docs_; +}; + +class GroupResultAccumulator { + private: + struct GroupResult { + std::string group_id; + std::vector 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 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> reverted_vector_list; + std::vector> 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 group_reverted_vectors; + std::vector 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( + std::move(group_list), std::move(reverted_vector_list), + std::move(reverted_sparse_values_list)); + } + + private: + std::unordered_map> docs_by_group_; +}; + +} // namespace + CombinedVectorColumnIndexer::CombinedVectorColumnIndexer( const std::vector &indexers, const std::vector &normal_indexers, @@ -54,9 +289,12 @@ CombinedVectorColumnIndexer::CombinedVectorColumnIndexer( Result CombinedVectorColumnIndexer::Search( const vector_column_params::VectorData &vector_data, const vector_column_params::QueryParams &query_params) { - core::IndexDocumentList doc_list; - std::vector reverted_vector_list; - std::vector 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 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 CombinedVectorColumnIndexer::Search( } auto index_results = result.value(); + + GroupVectorIndexResults *group_index_results = + dynamic_cast(index_results.get()); + if (group_index_results != nullptr) { + group_results.AddBlock(block_offsets_[i], group_index_results); + continue; + } + VectorIndexResults *vector_index_results = dynamic_cast(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( - 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 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 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 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 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( - 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 diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index 7fa0c16..dec1177 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -120,6 +120,16 @@ class ProximaEngineHelper { std::make_shared(rp); } } + if (db_query_params.group_by) { + engine_query_param->group_by_param = + std::make_shared(); + 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; } diff --git a/src/db/index/column/vector_column/vector_column_indexer.cc b/src/db/index/column/vector_column/vector_column_indexer.cc index 1859e24..eed5dda 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.cc +++ b/src/db/index/column/vector_column/vector_column_indexer.cc @@ -198,6 +198,15 @@ Result 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( + 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( is_sparse_, std::move(search_result.doc_list_), std::move(search_result.reverted_vector_list_), diff --git a/src/db/index/column/vector_column/vector_index_results.h b/src/db/index/column/vector_column/vector_index_results.h index 154ff6e..ed519f2 100644 --- a/src/db/index/column/vector_column/vector_index_results.h +++ b/src/db/index/column/vector_column/vector_index_results.h @@ -241,7 +241,9 @@ class GroupVectorIndexResults : public IndexResults { std::vector> &&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> &reverted_vector_list() { + return reverted_vector_list_; + } + + std::vector> &reverted_sparse_values_list() { + return reverted_sparse_values_list_; + } + private: const core::IndexDocument &document(size_t group_index, size_t doc_index) const { diff --git a/src/include/zvec/core/framework/index_context.h b/src/include/zvec/core/framework/index_context.h index c77fcf4..1cf0e03 100644 --- a/src/include/zvec/core/framework/index_context.h +++ b/src/include/zvec/core/framework/index_context.h @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index e9ed390..6c01d70 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -95,9 +95,13 @@ struct VectorDataBuffer { struct SearchResult { core::IndexDocumentList doc_list_; + core::IndexGroupDocumentList group_doc_list_; // use string to manage memory std::vector reverted_vector_list_{}; std::vector reverted_sparse_values_list_{}; + // Grouped reverted values, aligned with group_doc_list_. + std::vector> group_reverted_vector_list_{}; + std::vector> 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; diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 5d4e8a2..78bebd5 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include #include @@ -166,6 +167,13 @@ struct RefinerParam { std::shared_ptr reference_index = nullptr; }; +// --- GroupBy Parameters --- +struct GroupByParam { + uint32_t group_topk{0}; + uint32_t group_count{0}; + std::function 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 group_by_param = nullptr; virtual Pointer Clone() const = 0; }; diff --git a/tests/core/interface/index_group_by_test.cc b/tests/core/interface/index_group_by_test.cc new file mode 100644 index 0000000..bb41938 --- /dev/null +++ b/tests/core/interface/index_group_by_test.cc @@ -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 +#include +#include +#include +#include +#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> AllPks() { + auto pks = std::make_shared>(); + 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(); + 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>( + dimension); + for (size_t i = 0; i < kTrainCount; ++i) { + NumericalVector vec(dimension, static_cast(i)); + EXPECT_TRUE(holder->emplace(i, vec)); + } + + auto index_meta = + std::make_shared(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 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(); + 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 values; + std::vector 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(); + 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 values(test_case.dimension, static_cast(key)); + if (test_case.is_sparse) { + std::vector 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 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 *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(doc.key()); + const float *vector = nullptr; + if (has_reverted) { + vector = + reinterpret_cast((*group_vectors)[doc_idx].data()); + } else if (doc.vector() != nullptr) { + vector = reinterpret_cast(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(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 *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(sparse.sparse_indices().data()); + const float *values = nullptr; + if (has_reverted) { + values = reinterpret_cast( + (*group_sparse_values)[doc_idx].data()); + } else { + values = + reinterpret_cast(sparse.sparse_values().data()); + } + const float expected = static_cast(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 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 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 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); + } +} diff --git a/tests/db/index/CMakeLists.txt b/tests/db/index/CMakeLists.txt index d269994..03da6ff 100644 --- a/tests/db/index/CMakeLists.txt +++ b/tests/db/index/CMakeLists.txt @@ -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 diff --git a/tests/db/index/column/vector_column_group_by_test.cc b/tests/db/index/column/vector_column_group_by_test.cc new file mode 100644 index 0000000..4fedf3d --- /dev/null +++ b/tests/db/index/column/vector_column_group_by_test.cc @@ -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 +#include +#include +#include +#include +#include +#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 MakeGroupByParams( + uint32_t group_count = kGbNumGroups) { + return std::make_unique( + kGbGroupTopk, group_count, [](uint64_t key) -> std::string { + return std::to_string(key % kGbNumGroups); + }); +} + +std::unique_ptr MakeSegmentGroupByParams( + uint32_t group_topk, uint32_t group_count) { + return std::make_unique( + group_topk, group_count, + [](uint64_t key) -> std::string { return key < 2 ? "low" : "high"; }); +} + +std::vector AllPks() { + std::vector 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 dense; + std::vector sparse_indices; + std::vector 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(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 indices(kGbSparseCount); + std::vector values(kGbSparseCount); + for (uint32_t j = 0; j < kGbSparseCount; ++j) { + indices[j] = i * kGbSparseCount + j; + values[j] = static_cast(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 vec(tc.dimension, static_cast(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(results); + ASSERT_TRUE(group_results) + << "Expected GroupVectorIndexResults for " << tc.name; + ASSERT_EQ(kGbNumGroups, group_results->groups().size()) << tc.name; + + std::set 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_data.vector); + const float *vector = + reinterpret_cast(dense_vector.data); + const float expected = static_cast(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(300); + hnsw_linear_qp->set_is_linear(true); + + std::vector cases{ + {"dense_flat_graph", std::make_shared(MetricType::IP), + std::make_shared(IndexType::FLAT)}, + {"dense_hnsw_graph", + std::make_shared(MetricType::IP, 10, 100), + std::make_shared(300)}, + {"dense_hnsw_linear", + std::make_shared(MetricType::IP, 10, 100), + hnsw_linear_qp}, + {"dense_hnsw_bf_pks", + std::make_shared(MetricType::IP, 10, 100), + std::make_shared(300), + /*is_sparse=*/false, /*dimension=*/kGbDimension, + /*optional=*/false, /*with_bf_pks=*/true}, + {"dense_hnsw_fetch_vector", + std::make_shared(MetricType::IP, 10, 100), + std::make_shared(300), + /*is_sparse=*/false, /*dimension=*/kGbDimension, + /*optional=*/false, /*with_bf_pks=*/false, /*fetch_vector=*/true}, + {"dense_hnsw_fp16_fetch_vector", + std::make_shared(MetricType::IP, 10, 100, + QuantizeType::FP16), + std::make_shared(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(MetricType::IP); + FieldSchema schema("test", DataType::VECTOR_FP32, kGbDimension, false, + index_params); + + auto block0 = std::make_shared(block0_path, schema); + auto block1 = std::make_shared(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 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 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 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(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(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 cases{ + {"sparse_flat_graph", std::make_shared(MetricType::IP), + /*query_params=*/nullptr, + /*is_sparse=*/true}, + {"sparse_hnsw_graph", + std::make_shared(MetricType::IP, 10, 100), + /*query_params=*/nullptr, + /*is_sparse=*/true}, + }; + + for (const auto &tc : cases) { + RunOk(tc); + } +} + +TEST_F(GroupByIndexerTest, UnsupportedIndexTypes) { + std::vector cases{ + {"unsupported_ivf", std::make_shared(MetricType::IP, 4), + std::make_shared(4)}, + {"unsupported_diskann", + std::make_shared(MetricType::IP), + std::make_shared(), + /*is_sparse=*/false, /*dimension=*/kGbDimension, + /*optional=*/true}, + }; + + for (const auto &tc : cases) { + RunRejected(tc); + } +} + +#if defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic pop +#endif diff --git a/tests/db/index/column/vector_column_indexer_test.cc b/tests/db/index/column/vector_column_indexer_test.cc index b798e8d..b0bfd9c 100644 --- a/tests/db/index/column/vector_column_indexer_test.cc +++ b/tests/db/index/column/vector_column_indexer_test.cc @@ -2684,4 +2684,4 @@ TEST(VectorColumnIndexerTest, Refiner) { #if defined(__GNUC__) || defined(__GNUG__) #pragma GCC diagnostic pop -#endif \ No newline at end of file +#endif