lb_tantivy/aggregation/
collector.rs

1use super::agg_req::Aggregations;
2use super::agg_req_with_accessor::AggregationsWithAccessor;
3use super::agg_result::AggregationResults;
4use super::buf_collector::BufAggregationCollector;
5use super::intermediate_agg_result::IntermediateAggregationResults;
6use super::segment_agg_result::{
7    build_segment_agg_collector, AggregationLimitsGuard, SegmentAggregationCollector,
8};
9use crate::aggregation::agg_req_with_accessor::get_aggs_with_segment_accessor_and_validate;
10use crate::collector::{Collector, SegmentCollector};
11use crate::index::SegmentReader;
12use crate::{DocId, SegmentOrdinal, TantivyError};
13
14/// The default max bucket count, before the aggregation fails.
15pub const DEFAULT_BUCKET_LIMIT: u32 = 65000;
16
17/// The default memory limit in bytes before the aggregation fails. 500MB
18pub const DEFAULT_MEMORY_LIMIT: u64 = 500_000_000;
19
20/// Collector for aggregations.
21///
22/// The collector collects all aggregations by the underlying aggregation request.
23pub struct AggregationCollector {
24    agg: Aggregations,
25    limits: AggregationLimitsGuard,
26}
27
28impl AggregationCollector {
29    /// Create collector from aggregation request.
30    ///
31    /// Aggregation fails when the limits in `AggregationLimits` is exceeded. (memory limit and
32    /// bucket limit)
33    pub fn from_aggs(agg: Aggregations, limits: AggregationLimitsGuard) -> Self {
34        Self { agg, limits }
35    }
36}
37
38/// Collector for distributed aggregations.
39///
40/// The collector collects all aggregations by the underlying aggregation request.
41///
42/// # Purpose
43/// AggregationCollector returns `IntermediateAggregationResults` and not the final
44/// `AggregationResults`, so that results from different indices can be merged and then converted
45/// into the final `AggregationResults` via the `into_final_result()` method.
46pub struct DistributedAggregationCollector {
47    agg: Aggregations,
48    limits: AggregationLimitsGuard,
49}
50
51impl DistributedAggregationCollector {
52    /// Create collector from aggregation request.
53    ///
54    /// Aggregation fails when the limits in `AggregationLimits` is exceeded. (memory limit and
55    /// bucket limit)
56    pub fn from_aggs(agg: Aggregations, limits: AggregationLimitsGuard) -> Self {
57        Self { agg, limits }
58    }
59}
60
61impl Collector for DistributedAggregationCollector {
62    type Fruit = IntermediateAggregationResults;
63
64    type Child = AggregationSegmentCollector;
65
66    fn for_segment(
67        &self,
68        segment_local_id: crate::SegmentOrdinal,
69        reader: &crate::SegmentReader,
70    ) -> crate::Result<Self::Child> {
71        AggregationSegmentCollector::from_agg_req_and_reader(
72            &self.agg,
73            reader,
74            segment_local_id,
75            &self.limits,
76        )
77    }
78
79    fn requires_scoring(&self) -> bool {
80        false
81    }
82
83    fn merge_fruits(
84        &self,
85        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
86    ) -> crate::Result<Self::Fruit> {
87        merge_fruits(segment_fruits)
88    }
89}
90
91impl Collector for AggregationCollector {
92    type Fruit = AggregationResults;
93
94    type Child = AggregationSegmentCollector;
95
96    fn for_segment(
97        &self,
98        segment_local_id: crate::SegmentOrdinal,
99        reader: &crate::SegmentReader,
100    ) -> crate::Result<Self::Child> {
101        AggregationSegmentCollector::from_agg_req_and_reader(
102            &self.agg,
103            reader,
104            segment_local_id,
105            &self.limits,
106        )
107    }
108
109    fn requires_scoring(&self) -> bool {
110        false
111    }
112
113    fn merge_fruits(
114        &self,
115        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
116    ) -> crate::Result<Self::Fruit> {
117        let res = merge_fruits(segment_fruits)?;
118        res.into_final_result(self.agg.clone(), self.limits.clone())
119    }
120}
121
122fn merge_fruits(
123    mut segment_fruits: Vec<crate::Result<IntermediateAggregationResults>>,
124) -> crate::Result<IntermediateAggregationResults> {
125    if let Some(fruit) = segment_fruits.pop() {
126        let mut fruit = fruit?;
127        for next_fruit in segment_fruits {
128            fruit.merge_fruits(next_fruit?)?;
129        }
130        Ok(fruit)
131    } else {
132        Ok(IntermediateAggregationResults::default())
133    }
134}
135
136/// `AggregationSegmentCollector` does the aggregation collection on a segment.
137pub struct AggregationSegmentCollector {
138    aggs_with_accessor: AggregationsWithAccessor,
139    agg_collector: BufAggregationCollector,
140    error: Option<TantivyError>,
141}
142
143impl AggregationSegmentCollector {
144    /// Creates an `AggregationSegmentCollector from` an [`Aggregations`] request and a segment
145    /// reader. Also includes validation, e.g. checking field types and existence.
146    pub fn from_agg_req_and_reader(
147        agg: &Aggregations,
148        reader: &SegmentReader,
149        segment_ordinal: SegmentOrdinal,
150        limits: &AggregationLimitsGuard,
151    ) -> crate::Result<Self> {
152        let mut aggs_with_accessor =
153            get_aggs_with_segment_accessor_and_validate(agg, reader, segment_ordinal, limits)?;
154        let result =
155            BufAggregationCollector::new(build_segment_agg_collector(&mut aggs_with_accessor)?);
156        Ok(AggregationSegmentCollector {
157            aggs_with_accessor,
158            agg_collector: result,
159            error: None,
160        })
161    }
162}
163
164impl SegmentCollector for AggregationSegmentCollector {
165    type Fruit = crate::Result<IntermediateAggregationResults>;
166
167    #[inline]
168    fn collect(&mut self, doc: DocId, _score: crate::Score) {
169        if self.error.is_some() {
170            return;
171        }
172        if let Err(err) = self
173            .agg_collector
174            .collect(doc, &mut self.aggs_with_accessor)
175        {
176            self.error = Some(err);
177        }
178    }
179
180    /// The query pushes the documents to the collector via this method.
181    ///
182    /// Only valid for Collectors that ignore docs
183    fn collect_block(&mut self, docs: &[DocId]) {
184        if self.error.is_some() {
185            return;
186        }
187        if let Err(err) = self
188            .agg_collector
189            .collect_block(docs, &mut self.aggs_with_accessor)
190        {
191            self.error = Some(err);
192        }
193    }
194
195    fn harvest(mut self) -> Self::Fruit {
196        if let Some(err) = self.error {
197            return Err(err);
198        }
199        self.agg_collector.flush(&mut self.aggs_with_accessor)?;
200
201        let mut sub_aggregation_res = IntermediateAggregationResults::default();
202        Box::new(self.agg_collector).add_intermediate_aggregation_result(
203            &self.aggs_with_accessor,
204            &mut sub_aggregation_res,
205        )?;
206
207        Ok(sub_aggregation_res)
208    }
209}
lb_tantivy/aggregation/collector.rs

lb_tantivy/aggregation/
collector.rs
