From 32943eea248cdead364b98ce868f6019714c84da Mon Sep 17 00:00:00 2001 From: Yingfeng Zhang Date: Fri, 5 Jun 2026 00:01:10 +0800 Subject: [PATCH 1/6] Init --- src/executor/operator/physical_match_impl.cpp | 6 +- .../search/block_at_a_time_iterator.cppm | 110 ++++++++ .../search/block_at_a_time_iterator_impl.cpp | 261 ++++++++++++++++++ .../invertedindex/search/doc_iterator.cppm | 12 +- .../invertedindex/search/query_node_impl.cpp | 50 +++- 5 files changed, 421 insertions(+), 18 deletions(-) create mode 100644 src/storage/invertedindex/search/block_at_a_time_iterator.cppm create mode 100644 src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp diff --git a/src/executor/operator/physical_match_impl.cpp b/src/executor/operator/physical_match_impl.cpp index 7545e94180..ddeac2ac98 100644 --- a/src/executor/operator/physical_match_impl.cpp +++ b/src/executor/operator/physical_match_impl.cpp @@ -89,7 +89,8 @@ QueryIterators CreateQueryIterators(QueryBuilder &query_builder, case EarlyTermAlgo::kAuto: case EarlyTermAlgo::kNaive: case EarlyTermAlgo::kBatch: - case EarlyTermAlgo::kBMW: { + case EarlyTermAlgo::kBMW: + case EarlyTermAlgo::kBlockAtATime: { // ok break; } @@ -117,7 +118,8 @@ QueryIterators CreateQueryIterators(QueryBuilder &query_builder, case EarlyTermAlgo::kAuto: case EarlyTermAlgo::kNaive: case EarlyTermAlgo::kBatch: - case EarlyTermAlgo::kBMW: { + case EarlyTermAlgo::kBMW: + case EarlyTermAlgo::kBlockAtATime: { query_iterators.query_iter = get_iter(early_term_algo); break; } diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator.cppm b/src/storage/invertedindex/search/block_at_a_time_iterator.cppm new file mode 100644 index 0000000000..5c5504c533 --- /dev/null +++ b/src/storage/invertedindex/search/block_at_a_time_iterator.cppm @@ -0,0 +1,110 @@ +// Copyright(C) 2023 InfiniFlow, Inc. All rights reserved. +// +// 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 +// +// https://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. + +export module infinity_core:block_at_a_time_iterator; + +import :index_defines; +import :doc_iterator; +import :blockmax_leaf_iterator; +import :multi_doc_iterator; +import :simd_functions; + +import std.compat; +import internal_types; + +namespace infinity { + +// Block-at-a-time (BAAT) iterator for long queries with many common terms. +// Processes one term's posting list at a time, accumulating scores in a sparse +// accumulator, then extracts top-k results. +// +// Reference: Hornet engine BAAT traversal for agent workloads (long queries). +// BAAT is selected when query has >= 8 term children and average DF > 10% of total docs. +// In BAAT, terms are processed in descending IDF order (rare terms first) to +// quickly establish a high threshold, enabling aggressive block-level pruning +// for subsequent common terms. +export class BlockAtATimeIterator : public MultiDocIterator { +public: + explicit BlockAtATimeIterator(std::vector> &&iterators, u32 topn); + + ~BlockAtATimeIterator() override; + + DocIteratorType GetType() const override { return DocIteratorType::kBlockAtATimeIterator; } + + std::string Name() const override { return "BlockAtATimeIterator"; } + + void UpdateScoreThreshold(float threshold) override; + + bool Next(RowID doc_id) override; + + float BM25Score() const; + + float Score() override { return BM25Score(); } + + u32 MatchCount() const override; + +private: + // Per-term BM25 parameters (extracted from TermDocIterator) + struct TermBM25Params { + BlockMaxLeafIterator *leaf_iter = nullptr; + float f1 = 0.0f; + float f2 = 0.0f; + float bm25_common_score = 0.0f; + float score_upper_bound = 0.0f; + }; + + // Accumulator entry for a single document + struct AccumEntry { + float score = 0.0f; + u32 match_count = 0; + }; + + // Sort term indices by IDF descending (rare terms first) + void SortTermsByIDF(); + + // Process a single term's entire posting list, accumulating scores + void ProcessTerm(u32 term_idx); + + // Extract top-k results from the accumulator into a sorted list + void ExtractTopK(); + + // Top-k + u32 topn_; + + // Term BM25 parameters (pointer into children_) + std::vector term_params_; + + // Term order sorted by IDF descending + std::vector sorted_term_order_; + + // Sparse accumulator: doc_id -> (score, match_count) + std::unordered_map accumulator_; + + // Final top-k results (doc_id, score) sorted by score descending + std::vector> topk_results_; + + // Current iteration cursor into topk_results_ + size_t topk_cursor_ = 0; + + // BM25 score cache for current doc + mutable float bm25_score_cache_ = 0.0f; + bool score_cached_ = false; + + // Debug stats + u64 total_docs_processed_ = 0; + u64 total_blocks_skipped_ = 0; + u64 total_blocks_decoded_ = 0; +}; + +} // namespace infinity diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp new file mode 100644 index 0000000000..7848029f96 --- /dev/null +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -0,0 +1,261 @@ +// Copyright(C) 2023 InfiniFlow, Inc. All rights reserved. +// +// 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 +// +// https://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. + +module; + +#include +#include +#include + +module infinity_core:block_at_a_time_iterator.impl; + +import :block_at_a_time_iterator; +import :index_defines; +import :doc_iterator; +import :blockmax_leaf_iterator; +import :term_doc_iterator; +import :multi_doc_iterator; +import :logger; +import :infinity_exception; +import :simd_functions; + +import third_party; +import internal_types; + +namespace infinity { + +BlockAtATimeIterator::BlockAtATimeIterator(std::vector> &&iterators, const u32 topn) + : MultiDocIterator(std::move(iterators)), topn_(topn) { + bm25_score_upper_bound_ = 0.0f; + estimate_iterate_cost_ = {}; + size_t num_iterators = children_.size(); + term_params_.reserve(num_iterators); + for (size_t i = 0; i < num_iterators; i++) { + BlockMaxLeafIterator *leaf = dynamic_cast(children_[i].get()); + if (leaf == nullptr) { + UnrecoverableError("BlockAtATimeIterator only supports BlockMaxLeafIterator"); + } + TermBM25Params params; + params.leaf_iter = leaf; + // Extract BM25 parameters (f1, f2, bm25_common_score) from the TermDocIterator + const auto tdi = dynamic_cast(leaf); + if (tdi != nullptr) { + const auto [f1, f2, bm25_common_score] = tdi->Get_f1_f2_bm25_common_score(); + params.f1 = f1; + params.f2 = f2; + params.bm25_common_score = bm25_common_score; + } else { + // PhraseDocIterator fallback: use BlockMaxLeafIterator upper bound + params.f1 = 0.0f; + params.f2 = 0.0f; + params.bm25_common_score = leaf->BM25ScoreUpperBound(); + } + params.score_upper_bound = leaf->BM25ScoreUpperBound(); + bm25_score_upper_bound_ += params.score_upper_bound; + estimate_iterate_cost_ += leaf->GetEstimateIterateCost(); + term_params_.emplace_back(std::move(params)); + } + SortTermsByIDF(); + topk_results_.reserve(topn_); +} + +BlockAtATimeIterator::~BlockAtATimeIterator() { + if (SHOULD_LOG_TRACE()) { + LOG_TRACE(fmt::format("BlockAtATimeIterator stats: docs_processed={}, blocks_skipped={}, blocks_decoded={}", + total_docs_processed_, + total_blocks_skipped_, + total_blocks_decoded_)); + } +} + +void BlockAtATimeIterator::SortTermsByIDF() { + const size_t n = term_params_.size(); + sorted_term_order_.resize(n); + std::iota(sorted_term_order_.begin(), sorted_term_order_.end(), 0u); + + // Sort by IDF descending: rare terms first establish high threshold faster. + // IDF ≈ log(1 + (N - df + 0.5) / (df + 0.5)) + // score_upper_bound is proportional to IDF * (k1 + 1), so we sort by it. + std::sort(sorted_term_order_.begin(), sorted_term_order_.end(), [this](u32 a, u32 b) { + return term_params_[a].score_upper_bound > term_params_[b].score_upper_bound; + }); +} + +void BlockAtATimeIterator::UpdateScoreThreshold(const float threshold) { + if (threshold > threshold_) { + threshold_ = threshold; + // Propagate threshold to children for their internal BlockMax pruning + for (auto &p : term_params_) { + p.leaf_iter->UpdateScoreThreshold(threshold); + } + } +} + +bool BlockAtATimeIterator::Next(RowID doc_id) { + if (doc_id_ != INVALID_ROWID && doc_id_ >= doc_id) { + score_cached_ = false; + return true; + } + // If topk_results_ is not yet populated, run the full BAAT pipeline + if (topk_results_.empty()) { + // Initialize all term iterators + for (size_t i = 0; i < term_params_.size(); i++) { + term_params_[i].leaf_iter->Next(0); + } + + // Run block-at-a-time for all terms + for (u32 sort_idx = 0; sort_idx < sorted_term_order_.size(); ++sort_idx) { + const u32 term_idx = sorted_term_order_[sort_idx]; + ProcessTerm(term_idx); + } + + // Extract top-k from the accumulator + ExtractTopK(); + topk_cursor_ = 0; + } + + // Iterate through topk_results_ starting from the requested doc_id + while (topk_cursor_ < topk_results_.size()) { + const auto &[candidate_doc, candidate_score] = topk_results_[topk_cursor_]; + ++topk_cursor_; + if (candidate_doc >= doc_id && candidate_score > threshold_) { + doc_id_ = candidate_doc; + bm25_score_cache_ = candidate_score; + score_cached_ = true; + return true; + } + } + + doc_id_ = INVALID_ROWID; + return false; +} + +void BlockAtATimeIterator::ProcessTerm(const u32 term_idx) { + auto *leaf = term_params_[term_idx].leaf_iter; + + RowID target = 0; + while (true) { + // Advance the iterator. If its own threshold filters it out, we still need to + // process filtered docs because the accumulator may have partial scores from + // previously processed (rarer) terms. Set leaf threshold to 0 for full traversal. + leaf->UpdateScoreThreshold(0.0f); + + if (!leaf->Next(target)) { + break; // exhausted + } + + // --- Block-level pruning --- + RowID block_last = leaf->BlockLastDocID(); + RowID block_min = leaf->BlockMinPossibleDocID(); + + // If block_min is INVALID_ROWID (no block loaded), skip + if (block_min == INVALID_ROWID) { + break; + } + + float block_max_score = leaf->BlockMaxBM25Score(); + + // Compute the best possible additional score from remaining (unprocessed) terms + // for docs in accumulator. This is a conservative estimate - we use the + // current block's max score as the potential contribution. + if (threshold_ > 0.0f && block_max_score <= threshold_ * 0.01f) { + // Block max is negligible relative to threshold, skip entire block + ++total_blocks_skipped_; + target = block_last + 1; + continue; + } + + ++total_blocks_decoded_; + + // --- Process all docs in this block --- + RowID current_doc = leaf->DocID(); + while (current_doc <= block_last) { + ++total_docs_processed_; + const float score = leaf->Score(); + + // Accumulate + auto it = accumulator_.find(current_doc); + if (it != accumulator_.end()) { + it->second.score += score; + it->second.match_count += 1; + } else { + accumulator_.emplace(current_doc, AccumEntry{score, 1}); + } + + // Advance to next doc in this block + target = current_doc + 1; + if (!leaf->Next(target)) { + return; // exhausted + } + current_doc = leaf->DocID(); + if (current_doc > block_last) { + target = current_doc; + break; + } + } + } +} + +void BlockAtATimeIterator::ExtractTopK() { + topk_results_.clear(); + + // Minimum heap of size topn_ for efficient top-k extraction + auto cmp = [](const std::pair &a, const std::pair &b) { + return a.second > b.second; // min-heap by score + }; + std::priority_queue, std::vector>, decltype(cmp)> min_heap(cmp); + + for (const auto &[doc_id, entry] : accumulator_) { + if (entry.score <= threshold_) { + continue; + } + if (min_heap.size() < topn_) { + min_heap.emplace(doc_id, entry.score); + } else if (entry.score > min_heap.top().second) { + min_heap.pop(); + min_heap.emplace(doc_id, entry.score); + } + } + + // Extract from heap in descending score order + topk_results_.resize(min_heap.size()); + size_t idx = min_heap.size(); + while (!min_heap.empty()) { + --idx; + topk_results_[idx] = min_heap.top(); + min_heap.pop(); + } + + LOG_TRACE(fmt::format("BlockAtATimeIterator::ExtractTopK: accumulator size={}, topk size={}", accumulator_.size(), topk_results_.size())); +} + +float BlockAtATimeIterator::BM25Score() const { + if (score_cached_) { + return bm25_score_cache_; + } + return 0.0f; +} + +u32 BlockAtATimeIterator::MatchCount() const { + if (doc_id_ == INVALID_ROWID) { + return 0; + } + auto it = accumulator_.find(doc_id_); + if (it != accumulator_.end()) { + return it->second.match_count; + } + return 0; +} + +} // namespace infinity diff --git a/src/storage/invertedindex/search/doc_iterator.cppm b/src/storage/invertedindex/search/doc_iterator.cppm index 1bb14f1541..df018e2cce 100644 --- a/src/storage/invertedindex/search/doc_iterator.cppm +++ b/src/storage/invertedindex/search/doc_iterator.cppm @@ -21,11 +21,12 @@ import internal_types; namespace infinity { export enum class EarlyTermAlgo { - kAuto, // choose between kNaive, kBatch, kBMW - kNaive, // naive or - kBatch, // use batch_or if (sum_of_df > total_doc_num / 4) and term nodes under or node achieve a certain number - kBMW, // use bmw if it is "or iterator" on the top level and has only term children - kCompare, // compare bmw, batch, naive + kAuto, // choose between kNaive, kBatch, kBMW, kBlockAtATime + kNaive, // naive or + kBatch, // use batch_or if (sum_of_df > total_doc_num / 4) and term nodes under or node achieve a certain number + kBMW, // use bmw if it is "or iterator" on the top level and has only term children + kBlockAtATime, // use block-at-a-time for long queries (>= 8 terms with high DF) + kCompare, // compare bmw, batch, naive }; export enum class DocIteratorType : u8 { @@ -37,6 +38,7 @@ export enum class DocIteratorType : u8 { kMinimumShouldMatchIterator, kBatchOrIterator, kBMWIterator, + kBlockAtATimeIterator, kFilterIterator, kScoreThresholdIterator, kKeywordIterator, diff --git a/src/storage/invertedindex/search/query_node_impl.cpp b/src/storage/invertedindex/search/query_node_impl.cpp index 1c1bd3151c..d0c8707b28 100644 --- a/src/storage/invertedindex/search/query_node_impl.cpp +++ b/src/storage/invertedindex/search/query_node_impl.cpp @@ -25,6 +25,7 @@ import :parse_fulltext_options; import :keyword_iterator; import :must_first_iterator; import :batch_or_iterator; +import :block_at_a_time_iterator; import :blockmax_leaf_iterator; import :rank_feature_doc_iterator; import :rank_features_doc_iterator; @@ -616,6 +617,32 @@ std::unique_ptr OrQueryNode::CreateSearch(const CreateSearchParams } return df_sum && (df_sum * 5ull >= total_df); }; + // BAAT condition: long query (>= 8 terms) with many common terms (avg DF > 10% total docs) + auto term_children_need_baat = [&sub_doc_iters, ¶ms]() -> bool { + constexpr u32 kMinBaatTermCount = 8; + if (params.topn == 0u || sub_doc_iters.size() < kMinBaatTermCount) { + return false; + } + u64 total_df = 0u; + u64 df_sum = 0u; + u32 term_count = 0u; + for (const auto &iter : sub_doc_iters) { + if (iter->GetType() == DocIteratorType::kTermDocIterator) { + const auto tdi = static_cast(iter.get()); + if (term_count == 0u) { + total_df = tdi->GetTotalDF(); + } + df_sum += tdi->GetDocFreq(); + term_count++; + } + } + if (term_count < kMinBaatTermCount) { + return false; + } + // Use BAAT when average DF > 10% of total docs (common terms dominate) + const float avg_df = static_cast(df_sum) / static_cast(term_count); + return avg_df > static_cast(total_df) * 0.1f; + }; if (sub_doc_iters.empty() && keyword_iters.empty()) { return nullptr; } @@ -626,8 +653,11 @@ std::unique_ptr OrQueryNode::CreateSearch(const CreateSearchParams auto choose_algo = EarlyTermAlgo::kNaive; switch (params.early_term_algo) { case EarlyTermAlgo::kAuto: { - if (params.topn) { - // always prefer BMW + if (params.topn && term_children_need_baat()) { + // Long query with common terms: block-at-a-time mode + choose_algo = EarlyTermAlgo::kBlockAtATime; + } else if (params.topn) { + // Short query: use BMW (DAAT with block max pruning) choose_algo = EarlyTermAlgo::kBMW; } else if (term_children_need_batch()) { // topn == 0, case of filter @@ -635,19 +665,11 @@ std::unique_ptr OrQueryNode::CreateSearch(const CreateSearchParams } else { choose_algo = EarlyTermAlgo::kNaive; } - /* TODO: now always use BMW - if ((params.topn == 0u || sub_doc_iters.size() > term_num_threshold(params.topn)) && term_children_need_batch()) { - choose_algo = EarlyTermAlgo::kBatch; - } else if (params.topn == 0u) { - choose_algo = EarlyTermAlgo::kNaive; - } else { - choose_algo = EarlyTermAlgo::kBMW; - } - */ break; } case EarlyTermAlgo::kBMW: case EarlyTermAlgo::kBatch: + case EarlyTermAlgo::kBlockAtATime: case EarlyTermAlgo::kNaive: { choose_algo = params.early_term_algo; break; @@ -661,6 +683,12 @@ std::unique_ptr OrQueryNode::CreateSearch(const CreateSearchParams case EarlyTermAlgo::kBMW: { return GetIterResultT.template operator()(); } + case EarlyTermAlgo::kBlockAtATime: { + assert(all_are_term_or_phrase); + assert(params.topn > 0u); + // BAAT works with term doc iterators; they are already in sub_doc_iters + return std::make_unique(std::move(sub_doc_iters), params.topn); + } case EarlyTermAlgo::kNaive: { return GetIterResultT.template operator()(); } From 7f4e7057a0d8efd06de7ec0045d8ca9ea7ab5983 Mon Sep 17 00:00:00 2001 From: Yingfeng Zhang Date: Fri, 5 Jun 2026 00:17:14 +0800 Subject: [PATCH 2/6] Improve --- .../search/block_at_a_time_iterator.cppm | 59 +-- .../search/block_at_a_time_iterator_impl.cpp | 336 ++++++++++++------ 2 files changed, 257 insertions(+), 138 deletions(-) diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator.cppm b/src/storage/invertedindex/search/block_at_a_time_iterator.cppm index 5c5504c533..99b3aa156f 100644 --- a/src/storage/invertedindex/search/block_at_a_time_iterator.cppm +++ b/src/storage/invertedindex/search/block_at_a_time_iterator.cppm @@ -17,6 +17,7 @@ export module infinity_core:block_at_a_time_iterator; import :index_defines; import :doc_iterator; import :blockmax_leaf_iterator; +import :term_doc_iterator; import :multi_doc_iterator; import :simd_functions; @@ -26,14 +27,17 @@ import internal_types; namespace infinity { // Block-at-a-time (BAAT) iterator for long queries with many common terms. -// Processes one term's posting list at a time, accumulating scores in a sparse -// accumulator, then extracts top-k results. // -// Reference: Hornet engine BAAT traversal for agent workloads (long queries). -// BAAT is selected when query has >= 8 term children and average DF > 10% of total docs. -// In BAAT, terms are processed in descending IDF order (rare terms first) to -// quickly establish a high threshold, enabling aggressive block-level pruning -// for subsequent common terms. +// Processing order: terms sorted by IDF descending (rare first). +// - Rare terms → few docs, small accumulator → quickly establish top-k threshold. +// - Common terms → long posting lists → block-level pruning using threshold. +// +// For TermDocIterator children, per-block SIMD BM25 batch computation is used +// (reusing GetSIMD_FUNCTIONS().BatchBM25_func_ptr_) to compute all document +// scores in a block with a single AVX2-vectorized call. +// +// A running min-heap maintains the current top-k threshold during accumulation, +// updated after each term so subsequent terms can skip irrelevant blocks. export class BlockAtATimeIterator : public MultiDocIterator { public: explicit BlockAtATimeIterator(std::vector> &&iterators, u32 topn); @@ -54,10 +58,13 @@ public: u32 MatchCount() const override; + void PrintTree(std::ostream &os, const std::string &prefix, bool is_final) const override; + private: - // Per-term BM25 parameters (extracted from TermDocIterator) + // Per-term BM25 parameters (extracted from the child iterator) struct TermBM25Params { BlockMaxLeafIterator *leaf_iter = nullptr; + bool is_term_doc = false; // true → TermDocIterator, supports SIMD batch float f1 = 0.0f; float f2 = 0.0f; float bm25_common_score = 0.0f; @@ -70,38 +77,44 @@ private: u32 match_count = 0; }; - // Sort term indices by IDF descending (rare terms first) - void SortTermsByIDF(); + // Min-heap comparator by score (used for running top-k) + struct MinScore { + bool operator()(const std::pair &a, const std::pair &b) const { return a.second > b.second; } + }; - // Process a single term's entire posting list, accumulating scores + void SortTermsByIDF(); void ProcessTerm(u32 term_idx); - - // Extract top-k results from the accumulator into a sorted list + void ProcessBlockSIMD(TermDocIterator *tdi, RowID block_min, RowID block_last, u32 term_idx); + void AllocateBatchBuffers(); + void AccumulateDoc(RowID doc_id, float score, u32 match); + void UpdateThresholdFromAccumulator(); void ExtractTopK(); - // Top-k u32 topn_; - // Term BM25 parameters (pointer into children_) std::vector term_params_; - - // Term order sorted by IDF descending std::vector sorted_term_order_; - // Sparse accumulator: doc_id -> (score, match_count) + // Sparse accumulator: doc_id → (score, match_count) std::unordered_map accumulator_; - // Final top-k results (doc_id, score) sorted by score descending - std::vector> topk_results_; + // Running min-heap of top-k candidates; updated incrementally during accumulation + std::priority_queue, std::vector>, MinScore> running_topk_heap_; - // Current iteration cursor into topk_results_ + // Final sorted results + std::vector> topk_results_; size_t topk_cursor_ = 0; - // BM25 score cache for current doc mutable float bm25_score_cache_ = 0.0f; bool score_cached_ = false; - // Debug stats + // SIMD batch buffers (64-byte aligned, BAAT_BATCH_SIZE entries each) + void *aligned_buffer_ = nullptr; + u32 *batch_tf_ = nullptr; + u32 *batch_doc_len_ = nullptr; + u32 *batch_match_cnt_ = nullptr; + f32 *batch_score_sum_ = nullptr; + u64 total_docs_processed_ = 0; u64 total_blocks_skipped_ = 0; u64 total_blocks_decoded_ = 0; diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp index 7848029f96..261d2a0bd5 100644 --- a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -35,58 +35,73 @@ import internal_types; namespace infinity { +static constexpr u32 BAAT_BATCH_SIZE = MAX_DOC_PER_RECORD; // 128, already 8-aligned for AVX2 + BlockAtATimeIterator::BlockAtATimeIterator(std::vector> &&iterators, const u32 topn) : MultiDocIterator(std::move(iterators)), topn_(topn) { bm25_score_upper_bound_ = 0.0f; estimate_iterate_cost_ = {}; - size_t num_iterators = children_.size(); - term_params_.reserve(num_iterators); - for (size_t i = 0; i < num_iterators; i++) { - BlockMaxLeafIterator *leaf = dynamic_cast(children_[i].get()); + term_params_.reserve(children_.size()); + bool needs_simd_buffer = false; + for (auto &child : children_) { + auto *leaf = dynamic_cast(child.get()); if (leaf == nullptr) { UnrecoverableError("BlockAtATimeIterator only supports BlockMaxLeafIterator"); } TermBM25Params params; params.leaf_iter = leaf; - // Extract BM25 parameters (f1, f2, bm25_common_score) from the TermDocIterator - const auto tdi = dynamic_cast(leaf); - if (tdi != nullptr) { - const auto [f1, f2, bm25_common_score] = tdi->Get_f1_f2_bm25_common_score(); - params.f1 = f1; - params.f2 = f2; - params.bm25_common_score = bm25_common_score; + params.score_upper_bound = leaf->BM25ScoreUpperBound(); + bm25_score_upper_bound_ += params.score_upper_bound; + estimate_iterate_cost_ += leaf->GetEstimateIterateCost(); + if (auto *tdi = dynamic_cast(leaf)) { + params.is_term_doc = true; + needs_simd_buffer = true; + std::tie(params.f1, params.f2, params.bm25_common_score) = tdi->Get_f1_f2_bm25_common_score(); } else { - // PhraseDocIterator fallback: use BlockMaxLeafIterator upper bound params.f1 = 0.0f; params.f2 = 0.0f; params.bm25_common_score = leaf->BM25ScoreUpperBound(); } - params.score_upper_bound = leaf->BM25ScoreUpperBound(); - bm25_score_upper_bound_ += params.score_upper_bound; - estimate_iterate_cost_ += leaf->GetEstimateIterateCost(); term_params_.emplace_back(std::move(params)); } + if (needs_simd_buffer) { + AllocateBatchBuffers(); + } SortTermsByIDF(); topk_results_.reserve(topn_); } BlockAtATimeIterator::~BlockAtATimeIterator() { + std::free(aligned_buffer_); if (SHOULD_LOG_TRACE()) { - LOG_TRACE(fmt::format("BlockAtATimeIterator stats: docs_processed={}, blocks_skipped={}, blocks_decoded={}", + LOG_TRACE(fmt::format("BlockAtATimeIterator: docs={} blocks_skip={} blocks_dec={} thr={}", total_docs_processed_, total_blocks_skipped_, - total_blocks_decoded_)); + total_blocks_decoded_, + threshold_)); } } +void BlockAtATimeIterator::AllocateBatchBuffers() { + const u32 buf_size = BAAT_BATCH_SIZE * (sizeof(u32) + sizeof(u32) + sizeof(u32) + sizeof(f32)); + aligned_buffer_ = std::aligned_alloc(64, buf_size); + if (!aligned_buffer_) { + UnrecoverableError(fmt::format("{}: Out of memory!", __func__)); + } + auto *base = static_cast(aligned_buffer_); + batch_tf_ = reinterpret_cast(base); + base += BAAT_BATCH_SIZE * sizeof(u32); + batch_doc_len_ = reinterpret_cast(base); + base += BAAT_BATCH_SIZE * sizeof(u32); + batch_match_cnt_ = reinterpret_cast(base); + base += BAAT_BATCH_SIZE * sizeof(u32); + batch_score_sum_ = reinterpret_cast(base); +} + void BlockAtATimeIterator::SortTermsByIDF() { - const size_t n = term_params_.size(); - sorted_term_order_.resize(n); + sorted_term_order_.resize(term_params_.size()); std::iota(sorted_term_order_.begin(), sorted_term_order_.end(), 0u); - - // Sort by IDF descending: rare terms first establish high threshold faster. - // IDF ≈ log(1 + (N - df + 0.5) / (df + 0.5)) - // score_upper_bound is proportional to IDF * (k1 + 1), so we sort by it. + // Descending score_upper_bound → rare terms first std::sort(sorted_term_order_.begin(), sorted_term_order_.end(), [this](u32 a, u32 b) { return term_params_[a].score_upper_bound > term_params_[b].score_upper_bound; }); @@ -95,167 +110,258 @@ void BlockAtATimeIterator::SortTermsByIDF() { void BlockAtATimeIterator::UpdateScoreThreshold(const float threshold) { if (threshold > threshold_) { threshold_ = threshold; - // Propagate threshold to children for their internal BlockMax pruning for (auto &p : term_params_) { p.leaf_iter->UpdateScoreThreshold(threshold); } } } +// Rebuild running top-k heap from current accumulator, then update threshold_. +// Called after each term so subsequent terms benefit from block-level pruning. +void BlockAtATimeIterator::UpdateThresholdFromAccumulator() { + if (accumulator_.size() < topn_) { + return; + } + // Clear and rebuild + while (!running_topk_heap_.empty()) { + running_topk_heap_.pop(); + } + for (const auto &[doc_id, entry] : accumulator_) { + if (entry.score <= threshold_) + continue; + if (running_topk_heap_.size() < topn_) { + running_topk_heap_.emplace(doc_id, entry.score); + } else if (entry.score > running_topk_heap_.top().second) { + running_topk_heap_.pop(); + running_topk_heap_.emplace(doc_id, entry.score); + } + } + if (running_topk_heap_.size() >= topn_) { + threshold_ = running_topk_heap_.top().second; + } +} + bool BlockAtATimeIterator::Next(RowID doc_id) { if (doc_id_ != INVALID_ROWID && doc_id_ >= doc_id) { score_cached_ = false; return true; } - // If topk_results_ is not yet populated, run the full BAAT pipeline if (topk_results_.empty()) { - // Initialize all term iterators - for (size_t i = 0; i < term_params_.size(); i++) { - term_params_[i].leaf_iter->Next(0); + // Run BAAT pipeline once + while (!running_topk_heap_.empty()) { + running_topk_heap_.pop(); } - // Run block-at-a-time for all terms + for (auto &p : term_params_) { + p.leaf_iter->Next(0); + } + + // Process terms in IDF order; after each term, update threshold for (u32 sort_idx = 0; sort_idx < sorted_term_order_.size(); ++sort_idx) { - const u32 term_idx = sorted_term_order_[sort_idx]; - ProcessTerm(term_idx); + ProcessTerm(sorted_term_order_[sort_idx]); + // After processing a rare term, update threshold so subsequent + // common terms can skip entire blocks. + if (sort_idx < sorted_term_order_.size() - 1) { + UpdateThresholdFromAccumulator(); + } } - // Extract top-k from the accumulator ExtractTopK(); topk_cursor_ = 0; } - - // Iterate through topk_results_ starting from the requested doc_id while (topk_cursor_ < topk_results_.size()) { - const auto &[candidate_doc, candidate_score] = topk_results_[topk_cursor_]; - ++topk_cursor_; - if (candidate_doc >= doc_id && candidate_score > threshold_) { - doc_id_ = candidate_doc; - bm25_score_cache_ = candidate_score; + const auto &[cd, cs] = topk_results_[topk_cursor_++]; + if (cd >= doc_id && cs > threshold_) { + doc_id_ = cd; + bm25_score_cache_ = cs; score_cached_ = true; return true; } } - doc_id_ = INVALID_ROWID; return false; } +// ------------------------------------------------------------------ +// ProcessTerm: process one term's entire posting list +// ------------------------------------------------------------------ void BlockAtATimeIterator::ProcessTerm(const u32 term_idx) { - auto *leaf = term_params_[term_idx].leaf_iter; + auto ¶ms = term_params_[term_idx]; + auto *leaf = params.leaf_iter; + auto *tdi = params.is_term_doc ? dynamic_cast(leaf) : nullptr; + + // Do not use the term-level threshold filter — we want ALL docs from this term + // because even a low individual score adds to the accumulated total. + leaf->UpdateScoreThreshold(0.0f); + + // The global threshold_ is used only for block-level pruning below. + // A block whose max contribution can't meaningfully affect top-k candidates + // (even when combined with accumulated scores from earlier terms) is skipped. RowID target = 0; while (true) { - // Advance the iterator. If its own threshold filters it out, we still need to - // process filtered docs because the accumulator may have partial scores from - // previously processed (rarer) terms. Set leaf threshold to 0 for full traversal. - leaf->UpdateScoreThreshold(0.0f); - if (!leaf->Next(target)) { - break; // exhausted + break; } - - // --- Block-level pruning --- RowID block_last = leaf->BlockLastDocID(); RowID block_min = leaf->BlockMinPossibleDocID(); - - // If block_min is INVALID_ROWID (no block loaded), skip if (block_min == INVALID_ROWID) { break; } - float block_max_score = leaf->BlockMaxBM25Score(); - - // Compute the best possible additional score from remaining (unprocessed) terms - // for docs in accumulator. This is a conservative estimate - we use the - // current block's max score as the potential contribution. - if (threshold_ > 0.0f && block_max_score <= threshold_ * 0.01f) { - // Block max is negligible relative to threshold, skip entire block + // --- Block-level pruning --- + // threshold_ > 0 means we have top-k candidates from earlier (rarer) terms. + // A doc needs accumulated + this_block_max > threshold_ to enter top-k. + // Conservative: require block_max_score > threshold_ * 0.1f so that + // this term's contribution is at least 10% of threshold (the rest must + // come from earlier terms' accumulated scores already in the accumulator). + const float block_max_score = leaf->BlockMaxBM25Score(); + if (threshold_ > 0.0f && block_max_score * 10.0f <= threshold_) { ++total_blocks_skipped_; target = block_last + 1; continue; } - ++total_blocks_decoded_; - // --- Process all docs in this block --- - RowID current_doc = leaf->DocID(); - while (current_doc <= block_last) { - ++total_docs_processed_; - const float score = leaf->Score(); - - // Accumulate - auto it = accumulator_.find(current_doc); - if (it != accumulator_.end()) { - it->second.score += score; - it->second.match_count += 1; - } else { - accumulator_.emplace(current_doc, AccumEntry{score, 1}); + if (tdi != nullptr && batch_tf_ != nullptr) { + ProcessBlockSIMD(tdi, block_min, block_last, term_idx); + target = block_last + 1; + } else { + // Scalar fallback (PhraseDocIterator etc.) + RowID current = leaf->DocID(); + while (current <= block_last) { + ++total_docs_processed_; + const float score = leaf->Score(); + AccumulateDoc(current, score, 1); + target = current + 1; + if (!leaf->Next(target)) + return; + current = leaf->DocID(); } + } + } +} + +// ------------------------------------------------------------------ +// ProcessBlockSIMD: decode one entire block with BatchDecodeTo, +// compute all scores with SIMD BatchBM25, then accumulate. +// ------------------------------------------------------------------ +void BlockAtATimeIterator::ProcessBlockSIMD(TermDocIterator *tdi, const RowID block_min, const RowID block_last, const u32 term_idx) { + const u32 batch_len = static_cast(block_last - block_min) + 1; + const u32 aligned_len = ((batch_len + 7) / 8) * 8; // round up to 8 for AVX2 + + std::memset(batch_tf_, 0, aligned_len * sizeof(u32)); + std::memset(batch_doc_len_, 0, aligned_len * sizeof(u32)); + std::memset(batch_match_cnt_, 0, aligned_len * sizeof(u32)); + std::memset(batch_score_sum_, 0, aligned_len * sizeof(f32)); + + // BatchDecodeTo iterates SeekDoc from current position to block_last+1, + // filling tf and doc_len at offset (iter_doc_id - block_min) for each match. + tdi->BatchDecodeTo(block_min, block_last + 1, batch_tf_, batch_doc_len_); + + // SIMD batch BM25 for single keyword (children_num = 1) + const auto &tp = term_params_[term_idx]; + GetSIMD_FUNCTIONS() + .BatchBM25_func_ptr_(aligned_len, 1u, &tp.f1, &tp.f2, &tp.bm25_common_score, batch_tf_, batch_doc_len_, batch_match_cnt_, batch_score_sum_); + + // Accumulate + for (u32 i = 0; i < batch_len; ++i) { + if (batch_match_cnt_[i] > 0) { + const RowID doc_id = block_min + i; + const float score = batch_score_sum_[i]; + ++total_docs_processed_; + AccumulateDoc(doc_id, score, 1); + } + } +} + +// ------------------------------------------------------------------ +// AccumulateDoc: add score to the accumulator's entry for doc_id, +// and keep the running top-k heap up to date for progressive threshold. +// ------------------------------------------------------------------ +void BlockAtATimeIterator::AccumulateDoc(const RowID doc_id, const float score, const u32 match) { + auto it = accumulator_.find(doc_id); + float new_score; + if (it != accumulator_.end()) { + it->second.score += score; + it->second.match_count += match; + new_score = it->second.score; + } else { + accumulator_.emplace(doc_id, AccumEntry{score, match}); + new_score = score; + } - // Advance to next doc in this block - target = current_doc + 1; - if (!leaf->Next(target)) { - return; // exhausted + // Maintain running top-k heap for progressive threshold + if (new_score > threshold_) { + running_topk_heap_.emplace(doc_id, new_score); + // Lazy cleanup: only pop stale entries when heap is oversized + if (running_topk_heap_.size() >= topn_ * 3u) { + while (running_topk_heap_.size() > topn_ && running_topk_heap_.top().second <= threshold_) { + running_topk_heap_.pop(); } - current_doc = leaf->DocID(); - if (current_doc > block_last) { - target = current_doc; - break; + } + if (running_topk_heap_.size() >= topn_) { + // Pop stale low-score entries from the top + while (running_topk_heap_.top().second <= threshold_) { + running_topk_heap_.pop(); + if (running_topk_heap_.empty()) + break; + } + if (running_topk_heap_.size() >= topn_) { + threshold_ = running_topk_heap_.top().second; } } } } +// ------------------------------------------------------------------ +// ExtractTopK: rebuild clean heap from accumulator and sort into vector. +// ------------------------------------------------------------------ void BlockAtATimeIterator::ExtractTopK() { topk_results_.clear(); - - // Minimum heap of size topn_ for efficient top-k extraction - auto cmp = [](const std::pair &a, const std::pair &b) { - return a.second > b.second; // min-heap by score - }; - std::priority_queue, std::vector>, decltype(cmp)> min_heap(cmp); - + while (!running_topk_heap_.empty()) { + running_topk_heap_.pop(); + } for (const auto &[doc_id, entry] : accumulator_) { - if (entry.score <= threshold_) { - continue; - } - if (min_heap.size() < topn_) { - min_heap.emplace(doc_id, entry.score); - } else if (entry.score > min_heap.top().second) { - min_heap.pop(); - min_heap.emplace(doc_id, entry.score); + if (entry.score > threshold_) { + if (running_topk_heap_.size() < topn_) { + running_topk_heap_.emplace(doc_id, entry.score); + } else if (entry.score > running_topk_heap_.top().second) { + running_topk_heap_.pop(); + running_topk_heap_.emplace(doc_id, entry.score); + } } } - // Extract from heap in descending score order - topk_results_.resize(min_heap.size()); - size_t idx = min_heap.size(); - while (!min_heap.empty()) { - --idx; - topk_results_[idx] = min_heap.top(); - min_heap.pop(); + topk_results_.resize(running_topk_heap_.size()); + auto idx = running_topk_heap_.size(); + while (!running_topk_heap_.empty()) { + topk_results_[--idx] = running_topk_heap_.top(); + running_topk_heap_.pop(); } - - LOG_TRACE(fmt::format("BlockAtATimeIterator::ExtractTopK: accumulator size={}, topk size={}", accumulator_.size(), topk_results_.size())); } -float BlockAtATimeIterator::BM25Score() const { - if (score_cached_) { - return bm25_score_cache_; - } - return 0.0f; -} +float BlockAtATimeIterator::BM25Score() const { return score_cached_ ? bm25_score_cache_ : 0.0f; } u32 BlockAtATimeIterator::MatchCount() const { - if (doc_id_ == INVALID_ROWID) { + if (doc_id_ == INVALID_ROWID) return 0; - } auto it = accumulator_.find(doc_id_); - if (it != accumulator_.end()) { - return it->second.match_count; + return (it != accumulator_.end()) ? it->second.match_count : 0; +} + +void BlockAtATimeIterator::PrintTree(std::ostream &os, const std::string &prefix, bool is_final) const { + os << prefix; + os << (is_final ? "└──" : "├──"); + os << "BlockAtATimeIterator (topn: " << topn_ << ", threshold: " << threshold_ << ")"; + os << '\n'; + const std::string next_prefix = prefix + (is_final ? " " : "│ "); + for (u32 i = 0; i + 1 < children_.size(); ++i) { + children_[i]->PrintTree(os, next_prefix, false); + } + if (!children_.empty()) { + children_.back()->PrintTree(os, next_prefix, true); } - return 0; } } // namespace infinity From 16fff61dfa3b5032b475ed93fbaccf3fb147a3f1 Mon Sep 17 00:00:00 2001 From: Yingfeng Zhang Date: Fri, 5 Jun 2026 00:22:09 +0800 Subject: [PATCH 3/6] fix --- .../search/block_at_a_time_iterator_impl.cpp | 12 ++++++++++-- src/storage/invertedindex/search/query_node_impl.cpp | 5 ++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp index 261d2a0bd5..d84ea5b17a 100644 --- a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -167,11 +167,19 @@ bool BlockAtATimeIterator::Next(RowID doc_id) { } ExtractTopK(); + // Sort by doc_id ascending so linear cursor scan returns docs in increasing order. + std::sort(topk_results_.begin(), topk_results_.end(), + [](const auto &a, const auto &b) { return a.first < b.first; }); topk_cursor_ = 0; } while (topk_cursor_ < topk_results_.size()) { - const auto &[cd, cs] = topk_results_[topk_cursor_++]; - if (cd >= doc_id && cs > threshold_) { + const auto &[cd, cs] = topk_results_[topk_cursor_]; + if (cd < doc_id) { + ++topk_cursor_; + continue; + } + ++topk_cursor_; + if (cs > threshold_) { doc_id_ = cd; bm25_score_cache_ = cs; score_cached_ = true; diff --git a/src/storage/invertedindex/search/query_node_impl.cpp b/src/storage/invertedindex/search/query_node_impl.cpp index d0c8707b28..6444fdef6e 100644 --- a/src/storage/invertedindex/search/query_node_impl.cpp +++ b/src/storage/invertedindex/search/query_node_impl.cpp @@ -685,7 +685,10 @@ std::unique_ptr OrQueryNode::CreateSearch(const CreateSearchParams } case EarlyTermAlgo::kBlockAtATime: { assert(all_are_term_or_phrase); - assert(params.topn > 0u); + if (params.topn == 0u) [[unlikely]] { + // BAAT requires topn > 0 for top-k extraction; fall back to naive OR. + return GetIterResultT.template operator()(); + } // BAAT works with term doc iterators; they are already in sub_doc_iters return std::make_unique(std::move(sub_doc_iters), params.topn); } From 1259f36787aecdb407d09168075defdaded61496 Mon Sep 17 00:00:00 2001 From: Yingfeng Zhang Date: Fri, 5 Jun 2026 00:27:54 +0800 Subject: [PATCH 4/6] fix --- .../search/block_at_a_time_iterator_impl.cpp | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp index d84ea5b17a..8d61ebdd23 100644 --- a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -74,9 +74,8 @@ BlockAtATimeIterator::BlockAtATimeIterator(std::vectorUpdateScoreThreshold(0.0f); - // The global threshold_ is used only for block-level pruning below. - // A block whose max contribution can't meaningfully affect top-k candidates - // (even when combined with accumulated scores from earlier terms) is skipped. - RowID target = 0; while (true) { if (!leaf->Next(target)) { @@ -217,18 +212,6 @@ void BlockAtATimeIterator::ProcessTerm(const u32 term_idx) { break; } - // --- Block-level pruning --- - // threshold_ > 0 means we have top-k candidates from earlier (rarer) terms. - // A doc needs accumulated + this_block_max > threshold_ to enter top-k. - // Conservative: require block_max_score > threshold_ * 0.1f so that - // this term's contribution is at least 10% of threshold (the rest must - // come from earlier terms' accumulated scores already in the accumulator). - const float block_max_score = leaf->BlockMaxBM25Score(); - if (threshold_ > 0.0f && block_max_score * 10.0f <= threshold_) { - ++total_blocks_skipped_; - target = block_last + 1; - continue; - } ++total_blocks_decoded_; if (tdi != nullptr && batch_tf_ != nullptr) { From 34fef79989263263083d4aa4d8b64680cd31cce4 Mon Sep 17 00:00:00 2001 From: Yingfeng Zhang Date: Fri, 5 Jun 2026 00:33:00 +0800 Subject: [PATCH 5/6] fix --- .../search/block_at_a_time_iterator_impl.cpp | 35 +++++++++++++------ 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp index 8d61ebdd23..2d394e3eb7 100644 --- a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -74,8 +74,9 @@ BlockAtATimeIterator::BlockAtATimeIterator(std::vector running_topk_heap_.top().second) { + } else if (entry.score >= running_topk_heap_.top().second) { running_topk_heap_.pop(); running_topk_heap_.emplace(doc_id, entry.score); } @@ -178,7 +179,7 @@ bool BlockAtATimeIterator::Next(RowID doc_id) { continue; } ++topk_cursor_; - if (cs > threshold_) { + if (cs >= threshold_) { doc_id_ = cd; bm25_score_cache_ = cs; score_cached_ = true; @@ -203,7 +204,8 @@ void BlockAtATimeIterator::ProcessTerm(const u32 term_idx) { RowID target = 0; while (true) { - if (!leaf->Next(target)) { + // Step 1: Advance block cursor with NextShallow (skiplist only, no doc decode) + if (!leaf->NextShallow(target)) { break; } RowID block_last = leaf->BlockLastDocID(); @@ -212,8 +214,21 @@ void BlockAtATimeIterator::ProcessTerm(const u32 term_idx) { break; } + // Step 2: Block-level pruning. + // If this term's max possible contribution in this block cannot reach + // the current top-k threshold, skip the entire block without decoding. + if (threshold_ > 0.0f && leaf->BlockMaxBM25Score() <= threshold_) { + ++total_blocks_skipped_; + target = block_last + 1; + continue; + } ++total_blocks_decoded_; + // Step 3: Decode first document in the block (also positions doc cursor) + if (!leaf->Next(target)) { + break; + } + if (tdi != nullptr && batch_tf_ != nullptr) { ProcessBlockSIMD(tdi, block_min, block_last, term_idx); target = block_last + 1; @@ -283,17 +298,17 @@ void BlockAtATimeIterator::AccumulateDoc(const RowID doc_id, const float score, } // Maintain running top-k heap for progressive threshold - if (new_score > threshold_) { + if (new_score >= threshold_) { running_topk_heap_.emplace(doc_id, new_score); // Lazy cleanup: only pop stale entries when heap is oversized if (running_topk_heap_.size() >= topn_ * 3u) { - while (running_topk_heap_.size() > topn_ && running_topk_heap_.top().second <= threshold_) { + while (running_topk_heap_.size() > topn_ && running_topk_heap_.top().second < threshold_) { running_topk_heap_.pop(); } } if (running_topk_heap_.size() >= topn_) { // Pop stale low-score entries from the top - while (running_topk_heap_.top().second <= threshold_) { + while (running_topk_heap_.top().second < threshold_) { running_topk_heap_.pop(); if (running_topk_heap_.empty()) break; @@ -314,10 +329,10 @@ void BlockAtATimeIterator::ExtractTopK() { running_topk_heap_.pop(); } for (const auto &[doc_id, entry] : accumulator_) { - if (entry.score > threshold_) { + if (entry.score >= threshold_) { if (running_topk_heap_.size() < topn_) { running_topk_heap_.emplace(doc_id, entry.score); - } else if (entry.score > running_topk_heap_.top().second) { + } else if (entry.score >= running_topk_heap_.top().second) { running_topk_heap_.pop(); running_topk_heap_.emplace(doc_id, entry.score); } From d8186e08e4882395c065b4c739e7bb742b83b29d Mon Sep 17 00:00:00 2001 From: Yingfeng Zhang Date: Fri, 5 Jun 2026 01:03:03 +0800 Subject: [PATCH 6/6] fix --- .../invertedindex/search/block_at_a_time_iterator_impl.cpp | 6 +++--- .../invertedindex/search/blockmax_leaf_iterator.cppm | 5 +++++ src/storage/invertedindex/search/phrase_doc_iterator.cppm | 2 ++ src/storage/invertedindex/search/term_doc_iterator.cppm | 2 ++ 4 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp index 2d394e3eb7..203340c649 100644 --- a/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -198,9 +198,9 @@ void BlockAtATimeIterator::ProcessTerm(const u32 term_idx) { auto *leaf = params.leaf_iter; auto *tdi = params.is_term_doc ? dynamic_cast(leaf) : nullptr; - // Do not use the term-level threshold filter — we want ALL docs from this term - // because even a low individual score adds to the accumulated total. - leaf->UpdateScoreThreshold(0.0f); + // Unconditionally disable child-side threshold filtering so all docs are seen. + // UpdateScoreThreshold is monotonic so it cannot lower an already-raised value. + leaf->ForceSetScoreThreshold(0.0f); RowID target = 0; while (true) { diff --git a/src/storage/invertedindex/search/blockmax_leaf_iterator.cppm b/src/storage/invertedindex/search/blockmax_leaf_iterator.cppm index 6efbd0e3f0..d2f85afcdc 100644 --- a/src/storage/invertedindex/search/blockmax_leaf_iterator.cppm +++ b/src/storage/invertedindex/search/blockmax_leaf_iterator.cppm @@ -39,6 +39,11 @@ public: virtual bool NextShallow(RowID doc_id) = 0; virtual float BM25Score() = 0; + + // Unconditionally set the internal score threshold (unlike UpdateScoreThreshold + // which is monotonic and only raises). Used by BlockAtATimeIterator during + // BAAT traversal to disable child-side filtering so all documents are seen. + virtual void ForceSetScoreThreshold(float threshold) = 0; }; } // namespace infinity diff --git a/src/storage/invertedindex/search/phrase_doc_iterator.cppm b/src/storage/invertedindex/search/phrase_doc_iterator.cppm index bee9b1e661..887e87c8e1 100644 --- a/src/storage/invertedindex/search/phrase_doc_iterator.cppm +++ b/src/storage/invertedindex/search/phrase_doc_iterator.cppm @@ -64,6 +64,8 @@ public: threshold_ = threshold; } + void ForceSetScoreThreshold(float threshold) override { threshold_ = threshold; } + u32 MatchCount() const override { return DocID() != INVALID_ROWID; } void PrintTree(std::ostream &os, const std::string &prefix, bool is_final) const override; diff --git a/src/storage/invertedindex/search/term_doc_iterator.cppm b/src/storage/invertedindex/search/term_doc_iterator.cppm index 796ce6b5db..c838a3c732 100644 --- a/src/storage/invertedindex/search/term_doc_iterator.cppm +++ b/src/storage/invertedindex/search/term_doc_iterator.cppm @@ -81,6 +81,8 @@ public: threshold_ = threshold; } + void ForceSetScoreThreshold(float threshold) override { threshold_ = threshold; } + u32 MatchCount() const override { return DocID() != INVALID_ROWID; } void PrintTree(std::ostream &os, const std::string &prefix, bool is_final) const override;