return result from segment collector

Author: PSeitz

examples/custom_collector.rs

@@ -102,11 +102,12 @@ struct StatsSegmentCollector {
 impl SegmentCollector for StatsSegmentCollector {
     type Fruit = Option<Stats>;
 
-    fn collect(&mut self, doc: u32, _score: Score) {
+    fn collect(&mut self, doc: u32, _score: Score) -> crate::Result<()> {
         let value = self.fast_field_reader.get(doc) as f64;
         self.stats.count += 1;
         self.stats.sum += value;
         self.stats.squared_sum += value * value;
+        Ok(())
     }
 
     fn harvest(self) -> <Self as SegmentCollector>::Fruit {

src/aggregation/collector.rs

@@ -132,8 +132,9 @@ impl SegmentCollector for AggregationSegmentCollector {
     type Fruit = crate::Result<IntermediateAggregationResults>;
 
     #[inline]
-    fn collect(&mut self, doc: crate::DocId, _score: crate::Score) {
+    fn collect(&mut self, doc: crate::DocId, _score: crate::Score) -> crate::Result<()> {
         self.result.collect(doc, &self.aggs_with_accessor);
+        Ok(())
     }
 
     fn harvest(mut self) -> Self::Fruit {

src/collector/count_collector.rs

@@ -65,8 +65,9 @@ pub struct SegmentCountCollector {
 impl SegmentCollector for SegmentCountCollector {
     type Fruit = usize;
 
-    fn collect(&mut self, _: DocId, _: Score) {
+    fn collect(&mut self, _: DocId, _: Score) -> crate::Result<()> {
         self.count += 1;
+        Ok(())
     }
 
     fn harvest(self) -> usize {
@@ -92,18 +93,18 @@ mod tests {
         }
         {
             let mut count_collector = SegmentCountCollector::default();
-            count_collector.collect(0u32, 1.0);
+            count_collector.collect(0u32, 1.0).unwrap();
             assert_eq!(count_collector.harvest(), 1);
         }
         {
             let mut count_collector = SegmentCountCollector::default();
-            count_collector.collect(0u32, 1.0);
+            count_collector.collect(0u32, 1.0).unwrap();
             assert_eq!(count_collector.harvest(), 1);
         }
         {
             let mut count_collector = SegmentCountCollector::default();
-            count_collector.collect(0u32, 1.0);
-            count_collector.collect(1u32, 1.0);
+            count_collector.collect(0u32, 1.0).unwrap();
+            count_collector.collect(1u32, 1.0).unwrap();
             assert_eq!(count_collector.harvest(), 2);
         }
     }

src/collector/custom_score_top_collector.rs

@@ -90,9 +90,10 @@ where
 {
     type Fruit = Vec<(TScore, DocAddress)>;
 
-    fn collect(&mut self, doc: DocId, _score: Score) {
+    fn collect(&mut self, doc: DocId, _score: Score) -> crate::Result<()> {
         let score = self.segment_scorer.score(doc);
         self.segment_collector.collect(doc, score);
+        Ok(())
     }
 
     fn harvest(self) -> Vec<(TScore, DocAddress)> {

src/collector/docset_collector.rs

@@ -50,8 +50,9 @@ pub struct DocSetChildCollector {
 impl SegmentCollector for DocSetChildCollector {
     type Fruit = (u32, HashSet<DocId>);
 
-    fn collect(&mut self, doc: crate::DocId, _score: Score) {
+    fn collect(&mut self, doc: crate::DocId, _score: Score) -> crate::Result<()> {
         self.docs.insert(doc);
+        Ok(())
     }
 
     fn harvest(self) -> (u32, HashSet<DocId>) {

src/collector/facet_collector.rs

@@ -333,7 +333,7 @@ impl Collector for FacetCollector {
 impl SegmentCollector for FacetSegmentCollector {
     type Fruit = FacetCounts;
 
-    fn collect(&mut self, doc: DocId, _: Score) {
+    fn collect(&mut self, doc: DocId, _: Score) -> crate::Result<()> {
         self.reader.facet_ords(doc, &mut self.facet_ords_buf);
         let mut previous_collapsed_ord: usize = usize::MAX;
         for &facet_ord in &self.facet_ords_buf {
@@ -345,6 +345,7 @@ impl SegmentCollector for FacetSegmentCollector {
             };
             previous_collapsed_ord = collapsed_ord;
         }
+        Ok(())
     }
 
     /// Returns the results of the collection.

src/collector/filter_collector_wrapper.rs

@@ -173,11 +173,12 @@ where
 {
     type Fruit = TSegmentCollector::Fruit;
 
-    fn collect(&mut self, doc: u32, score: Score) {
+    fn collect(&mut self, doc: u32, score: Score) -> crate::Result<()> {
         let value = self.fast_field_reader.get(doc);
         if (self.predicate)(value) {
-            self.segment_collector.collect(doc, score)
+            self.segment_collector.collect(doc, score)?;
         }
+        Ok(())
     }
 
     fn harvest(self) -> <TSegmentCollector as SegmentCollector>::Fruit {

src/collector/histogram_collector.rs

@@ -91,9 +91,10 @@ pub struct SegmentHistogramCollector {
 impl SegmentCollector for SegmentHistogramCollector {
     type Fruit = Vec<u64>;
 
-    fn collect(&mut self, doc: DocId, _score: Score) {
+    fn collect(&mut self, doc: DocId, _score: Score) -> crate::Result<()> {
         let value = self.ff_reader.get(doc);
         self.histogram_computer.add_value(value);
+        Ok(())
     }
 
     fn harvest(self) -> Self::Fruit {

src/collector/mod.rs

@@ -175,12 +175,12 @@ pub trait Collector: Sync + Send {
         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);
+                    segment_collector.collect(doc, score).unwrap(); // TODO
                 }
             })?;
         } else {
             weight.for_each(reader, &mut |doc, score| {
-                segment_collector.collect(doc, score);
+                segment_collector.collect(doc, score).unwrap(); // TODO
             })?;
         }
         Ok(segment_collector.harvest())
@@ -190,10 +190,11 @@ pub trait Collector: Sync + Send {
 impl<TSegmentCollector: SegmentCollector> SegmentCollector for Option<TSegmentCollector> {
     type Fruit = Option<TSegmentCollector::Fruit>;
 
-    fn collect(&mut self, doc: DocId, score: Score) {
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
         if let Some(segment_collector) = self {
-            segment_collector.collect(doc, score);
+            segment_collector.collect(doc, score)?;
         }
+        Ok(())
     }
 
     fn harvest(self) -> Self::Fruit {
@@ -253,7 +254,7 @@ pub trait SegmentCollector: 'static {
     type Fruit: Fruit;
 
     /// The query pushes the scored document to the collector via this method.
-    fn collect(&mut self, doc: DocId, score: Score);
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()>;
 
     /// Extract the fruit of the collection from the `SegmentCollector`.
     fn harvest(self) -> Self::Fruit;
@@ -308,9 +309,10 @@ where
 {
     type Fruit = (Left::Fruit, Right::Fruit);
 
-    fn collect(&mut self, doc: DocId, score: Score) {
-        self.0.collect(doc, score);
-        self.1.collect(doc, score);
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
+        self.0.collect(doc, score)?;
+        self.1.collect(doc, score)?;
+        Ok(())
     }
 
     fn harvest(self) -> <Self as SegmentCollector>::Fruit {
@@ -372,10 +374,11 @@ where
 {
     type Fruit = (One::Fruit, Two::Fruit, Three::Fruit);
 
-    fn collect(&mut self, doc: DocId, score: Score) {
-        self.0.collect(doc, score);
-        self.1.collect(doc, score);
-        self.2.collect(doc, score);
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
+        self.0.collect(doc, score)?;
+        self.1.collect(doc, score)?;
+        self.2.collect(doc, score)?;
+        Ok(())
     }
 
     fn harvest(self) -> <Self as SegmentCollector>::Fruit {
@@ -446,11 +449,12 @@ where
 {
     type Fruit = (One::Fruit, Two::Fruit, Three::Fruit, Four::Fruit);
 
-    fn collect(&mut self, doc: DocId, score: Score) {
-        self.0.collect(doc, score);
-        self.1.collect(doc, score);
-        self.2.collect(doc, score);
-        self.3.collect(doc, score);
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
+        self.0.collect(doc, score)?;
+        self.1.collect(doc, score)?;
+        self.2.collect(doc, score)?;
+        self.3.collect(doc, score)?;
+        Ok(())
     }
 
     fn harvest(self) -> <Self as SegmentCollector>::Fruit {

src/collector/multi_collector.rs

@@ -52,8 +52,9 @@ impl<TCollector: Collector> Collector for CollectorWrapper<TCollector> {
 impl SegmentCollector for Box<dyn BoxableSegmentCollector> {
     type Fruit = Box<dyn Fruit>;
 
-    fn collect(&mut self, doc: u32, score: Score) {
-        self.as_mut().collect(doc, score);
+    fn collect(&mut self, doc: u32, score: Score) -> crate::Result<()> {
+        self.as_mut().collect(doc, score)?;
+        Ok(())
     }
 
     fn harvest(self) -> Box<dyn Fruit> {
@@ -62,7 +63,7 @@ impl SegmentCollector for Box<dyn BoxableSegmentCollector> {
 }
 
 pub trait BoxableSegmentCollector {
-    fn collect(&mut self, doc: u32, score: Score);
+    fn collect(&mut self, doc: u32, score: Score) -> crate::Result<()>;
     fn harvest_from_box(self: Box<Self>) -> Box<dyn Fruit>;
 }
 
@@ -71,8 +72,8 @@ pub struct SegmentCollectorWrapper<TSegmentCollector: SegmentCollector>(TSegment
 impl<TSegmentCollector: SegmentCollector> BoxableSegmentCollector
     for SegmentCollectorWrapper<TSegmentCollector>
 {
-    fn collect(&mut self, doc: u32, score: Score) {
-        self.0.collect(doc, score);
+    fn collect(&mut self, doc: u32, score: Score) -> crate::Result<()> {
+        self.0.collect(doc, score)
     }
 
     fn harvest_from_box(self: Box<Self>) -> Box<dyn Fruit> {
@@ -228,10 +229,11 @@ pub struct MultiCollectorChild {
 impl SegmentCollector for MultiCollectorChild {
     type Fruit = MultiFruit;
 
-    fn collect(&mut self, doc: DocId, score: Score) {
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
         for child in &mut self.children {
-            child.collect(doc, score);
+            child.collect(doc, score)?;
         }
+        Ok(())
     }
 
     fn harvest(self) -> MultiFruit {

src/collector/tests.rs

@@ -138,9 +138,10 @@ impl Collector for TestCollector {
 impl SegmentCollector for TestSegmentCollector {
     type Fruit = TestFruit;
 
-    fn collect(&mut self, doc: DocId, score: Score) {
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
         self.fruit.docs.push(DocAddress::new(self.segment_id, doc));
         self.fruit.scores.push(score);
+        Ok(())
     }
 
     fn harvest(self) -> <Self as SegmentCollector>::Fruit {
@@ -198,9 +199,10 @@ impl Collector for FastFieldTestCollector {
 impl SegmentCollector for FastFieldSegmentCollector {
     type Fruit = Vec<u64>;
 
-    fn collect(&mut self, doc: DocId, _score: Score) {
+    fn collect(&mut self, doc: DocId, _score: Score) -> crate::Result<()> {
         let val = self.reader.get(doc);
         self.vals.push(val);
+        Ok(())
     }
 
     fn harvest(self) -> Vec<u64> {
@@ -255,9 +257,10 @@ impl Collector for BytesFastFieldTestCollector {
 impl SegmentCollector for BytesFastFieldSegmentCollector {
     type Fruit = Vec<u8>;
 
-    fn collect(&mut self, doc: u32, _score: Score) {
+    fn collect(&mut self, doc: u32, _score: Score) -> crate::Result<()> {
         let data = self.reader.get_bytes(doc);
         self.vals.extend(data);
+        Ok(())
     }
 
     fn harvest(self) -> <Self as SegmentCollector>::Fruit {

src/collector/top_score_collector.rs

@@ -699,8 +699,9 @@ pub struct TopScoreSegmentCollector(TopSegmentCollector<Score>);
 impl SegmentCollector for TopScoreSegmentCollector {
     type Fruit = Vec<(Score, DocAddress)>;
 
-    fn collect(&mut self, doc: DocId, score: Score) {
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
         self.0.collect(doc, score);
+        Ok(())
     }
 
     fn harvest(self) -> Vec<(Score, DocAddress)> {

src/collector/tweak_score_top_collector.rs

@@ -93,9 +93,10 @@ where
 {
     type Fruit = Vec<(TScore, DocAddress)>;
 
-    fn collect(&mut self, doc: DocId, score: Score) {
+    fn collect(&mut self, doc: DocId, score: Score) -> crate::Result<()> {
         let score = self.segment_scorer.score(doc, score);
         self.segment_collector.collect(doc, score);
+        Ok(())
     }
 
     fn harvest(self) -> Vec<(TScore, DocAddress)> {