use crate::error::{LaurusError, Result};
use crate::lexical::index::inverted::reader::InvertedIndexReader;
use crate::lexical::index::inverted::searcher::Deadline;
use crate::lexical::query::Query;
use crate::lexical::query::boolean::{BooleanQuery, Occur};
use crate::lexical::query::collector::Collector;
use crate::lexical::query::term::TermQuery;
use crate::lexical::reader::LexicalIndexReader;
struct BmwClause {
scorer: crate::lexical::query::scorer::LeafScorer,
matcher: crate::lexical::query::matcher::LeafMatcher,
field_name: Option<String>,
}
pub(crate) struct BlockMaxOrExecutor<'r> {
clauses: Vec<BmwClause>,
inverted_reader: Option<&'r InvertedIndexReader>,
}
impl<'r> BlockMaxOrExecutor<'r> {
pub fn new(boolean_query: &BooleanQuery, reader: &'r dyn LexicalIndexReader) -> Result<Self> {
let mut clauses = Vec::with_capacity(boolean_query.clauses().len());
for clause in boolean_query.clauses() {
let scorer = clause.query.scorer(reader)?;
if scorer.next_block_boundary(0).is_none() {
return Err(LaurusError::InvalidOperation(
"BMW fast path requires per-block scorer for every clause".into(),
));
}
let matcher = clause.query.matcher(reader)?;
let field_name = field_name_of(clause.query.as_ref());
clauses.push(BmwClause {
scorer: crate::lexical::query::scorer::LeafScorer::from_box(scorer),
matcher: crate::lexical::query::matcher::LeafMatcher::from_box(matcher),
field_name,
});
}
let inverted_reader = reader.as_any().downcast_ref::<InvertedIndexReader>();
Ok(BlockMaxOrExecutor {
clauses,
inverted_reader,
})
}
pub fn run<C: Collector>(mut self, mut collector: C, deadline: Option<Deadline>) -> Result<C> {
let mut active: Vec<usize> = Vec::with_capacity(self.clauses.len());
for (i, c) in self.clauses.iter().enumerate() {
if !c.matcher.is_exhausted() && c.matcher.doc_id() != u64::MAX {
active.push(i);
}
}
let mut scanned: u64 = 0;
loop {
if let Some(d) = deadline {
d.check(scanned)?;
}
scanned = scanned.wrapping_add(1);
if active.is_empty() {
break;
}
active.sort_by_key(|&i| self.clauses[i].matcher.doc_id());
let min_comp = collector.min_competitive();
let mut sum = 0.0_f32;
let mut pivot_k: Option<usize> = None;
for (j, &i) in active.iter().enumerate() {
let doc_id = self.clauses[i].matcher.doc_id();
sum += self.clauses[i].scorer.current_block_max_score(doc_id);
if sum > min_comp {
pivot_k = Some(j);
break;
}
}
match pivot_k {
None => {
let mut cum_sum = 0.0_f32;
for &i in &active {
let d = self.clauses[i].matcher.doc_id();
cum_sum += self.clauses[i].scorer.block_max_score_at(d);
}
if cum_sum <= min_comp {
break;
}
let lead = active[0];
let lead_doc = self.clauses[lead].matcher.doc_id();
self.advance_clause_past_block(lead, lead_doc)?;
}
Some(k) => {
let pivot_doc = self.clauses[active[k]].matcher.doc_id();
let head_doc = self.clauses[active[0]].matcher.doc_id();
if head_doc == pivot_doc {
let mut total_score = 0.0_f32;
for &i in &active {
if self.clauses[i].matcher.doc_id() == pivot_doc {
let tf = self.clauses[i].matcher.term_freq() as f32;
let fl = self.field_length_for(i, pivot_doc);
total_score += self.clauses[i].scorer.score(pivot_doc, tf, fl);
} else {
break;
}
}
collector.collect(pivot_doc, total_score)?;
if !collector.needs_more() {
break;
}
for &i in active.iter() {
if self.clauses[i].matcher.doc_id() == pivot_doc {
self.clauses[i].matcher.next()?;
} else {
break;
}
}
} else {
for &i in &active[..k] {
self.clauses[i].matcher.skip_to(pivot_doc)?;
}
}
}
}
active.retain(|&i| {
!self.clauses[i].matcher.is_exhausted()
&& self.clauses[i].matcher.doc_id() != u64::MAX
});
}
Ok(collector)
}
fn advance_clause_past_block(&mut self, clause_idx: usize, at_doc: u64) -> Result<()> {
let target = self.clauses[clause_idx].scorer.next_block_boundary(at_doc);
match target {
Some(t) if t == u64::MAX => {
self.clauses[clause_idx].matcher.skip_to(u64::MAX)?;
}
Some(t) => {
self.clauses[clause_idx].matcher.skip_to(t)?;
}
None => {
self.clauses[clause_idx].matcher.next()?;
}
}
Ok(())
}
fn field_length_for(&self, clause_idx: usize, doc_id: u64) -> Option<f32> {
let field = self.clauses[clause_idx].field_name.as_deref()?;
let reader = self.inverted_reader?;
reader
.field_length(doc_id, field)
.ok()
.flatten()
.map(|n| n as f32)
}
}
fn field_name_of(query: &dyn Query) -> Option<String> {
query
.as_any()
.downcast_ref::<TermQuery>()
.map(|t| t.field().to_string())
}
pub(crate) fn is_bmw_eligible(query: &dyn Query) -> Option<&BooleanQuery> {
let bq = query.as_any().downcast_ref::<BooleanQuery>()?;
if bq.minimum_should_match() > 0 {
return None;
}
if bq.clauses().len() < 2 {
return None;
}
if bq
.clauses()
.iter()
.any(|c| !matches!(c.occur, Occur::Should))
{
return None;
}
if bq
.clauses()
.iter()
.any(|c| c.query.as_any().downcast_ref::<TermQuery>().is_none())
{
return None;
}
Some(bq)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::lexical::query::boolean::BooleanQueryBuilder;
use crate::lexical::query::term::TermQuery;
#[test]
fn eligibility_rejects_non_should_only_or_thin_queries() {
let single = BooleanQueryBuilder::new()
.should(Box::new(TermQuery::new("text", "x")))
.build();
assert!(is_bmw_eligible(&single).is_none());
let mixed = BooleanQueryBuilder::new()
.must(Box::new(TermQuery::new("text", "x")))
.should(Box::new(TermQuery::new("text", "y")))
.build();
assert!(is_bmw_eligible(&mixed).is_none());
let msm = BooleanQueryBuilder::new()
.should(Box::new(TermQuery::new("text", "x")))
.should(Box::new(TermQuery::new("text", "y")))
.minimum_should_match(2)
.build();
assert!(is_bmw_eligible(&msm).is_none());
let ok = BooleanQueryBuilder::new()
.should(Box::new(TermQuery::new("text", "x")))
.should(Box::new(TermQuery::new("text", "y")))
.build();
assert!(is_bmw_eligible(&ok).is_some());
}
}