Skip to content

Commit 3e980ea

Browse files
author
David
committed
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.
1 parent 350fedb commit 3e980ea

6 files changed

Lines changed: 811 additions & 18 deletions

File tree

include/knowhere/sparse_utils.h

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include <cstring>
2222
#include <functional>
2323
#include <type_traits>
24+
#include <unordered_set>
2425
#include <vector>
2526

2627
#include "knowhere/expected.h"
@@ -319,6 +320,58 @@ class MaxMinHeap {
319320
std::vector<SparseIdVal<T>> pool_;
320321
}; // class MaxMinHeap
321322

323+
// A MaxMinHeap variant for iterative (batched) retrieval. It keeps the top
324+
// `capacity` elements like MaxMinHeap, but additionally rejects any element
325+
// scoring above `ceiling` and -- at exactly `ceiling` -- any id in `excluded`.
326+
//
327+
// This lets a caller retrieve the score-descending band that follows a batch it
328+
// has already returned, without re-returning it: `ceiling` is the previous
329+
// batch's minimum score and `excluded` is the set of ids already returned at
330+
// exactly that score. On the first batch, pass `ceiling = +inf` and an empty
331+
// set, which makes it behave exactly like MaxMinHeap.
332+
template <typename T>
333+
class BoundedMaxMinHeap {
334+
public:
335+
BoundedMaxMinHeap(int capacity, T ceiling, const std::unordered_set<table_t>& excluded)
336+
: heap_(capacity), ceiling_(ceiling), excluded_(excluded) {
337+
}
338+
void
339+
push(table_t id, T val) {
340+
if (val > ceiling_) {
341+
return;
342+
}
343+
if (val == ceiling_ && excluded_.count(id) != 0) {
344+
return;
345+
}
346+
heap_.push(id, val);
347+
}
348+
table_t
349+
pop() {
350+
return heap_.pop();
351+
}
352+
[[nodiscard]] size_t
353+
size() const {
354+
return heap_.size();
355+
}
356+
[[nodiscard]] bool
357+
empty() const {
358+
return heap_.empty();
359+
}
360+
SparseIdVal<T>
361+
top() const {
362+
return heap_.top();
363+
}
364+
[[nodiscard]] bool
365+
full() const {
366+
return heap_.full();
367+
}
368+
369+
private:
370+
MaxMinHeap<T> heap_;
371+
const T ceiling_;
372+
const std::unordered_set<table_t>& excluded_;
373+
}; // class BoundedMaxMinHeap
374+
322375
// A std::vector like container but uses fixed size free memory(typically from
323376
// mmap) as backing store and can only be appended at the end.
324377
//

src/index/index_node.cc

Lines changed: 282 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
#include "knowhere/index/index_node.h"
1313

1414
#include <cmath>
15+
#include <cstring>
1516
#include <queue>
17+
#include <unordered_map>
1618
#include <unordered_set>
1719

1820
#include "knowhere/context.h"
@@ -27,6 +29,191 @@
2729

2830
namespace knowhere {
2931

32+
namespace {
33+
34+
// Chunk-level streaming iterator for emb_list (ArrayOfVector) indexes.
35+
//
36+
// An emb_list index groups consecutive paragraph vectors into chunks (emb_lists); a
37+
// query is itself a group of `m` vectors. This iterator is a grouping layer over the
38+
// `m` per-query-vector iterators that the underlying index's AnnIterator already
39+
// returns for one query group: it consumes paragraph-level hits, resolves each to its
40+
// chunk, computes the exact MAX_SIM score of the whole chunk on first sighting, and
41+
// emits chunk-level (chunk_id, score) pairs in approximately-descending order.
42+
//
43+
// Scoring is always exact -- a chunk's paragraphs are contiguous and few, so the full
44+
// chunk is brute-force scored the moment any one paragraph is touched. The `ub` bound
45+
// governs emission ordering only, and is soft because best-first ANN traversal can
46+
// move "uphill" (see SPEC 6.1).
47+
class EmbListIterator : public IndexNode::iterator {
48+
public:
49+
// Computes the exact MAX_SIM score of a chunk from its paragraph vector ids;
50+
// std::nullopt signals a scoring failure for that chunk.
51+
using ChunkScorer = std::function<std::optional<float>(const std::vector<int64_t>&)>;
52+
53+
EmbListIterator(std::vector<IndexNode::IteratorPtr>&& sub_iters, const EmbListOffset* el_offset,
54+
ChunkScorer score_chunk, bool larger_is_closer)
55+
: sub_iters_(std::move(sub_iters)),
56+
el_offset_(el_offset),
57+
score_chunk_(std::move(score_chunk)),
58+
larger_is_closer_(larger_is_closer) {
59+
}
60+
61+
std::pair<int64_t, float>
62+
Next() override {
63+
prepare();
64+
if (!has_next_) {
65+
throw std::runtime_error("No more elements");
66+
}
67+
prepared_ = false;
68+
return next_chunk_;
69+
}
70+
71+
[[nodiscard]] bool
72+
HasNext() override {
73+
prepare();
74+
return has_next_;
75+
}
76+
77+
private:
78+
struct ScoredChunk {
79+
int64_t id;
80+
float score;
81+
// sign-normalised score so that pending_ is always a max-heap on the best chunk
82+
float key;
83+
84+
bool
85+
operator<(const ScoredChunk& other) const {
86+
return key < other.key;
87+
}
88+
};
89+
90+
// Sign-normalise so that "more promising" is always "larger": a similarity for
91+
// IP/COSINE, a negated distance for L2.
92+
float
93+
to_key(float val) const {
94+
return larger_is_closer_ ? val : -val;
95+
}
96+
97+
// Pull the next paragraph hit from sub-iterator `i` into head_[i], or clear it.
98+
void
99+
refill_head(size_t i) {
100+
if (sub_iters_[i] != nullptr && sub_iters_[i]->HasNext()) {
101+
head_[i] = sub_iters_[i]->Next();
102+
} else {
103+
head_[i] = std::nullopt;
104+
}
105+
}
106+
107+
// ub_key_ = sum of the per-sub-iterator head keys: a soft bound that no
108+
// not-yet-scored chunk is expected to outrank.
109+
void
110+
recompute_ub() {
111+
float ub = 0.0f;
112+
for (const auto& h : head_) {
113+
if (h.has_value()) {
114+
ub += to_key(h->second);
115+
}
116+
}
117+
ub_key_ = ub;
118+
}
119+
120+
// Advance the underlying traversals until pending_'s best chunk is safe to emit,
121+
// or the iterator is exhausted. Caches the outcome in has_next_ / next_chunk_.
122+
void
123+
prepare() {
124+
if (prepared_) {
125+
return;
126+
}
127+
if (!started_) {
128+
head_.resize(sub_iters_.size());
129+
for (size_t i = 0; i < sub_iters_.size(); i++) {
130+
refill_head(i);
131+
}
132+
recompute_ub();
133+
started_ = true;
134+
}
135+
136+
while (true) {
137+
// pick the most promising sub-traversal to advance next
138+
bool any_head = false;
139+
size_t best_i = 0;
140+
float best_key = 0.0f;
141+
for (size_t i = 0; i < head_.size(); i++) {
142+
if (head_[i].has_value()) {
143+
const float k = to_key(head_[i]->second);
144+
if (!any_head || k > best_key) {
145+
any_head = true;
146+
best_key = k;
147+
best_i = i;
148+
}
149+
}
150+
}
151+
152+
if (!pending_.empty()) {
153+
// Once no sub-iterator can advance, ub no longer constrains anything,
154+
// so drain pending_ in exact-score order.
155+
if (!any_head || pending_.top().key >= ub_key_) {
156+
const auto& top = pending_.top();
157+
next_chunk_ = {top.id, top.score};
158+
pending_.pop();
159+
has_next_ = true;
160+
prepared_ = true;
161+
return;
162+
}
163+
} else if (!any_head) {
164+
has_next_ = false;
165+
prepared_ = true;
166+
return;
167+
}
168+
169+
// advance the chosen sub-traversal by one paragraph
170+
const int64_t para_id = head_[best_i]->first;
171+
refill_head(best_i);
172+
recompute_ub();
173+
174+
if (para_id < 0) {
175+
continue;
176+
}
177+
const size_t chunk_id = el_offset_->get_el_id(static_cast<size_t>(para_id));
178+
if (chunk_id >= el_offset_->num_el()) {
179+
continue;
180+
}
181+
const int64_t cid = static_cast<int64_t>(chunk_id);
182+
if (scored_.count(cid) != 0) {
183+
continue;
184+
}
185+
const auto vids = el_offset_->get_vids(chunk_id);
186+
const auto score_or = score_chunk_(vids);
187+
if (!score_or.has_value()) {
188+
// Defensive: mark the chunk scored so it is not retried, but skip
189+
// emitting it rather than aborting the whole iterator.
190+
scored_.emplace(cid, 0.0f);
191+
continue;
192+
}
193+
const float score = score_or.value();
194+
scored_.emplace(cid, score);
195+
pending_.push(ScoredChunk{cid, score, to_key(score)});
196+
}
197+
}
198+
199+
std::vector<IndexNode::IteratorPtr> sub_iters_;
200+
const EmbListOffset* el_offset_;
201+
ChunkScorer score_chunk_;
202+
const bool larger_is_closer_;
203+
204+
bool started_ = false;
205+
bool prepared_ = false;
206+
bool has_next_ = false;
207+
std::pair<int64_t, float> next_chunk_;
208+
209+
std::vector<std::optional<std::pair<int64_t, float>>> head_;
210+
float ub_key_ = 0.0f;
211+
std::unordered_map<int64_t, float> scored_;
212+
std::priority_queue<ScoredChunk> pending_;
213+
};
214+
215+
} // namespace
216+
30217
// NOLINTBEGIN(google-default-arguments)
31218
expected<DataSetPtr>
32219
IndexNode::RangeSearch(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset,
@@ -419,15 +606,103 @@ IndexNode::RangeSearchEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr<Co
419606
expected<std::vector<IndexNode::IteratorPtr>>
420607
IndexNode::AnnIteratorEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset,
421608
bool use_knowhere_search_pool, milvus::OpContext* op_context) const {
422-
auto config = static_cast<const knowhere::BaseConfig&>(*cfg);
423-
auto el_metric_type_or = get_el_metric_type(config.metric_type.value());
424-
auto metric_is_emb_list = el_metric_type_or.has_value();
425-
if (metric_is_emb_list) {
426-
LOG_KNOWHERE_WARNING_ << "Ann iterator is not supported for emb_list";
609+
auto& config = static_cast<BaseConfig&>(*cfg);
610+
auto metric_type = config.metric_type.value();
611+
if (!get_el_metric_type(metric_type).has_value()) {
612+
// not an emb_list metric: regular per-vector iterator
613+
return AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context);
614+
}
615+
if (emb_list_offset_ == nullptr) {
616+
LOG_KNOWHERE_WARNING_ << "emb_list metric type, but index has no emb_list offset";
617+
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error, "index is not an emb_list index");
618+
}
619+
620+
// the query dataset is itself grouped into emb_lists
621+
const size_t* lims = dataset->Get<const size_t*>(knowhere::meta::EMB_LIST_OFFSET);
622+
if (lims == nullptr) {
623+
LOG_KNOWHERE_WARNING_ << "emb_list metric type, but query dataset has no emb_list offset";
427624
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error,
428-
"ann iterator is not supported for emb_list");
625+
"missing emb_list offset in query dataset");
626+
}
627+
auto num_q_vecs = static_cast<size_t>(dataset->GetRows());
628+
if (num_q_vecs == 0) {
629+
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error, "empty query dataset");
630+
}
631+
EmbListOffset query_el_offset(lims, num_q_vecs);
632+
auto num_q_el = query_el_offset.num_el();
633+
634+
auto sub_metric_type_or = get_sub_metric_type(metric_type);
635+
if (!sub_metric_type_or.has_value()) {
636+
LOG_KNOWHERE_WARNING_ << "Invalid emb_list metric type: " << metric_type;
637+
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error, "invalid emb_list metric type");
638+
}
639+
auto sub_metric_type = sub_metric_type_or.value();
640+
bool larger_is_closer = true;
641+
if (sub_metric_type == metric::L2 || sub_metric_type == metric::HAMMING || sub_metric_type == metric::JACCARD) {
642+
larger_is_closer = false;
643+
}
644+
bool is_cosine = sub_metric_type == metric::COSINE ? true : false;
645+
646+
auto query_code_size_or = GetQueryCodeSize(dataset);
647+
if (!query_code_size_or.has_value()) {
648+
LOG_KNOWHERE_ERROR_ << "could not get query code size for emb_list iterator";
649+
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error, "could not get query code size");
650+
}
651+
auto query_code_size = query_code_size_or.value();
652+
auto dim = dataset->GetDim();
653+
const char* query_tensor = static_cast<const char*>(dataset->GetTensor());
654+
655+
// The underlying per-vector iterator dispatches on the sub-metric (IP / COSINE /
656+
// L2); rewrite the config so it does not see the emb_list MAX_SIM_* metric.
657+
config.metric_type = sub_metric_type;
658+
auto sub_iters_or = AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context);
659+
if (!sub_iters_or.has_value()) {
660+
return sub_iters_or;
661+
}
662+
auto sub_iters = sub_iters_or.value();
663+
if (sub_iters.size() != num_q_vecs) {
664+
LOG_KNOWHERE_ERROR_ << "unexpected sub-iterator count: " << sub_iters.size() << " vs " << num_q_vecs;
665+
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error, "unexpected sub-iterator count");
666+
}
667+
668+
std::vector<IteratorPtr> result(num_q_el);
669+
try {
670+
for (size_t i = 0; i < num_q_el; i++) {
671+
auto start = query_el_offset.offset[i];
672+
auto end = query_el_offset.offset[i + 1];
673+
auto nq = end - start;
674+
675+
std::vector<IteratorPtr> group_iters(sub_iters.begin() + start, sub_iters.begin() + end);
676+
677+
// own a private copy of this query emb_list's vectors so the scorer stays
678+
// valid for the whole (lazy) lifetime of the iterator
679+
auto group_buf = std::make_unique<char[]>(nq * query_code_size);
680+
std::memcpy(group_buf.get(), query_tensor + start * query_code_size, nq * query_code_size);
681+
auto group_query = GenDataSet(static_cast<int64_t>(nq), dim, group_buf.release());
682+
group_query->SetIsOwner(true);
683+
684+
EmbListIterator::ChunkScorer scorer = [this, bitset, group_query, nq, is_cosine, larger_is_closer](
685+
const std::vector<int64_t>& vids) -> std::optional<float> {
686+
if (vids.empty()) {
687+
return std::nullopt;
688+
}
689+
// exact MAX_SIM: brute-force this query emb_list's vectors against
690+
// every paragraph of the candidate chunk, then aggregate
691+
auto dist_or = CalcDistByIDs(group_query, bitset, vids.data(), vids.size(), is_cosine);
692+
if (!dist_or.has_value()) {
693+
return std::nullopt;
694+
}
695+
return get_sum_max_sim(dist_or.value()->GetDistance(), nq, vids.size(), larger_is_closer);
696+
};
697+
698+
result[i] = std::make_shared<EmbListIterator>(std::move(group_iters), emb_list_offset_.get(),
699+
std::move(scorer), larger_is_closer);
700+
}
701+
} catch (const std::exception& e) {
702+
LOG_KNOWHERE_WARNING_ << "emb_list iterator error: " << e.what();
703+
return expected<std::vector<IteratorPtr>>::Err(Status::emb_list_inner_error, e.what());
429704
}
430-
return AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context);
705+
return result;
431706
}
432707
// NOLINTEND(google-default-arguments)
433708

0 commit comments

Comments
 (0)