Deduplicate multivalued documents within histogram buckets

src/aggregation/bucket/histogram/histogram.rs

@@ -440,7 +440,7 @@ impl<B: BucketIdSlot> HistogramBuckets<B> {
 /// The collector puts values from the fast field into the correct buckets and does a conversion to
 /// the correct datatype.
 #[derive(Debug)]
-pub struct SegmentHistogramCollector<B> {
+pub struct SegmentHistogramCollector<B, const SOURCE_CONTAINS_MULTIVALUES: bool> {
     /// The buckets containing the aggregation data.
     /// One Histogram bucket per parent bucket id.
     parent_buckets: Vec<HistogramBuckets<B>>,
@@ -453,7 +453,9 @@ pub struct SegmentHistogramCollector<B> {
     dense_range: Option<DenseRange>,
 }
 
-impl<B: BucketIdSlot> SegmentAggregationCollector for SegmentHistogramCollector<B> {
+impl<B: BucketIdSlot, const SOURCE_CONTAINS_MULTIVALUES: bool> SegmentAggregationCollector
+    for SegmentHistogramCollector<B, SOURCE_CONTAINS_MULTIVALUES>
+{
     fn add_intermediate_aggregation_result(
         &mut self,
         agg_data: &AggregationsSegmentCtx,
@@ -489,25 +491,40 @@ impl<B: BucketIdSlot> SegmentAggregationCollector for SegmentHistogramCollector<
         let offset = req.offset;
         let get_bucket_pos = |val| get_bucket_pos_f64(val, interval, offset) as i64;
 
-        agg_data
-            .column_block_accessor
-            .fetch_block(docs, &*req.accessor);
-        // special path for nested buckets
-        if let Some(sub_agg) = &mut self.sub_agg {
-            for (doc, val) in agg_data.column_block_accessor.iter_docid_vals(docs) {
+        let accessor = &mut agg_data.column_block_accessor;
+        if SOURCE_CONTAINS_MULTIVALUES {
+            accessor.fetch_block_with_missing_unique_per_doc(docs, &*req.accessor, None, false);
+        } else {
+            accessor.fetch_block(docs, &*req.accessor);
+        }
+        // Known single-valued sources compile out deduplication; otherwise check the loaded batch.
+        let multivalued = SOURCE_CONTAINS_MULTIVALUES && accessor.is_batch_multivalued();
+        // Document IDs are needed for child collection and multivalued deduplication.
+        if self.sub_agg.is_some() || multivalued {
+            let mut previous = None;
+            for (doc, val) in accessor.iter_docid_vals(docs) {
                 let val = f64_from_fastfield_u64(val, self.column_type);
                 if bounds.contains(val) {
-                    let bucket = store.get_or_create(
-                        get_bucket_pos(val),
-                        &mut self.bucket_id_provider,
-                        |pos| get_bucket_key_from_pos(pos as f64, interval, offset),
-                    );
+                    let bucket_pos = get_bucket_pos(val);
+                    if multivalued {
+                        // The fetcher makes equal bucket hits consecutive within each document.
+                        if previous == Some((doc, bucket_pos)) {
+                            continue;
+                        }
+                        previous = Some((doc, bucket_pos));
+                    }
+                    let bucket =
+                        store.get_or_create(bucket_pos, &mut self.bucket_id_provider, |pos| {
+                            get_bucket_key_from_pos(pos as f64, interval, offset)
+                        });
                     bucket.doc_count += 1;
-                    sub_agg.push(bucket.bucket_id.to_bucket_id(), doc);
+                    if let Some(sub_agg) = &mut self.sub_agg {
+                        sub_agg.push(bucket.bucket_id.to_bucket_id(), doc);
+                    }
                 }
             }
         } else {
-            for val in agg_data.column_block_accessor.iter_vals() {
+            for val in accessor.iter_vals() {
                 let val = f64_from_fastfield_u64(val, self.column_type);
                 if bounds.contains(val) {
                     let bucket = store.get_or_create(
@@ -565,7 +582,9 @@ impl<B: BucketIdSlot> SegmentAggregationCollector for SegmentHistogramCollector<
     }
 }
 
-impl<B: BucketIdSlot> SegmentHistogramCollector<B> {
+impl<B: BucketIdSlot, const SOURCE_CONTAINS_MULTIVALUES: bool>
+    SegmentHistogramCollector<B, SOURCE_CONTAINS_MULTIVALUES>
+{
     fn get_memory_consumption(&self, parent_bucket_id: BucketId) -> u64 {
         self.parent_buckets[parent_bucket_id as usize].memory_consumption()
     }
@@ -627,7 +646,7 @@ impl<B: BucketIdSlot> SegmentHistogramCollector<B> {
     }
 }
 
-impl SegmentHistogramCollector<()> {
+impl SegmentHistogramCollector<(), false> {
     /// Builds a histogram collector whose parent `t` is a dense histogram filled from
     /// `counts[t * num_time_buckets .. (t + 1) * num_time_buckets]` (row-major), consolidating each
     /// cell's count lanes. Used by the flattened terms×histogram collector to turn its flat 2D
@@ -725,21 +744,38 @@ pub(crate) fn prepare_histogram_dense_range(
     Ok(dense_range.map(|range| (req_data, range)))
 }
 
-/// Builds a boxed histogram (or date histogram) segment collector, picking the bucket-id storage
-/// based on whether there are sub aggregations: `()` (no id stored) when there are none, otherwise
-/// [`BucketId`].
+/// Builds a histogram (or date histogram) collector specialized for source cardinality and
+/// bucket-id storage: `()` when there are no sub aggregations, otherwise [`BucketId`].
 pub(crate) fn build_segment_histogram_collector(
     agg_data: &mut AggregationsSegmentCtx,
     node: &AggRefNode,
 ) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
-    if node.children.is_empty() {
-        Ok(Box::new(
-            SegmentHistogramCollector::<()>::from_req_and_validate(agg_data, node)?,
-        ))
+    let accessor = &agg_data.per_request.histogram_req_data[node.idx_in_req_data].accessor;
+    // Computed sources may change cardinality between blocks.
+    let source_contains_multivalues = accessor.as_column().map_or(true, |column| {
+        column.index.get_cardinality().is_multivalue()
+    });
+    if source_contains_multivalues {
+        build_histogram_collector_with_cardinality::<true>(agg_data, node)
     } else {
-        Ok(Box::new(
-            SegmentHistogramCollector::<BucketId>::from_req_and_validate(agg_data, node)?,
-        ))
+        build_histogram_collector_with_cardinality::<false>(agg_data, node)
+    }
+}
+
+fn build_histogram_collector_with_cardinality<const SOURCE_CONTAINS_MULTIVALUES: bool>(
+    agg_data: &mut AggregationsSegmentCtx,
+    node: &AggRefNode,
+) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
+    if node.children.is_empty() {
+        Ok(Box::new(SegmentHistogramCollector::<
+            (),
+            SOURCE_CONTAINS_MULTIVALUES,
+        >::from_req_and_validate(agg_data, node)?))
+    } else {
+        Ok(Box::new(SegmentHistogramCollector::<
+            BucketId,
+            SOURCE_CONTAINS_MULTIVALUES,
+        >::from_req_and_validate(agg_data, node)?))
     }
 }
 
@@ -958,6 +994,136 @@ mod tests {
     };
     use crate::query::AllQuery;
 
+    #[test]
+    fn histogram_counts_each_document_once_per_bucket() -> crate::Result<()> {
+        use std::collections::{BTreeMap, BTreeSet};
+
+        use crate::schema::{Schema, FAST};
+        use crate::{DateTime, Index};
+
+        let rows = [
+            vec![12, 15],         // Distinct values in one bucket.
+            vec![18, 11],         // Two descending values in one bucket.
+            vec![25, 12],         // Two descending values in different buckets.
+            vec![12, 25, 15],     // Return to an earlier bucket.
+            vec![12, 12, 25, 25], // Repeated raw values.
+            vec![-18, -15],       // Negative bucket positions.
+            vec![],
+            vec![12],
+        ];
+        // One block uses sparse storage; multiple blocks allow densification.
+        for repetitions in [1, 16] {
+            let mut schema = Schema::builder();
+            let value = schema.add_i64_field("value", FAST);
+            let date = schema.add_date_field("date", FAST);
+            let score = schema.add_u64_field("score", FAST);
+            let index = Index::create_in_ram(schema.build());
+            let mut writer = index.writer_with_num_threads(1, 20_000_000)?;
+            let documents: Vec<_> = rows.iter().cycle().take(rows.len() * repetitions).collect();
+            for (doc, values) in documents.iter().enumerate() {
+                let mut document = doc!(score => doc as u64 + 1);
+                for &val in values.iter() {
+                    document.add_i64(value, val);
+                    document.add_date(date, DateTime::from_timestamp_secs(val));
+                }
+                writer.add_document(document)?;
+            }
+            writer.commit()?;
+
+            for (date_histogram, bounded) in [(false, false), (false, true), (true, false)] {
+                let offset = if bounded { 2 } else { 0 };
+                for with_children in [false, true] {
+                    let mut histogram = if date_histogram {
+                        json!({"date_histogram": {
+                            "field": "date", "fixed_interval": "10s", "min_doc_count": 1
+                        }})
+                    } else {
+                        json!({"histogram": {
+                            "field": "value", "interval": 10, "offset": offset, "min_doc_count": 1
+                        }})
+                    };
+                    if bounded {
+                        histogram["histogram"]["hard_bounds"] = json!({"min": 12, "max": 25});
+                    }
+                    if with_children {
+                        histogram["aggs"] = json!({
+                            "score": {"sum": {"field": "score"}},
+                            "values": {"sum": {"field": "value"}}
+                        });
+                    }
+                    let request = serde_json::from_value(json!({"histogram": histogram}))?;
+                    let result = exec_request(request, &index)?;
+                    let mut expected = BTreeMap::<i64, (u64, u64, i64)>::new();
+                    for (doc, values) in documents.iter().enumerate() {
+                        let keys: BTreeSet<_> = values
+                            .iter()
+                            .filter(|&&val| !bounded || (12..=25).contains(&val))
+                            .map(|&val| (val - offset).div_euclid(10) * 10 + offset)
+                            .collect();
+                        for key in keys {
+                            let entry = expected.entry(key).or_default();
+                            entry.0 += 1;
+                            entry.1 += doc as u64 + 1;
+                            entry.2 += values.iter().sum::<i64>();
+                        }
+                    }
+                    let buckets = result["histogram"]["buckets"].as_array().unwrap();
+                    assert_eq!(buckets.len(), expected.len());
+                    for (bucket, (key, (count, score_sum, value_sum))) in
+                        buckets.iter().zip(expected)
+                    {
+                        let key = if date_histogram { key * 1000 } else { key };
+                        assert_eq!(bucket["key"], key as f64);
+                        assert_eq!(bucket["doc_count"], count);
+                        if with_children {
+                            assert_eq!(bucket["score"]["value"], score_sum as f64);
+                            // Children retain every raw value, including duplicates.
+                            assert_eq!(bucket["values"]["value"], value_sum as f64);
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn histogram_counts_repeated_documents_in_separate_calls() -> crate::Result<()> {
+        use crate::aggregation::agg_data::{
+            build_aggregations_data_from_req, build_segment_agg_collectors_root,
+        };
+        use crate::schema::{Schema, FAST};
+        use crate::Index;
+
+        let mut schema = Schema::builder();
+        let value = schema.add_u64_field("value", FAST);
+        let index = Index::create_in_ram(schema.build());
+        let mut writer = index.writer_with_num_threads(1, 20_000_000)?;
+        writer.add_document(doc!(value => 12u64, value => 15u64))?;
+        writer.commit()?;
+        let reader = index.reader()?;
+        let searcher = reader.searcher();
+        let request = serde_json::from_value(json!({
+            "histogram": {"histogram": {"field": "value", "interval": 10}}
+        }))?;
+        let mut ctx = build_aggregations_data_from_req(
+            &request,
+            searcher.segment_reader(0),
+            0,
+            Default::default(),
+        )?;
+        let mut collector = build_segment_agg_collectors_root(&mut ctx)?;
+        collector.prepare_max_bucket(0, &ctx)?;
+        collector.collect(0, &[0], &mut ctx)?;
+        collector.collect(0, &[0], &mut ctx)?;
+        collector.flush(&mut ctx)?;
+        let mut result = IntermediateAggregationResults::default();
+        collector.add_intermediate_aggregation_result(&ctx, &mut result, 0)?;
+        let result = serde_json::to_value(result.into_final_result(request, Default::default())?)?;
+        assert_eq!(result["histogram"]["buckets"][0]["doc_count"], 2);
+        Ok(())
+    }
+
     #[test]
     fn histogram_test_crooked_values() -> crate::Result<()> {
         let values = vec![-12.0, 12.31, 14.33, 16.23];

src/aggregation/bucket/range.rs

@@ -157,7 +157,7 @@ pub(crate) struct SegmentRangeAndBucketEntry {
 
 /// The collector puts values from the fast field into the correct buckets and does a conversion to
 /// the correct datatype.
-pub struct SegmentRangeCollector<B: SubAggBuffer, const IS_MULTI_VALUED: bool> {
+pub struct SegmentRangeCollector<B: SubAggBuffer, const SOURCE_CONTAINS_MULTIVALUES: bool> {
     /// The buckets containing the aggregation data.
     /// One for each ParentBucketId
     parent_buckets: Vec<Vec<SegmentRangeAndBucketEntry>>,
@@ -179,8 +179,8 @@ pub struct SegmentRangeCollector<B: SubAggBuffer, const IS_MULTI_VALUED: bool> {
     limits: AggregationLimitsGuard,
 }
 
-impl<B: SubAggBuffer, const IS_MULTI_VALUED: bool> Debug
-    for SegmentRangeCollector<B, IS_MULTI_VALUED>
+impl<B: SubAggBuffer, const SOURCE_CONTAINS_MULTIVALUES: bool> Debug
+    for SegmentRangeCollector<B, SOURCE_CONTAINS_MULTIVALUES>
 {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("SegmentRangeCollector")
@@ -232,8 +232,8 @@ impl SegmentRangeBucketEntry {
     }
 }
 
-impl<B: SubAggBuffer, const IS_MULTI_VALUED: bool> SegmentAggregationCollector
-    for SegmentRangeCollector<B, IS_MULTI_VALUED>
+impl<B: SubAggBuffer, const SOURCE_CONTAINS_MULTIVALUES: bool> SegmentAggregationCollector
+    for SegmentRangeCollector<B, SOURCE_CONTAINS_MULTIVALUES>
 {
     fn add_intermediate_aggregation_result(
         &mut self,
@@ -283,7 +283,7 @@ impl<B: SubAggBuffer, const IS_MULTI_VALUED: bool> SegmentAggregationCollector
         agg_data: &mut AggregationsSegmentCtx,
     ) -> crate::Result<()> {
         let accessor = &mut agg_data.column_block_accessor;
-        if IS_MULTI_VALUED {
+        if SOURCE_CONTAINS_MULTIVALUES {
             accessor.fetch_block_with_missing_unique_per_doc(
                 docs,
                 &*self.req_data.accessor,
@@ -296,7 +296,8 @@ impl<B: SubAggBuffer, const IS_MULTI_VALUED: bool> SegmentAggregationCollector
 
         let buckets = &mut self.parent_buckets[parent_bucket_id as usize];
 
-        let multivalued = IS_MULTI_VALUED && accessor.is_batch_multivalued();
+        // Known single-valued sources compile out deduplication; otherwise check the loaded batch.
+        let multivalued = SOURCE_CONTAINS_MULTIVALUES && accessor.is_batch_multivalued();
         let mut previous = None;
         for (doc, val) in accessor.iter_docid_vals(docs) {
             let bucket_pos = get_bucket_pos(val, buckets);
@@ -359,17 +360,17 @@ pub(crate) fn build_segment_range_collector(
 ) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
     let accessor = &agg_data.per_request.range_req_data[node.idx_in_req_data].accessor;
     // Computed sources may change cardinality between blocks.
-    let multivalued = accessor.as_column().map_or(true, |column| {
+    let source_contains_multivalues = accessor.as_column().map_or(true, |column| {
         column.index.get_cardinality().is_multivalue()
     });
-    if multivalued {
+    if source_contains_multivalues {
         build_range_collector_with_cardinality::<true>(agg_data, node)
     } else {
         build_range_collector_with_cardinality::<false>(agg_data, node)
     }
 }
 
-fn build_range_collector_with_cardinality<const IS_MULTI_VALUED: bool>(
+fn build_range_collector_with_cardinality<const SOURCE_CONTAINS_MULTIVALUES: bool>(
     agg_data: &mut AggregationsSegmentCtx,
     node: &AggRefNode,
 ) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
@@ -393,7 +394,7 @@ fn build_range_collector_with_cardinality<const IS_MULTI_VALUED: bool>(
     if is_low_card {
         Ok(Box::new(SegmentRangeCollector::<
             LowCardSubAggBuffer,
-            IS_MULTI_VALUED,
+            SOURCE_CONTAINS_MULTIVALUES,
         > {
             sub_agg: sub_agg.map(LowCardBufferedSubAggs::new),
             req_data,
@@ -404,7 +405,7 @@ fn build_range_collector_with_cardinality<const IS_MULTI_VALUED: bool>(
     } else {
         Ok(Box::new(SegmentRangeCollector::<
             HighCardSubAggBuffer,
-            IS_MULTI_VALUED,
+            SOURCE_CONTAINS_MULTIVALUES,
         > {
             sub_agg: sub_agg.map(BufferedSubAggs::new),
             req_data,
@@ -415,7 +416,9 @@ fn build_range_collector_with_cardinality<const IS_MULTI_VALUED: bool>(
     }
 }
 
-impl<B: SubAggBuffer, const IS_MULTI_VALUED: bool> SegmentRangeCollector<B, IS_MULTI_VALUED> {
+impl<B: SubAggBuffer, const SOURCE_CONTAINS_MULTIVALUES: bool>
+    SegmentRangeCollector<B, SOURCE_CONTAINS_MULTIVALUES>
+{
     pub(crate) fn create_new_buckets(&mut self) -> crate::Result<Vec<SegmentRangeAndBucketEntry>> {
         let req_data = &self.req_data;
         let field_type = req_data.accessor.column_type();

src/aggregation/bucket/term_agg/flattened_term_histogram.rs

@@ -423,7 +423,7 @@ impl<R: BucketResolver, const LANES: usize> SegmentAggregationCollector
                 })
                 .collect(),
         };
-        let mut histogram = SegmentHistogramCollector::<()>::from_dense_rows(
+        let mut histogram = SegmentHistogramCollector::<(), false>::from_dense_rows(
             self.hist_req_data.clone(),
             self.base_pos,
             num_time_buckets,

src/aggregation/value_source/block_accessor.rs

@@ -235,9 +235,13 @@ impl ColumnBlockAccessor {
             && (self.cardinality.is_full() || self.docid_cache == docs)
     }
 
+    /// Whether any document has multiple values in the loaded batch.
+    /// Values must be grouped by document.
     #[inline]
     pub(crate) fn is_batch_multivalued(&self) -> bool {
+        // Full/Optional cannot repeat documents; Full may leave docid_cache stale.
         self.cardinality.is_multivalue()
+            && self.docid_cache.windows(2).any(|pair| pair[0] == pair[1])
     }
 
     #[inline]
@@ -372,6 +376,52 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_is_batch_multivalued_checks_loaded_values() {
+        let docs = [0, 1, 2];
+        let mut accessor = ColumnBlockAccessor::default();
+        for (entries, expected) in [
+            (vec![], false),
+            (vec![(0, 12)], false),
+            (vec![(0, 12), (1, 15), (2, 25)], false),
+            (vec![(0, 12), (2, 25)], false),
+            (vec![(0, 12), (0, 15)], true),
+            (vec![(0, 12), (2, 25), (2, 28)], true),
+        ] {
+            let source = TestValueSource {
+                cardinality: Cardinality::Multivalued,
+                entries,
+            };
+            accessor.fetch_block(&docs, &source);
+            assert_eq!(accessor.is_batch_multivalued(), expected);
+        }
+
+        let source = TestValueSource {
+            cardinality: Cardinality::Multivalued,
+            entries: vec![(0, 12), (0, 12)],
+        };
+        accessor.fetch_block(&docs, &source);
+        assert!(accessor.is_batch_multivalued());
+        accessor.fetch_block_with_missing_unique_per_doc(&docs, &source, None, false);
+        assert!(!accessor.is_batch_multivalued());
+    }
+
+    #[test]
+    fn test_is_batch_multivalued_ignores_stale_full_docids() {
+        let mut accessor = ColumnBlockAccessor::default();
+        let source = TestValueSource {
+            cardinality: Cardinality::Multivalued,
+            entries: vec![(0, 12), (0, 15)],
+        };
+        accessor.fetch_block(&[0], &source);
+        assert!(accessor.is_batch_multivalued());
+
+        let column = full_column(&[25]);
+        accessor.fetch_full_column_block(&[0], &*column.values);
+        assert_eq!(accessor.docids(), &[0, 0]);
+        assert!(!accessor.is_batch_multivalued());
+    }
+
     #[test]
     fn test_as_column_distinguishes_the_two_kinds() {
         let column: Arc<dyn ValueSource> = Arc::new((full_column(&[5, 6, 7]), ColumnType::U64));