diff --git a/src/aggregation/bucket/histogram/histogram.rs b/src/aggregation/bucket/histogram/histogram.rs index 5279310de..4e56a27b6 100644 --- a/src/aggregation/bucket/histogram/histogram.rs +++ b/src/aggregation/bucket/histogram/histogram.rs @@ -492,8 +492,32 @@ impl SegmentAggregationCollector for SegmentHistogramCollector< agg_data .column_block_accessor .fetch_block(docs, &*req.accessor); - // special path for nested buckets - if let Some(sub_agg) = &mut self.sub_agg { + if agg_data.column_block_accessor.is_multivalued() { + // Like Elasticsearch, a doc is counted once per bucket, even if several of its values + // fall in that bucket. Once the values of a doc are sorted, its bucket positions are + // non-decreasing, so comparing with the previous (doc, bucket_pos) is enough. + agg_data.column_block_accessor.sort_values_within_docs(); + let mut previous_doc_bucket_pos: Option<(crate::DocId, i64)> = None; + for (doc, val) in agg_data.column_block_accessor.iter_docid_vals(docs) { + let val = f64_from_fastfield_u64(val, self.column_type); + if !bounds.contains(val) { + continue; + } + let bucket_pos = get_bucket_pos(val); + if previous_doc_bucket_pos == Some((doc, bucket_pos)) { + continue; + } + previous_doc_bucket_pos = 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; + if let Some(sub_agg) = &mut self.sub_agg { + sub_agg.push(bucket.bucket_id.to_bucket_id(), doc); + } + } + } else if let Some(sub_agg) = &mut self.sub_agg { + // special path for nested buckets for (doc, val) in agg_data.column_block_accessor.iter_docid_vals(docs) { let val = f64_from_fastfield_u64(val, self.column_type); if bounds.contains(val) { @@ -1882,4 +1906,52 @@ mod tests { assert_eq!(first_bucket["key"], 0.0); assert_eq!(first_bucket["avg_metric"]["value"], 10.0); } + + /// Like Elasticsearch, a doc is counted once per bucket, even if several of its values fall + /// in it. Values are deliberately unsorted within the doc. + #[test] + fn histogram_multivalued_doc_counted_once_per_bucket() { + histogram_multivalued_doc_counted_once_per_bucket_aux(false); + histogram_multivalued_doc_counted_once_per_bucket_aux(true); + } + + fn histogram_multivalued_doc_counted_once_per_bucket_aux(with_sub_agg: bool) { + let mut schema_builder = crate::schema::Schema::builder(); + let bucket_field = schema_builder.add_f64_field("bucket_field", crate::schema::FAST); + let metric_field = schema_builder.add_f64_field("metric_field", crate::schema::FAST); + let index = crate::Index::create_in_ram(schema_builder.build()); + let mut index_writer: crate::IndexWriter = index.writer_for_tests().unwrap(); + index_writer + .add_document(doc!( + bucket_field => 3.0, + bucket_field => 25.0, + bucket_field => 1.0, + metric_field => 10.0 + )) + .unwrap(); + index_writer + .add_document(doc!(bucket_field => 2.0, metric_field => 20.0)) + .unwrap(); + index_writer.commit().unwrap(); + + let mut agg_req_json = json!({ + "histo": { "histogram": { "field": "bucket_field", "interval": 10.0 } } + }); + if with_sub_agg { + agg_req_json["histo"]["aggs"] = + json!({ "sum_metric": { "sum": { "field": "metric_field" } } }); + } + let agg_req: Aggregations = serde_json::from_value(agg_req_json).unwrap(); + + let res = exec_request(agg_req, &index).unwrap(); + let buckets = &res["histo"]["buckets"]; + assert_eq!(buckets[0]["key"], 0.0); + assert_eq!(buckets[0]["doc_count"], 2); + assert_eq!(buckets[2]["key"], 20.0); + assert_eq!(buckets[2]["doc_count"], 1); + if with_sub_agg { + assert_eq!(buckets[0]["sum_metric"]["value"], 30.0); + assert_eq!(buckets[2]["sum_metric"]["value"], 10.0); + } + } } diff --git a/src/aggregation/bucket/range.rs b/src/aggregation/bucket/range.rs index e2b581367..f00bf6b53 100644 --- a/src/aggregation/bucket/range.rs +++ b/src/aggregation/bucket/range.rs @@ -282,10 +282,21 @@ impl SegmentAggregationCollector for SegmentRangeCollector { .column_block_accessor .fetch_block(docs, &*self.req_data.accessor); + // Like Elasticsearch, a doc is counted once per bucket, even if several of its values + // fall in that bucket. Ranges do not overlap, so once the values of a doc are sorted, its + // bucket positions are non-decreasing, and comparing with the previous (doc, bucket_pos) + // is enough. Outside of multivalued blocks, a doc never appears twice. + agg_data.column_block_accessor.sort_values_within_docs(); + let buckets = &mut self.parent_buckets[parent_bucket_id as usize]; + let mut previous_doc_bucket_pos: Option<(crate::DocId, usize)> = None; for (doc, val) in agg_data.column_block_accessor.iter_docid_vals(docs) { let bucket_pos = get_bucket_pos(val, buckets); + if previous_doc_bucket_pos == Some((doc, bucket_pos)) { + continue; + } + previous_doc_bucket_pos = Some((doc, bucket_pos)); let bucket = &mut buckets[bucket_pos]; bucket.bucket.doc_count += 1; if let Some(sub_agg) = self.sub_agg.as_mut() { @@ -950,4 +961,47 @@ mod tests { assert_eq!(search(u64::MAX - 1), 2); // Since the end range is never included, // the max value } + + /// Like Elasticsearch, a doc is counted once per range, even if several of its values fall + /// in it. Values are deliberately unsorted within the doc. + #[test] + fn range_multivalued_doc_counted_once_per_bucket() { + let mut schema_builder = crate::schema::Schema::builder(); + let range_field = schema_builder.add_f64_field("range_field", crate::schema::FAST); + let metric_field = schema_builder.add_f64_field("metric_field", crate::schema::FAST); + let index = crate::Index::create_in_ram(schema_builder.build()); + let mut index_writer: crate::IndexWriter = index.writer_for_tests().unwrap(); + index_writer + .add_document(doc!( + range_field => 3.0, + range_field => 25.0, + range_field => 1.0, + metric_field => 10.0 + )) + .unwrap(); + index_writer + .add_document(doc!(range_field => 2.0, metric_field => 20.0)) + .unwrap(); + index_writer.commit().unwrap(); + + let agg_req: Aggregations = serde_json::from_value(json!({ + "my_range": { + "range": { + "field": "range_field", + "ranges": [{ "to": 10.0 }, { "from": 20.0, "to": 30.0 }] + }, + "aggs": { "sum_metric": { "sum": { "field": "metric_field" } } } + } + })) + .unwrap(); + + let res = exec_request(agg_req, &index).unwrap(); + let buckets = &res["my_range"]["buckets"]; + assert_eq!(buckets[0]["key"], "*-10"); + assert_eq!(buckets[0]["doc_count"], 2); + assert_eq!(buckets[0]["sum_metric"]["value"], 30.0); + assert_eq!(buckets[2]["key"], "20-30"); + assert_eq!(buckets[2]["doc_count"], 1); + assert_eq!(buckets[2]["sum_metric"]["value"], 10.0); + } } diff --git a/src/aggregation/value_source/block_accessor.rs b/src/aggregation/value_source/block_accessor.rs index ccc6b9951..e352403dc 100644 --- a/src/aggregation/value_source/block_accessor.rs +++ b/src/aggregation/value_source/block_accessor.rs @@ -1,6 +1,8 @@ use std::cmp::Ordering; use columnar::{Cardinality, ColumnValues, RowId}; +use common::GroupByIteratorExtended; +use itertools::Itertools; use crate::aggregation::value_source::ValueSource; use crate::DocId; @@ -56,7 +58,7 @@ impl ColumnBlockAccessor { /// Fetches a block from a column known to be full (hence we pass the ColumnValue Object /// directly). /// - /// docs needs to be strictly increasing. + /// `docs` can be in any order and contain duplicates. #[inline] pub(crate) fn fetch_full_column_block( &mut self, @@ -165,6 +167,30 @@ impl ColumnBlockAccessor { } } + /// Sorts the values of each document of a multivalued block, so that the values of a given doc + /// come in ascending `u64` order. No-op for non-multivalued blocks. + /// + /// Precondition: the entries of a given doc are contiguous in the docid cache. This holds + /// for blocks loaded by `fetch_block` (see [`ValueSource::load_block`]), but not after + /// `fetch_block_with_missing` with `ordered == false`. + pub(crate) fn sort_values_within_docs(&mut self) { + if !self.cardinality.is_multivalue() { + return; + } + let mut start = 0; + while start < self.docid_cache.len() { + let doc = self.docid_cache[start]; + let num_values_for_doc: usize = 1 + self.docid_cache[start + 1..] + .iter() + .take_while(|next_doc| **next_doc == doc) + .count(); + if num_values_for_doc > 1 { + self.val_cache[start..][..num_values_for_doc].sort_unstable(); + } + start += num_values_for_doc; + } + } + /// Removes duplicate (doc_id, value) pairs from the caches. /// /// After `fetch_block`, entries are sorted by doc_id, but values within diff --git a/src/aggregation/value_source/mod.rs b/src/aggregation/value_source/mod.rs index 12fec6a41..693be97e1 100644 --- a/src/aggregation/value_source/mod.rs +++ b/src/aggregation/value_source/mod.rs @@ -32,7 +32,8 @@ pub trait ValueSource: std::fmt::Debug { /// - `Full`: `values.len() == docs.len()` and `values[i]` is the value of `docs[i]`. `docids` /// is left unspecified and must not be read by the caller. /// - `Optional` / `Multivalued`: `docids.len() == values.len()` and `values[i]` is a value of - /// `docids[i]`. `docids` only contains docs from `docs`. A doc is repeated once per value. + /// `docids[i]`. `docids` only contains docs from `docs`. A doc is repeated once per value, + /// and the entries of a given doc are contiguous. /// /// `row_ids` is scratch the implementation may use freely. fn load_block( @@ -90,7 +91,7 @@ impl> + std::fmt::Debug> ValueSource for (ColumnRe } } -/// `docs` has to be sorted ascending and free of duplicates. +/// `docs` can be in any order and contain duplicates. #[inline] fn load_full_column_values( docs: &[DocId], @@ -112,15 +113,15 @@ fn load_full_column_values( /// Returns true if `docs` is a contiguous ascending run `[d, d + 1, ..., d + n - 1]`. /// -/// `docs` has to be sorted ascending and free of duplicates. +/// Accepts any input: sub-aggregations can receive duplicated or unordered docs. #[inline] fn is_contiguous(docs: &[u32]) -> bool { let (Some(&first), Some(&last)) = (docs.first(), docs.last()) else { return false; }; - debug_assert!( - docs.windows(2).all(|w| w[0] < w[1]), - "fetch_block requires docs sorted ascending without duplicates" - ); - (last - first) as usize + 1 == docs.len() + if last < first || (last - first) as usize + 1 != docs.len() { + return false; + } + // The span check alone is fooled by duplicates or unordered docs, e.g. `[0, 0, 2]`. + docs.windows(2).all(|pair| pair[0] + 1 == pair[1]) }