use crate::core::{DocId, NO_MORE_DOCS, Result, ScoreMode, Scorer, TwoPhaseIterator};
use crate::query::{BoundQuery, Query, ScorerSupplier};
use crate::search::searcher::Searcher;
use crate::segment::reader::SegmentReader;
pub struct NestedQuery {
pub path: String,
pub(crate) inner: Box<dyn Query>,
pub inner_hits: Option<crate::query::ast::InnerHitsConfig>,
}
impl Query for NestedQuery {
fn bind(&self, searcher: &Searcher, score_mode: ScoreMode) -> Result<Box<dyn BoundQuery>> {
let inner_weight = self.inner.bind(searcher, score_mode)?;
let inner_hits_weight = if self.inner_hits.is_some() {
Some(self.inner.bind(searcher, score_mode)?)
} else {
None
};
Ok(Box::new(BoundNestedQuery {
_path: self.path.clone(),
inner: inner_weight,
_inner_hits_config: self.inner_hits.clone(),
_inner_hits_weight: inner_hits_weight,
}))
}
}
pub struct InnerHitSpec {
pub name: String,
pub path: String,
pub config: crate::query::ast::InnerHitsConfig,
pub(crate) weight: Box<dyn BoundQuery>,
}
struct BoundNestedQuery {
_path: String,
inner: Box<dyn BoundQuery>,
_inner_hits_config: Option<crate::query::ast::InnerHitsConfig>,
_inner_hits_weight: Option<Box<dyn BoundQuery>>,
}
impl BoundQuery for BoundNestedQuery {
fn scorer_supplier(&self, reader: &SegmentReader) -> Result<Option<Box<dyn ScorerSupplier>>> {
let parent_bitset = match reader.parent_bitset() {
Some(b) => b,
None => return Ok(None), };
let inner_supplier = match self.inner.scorer_supplier(reader)? {
Some(s) => s,
None => return Ok(None),
};
let inner_scorer = inner_supplier.scorer()?;
let mut parent_matches: Vec<(u32, f32)> = Vec::new();
let mut scorer = inner_scorer;
let mut last_parent = u32::MAX;
while scorer.doc_id() != NO_MORE_DOCS {
let nested_doc = scorer.doc_id().as_u32();
let score = scorer.score();
let parent_doc = find_parent(&parent_bitset, nested_doc);
if parent_doc != last_parent {
parent_matches.push((parent_doc, score));
last_parent = parent_doc;
} else if let Some(last) = parent_matches.last_mut() {
if score > last.1 {
last.1 = score;
}
}
scorer.next();
}
if parent_matches.is_empty() {
return Ok(None);
}
parent_matches.dedup_by_key(|m| m.0);
Ok(Some(Box::new(NestedScorerSupplier {
matches: parent_matches,
})))
}
}
fn find_parent(parent_bitset: &[bool], nested_doc: u32) -> u32 {
let mut i = nested_doc as usize;
while i > 0 {
if i < parent_bitset.len() && parent_bitset[i] {
return i as u32;
}
i -= 1;
}
if !parent_bitset.is_empty() && parent_bitset[0] {
return 0;
}
0
}
struct NestedScorerSupplier {
matches: Vec<(u32, f32)>,
}
impl ScorerSupplier for NestedScorerSupplier {
fn cost(&self) -> u64 {
self.matches.len() as u64
}
fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
Ok(Box::new(NestedScorer {
matches: self.matches,
pos: 0,
}))
}
}
struct NestedScorer {
matches: Vec<(u32, f32)>,
pos: usize,
}
impl Scorer for NestedScorer {
fn doc_id(&self) -> DocId {
if self.pos < self.matches.len() {
DocId::new(self.matches[self.pos].0)
} else {
NO_MORE_DOCS
}
}
fn next(&mut self) -> DocId {
if self.pos < self.matches.len() {
self.pos += 1;
}
self.doc_id()
}
fn advance(&mut self, target: DocId) -> DocId {
while self.pos < self.matches.len() && self.matches[self.pos].0 < target.as_u32() {
self.pos += 1;
}
self.doc_id()
}
fn score(&mut self) -> f32 {
if self.pos < self.matches.len() {
self.matches[self.pos].1
} else {
0.0
}
}
fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
None
}
}