diff --git a/benches/agg_bench.rs b/benches/agg_bench.rs index d89c9b5bc..e5bb87170 100644 --- a/benches/agg_bench.rs +++ b/benches/agg_bench.rs @@ -73,6 +73,8 @@ fn bench_agg(mut group: InputGroup) { register!(group, terms_zipf_1000); register!(group, terms_zipf_1000_with_histogram); register!(group, terms_zipf_1000_with_avg_sub_agg); + register!(group, terms_zipf_90); + register!(group, terms_zipf_90_with_sum_sub_agg); register!(group, terms_many_json_mixed_type_with_avg_sub_agg); @@ -495,6 +497,27 @@ fn terms_zipf_1000(index: &Index) { execute_agg(index, agg_req); } +fn terms_zipf_90(index: &Index) { + let agg_req = json!({ + "my_texts": { "terms": { "field": "text_90_terms_zipf", "size": 100 } }, + }); + execute_agg(index, agg_req); +} + +// 90-term (low-cardinality Vec path) terms agg with a metric sub-agg. The skewed distribution keeps +// a dominant bucket crossing the sub-agg flush threshold, exercising the buffer flush path. +fn terms_zipf_90_with_sum_sub_agg(index: &Index) { + let agg_req = json!({ + "my_texts": { + "terms": { "field": "text_90_terms_zipf", "size": 100 }, + "aggs": { + "sum_score": { "sum": { "field": "score" } } + } + }, + }); + execute_agg(index, agg_req); +} + fn terms_many_json_mixed_type_with_avg_sub_agg(index: &Index) { let agg_req = json!({ "my_texts": { @@ -762,6 +785,8 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result { let text_field_few_terms = schema_builder.add_text_field("text_few_terms", STRING | FAST); let text_field_few_terms_status = schema_builder.add_text_field("text_few_terms_status", STRING | FAST); + let text_field_90_terms_zipf = + schema_builder.add_text_field("text_90_terms_zipf", STRING | FAST); let text_field_1000_terms_zipf = schema_builder.add_text_field("text_1000_terms_zipf", STRING | FAST); let score_fieldtype = tantivy::schema::NumericOptions::default().set_fast(); @@ -800,6 +825,12 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result { let terms_1000: Vec = (1..=1000).map(|i| format!("term_{i}")).collect(); let zipf_1000 = rand_distr::Zipf::new(1000.0, 1.1f64).unwrap(); + // 90 terms (< MAX_NUM_TERMS_FOR_VEC), skewed via Zipf so a dominant bucket keeps crossing the + // sub-agg flush threshold while minority buckets stay small. Exercises the low-cardinality + // sub-agg buffer flush path (see issue #2992). + let terms_90: Vec = (1..=90).map(|i| format!("term_{i}")).collect(); + let zipf_90 = rand_distr::Zipf::new(90.0, 1.1f64).unwrap(); + { let mut rng = StdRng::from_seed([1u8; 32]); let mut index_writer = index.writer_with_num_threads(1, 200_000_000)?; @@ -815,6 +846,8 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result { let idx_b = zipf_1000.sample(&mut rng) as usize - 1; let term_1000_a = &terms_1000[idx_a]; let term_1000_b = &terms_1000[idx_b]; + let term_90_a = &terms_90[zipf_90.sample(&mut rng) as usize - 1]; + let term_90_b = &terms_90[zipf_90.sample(&mut rng) as usize - 1]; index_writer.add_document(doc!( json_field => json!({"mixed_type": 10.0}), json_field => json!({"mixed_type": 10.0}), @@ -830,6 +863,8 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result { text_field_few_terms => "cool", text_field_few_terms_status => log_level_sample_a, text_field_few_terms_status => log_level_sample_b, + text_field_90_terms_zipf => term_90_a.as_str(), + text_field_90_terms_zipf => term_90_b.as_str(), text_field_1000_terms_zipf => term_1000_a.as_str(), text_field_1000_terms_zipf => term_1000_b.as_str(), score_field => 1u64, @@ -866,6 +901,7 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result { text_field_many_terms => many_terms_data.choose(&mut rng).unwrap().to_string(), text_field_few_terms => few_terms_data.choose(&mut rng).unwrap().to_string(), text_field_few_terms_status => status_field_data[log_level_distribution.sample(&mut rng)].0, + text_field_90_terms_zipf => terms_90[zipf_90.sample(&mut rng) as usize - 1].as_str(), text_field_1000_terms_zipf => terms_1000[zipf_1000.sample(&mut rng) as usize - 1].as_str(), score_field => val as u64, score_field_f64 => lg_norm.sample(&mut rng), diff --git a/src/aggregation/agg_tests.rs b/src/aggregation/agg_tests.rs index 09d9752e9..49b14898e 100644 --- a/src/aggregation/agg_tests.rs +++ b/src/aggregation/agg_tests.rs @@ -664,6 +664,82 @@ fn test_aggregation_flushing_variants() { test_aggregation_flushing(true, true).unwrap(); } +// Regression test for https://github.com/quickwit-oss/tantivy/issues/2992 +// +// A skewed terms bucket over <100 terms uses the low-cardinality (Vec) sub-agg buffer. A dominant +// term keeps crossing the periodic flush threshold (every 2048 docs), which used to drop the +// cached doc ids of the minority buckets before collecting them — corrupting their metric +// sub-aggregations while their doc counts stayed exact. +#[test] +fn test_terms_sub_agg_flushing_skewed_buckets() -> crate::Result<()> { + use std::collections::HashMap; + + // 89 minority terms + 1 dominant term = 90 distinct terms, staying below + // MAX_NUM_TERMS_FOR_VEC (100) so the low-cardinality Vec sub-agg buffer is used. + const NUM_MINORITY_TERMS: usize = 89; + + let mut values: Vec<(f64, String)> = Vec::new(); + let mut minority_idx = 0usize; + // Enough docs to cross the 2048 flush threshold multiple times. The minority docs are + // sprinkled among the dominant ones so they land in different flush windows. + for i in 0..5000u64 { + if i % 25 == 0 { + let term = format!("minority_{:02}", minority_idx % NUM_MINORITY_TERMS); + minority_idx += 1; + values.push(((i % 13 + 1) as f64, term)); + } else { + values.push((7.0, "dominant".to_string())); + } + } + + let mut truth: HashMap = HashMap::new(); + for (score, term) in &values { + let entry = truth.entry(term.clone()).or_insert((0, 0.0)); + entry.0 += 1; + entry.1 += *score; + } + // Sanity check on the shape of the generated data. + assert_eq!(truth.len(), NUM_MINORITY_TERMS + 1); + + let index = get_test_index_from_values_and_terms(false, &[values])?; + let reader = index.reader()?; + + let agg_req: Aggregations = serde_json::from_value(json!({ + "my_terms": { + "terms": { "field": "string_id", "size": 100 }, + "aggs": { + "sum_score": { "sum": { "field": "score" } } + } + } + })) + .unwrap(); + + let collector = get_collector(agg_req); + let searcher = reader.searcher(); + let agg_res = searcher.search(&AllQuery, &collector)?; + let res: Value = serde_json::from_str(&serde_json::to_string(&agg_res)?)?; + + let buckets = res["my_terms"]["buckets"].as_array().unwrap(); + // size 100 >= 90 distinct terms, so every bucket is returned. + assert_eq!(buckets.len(), truth.len()); + for bucket in buckets { + let key = bucket["key"].as_str().unwrap(); + let (true_count, true_sum) = truth[key]; + assert_eq!( + bucket["doc_count"].as_u64().unwrap(), + true_count, + "doc_count mismatch for {key}" + ); + assert_eq!( + bucket["sum_score"]["value"].as_f64().unwrap(), + true_sum, + "sum sub-agg mismatch for {key}" + ); + } + + Ok(()) +} + #[test] fn test_aggregation_level1_simple() -> crate::Result<()> { let index = get_test_index_2_segments(true)?; diff --git a/src/aggregation/buffered_sub_aggs.rs b/src/aggregation/buffered_sub_aggs.rs index 0e8c76706..c3555c832 100644 --- a/src/aggregation/buffered_sub_aggs.rs +++ b/src/aggregation/buffered_sub_aggs.rs @@ -2,7 +2,6 @@ use std::fmt::Debug; use super::segment_agg_result::SegmentAggregationCollector; use crate::aggregation::agg_data::AggregationsSegmentCtx; -use crate::aggregation::bucket::MAX_NUM_TERMS_FOR_VEC; use crate::aggregation::BucketId; use crate::DocId; @@ -45,7 +44,6 @@ pub trait SubAggBuffer: Debug { &mut self, sub_agg: &mut Box, agg_data: &mut AggregationsSegmentCtx, - force: bool, ) -> crate::Result<()>; } @@ -76,7 +74,7 @@ impl BufferedSubAggs { ) -> crate::Result<()> { if self.num_docs >= FLUSH_THRESHOLD { self.buffer - .flush_local(&mut self.sub_agg_collector, agg_data, false)?; + .flush_local(&mut self.sub_agg_collector, agg_data)?; self.num_docs = 0; } Ok(()) @@ -86,7 +84,7 @@ impl BufferedSubAggs { pub fn flush(&mut self, agg_data: &mut AggregationsSegmentCtx) -> crate::Result<()> { if self.num_docs != 0 { self.buffer - .flush_local(&mut self.sub_agg_collector, agg_data, true)?; + .flush_local(&mut self.sub_agg_collector, agg_data)?; self.num_docs = 0; } self.sub_agg_collector.flush(agg_data)?; @@ -150,7 +148,6 @@ impl SubAggBuffer for HighCardSubAggBuffer { &mut self, sub_agg: &mut Box, agg_data: &mut AggregationsSegmentCtx, - _force: bool, ) -> crate::Result<()> { let mut max_bucket = 0u32; for partition in self.partitions.iter() { @@ -210,34 +207,11 @@ impl SubAggBuffer for LowCardSubAggBuffer { &mut self, sub_agg: &mut Box, agg_data: &mut AggregationsSegmentCtx, - force: bool, ) -> crate::Result<()> { // Pre-aggregated: call collect per bucket. let max_bucket = (self.per_bucket_docs.len() as BucketId).saturating_sub(1); sub_agg.prepare_max_bucket(max_bucket, agg_data)?; - // The threshold above which we flush buckets individually. - // Note: We need to make sure that we don't lock ourselves into a situation where we hit - // the FLUSH_THRESHOLD, but never flush any buckets. (except the final flush) - let mut bucket_treshold = FLUSH_THRESHOLD / (self.per_bucket_docs.len().max(1) * 2); - const _: () = { - // MAX_NUM_TERMS_FOR_VEC threshold is used for term aggregations - // Note: There may be other flexible values, for other aggregations, but we can use the - // const value here as a upper bound. (better than nothing) - let bucket_treshold_limit = FLUSH_THRESHOLD / (MAX_NUM_TERMS_FOR_VEC as usize * 2); - assert!( - bucket_treshold_limit > 0, - "Bucket threshold must be greater than 0" - ); - }; - if force { - bucket_treshold = 0; - } - for (bucket_id, docs) in self - .per_bucket_docs - .iter() - .enumerate() - .filter(|(_, docs)| docs.len() > bucket_treshold) - { + for (bucket_id, docs) in self.per_bucket_docs.iter().enumerate() { sub_agg.collect(bucket_id as BucketId, docs, agg_data)?; }