diff --git a/src/aggregation/bucket/composite/collector.rs b/src/aggregation/bucket/composite/collector.rs index 2fb4e53fb..09bd938f5 100644 --- a/src/aggregation/bucket/composite/collector.rs +++ b/src/aggregation/bucket/composite/collector.rs @@ -264,7 +264,7 @@ impl SegmentCompositeCollector { dict.insert( key, IntermediateCompositeBucketEntry { - doc_count: agg.count, + doc_count: agg.count as u64, sub_aggregation: sub_aggregation_res, }, ); diff --git a/src/aggregation/bucket/term_agg.rs b/src/aggregation/bucket/term_agg.rs index 6947f8f26..43a298534 100644 --- a/src/aggregation/bucket/term_agg.rs +++ b/src/aggregation/bucket/term_agg.rs @@ -977,12 +977,12 @@ where bucket.bucket_id, )?; Ok(IntermediateTermBucketEntry { - doc_count: bucket.count, + doc_count: bucket.count as u64, sub_aggregation: sub_aggregation_res, }) } else { Ok(IntermediateTermBucketEntry { - doc_count: bucket.count, + doc_count: bucket.count as u64, sub_aggregation: Default::default(), }) } diff --git a/src/aggregation/bucket/term_missing_agg.rs b/src/aggregation/bucket/term_missing_agg.rs index 47c3989c6..b2e1dcf71 100644 --- a/src/aggregation/bucket/term_missing_agg.rs +++ b/src/aggregation/bucket/term_missing_agg.rs @@ -98,7 +98,7 @@ impl SegmentAggregationCollector for TermMissingAgg { let missing_count = &self.missing_count_per_bucket[parent_bucket_id as usize]; let mut missing_entry = IntermediateTermBucketEntry { - doc_count: missing_count.missing_count, + doc_count: missing_count.missing_count as u64, sub_aggregation: Default::default(), }; if let Some(sub_agg) = &mut self.sub_agg { diff --git a/src/aggregation/intermediate_agg_result.rs b/src/aggregation/intermediate_agg_result.rs index 6c6eeeb7e..a93b236f7 100644 --- a/src/aggregation/intermediate_agg_result.rs +++ b/src/aggregation/intermediate_agg_result.rs @@ -915,7 +915,7 @@ impl IntermediateRangeBucketEntry { #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] pub struct IntermediateTermBucketEntry { /// The number of documents in the bucket. - pub doc_count: u32, + pub doc_count: u64, /// The sub_aggregation in this bucket. pub sub_aggregation: IntermediateAggregationResults, } @@ -1229,6 +1229,24 @@ mod tests { assert_eq!(tree_left, tree_expected); } + #[test] + fn test_term_bucket_doc_count_no_u32_overflow() { + // Two segments each contributing (u32::MAX - 100) docs to the same term. Summing them + // overflowed when doc_count was u32. + let per_segment = u32::MAX as u64 - 100; + let mut entry = IntermediateTermBucketEntry { + doc_count: per_segment, + sub_aggregation: Default::default(), + }; + entry + .merge_fruits(IntermediateTermBucketEntry { + doc_count: per_segment, + sub_aggregation: Default::default(), + }) + .unwrap(); + assert_eq!(entry.doc_count, per_segment * 2); + } + #[test] fn test_merge_fruits_tree_empty() { let mut tree_left = get_intermediate_tree_with_ranges(&[