|
12 | 12 | #include "knowhere/index/index_node.h" |
13 | 13 |
|
14 | 14 | #include <cmath> |
| 15 | +#include <cstring> |
15 | 16 | #include <queue> |
| 17 | +#include <unordered_map> |
16 | 18 | #include <unordered_set> |
17 | 19 |
|
18 | 20 | #include "knowhere/context.h" |
|
27 | 29 |
|
28 | 30 | namespace knowhere { |
29 | 31 |
|
| 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 | + |
30 | 217 | // NOLINTBEGIN(google-default-arguments) |
31 | 218 | expected<DataSetPtr> |
32 | 219 | 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 |
419 | 606 | expected<std::vector<IndexNode::IteratorPtr>> |
420 | 607 | IndexNode::AnnIteratorEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset, |
421 | 608 | 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"; |
427 | 624 | 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()); |
429 | 704 | } |
430 | | - return AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, op_context); |
| 705 | + return result; |
431 | 706 | } |
432 | 707 | // NOLINTEND(google-default-arguments) |
433 | 708 |
|
|
0 commit comments