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..99b3aa156f --- /dev/null +++ b/src/storage/invertedindex/search/block_at_a_time_iterator.cppm @@ -0,0 +1,123 @@ +// 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 :term_doc_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. +// +// 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); + + ~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; + + void PrintTree(std::ostream &os, const std::string &prefix, bool is_final) const override; + +private: + // 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; + float score_upper_bound = 0.0f; + }; + + // Accumulator entry for a single document + struct AccumEntry { + float score = 0.0f; + u32 match_count = 0; + }; + + // 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; } + }; + + void SortTermsByIDF(); + void ProcessTerm(u32 term_idx); + 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(); + + u32 topn_; + + std::vector term_params_; + std::vector sorted_term_order_; + + // Sparse accumulator: doc_id → (score, match_count) + std::unordered_map accumulator_; + + // Running min-heap of top-k candidates; updated incrementally during accumulation + std::priority_queue, std::vector>, MinScore> running_topk_heap_; + + // Final sorted results + std::vector> topk_results_; + size_t topk_cursor_ = 0; + + mutable float bm25_score_cache_ = 0.0f; + bool score_cached_ = false; + + // 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; +}; + +} // 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..203340c649 --- /dev/null +++ b/src/storage/invertedindex/search/block_at_a_time_iterator_impl.cpp @@ -0,0 +1,373 @@ +// 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 { + +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_ = {}; + 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; + 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 { + params.f1 = 0.0f; + params.f2 = 0.0f; + params.bm25_common_score = leaf->BM25ScoreUpperBound(); + } + 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: docs={} blk_skip={} blk_dec={} thr={}", + total_docs_processed_, + total_blocks_skipped_, + 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() { + sorted_term_order_.resize(term_params_.size()); + std::iota(sorted_term_order_.begin(), sorted_term_order_.end(), 0u); + // 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; + }); +} + +void BlockAtATimeIterator::UpdateScoreThreshold(const float threshold) { + if (threshold > threshold_) { + threshold_ = threshold; + 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_.empty()) { + // Run BAAT pipeline once + while (!running_topk_heap_.empty()) { + running_topk_heap_.pop(); + } + + 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) { + 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(); + } + } + + 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) { + ++topk_cursor_; + continue; + } + ++topk_cursor_; + if (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 ¶ms = term_params_[term_idx]; + auto *leaf = params.leaf_iter; + auto *tdi = params.is_term_doc ? dynamic_cast(leaf) : nullptr; + + // 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) { + // Step 1: Advance block cursor with NextShallow (skiplist only, no doc decode) + if (!leaf->NextShallow(target)) { + break; + } + RowID block_last = leaf->BlockLastDocID(); + RowID block_min = leaf->BlockMinPossibleDocID(); + if (block_min == INVALID_ROWID) { + 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; + } 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; + } + + // 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(); + } + } + 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(); + while (!running_topk_heap_.empty()) { + running_topk_heap_.pop(); + } + for (const auto &[doc_id, entry] : accumulator_) { + 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); + } + } + } + + 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(); + } +} + +float BlockAtATimeIterator::BM25Score() const { return score_cached_ ? bm25_score_cache_ : 0.0f; } + +u32 BlockAtATimeIterator::MatchCount() const { + if (doc_id_ == INVALID_ROWID) + return 0; + auto it = accumulator_.find(doc_id_); + 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); + } +} + +} // namespace infinity 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/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/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/query_node_impl.cpp b/src/storage/invertedindex/search/query_node_impl.cpp index 1c1bd3151c..6444fdef6e 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,15 @@ std::unique_ptr OrQueryNode::CreateSearch(const CreateSearchParams case EarlyTermAlgo::kBMW: { return GetIterResultT.template operator()(); } + case EarlyTermAlgo::kBlockAtATime: { + assert(all_are_term_or_phrase); + 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); + } case EarlyTermAlgo::kNaive: { return GetIterResultT.template operator()(); } 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;