From 04198175de9be487d2f0ed1187e6eb4f36771c0b Mon Sep 17 00:00:00 2001 From: David Date: Mon, 8 Jun 2026 20:35:17 +0000 Subject: [PATCH] feat: emb_list (ArrayOfVector) and sparse AnnIterators for streaming retrieval Adds knowhere AnnIterator support enabling bounded-memory streaming (search_iterator) over two index families: - emb_list / ArrayOfVector (EMB_LIST_HNSW, MAX_SIM): a chunk-level streaming iterator with one sub-iterator per query vector, aggregating MAX_SIM per row. - sparse inverted index: a streaming bounded-WAND AnnIterator. Engine-level support for milvus search_iterator over emb_list / hybrid / sparse (milvus-io/milvus#49906). Tests: tests/ut/test_emb_list.cc, test_sparse.cc. Signed-off-by: David --- include/knowhere/sparse_utils.h | 53 +++++ src/index/index_node.cc | 289 ++++++++++++++++++++++- src/index/sparse/sparse_index_node.cc | 131 +++++++++- src/index/sparse/sparse_inverted_index.h | 57 ++++- tests/ut/test_emb_list.cc | 184 +++++++++++++++ tests/ut/test_sparse.cc | 115 +++++++++ 6 files changed, 811 insertions(+), 18 deletions(-) diff --git a/include/knowhere/sparse_utils.h b/include/knowhere/sparse_utils.h index 0c76e4c3c..52c687a35 100644 --- a/include/knowhere/sparse_utils.h +++ b/include/knowhere/sparse_utils.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "knowhere/expected.h" @@ -319,6 +320,58 @@ class MaxMinHeap { std::vector> pool_; }; // class MaxMinHeap +// A MaxMinHeap variant for iterative (batched) retrieval. It keeps the top +// `capacity` elements like MaxMinHeap, but additionally rejects any element +// scoring above `ceiling` and -- at exactly `ceiling` -- any id in `excluded`. +// +// This lets a caller retrieve the score-descending band that follows a batch it +// has already returned, without re-returning it: `ceiling` is the previous +// batch's minimum score and `excluded` is the set of ids already returned at +// exactly that score. On the first batch, pass `ceiling = +inf` and an empty +// set, which makes it behave exactly like MaxMinHeap. +template +class BoundedMaxMinHeap { + public: + BoundedMaxMinHeap(int capacity, T ceiling, const std::unordered_set& excluded) + : heap_(capacity), ceiling_(ceiling), excluded_(excluded) { + } + void + push(table_t id, T val) { + if (val > ceiling_) { + return; + } + if (val == ceiling_ && excluded_.count(id) != 0) { + return; + } + heap_.push(id, val); + } + table_t + pop() { + return heap_.pop(); + } + [[nodiscard]] size_t + size() const { + return heap_.size(); + } + [[nodiscard]] bool + empty() const { + return heap_.empty(); + } + SparseIdVal + top() const { + return heap_.top(); + } + [[nodiscard]] bool + full() const { + return heap_.full(); + } + + private: + MaxMinHeap heap_; + const T ceiling_; + const std::unordered_set& excluded_; +}; // class BoundedMaxMinHeap + // A std::vector like container but uses fixed size free memory(typically from // mmap) as backing store and can only be appended at the end. // diff --git a/src/index/index_node.cc b/src/index/index_node.cc index 1b3b2f432..60297f780 100644 --- a/src/index/index_node.cc +++ b/src/index/index_node.cc @@ -12,7 +12,9 @@ #include "knowhere/index/index_node.h" #include +#include #include +#include #include #include "knowhere/context.h" @@ -27,6 +29,191 @@ namespace knowhere { +namespace { + +// Chunk-level streaming iterator for emb_list (ArrayOfVector) indexes. +// +// An emb_list index groups consecutive paragraph vectors into chunks (emb_lists); a +// query is itself a group of `m` vectors. This iterator is a grouping layer over the +// `m` per-query-vector iterators that the underlying index's AnnIterator already +// returns for one query group: it consumes paragraph-level hits, resolves each to its +// chunk, computes the exact MAX_SIM score of the whole chunk on first sighting, and +// emits chunk-level (chunk_id, score) pairs in approximately-descending order. +// +// Scoring is always exact -- a chunk's paragraphs are contiguous and few, so the full +// chunk is brute-force scored the moment any one paragraph is touched. The `ub` bound +// governs emission ordering only, and is soft because best-first ANN traversal can +// move "uphill" (see SPEC 6.1). +class EmbListIterator : public IndexNode::iterator { + public: + // Computes the exact MAX_SIM score of a chunk from its paragraph vector ids; + // std::nullopt signals a scoring failure for that chunk. + using ChunkScorer = std::function(const std::vector&)>; + + EmbListIterator(std::vector&& sub_iters, const EmbListOffset* el_offset, + ChunkScorer score_chunk, bool larger_is_closer) + : sub_iters_(std::move(sub_iters)), + el_offset_(el_offset), + score_chunk_(std::move(score_chunk)), + larger_is_closer_(larger_is_closer) { + } + + std::pair + Next() override { + prepare(); + if (!has_next_) { + throw std::runtime_error("No more elements"); + } + prepared_ = false; + return next_chunk_; + } + + [[nodiscard]] bool + HasNext() override { + prepare(); + return has_next_; + } + + private: + struct ScoredChunk { + int64_t id; + float score; + // sign-normalised score so that pending_ is always a max-heap on the best chunk + float key; + + bool + operator<(const ScoredChunk& other) const { + return key < other.key; + } + }; + + // Sign-normalise so that "more promising" is always "larger": a similarity for + // IP/COSINE, a negated distance for L2. + float + to_key(float val) const { + return larger_is_closer_ ? val : -val; + } + + // Pull the next paragraph hit from sub-iterator `i` into head_[i], or clear it. + void + refill_head(size_t i) { + if (sub_iters_[i] != nullptr && sub_iters_[i]->HasNext()) { + head_[i] = sub_iters_[i]->Next(); + } else { + head_[i] = std::nullopt; + } + } + + // ub_key_ = sum of the per-sub-iterator head keys: a soft bound that no + // not-yet-scored chunk is expected to outrank. + void + recompute_ub() { + float ub = 0.0f; + for (const auto& h : head_) { + if (h.has_value()) { + ub += to_key(h->second); + } + } + ub_key_ = ub; + } + + // Advance the underlying traversals until pending_'s best chunk is safe to emit, + // or the iterator is exhausted. Caches the outcome in has_next_ / next_chunk_. + void + prepare() { + if (prepared_) { + return; + } + if (!started_) { + head_.resize(sub_iters_.size()); + for (size_t i = 0; i < sub_iters_.size(); i++) { + refill_head(i); + } + recompute_ub(); + started_ = true; + } + + while (true) { + // pick the most promising sub-traversal to advance next + bool any_head = false; + size_t best_i = 0; + float best_key = 0.0f; + for (size_t i = 0; i < head_.size(); i++) { + if (head_[i].has_value()) { + const float k = to_key(head_[i]->second); + if (!any_head || k > best_key) { + any_head = true; + best_key = k; + best_i = i; + } + } + } + + if (!pending_.empty()) { + // Once no sub-iterator can advance, ub no longer constrains anything, + // so drain pending_ in exact-score order. + if (!any_head || pending_.top().key >= ub_key_) { + const auto& top = pending_.top(); + next_chunk_ = {top.id, top.score}; + pending_.pop(); + has_next_ = true; + prepared_ = true; + return; + } + } else if (!any_head) { + has_next_ = false; + prepared_ = true; + return; + } + + // advance the chosen sub-traversal by one paragraph + const int64_t para_id = head_[best_i]->first; + refill_head(best_i); + recompute_ub(); + + if (para_id < 0) { + continue; + } + const size_t chunk_id = el_offset_->get_el_id(static_cast(para_id)); + if (chunk_id >= el_offset_->num_el()) { + continue; + } + const int64_t cid = static_cast(chunk_id); + if (scored_.count(cid) != 0) { + continue; + } + const auto vids = el_offset_->get_vids(chunk_id); + const auto score_or = score_chunk_(vids); + if (!score_or.has_value()) { + // Defensive: mark the chunk scored so it is not retried, but skip + // emitting it rather than aborting the whole iterator. + scored_.emplace(cid, 0.0f); + continue; + } + const float score = score_or.value(); + scored_.emplace(cid, score); + pending_.push(ScoredChunk{cid, score, to_key(score)}); + } + } + + std::vector sub_iters_; + const EmbListOffset* el_offset_; + ChunkScorer score_chunk_; + const bool larger_is_closer_; + + bool started_ = false; + bool prepared_ = false; + bool has_next_ = false; + std::pair next_chunk_; + + std::vector>> head_; + float ub_key_ = 0.0f; + std::unordered_map scored_; + std::priority_queue pending_; +}; + +} // namespace + // NOLINTBEGIN(google-default-arguments) expected IndexNode::RangeSearch(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, @@ -419,15 +606,103 @@ IndexNode::RangeSearchEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr> IndexNode::AnnIteratorEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, bool use_knowhere_search_pool, milvus::OpContext* op_context) const { - auto config = static_cast(*cfg); - auto el_metric_type_or = get_el_metric_type(config.metric_type.value()); - auto metric_is_emb_list = el_metric_type_or.has_value(); - if (metric_is_emb_list) { - LOG_KNOWHERE_WARNING_ << "Ann iterator is not supported for emb_list"; + auto& config = static_cast(*cfg); + auto metric_type = config.metric_type.value(); + if (!get_el_metric_type(metric_type).has_value()) { + // not an emb_list metric: regular per-vector iterator + return AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context); + } + if (emb_list_offset_ == nullptr) { + LOG_KNOWHERE_WARNING_ << "emb_list metric type, but index has no emb_list offset"; + return expected>::Err(Status::emb_list_inner_error, "index is not an emb_list index"); + } + + // the query dataset is itself grouped into emb_lists + const size_t* lims = dataset->Get(knowhere::meta::EMB_LIST_OFFSET); + if (lims == nullptr) { + LOG_KNOWHERE_WARNING_ << "emb_list metric type, but query dataset has no emb_list offset"; return expected>::Err(Status::emb_list_inner_error, - "ann iterator is not supported for emb_list"); + "missing emb_list offset in query dataset"); + } + auto num_q_vecs = static_cast(dataset->GetRows()); + if (num_q_vecs == 0) { + return expected>::Err(Status::emb_list_inner_error, "empty query dataset"); + } + EmbListOffset query_el_offset(lims, num_q_vecs); + auto num_q_el = query_el_offset.num_el(); + + auto sub_metric_type_or = get_sub_metric_type(metric_type); + if (!sub_metric_type_or.has_value()) { + LOG_KNOWHERE_WARNING_ << "Invalid emb_list metric type: " << metric_type; + return expected>::Err(Status::emb_list_inner_error, "invalid emb_list metric type"); + } + auto sub_metric_type = sub_metric_type_or.value(); + bool larger_is_closer = true; + if (sub_metric_type == metric::L2 || sub_metric_type == metric::HAMMING || sub_metric_type == metric::JACCARD) { + larger_is_closer = false; + } + bool is_cosine = sub_metric_type == metric::COSINE ? true : false; + + auto query_code_size_or = GetQueryCodeSize(dataset); + if (!query_code_size_or.has_value()) { + LOG_KNOWHERE_ERROR_ << "could not get query code size for emb_list iterator"; + return expected>::Err(Status::emb_list_inner_error, "could not get query code size"); + } + auto query_code_size = query_code_size_or.value(); + auto dim = dataset->GetDim(); + const char* query_tensor = static_cast(dataset->GetTensor()); + + // The underlying per-vector iterator dispatches on the sub-metric (IP / COSINE / + // L2); rewrite the config so it does not see the emb_list MAX_SIM_* metric. + config.metric_type = sub_metric_type; + auto sub_iters_or = AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context); + if (!sub_iters_or.has_value()) { + return sub_iters_or; + } + auto sub_iters = sub_iters_or.value(); + if (sub_iters.size() != num_q_vecs) { + LOG_KNOWHERE_ERROR_ << "unexpected sub-iterator count: " << sub_iters.size() << " vs " << num_q_vecs; + return expected>::Err(Status::emb_list_inner_error, "unexpected sub-iterator count"); + } + + std::vector result(num_q_el); + try { + for (size_t i = 0; i < num_q_el; i++) { + auto start = query_el_offset.offset[i]; + auto end = query_el_offset.offset[i + 1]; + auto nq = end - start; + + std::vector group_iters(sub_iters.begin() + start, sub_iters.begin() + end); + + // own a private copy of this query emb_list's vectors so the scorer stays + // valid for the whole (lazy) lifetime of the iterator + auto group_buf = std::make_unique(nq * query_code_size); + std::memcpy(group_buf.get(), query_tensor + start * query_code_size, nq * query_code_size); + auto group_query = GenDataSet(static_cast(nq), dim, group_buf.release()); + group_query->SetIsOwner(true); + + EmbListIterator::ChunkScorer scorer = [this, bitset, group_query, nq, is_cosine, larger_is_closer]( + const std::vector& vids) -> std::optional { + if (vids.empty()) { + return std::nullopt; + } + // exact MAX_SIM: brute-force this query emb_list's vectors against + // every paragraph of the candidate chunk, then aggregate + auto dist_or = CalcDistByIDs(group_query, bitset, vids.data(), vids.size(), is_cosine); + if (!dist_or.has_value()) { + return std::nullopt; + } + return get_sum_max_sim(dist_or.value()->GetDistance(), nq, vids.size(), larger_is_closer); + }; + + result[i] = std::make_shared(std::move(group_iters), emb_list_offset_.get(), + std::move(scorer), larger_is_closer); + } + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "emb_list iterator error: " << e.what(); + return expected>::Err(Status::emb_list_inner_error, e.what()); } - return AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context); + return result; } // NOLINTEND(google-default-arguments) diff --git a/src/index/sparse/sparse_index_node.cc b/src/index/sparse/sparse_index_node.cc index ee7dfc72d..bcc366704 100644 --- a/src/index/sparse/sparse_index_node.cc +++ b/src/index/sparse/sparse_index_node.cc @@ -12,6 +12,8 @@ #include #include +#include +#include #include "index/sparse/sparse_inverted_index.h" #include "index/sparse/sparse_inverted_index_config.h" @@ -32,6 +34,14 @@ namespace knowhere { +namespace { +// Internal batch size for the streaming bounded-WAND iterator: the number of +// results one posting-list traversal retrieves before the cursor advances. +// Bounded-WAND re-traverses per batch, so cost is ~O(L^2 / batch) (SPEC 13 D2); +// a few hundred keeps per-batch memory small while amortising the re-traversal. +constexpr size_t kSparseIteratorBatchSize = 256; +} // namespace + // Inverted Index impl for sparse vectors. // // Not overriding RangeSearch, will use the default implementation in IndexNode. @@ -184,11 +194,107 @@ class SparseInvertedIndexNode : public IndexNode { bool first_return_ = true; }; + // Streaming bounded-WAND iterator (SPEC 6.2, v1). Each refill runs one bounded + // top-`batch` retrieval (WAND / MaxScore / TAAT) just below a + // (last_score, tie-band) cursor, so results stream out in descending score in + // O(batch) memory -- unlike PrecomputedDistanceIterator, whose first Next() + // materialises the whole result set. The cursor is the only state carried + // between batches; per-batch posting traversal is re-done from scratch. + class BoundedWandIterator : public IndexNode::iterator { + public: + BoundedWandIterator(const sparse::BaseInvertedIndex* index, sparse::SparseRow&& query, + const sparse::DocValueComputer& computer, const BitsetView& bitset, + const sparse::InvertedIndexApproxSearchParams& approx_params, size_t batch_size) + : index_(index), + query_(std::move(query)), + computer_(computer), + bitset_(bitset), + approx_params_(approx_params), + batch_size_(batch_size) { + } + + std::pair + Next() override { + refill_if_needed(); + if (pos_ >= buffer_.size()) { + throw std::runtime_error("index out of range while iterating BoundedWandIterator"); + } + const auto& res = buffer_[pos_++]; + return {res.id, res.val}; + } + + [[nodiscard]] bool + HasNext() override { + refill_if_needed(); + return pos_ < buffer_.size(); + } + + private: + // Refill the batch buffer once drained, advancing the (last_score, tie-band) + // cursor so the next batch resumes exactly where this one stopped. + void + refill_if_needed() { + if (pos_ < buffer_.size() || exhausted_) { + return; + } + buffer_ = index_->IterativeSearch(query_, batch_size_, last_score_, tie_band_, bitset_, computer_, + approx_params_); + pos_ = 0; + // a short batch means no more candidates exist below the cursor + if (buffer_.size() < batch_size_) { + exhausted_ = true; + } + if (buffer_.empty()) { + return; + } + // buffer_ is descending; its minimum score is the next batch's ceiling + const float batch_min = buffer_.back().val; + std::unordered_set next_tie_band; + for (const auto& d : buffer_) { + if (d.val == batch_min) { + next_tie_band.insert(static_cast(d.id)); + } + } + // a tie band straddling the batch boundary must keep excluding the ids + // the previous batch already returned at this exact score + if (batch_min == last_score_) { + next_tie_band.insert(tie_band_.begin(), tie_band_.end()); + } + last_score_ = batch_min; + tie_band_ = std::move(next_tie_band); + } + + const sparse::BaseInvertedIndex* index_; + sparse::SparseRow query_; + const sparse::DocValueComputer computer_; + const BitsetView bitset_; + sparse::InvertedIndexApproxSearchParams approx_params_; + const size_t batch_size_; + + std::vector buffer_; + size_t pos_ = 0; + float last_score_ = std::numeric_limits::max(); + std::unordered_set tie_band_; + bool exhausted_ = false; + }; + public: - // TODO: for now inverted index and wand use the same impl for AnnIterator. [[nodiscard]] expected> AnnIterator(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, bool use_knowhere_search_pool, milvus::OpContext* op_context) const override { + return CreateAnnIterators(dataset, std::move(config), bitset, use_knowhere_search_pool, /*streaming=*/true); + } + + protected: + // Shared body for AnnIterator. `streaming` selects the bounded-WAND iterator + // (SPEC 6.2: streaming, O(batch) memory); otherwise the materialising + // PrecomputedDistanceIterator is used. The concurrent (CC) index keeps the + // materialising path -- the streaming iterator's deferred per-batch traversal + // is not yet hardened against concurrent Add. + // TODO: for now inverted index and wand share this impl for AnnIterator. + [[nodiscard]] expected> + CreateAnnIterators(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, + bool use_knowhere_search_pool, bool streaming) const { if (!index_) { LOG_KNOWHERE_WARNING_ << "creating iterator on empty index"; return expected>>::Err(Status::empty_index, @@ -206,12 +312,24 @@ class SparseInvertedIndexNode : public IndexNode { auto computer = computer_or.value(); auto drop_ratio_search = cfg.drop_ratio_search.value_or(0.0f); + sparse::InvertedIndexApproxSearchParams approx_params = { + .refine_factor = 1, + .drop_ratio_search = drop_ratio_search, + .dim_max_score_ratio = cfg.dim_max_score_ratio.value(), + }; + // TODO: set approximated to false for now since the refinement is too slow after forward index is removed. const bool approximated = false; auto vec = std::vector>(nq, nullptr); try { for (int i = 0; i < nq; ++i) { + if (streaming) { + sparse::SparseRow query_copy(queries[i]); + vec[i] = std::make_shared(index_, std::move(query_copy), computer, bitset, + approx_params, kSparseIteratorBatchSize); + continue; + } // Heavy computations with `compute_dist_func` will be deferred until the first call to // 'Iterator->Next()'. auto compute_dist_func = [=]() -> std::vector { @@ -248,6 +366,7 @@ class SparseInvertedIndexNode : public IndexNode { return vec; } + public: [[nodiscard]] expected GetVectorByIds(const DataSetPtr dataset, milvus::OpContext* op_context) const override { return expected::Err(Status::not_implemented, "GetVectorByIds not implemented"); @@ -528,12 +647,14 @@ class SparseInvertedIndexNodeCC : public SparseInvertedIndexNode { bool use_knowhere_search_pool, milvus::OpContext* op_context) const override { ReadPermission permission(*this); // Always uses PrecomputedDistanceIterator for SparseInvertedIndexNodeCC: - // If we want to use RefineIterator, it needs to get another ReadPermission when calling - // index_->GetRawDistance(). If an Add task is added in between, there will be a deadlock. + // the streaming bounded-WAND iterator defers posting-list traversal to each + // Next() call, which would run without a ReadPermission and could race a + // concurrent Add. Using RefineIterator would likewise need another + // ReadPermission inside index_->GetRawDistance() and could deadlock. auto config = static_cast(*cfg); config.drop_ratio_search = 0.0f; - return SparseInvertedIndexNode::AnnIterator(dataset, std::move(cfg), bitset, - use_knowhere_search_pool, op_context); + return SparseInvertedIndexNode::CreateAnnIterators(dataset, std::move(cfg), bitset, + use_knowhere_search_pool, /*streaming=*/false); } expected diff --git a/src/index/sparse/sparse_inverted_index.h b/src/index/sparse/sparse_inverted_index.h index 63dcf8471..bd6fb8750 100644 --- a/src/index/sparse/sparse_inverted_index.h +++ b/src/index/sparse/sparse_inverted_index.h @@ -100,6 +100,15 @@ class BaseInvertedIndex { Search(const SparseRow& query, size_t k, float* distances, label_t* labels, const BitsetView& bitset, const DocValueComputer& computer, InvertedIndexApproxSearchParams& approx_params) const = 0; + // One score-descending batch for the streaming bounded-WAND iterator (SPEC 6.2): + // up to `batch_size` results scoring at most `score_ceiling`, excluding ids in + // `excluded` at exactly `score_ceiling`. First batch: score_ceiling = +inf, + // empty `excluded`. Fewer than `batch_size` results => exhausted below the cursor. + virtual std::vector + IterativeSearch(const SparseRow& query, size_t batch_size, float score_ceiling, + const std::unordered_set& excluded, const BitsetView& bitset, + const DocValueComputer& computer, InvertedIndexApproxSearchParams& approx_params) const = 0; + virtual std::vector GetAllDistances(const SparseRow& query, float drop_ratio_search, const BitsetView& bitset, const DocValueComputer& computer) const = 0; @@ -855,6 +864,42 @@ class InvertedIndex : public BaseInvertedIndex { } } + // One score-descending batch for the streaming bounded-WAND iterator: the top + // `batch_size` results scoring at most `score_ceiling`, excluding ids in + // `excluded` that sit at exactly `score_ceiling`. The same DAAT_WAND / + // DAAT_MAXSCORE / TAAT traversal as Search is reused; the BoundedMaxMinHeap + // applies the cursor. Fewer than `batch_size` results => exhausted below the + // cursor. See SPEC 6.2 / Appendix A.2. + std::vector + IterativeSearch(const SparseRow& query, size_t batch_size, float score_ceiling, + const std::unordered_set& excluded, const BitsetView& bitset, + const DocValueComputer& computer, + InvertedIndexApproxSearchParams& approx_params) const override { + std::vector results; + if (batch_size == 0 || query.size() == 0) { + return results; + } + auto q_vec = parse_query(query, approx_params.drop_ratio_search); + if (q_vec.empty()) { + return results; + } + BoundedMaxMinHeap heap(static_cast(batch_size), score_ceiling, excluded); + if constexpr (algo == InvertedIndexAlgo::DAAT_WAND) { + search_daat_wand(q_vec, heap, bitset, computer, approx_params.dim_max_score_ratio); + } else if constexpr (algo == InvertedIndexAlgo::DAAT_MAXSCORE) { + search_daat_maxscore(q_vec, heap, bitset, computer, approx_params.dim_max_score_ratio); + } else { + search_taat_naive(q_vec, heap, bitset, computer); + } + results.resize(heap.size()); + for (int i = static_cast(heap.size()) - 1; i >= 0; --i) { + auto top = heap.top(); + results[i] = DistId(static_cast(top.id), top.val); + heap.pop(); + } + return results; + } + // Returned distances are inaccurate based on the drop_ratio. std::vector GetAllDistances(const SparseRow& query, float drop_ratio_search, const BitsetView& bitset, @@ -1094,9 +1139,9 @@ class InvertedIndex : public BaseInvertedIndex { // find the top-k candidates using brute force search, k as specified by the capacity of the heap. // any value in q_vec that is smaller than q_threshold and any value with dimension >= n_cols() will be ignored. // TODO: may switch to row-wise brute force if filter rate is high. Benchmark needed. - template + template void - search_taat_naive(const std::vector>& q_vec, MaxMinHeap& heap, DocIdFilter& filter, + search_taat_naive(const std::vector>& q_vec, HeapType& heap, DocIdFilter& filter, const DocValueComputer& computer) const { auto scores = compute_all_distances(q_vec, computer); for (size_t i = 0; i < n_rows_internal_; ++i) { @@ -1106,9 +1151,9 @@ class InvertedIndex : public BaseInvertedIndex { } } - template + template void - search_daat_wand(const std::vector>& q_vec, MaxMinHeap& heap, DocIdFilter& filter, + search_daat_wand(const std::vector>& q_vec, HeapType& heap, DocIdFilter& filter, const DocValueComputer& computer, float dim_max_score_ratio) const { std::vector> cursors = make_cursors(q_vec, computer, filter, dim_max_score_ratio); std::vector*> cursor_ptrs(cursors.size()); @@ -1171,9 +1216,9 @@ class InvertedIndex : public BaseInvertedIndex { } } - template + template void - search_daat_maxscore(std::vector>& q_vec, MaxMinHeap& heap, DocIdFilter& filter, + search_daat_maxscore(std::vector>& q_vec, HeapType& heap, DocIdFilter& filter, const DocValueComputer& computer, float dim_max_score_ratio) const { std::sort(q_vec.begin(), q_vec.end(), [this](auto& a, auto& b) { return a.second * max_score_in_dim_spans_[a.first] > b.second * max_score_in_dim_spans_[b.first]; diff --git a/tests/ut/test_emb_list.cc b/tests/ut/test_emb_list.cc index 258206c90..fde72a2c6 100644 --- a/tests/ut/test_emb_list.cc +++ b/tests/ut/test_emb_list.cc @@ -2178,3 +2178,187 @@ TEST_CASE("Test brute force anniterator on chunk", "[on_chunk]") { } } } + +namespace { + +// Spearman rank correlation between the iterator's emission order and the order +// implied by the (exact) emitted scores. 1.0 == perfectly ordered emission. +double +emission_spearman_rho(const std::vector>& emitted, bool larger_is_closer) { + const size_t n = emitted.size(); + if (n < 2) { + return 1.0; + } + std::vector by_score(n); + for (size_t i = 0; i < n; i++) { + by_score[i] = i; + } + std::stable_sort(by_score.begin(), by_score.end(), [&](size_t a, size_t b) { + return larger_is_closer ? emitted[a].second > emitted[b].second : emitted[a].second < emitted[b].second; + }); + std::vector ideal_rank(n); + for (size_t r = 0; r < n; r++) { + ideal_rank[by_score[r]] = r; + } + double sum_d2 = 0.0; + for (size_t i = 0; i < n; i++) { + const double d = static_cast(i) - static_cast(ideal_rank[i]); + sum_d2 += d * d; + } + const double dn = static_cast(n); + return 1.0 - (6.0 * sum_d2) / (dn * (dn * dn - 1.0)); +} + +// Drain an emb_list iterator up to `limit` results, asserting there are no duplicate +// chunk ids along the way. +std::vector> +drain_emb_list_iterator(const knowhere::IndexNode::IteratorPtr& it, size_t limit) { + std::vector> out; + std::unordered_set seen; + while (it->HasNext() && out.size() < limit) { + auto [id, score] = it->Next(); + REQUIRE(seen.insert(id).second); // zero duplicate ids + out.emplace_back(id, score); + } + return out; +} + +} // namespace + +TEST_CASE("EmbList HNSW AnnIterator", "[emb_list][iterator]") { + const std::vector DISTANCE_TYPES = {"MAX_SIM_IP", "MAX_SIM_COSINE", "MAX_SIM_L2"}; + const int32_t dim = 16; + const int32_t nb = 800; + const int32_t each_el_len = 8; + const int32_t num_chunks = nb / each_el_len; + const int32_t NQ = 10; + const int32_t TOPK = 10; + const auto version = GenTestEmbListVersionList(); + + knowhere::Json base_conf; + base_conf[knowhere::meta::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW; + base_conf[knowhere::meta::DIM] = dim; + base_conf[knowhere::meta::ROWS] = nb; + base_conf[knowhere::meta::TOPK] = TOPK; + base_conf[knowhere::indexparam::HNSW_M] = 32; + base_conf[knowhere::indexparam::EFCONSTRUCTION] = 200; + base_conf[knowhere::indexparam::EF] = 200; + base_conf[knowhere::indexparam::RETRIEVAL_ANN_RATIO] = 3.0f; + + SECTION("complete, deduplicated, ordered, exact-scoring chunk-level iteration") { + for (const auto& metric : DISTANCE_TYPES) { + const bool larger_is_closer = (metric != "MAX_SIM_L2"); + knowhere::Json conf = base_conf; + conf[knowhere::meta::METRIC_TYPE] = metric; + + auto base_ds = GenEmbListDataSet(nb, dim, 42, each_el_len); + // query emb_lists of mixed size (1, 3, 5, ...) -- exercises m > 1 + auto query_ds = GenQueryEmbListDataSet(NQ, dim, 1234); + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW, version) + .value(); + REQUIRE(index.Build(base_ds, conf) == knowhere::Status::success); + + // reference: the two-stage emb_list Search + auto search_or = index.Search(query_ds, conf, nullptr); + REQUIRE(search_or.has_value()); + const auto search_res = search_or.value(); + const int64_t num_q_el = search_res->GetRows(); + const int64_t search_k = search_res->GetDim(); + const int64_t* search_ids = search_res->GetIds(); + const float* search_dists = search_res->GetDistance(); + + auto iters_or = index.AnnIterator(query_ds, conf, nullptr); + REQUIRE(iters_or.has_value()); + auto iters = iters_or.value(); + REQUIRE(static_cast(iters.size()) == num_q_el); + + for (int64_t q = 0; q < num_q_el; q++) { + // chunk ids + exact scores reported by the two-stage Search + std::unordered_set search_set; + std::unordered_map search_score; + for (int64_t j = 0; j < search_k; j++) { + const int64_t id = search_ids[q * search_k + j]; + if (id < 0) { + continue; + } + search_set.insert(id); + search_score.emplace(id, search_dists[q * search_k + j]); + } + + // (b) drain the iterator fully -- the helper asserts zero duplicates + const auto full = drain_emb_list_iterator(iters[q], static_cast(num_chunks) + 1); + REQUIRE(!iters[q]->HasNext()); + + // unfiltered, every chunk is reachable: the iterator emits the whole + // chunk set exactly once each + std::unordered_set iter_set; + for (const auto& [id, score] : full) { + REQUIRE(id >= 0); + REQUIRE(id < num_chunks); + iter_set.insert(id); + } + REQUIRE(static_cast(iter_set.size()) == num_chunks); + + // (a) the iterator's chunk-id set is a consistent superset of the + // Search result set -- guards against id-space divergence between the + // iterator path and the Search path + for (const int64_t id : search_set) { + REQUIRE(iter_set.count(id) != 0); + } + + // (d) emission order is approximately descending by score + std::vector> head( + full.begin(), full.begin() + std::min(full.size(), static_cast(TOPK))); + REQUIRE(emission_spearman_rho(head, larger_is_closer) >= 0.98); + + // (c) a chunk also found by Search carries the same exact MAX_SIM score + for (const auto& [id, score] : head) { + auto it = search_score.find(id); + if (it != search_score.end()) { + REQUIRE(score == Catch::Approx(it->second).epsilon(1e-4)); + } + } + } + } + } + + SECTION("single-chunk and full-capacity edge cases") { + // single chunk: the whole index is one emb_list + // full capacity: chunks of 20 paragraphs (the emb_list max_capacity) + const std::vector> edges = {{8, 8}, {200, 20}}; + for (const auto& metric : DISTANCE_TYPES) { + for (const auto& [edge_nb, edge_el_len] : edges) { + knowhere::Json conf = base_conf; + conf[knowhere::meta::METRIC_TYPE] = metric; + conf[knowhere::meta::ROWS] = edge_nb; + + auto base_ds = GenEmbListDataSet(edge_nb, dim, 7, edge_el_len); + auto query_ds = GenQueryEmbListDataSet(NQ, dim, 99); + const int32_t edge_chunks = edge_nb / edge_el_len; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW, version) + .value(); + REQUIRE(index.Build(base_ds, conf) == knowhere::Status::success); + + auto iters_or = index.AnnIterator(query_ds, conf, nullptr); + REQUIRE(iters_or.has_value()); + auto iters = iters_or.value(); + + for (auto& it : iters) { + // draining without a cap must terminate, never duplicate a chunk, + // and emit exactly the chunks the index contains + const auto emitted = drain_emb_list_iterator(it, static_cast(edge_chunks) + 1); + REQUIRE(!it->HasNext()); + REQUIRE(static_cast(emitted.size()) == edge_chunks); + for (const auto& [id, score] : emitted) { + REQUIRE(id >= 0); + REQUIRE(id < edge_chunks); + } + } + } + } + } +} diff --git a/tests/ut/test_sparse.cc b/tests/ut/test_sparse.cc index 6d19974ce..67cd6f9b3 100644 --- a/tests/ut/test_sparse.cc +++ b/tests/ut/test_sparse.cc @@ -10,8 +10,15 @@ // or implied. See the License for the specific language governing permissions and limitations under the License. #include +#include +#include #include +#include +#include +#include +#include +#include "catch2/catch_approx.hpp" #include "catch2/catch_test_macros.hpp" #include "catch2/generators/catch_generators.hpp" #include "knowhere/bitsetview.h" @@ -580,3 +587,111 @@ TEST_CASE("Test Mem Sparse Index CC", "[float metrics]") { } } } + +// Bounded-WAND streaming iterator (SPEC 6.2): SparseInvertedIndexNode::AnnIterator +// returns a streaming iterator that yields (doc_id, score) in descending score in +// O(batch) memory, instead of the materialising PrecomputedDistanceIterator. +TEST_CASE("Test Sparse Bounded-WAND Iterator", "[sparse][iterator]") { + const int32_t nb = 2000; + const int32_t nq = 10; + const int32_t dim = 300; + const int32_t L = 100; + auto version = GenTestVersionList(); + + auto metric = GENERATE(knowhere::metric::IP, knowhere::metric::BM25); + auto algo = GENERATE("TAAT_NAIVE", "DAAT_WAND", "DAAT_MAXSCORE"); + auto name = GENERATE(knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX, knowhere::IndexEnum::INDEX_SPARSE_WAND); + CAPTURE(metric, algo, name); + + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::METRIC_TYPE] = metric; + json[knowhere::meta::TOPK] = L; + json[knowhere::meta::BM25_K1] = 1.2; + json[knowhere::meta::BM25_B] = 0.75; + json[knowhere::meta::BM25_AVGDL] = 100; + json[knowhere::indexparam::DROP_RATIO_SEARCH] = 0.0; + json[knowhere::indexparam::INVERTED_INDEX_ALGO] = algo; + + auto gen_ds = [&](int32_t rows, float sparsity, int seed) -> knowhere::DataSetPtr { + if (metric == knowhere::metric::BM25) { + return GenSparseDataSetWithMaxVal(rows, dim, sparsity, 256, true); + } + return GenSparseDataSet(rows, dim, sparsity, seed); + }; + auto train_ds = gen_ds(nb, 0.95f, 42); + auto query_ds = gen_ds(nq, 0.97f, 7); + + auto idx = knowhere::IndexFactory::Instance().Create(name, version).value(); + REQUIRE(idx.Build(train_ds, json) == knowhere::Status::success); + + SECTION("streams Search top-L: complete, deduplicated, descending, exact scores") { + auto search_or = idx.Search(query_ds, json, nullptr); + REQUIRE(search_or.has_value()); + const auto search_res = search_or.value(); + const int64_t search_k = search_res->GetDim(); + const int64_t* search_ids = search_res->GetIds(); + const float* search_dists = search_res->GetDistance(); + + auto iters_or = idx.AnnIterator(query_ds, json, nullptr); + REQUIRE(iters_or.has_value()); + auto iters = iters_or.value(); + REQUIRE(iters.size() == static_cast(nq)); + + for (int q = 0; q < nq; ++q) { + // (b) drain the iterator fully -- zero duplicate ids; (d) non-increasing score + std::vector> emitted; + std::unordered_set seen; + std::unordered_map iter_score; + float prev = std::numeric_limits::max(); + auto& it = iters[q]; + while (it->HasNext()) { + auto [id, score] = it->Next(); + REQUIRE(seen.insert(id).second); + REQUIRE(score <= prev); + prev = score; + emitted.emplace_back(id, score); + iter_score.emplace(id, score); + } + + // (a) + (c) every doc the two-stage Search returns is in the stream, + // carrying the same exact score + int64_t valid = 0; + for (int64_t j = 0; j < search_k; ++j) { + const int64_t sid = search_ids[q * search_k + j]; + if (sid < 0) { + continue; + } + ++valid; + auto f = iter_score.find(sid); + REQUIRE(f != iter_score.end()); + REQUIRE(f->second == Catch::Approx(search_dists[q * search_k + j]).epsilon(1e-4)); + } + // iterating to L yields the same score-ranked prefix as Search(topK=L) + REQUIRE(static_cast(emitted.size()) >= valid); + for (int64_t j = 0; j < valid; ++j) { + REQUIRE(emitted[j].second == Catch::Approx(search_dists[q * search_k + j]).epsilon(1e-4)); + } + } + } + + SECTION("degenerate: empty query yields an immediately exhausted iterator") { + std::vector> rows = {{}, {{1, 1.0f}, {7, 2.0f}, {42, 3.0f}}}; + auto degenerate_query = GenSparseDataSet(rows, dim); + auto iters_or = idx.AnnIterator(degenerate_query, json, nullptr); + REQUIRE(iters_or.has_value()); + auto iters = iters_or.value(); + REQUIRE(iters.size() == 2); + // an empty query row has nothing to score + REQUIRE(!iters[0]->HasNext()); + // a non-empty row still streams a well-formed (deduplicated, descending) result + std::unordered_set seen; + float prev = std::numeric_limits::max(); + while (iters[1]->HasNext()) { + auto [id, score] = iters[1]->Next(); + REQUIRE(seen.insert(id).second); + REQUIRE(score <= prev); + prev = score; + } + } +}