Skip to content

Commit

Permalink
introduce optional collect_block in segmentcollector
Browse files Browse the repository at this point in the history
add collect_block in segment_collector to handle groups of documents as performance optimization
  • Loading branch information
PSeitz committed May 17, 2022
1 parent 44ea731 commit ac5f3bf
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 3 deletions.
5 changes: 5 additions & 0 deletions src/collector/count_collector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,11 @@ impl SegmentCollector for SegmentCountCollector {
Ok(())
}

fn collect_block(&mut self, docs: &[(DocId, Score)]) -> crate::Result<()> {
self.count += docs.len();
Ok(())
}

fn harvest(self) -> usize {
self.count
}
Expand Down
47 changes: 45 additions & 2 deletions src/collector/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -172,19 +172,33 @@ pub trait Collector: Sync + Send {
) -> crate::Result<<Self::Child as SegmentCollector>::Fruit> {
let mut segment_collector = self.for_segment(segment_ord as u32, reader)?;

let mut cache_pos = 0;
let mut cache = [(0, 0.0); 64];

if let Some(alive_bitset) = reader.alive_bitset() {
weight.for_each(reader, &mut |doc, score| {
if alive_bitset.is_alive(doc) {
segment_collector.collect(doc, score)?;
cache[cache_pos] = (doc, score);
cache_pos += 1;
if cache_pos == 64 {
segment_collector.collect_block(&cache)?;
cache_pos = 0;
}
}
Ok(())
})?;
} else {
weight.for_each(reader, &mut |doc, score| {
segment_collector.collect(doc, score)?;
cache[cache_pos] = (doc, score);
cache_pos += 1;
if cache_pos == 64 {
segment_collector.collect_block(&cache)?;
cache_pos = 0;
}
Ok(())
})?;
}
segment_collector.collect_block(&cache[..cache_pos])?;
Ok(segment_collector.harvest())
}
}
Expand Down Expand Up @@ -258,6 +272,14 @@ pub trait SegmentCollector: 'static {
/// The query pushes the scored document to the collector via this method.
fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()>;

/// The query pushes the scored document to the collector via this method.
fn collect_block(&mut self, docs: &[(DocId, Score)]) -> crate::Result<()> {
for (doc, score) in docs {
self.collect(*doc, *score)?;
}
Ok(())
}

/// Extract the fruit of the collection from the `SegmentCollector`.
fn harvest(self) -> Self::Fruit;
}
Expand Down Expand Up @@ -317,6 +339,12 @@ where
Ok(())
}

fn collect_block(&mut self, docs: &[(DocId, Score)]) -> crate::Result<()> {
self.0.collect_block(docs)?;
self.1.collect_block(docs)?;
Ok(())
}

fn harvest(self) -> <Self as SegmentCollector>::Fruit {
(self.0.harvest(), self.1.harvest())
}
Expand Down Expand Up @@ -383,6 +411,13 @@ where
Ok(())
}

fn collect_block(&mut self, docs: &[(DocId, Score)]) -> crate::Result<()> {
self.0.collect_block(docs)?;
self.1.collect_block(docs)?;
self.2.collect_block(docs)?;
Ok(())
}

fn harvest(self) -> <Self as SegmentCollector>::Fruit {
(self.0.harvest(), self.1.harvest(), self.2.harvest())
}
Expand Down Expand Up @@ -459,6 +494,14 @@ where
Ok(())
}

fn collect_block(&mut self, docs: &[(DocId, Score)]) -> crate::Result<()> {
self.0.collect_block(docs)?;
self.1.collect_block(docs)?;
self.2.collect_block(docs)?;
self.3.collect_block(docs)?;
Ok(())
}

fn harvest(self) -> <Self as SegmentCollector>::Fruit {
(
self.0.harvest(),
Expand Down
3 changes: 2 additions & 1 deletion src/collector/top_collector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,8 @@ pub(crate) struct TopCollector<T> {
}

impl<T> TopCollector<T>
where T: PartialOrd + Clone
where
T: PartialOrd + Clone,
{
/// Creates a top collector, with a number of documents equal to "limit".
///
Expand Down
8 changes: 8 additions & 0 deletions src/collector/top_score_collector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -704,6 +704,14 @@ impl SegmentCollector for TopScoreSegmentCollector {
Ok(())
}

#[inline]
fn collect_block(&mut self, docs: &[(DocId, Score)]) -> crate::Result<()> {
for (doc, score) in docs {
self.0.collect(*doc, *score);
}
Ok(())
}

fn harvest(self) -> Vec<(Score, DocAddress)> {
self.0.harvest()
}
Expand Down

0 comments on commit ac5f3bf

Please sign in to comment.