remove dynamic dispatch in collect_segment

This commit is contained in:
Pascal Seitz
2023-03-02 20:15:47 +08:00
parent ca20bfa776
commit 6037cdfe7e
3 changed files with 19 additions and 11 deletions

View File

@@ -106,7 +106,7 @@ mod tweak_score_top_collector;
pub use self::tweak_score_top_collector::{ScoreSegmentTweaker, ScoreTweaker};
mod facet_collector;
pub use self::facet_collector::{FacetCollector, FacetCounts};
use crate::query::Weight;
use crate::query::{for_each_docset, for_each_scorer, Weight};
mod docset_collector;
pub use self::docset_collector::DocSetCollector;
@@ -173,28 +173,32 @@ pub trait Collector: Sync + Send {
match (reader.alive_bitset(), self.requires_scoring()) {
(Some(alive_bitset), true) => {
weight.for_each(reader, &mut |doc, score| {
let mut scorer = weight.scorer(reader, 1.0)?;
for_each_scorer(scorer.as_mut(), |doc, score| {
if alive_bitset.is_alive(doc) {
segment_collector.collect(doc, score);
}
})?;
});
}
(Some(alive_bitset), false) => {
weight.for_each_no_score(reader, &mut |doc| {
let mut docset = weight.scorer(reader, 1.0)?;
for_each_docset(docset.as_mut(), |doc| {
if alive_bitset.is_alive(doc) {
segment_collector.collect(doc, 0.0);
}
})?;
});
}
(None, true) => {
weight.for_each(reader, &mut |doc, score| {
let mut scorer = weight.scorer(reader, 1.0)?;
for_each_scorer(scorer.as_mut(), |doc, score| {
segment_collector.collect(doc, score);
})?;
});
}
(None, false) => {
weight.for_each_no_score(reader, &mut |doc| {
let mut docset = weight.scorer(reader, 1.0)?;
for_each_docset(docset.as_mut(), |doc| {
segment_collector.collect(doc, 0.0);
})?;
});
}
}

View File

@@ -65,6 +65,7 @@ pub use self::union::Union;
#[cfg(test)]
pub use self::vec_docset::VecDocSet;
pub use self::weight::Weight;
pub(crate) use self::weight::{for_each_docset, for_each_pruning_scorer, for_each_scorer};
#[cfg(test)]
mod tests {

View File

@@ -5,9 +5,10 @@ use crate::{DocId, DocSet, Score, TERMINATED};
/// Iterates through all of the documents and scores matched by the DocSet
/// `DocSet`.
#[inline]
pub(crate) fn for_each_scorer<TScorer: Scorer + ?Sized>(
scorer: &mut TScorer,
callback: &mut dyn FnMut(DocId, Score),
mut callback: impl FnMut(DocId, Score),
) {
let mut doc = scorer.doc();
while doc != TERMINATED {
@@ -18,7 +19,8 @@ pub(crate) fn for_each_scorer<TScorer: Scorer + ?Sized>(
/// Iterates through all of the documents matched by the DocSet
/// `DocSet`.
pub(crate) fn for_each_docset<T: DocSet + ?Sized>(docset: &mut T, callback: &mut dyn FnMut(DocId)) {
#[inline]
pub(crate) fn for_each_docset<T: DocSet + ?Sized>(docset: &mut T, mut callback: impl FnMut(DocId)) {
let mut doc = docset.doc();
while doc != TERMINATED {
callback(doc);
@@ -36,6 +38,7 @@ pub(crate) fn for_each_docset<T: DocSet + ?Sized>(docset: &mut T, callback: &mut
///
/// More importantly, it makes it possible for scorers to implement
/// important optimization (e.g. BlockWAND for union).
#[inline]
pub(crate) fn for_each_pruning_scorer<TScorer: Scorer + ?Sized>(
scorer: &mut TScorer,
mut threshold: Score,