From d134c0e8306cd6dbb24e58afdf73e525dacae16e Mon Sep 17 00:00:00 2001 From: DeviousCardi <115358213+DeviousCardi@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:06:17 +0530 Subject: [PATCH 1/3] do not prune intermediate aggregation results at the root Pruning to `segment_size` bounds what a leaf sends upstream, so it is worth paying there. The root has no such consumer: finalization already prunes to the requested `size`, and cutting the candidate set a second time discards counts that finalization could still have used, which surfaces as an inflated `doc_count_error_upper_bound`. Give `QuickwitCollector` a `MergeLevel`, set by whoever builds it, and prune intermediate results only when merging at a leaf. The effect is visible once the root merges two or more leaf responses; a root with a single leaf response returns it untouched and never reached the pruning either way. --- quickwit/quickwit-search/src/collector.rs | 78 +++++++++++++++++++---- quickwit/quickwit-search/src/leaf.rs | 18 ++++-- quickwit/quickwit-search/src/root.rs | 12 ++-- 3 files changed, 86 insertions(+), 22 deletions(-) diff --git a/quickwit/quickwit-search/src/collector.rs b/quickwit/quickwit-search/src/collector.rs index 08ed20a8486..b12fc2a2f0d 100644 --- a/quickwit/quickwit-search/src/collector.rs +++ b/quickwit/quickwit-search/src/collector.rs @@ -682,7 +682,7 @@ impl QuickwitIncrementalAggregations { } } - fn finalize(self) -> tantivy::Result>> { + fn finalize(self, merge_level: MergeLevel) -> tantivy::Result>> { match self { QuickwitIncrementalAggregations::FindTraceIdsAggregation(collector, mut state) => { let merged_fruit = if state.len() > 1 { @@ -697,6 +697,7 @@ impl QuickwitIncrementalAggregations { merge_intermediate_aggregation_result( &Some(QuickwitAggregations::TantivyAggregations(aggregation)), state.iter().map(|vec| vec.as_slice()), + merge_level, ) } QuickwitIncrementalAggregations::NoAggregation => Ok(None), @@ -704,6 +705,18 @@ impl QuickwitIncrementalAggregations { } } +/// Where in the search tree a collector is merging results. +/// +/// Leaf merges are followed by more merging upstream, so their intermediate results are +/// pruned back to `segment_size` to bound what crosses the network. A root merge has no +/// such consumer: finalization prunes it to the requested `size` anyway, and pruning it +/// twice only widens `doc_count_error_upper_bound`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MergeLevel { + Leaf, + Root, +} + /// The quickwit collector is the tantivy Collector used in Quickwit. /// /// It defines the data that should be accumulated about the documents matching @@ -716,6 +729,7 @@ pub(crate) struct QuickwitCollector { pub sort_by: SortByPair, pub aggregation: Option, pub agg_context_params: AggContextParams, + pub merge_level: MergeLevel, search_after: Option, } @@ -846,6 +860,7 @@ impl Collector for QuickwitCollector { sort_order1, sort_order2, num_hits, + self.merge_level, )?; // ... and drop the first [..start_offsets) hits. // note that self.start_offset is 0 when merging from leaf_search, and is only set when @@ -870,6 +885,7 @@ fn map_error(error: postcard::Error) -> TantivyError { fn merge_intermediate_aggregation_result<'a>( aggregations_opt: &Option, intermediate_aggregation_results: impl Iterator, + merge_level: MergeLevel, ) -> tantivy::Result>> { let merged_intermediate_aggregation_result = match aggregations_opt { Some(QuickwitAggregations::FindTraceIdsAggregation(collector)) => { @@ -903,7 +919,14 @@ fn merge_intermediate_aggregation_result<'a>( let mut merged = merged_opt.unwrap_or_default(); // Leaf results can be merged again at the root or by a federated query. Keep the // intermediate candidate set (`segment_size`) rather than applying final pruning. - merged.prune_intermediate_results(aggregations, PruneMode::Intermediate)?; + // + // At the root there is no further merge to bound: finalization already prunes to + // the requested `size`, and cutting the candidate set here as well only discards + // counts that finalization could still have used, inflating + // `doc_count_error_upper_bound`. + if merge_level == MergeLevel::Leaf { + merged.prune_intermediate_results(aggregations, PruneMode::Intermediate)?; + } let serialized = postcard::to_allocvec(&merged).map_err(map_error)?; Some(serialized) } @@ -920,6 +943,7 @@ fn merge_leaf_responses( sort_order1: SortOrder, sort_order2: SortOrder, max_hits: usize, + merge_level: MergeLevel, ) -> tantivy::Result { // Optimization: No merging needed if there is only one result. if leaf_responses.len() == 1 { @@ -937,6 +961,7 @@ fn merge_leaf_responses( leaf_responses.iter().filter_map(|leaf_response| { leaf_response.intermediate_aggregation_result.as_deref() }), + merge_level, )?; let num_attempted_splits = leaf_responses .iter() @@ -1050,6 +1075,7 @@ pub(crate) fn make_collector_for_split( sort_by, aggregation, agg_context_params, + merge_level: MergeLevel::Leaf, search_after: search_request.search_after.clone(), }) } @@ -1058,6 +1084,7 @@ pub(crate) fn make_collector_for_split( pub(crate) fn make_merge_collector( search_request: &SearchRequest, agg_limits: AggregationLimitsGuard, + merge_level: MergeLevel, ) -> crate::Result { // Note: at this point the tokenizer manager is not used anymore by aggregations (filter query), // so we can create an empty one. So if it will ever be used, it would panic. @@ -1078,6 +1105,7 @@ pub(crate) fn make_merge_collector( sort_by, aggregation, agg_context_params, + merge_level, search_after: search_request.search_after.clone(), }) } @@ -1204,6 +1232,7 @@ pub(crate) struct IncrementalCollector { num_successful_splits: u64, start_offset: usize, resource_stats: Option, + merge_level: MergeLevel, } impl IncrementalCollector { @@ -1214,6 +1243,7 @@ impl IncrementalCollector { .as_ref() .map(QuickwitAggregations::maybe_incremental_aggregator) .unwrap_or(QuickwitIncrementalAggregations::NoAggregation); + let merge_level = collector.merge_level; let (order1, order2) = collector.sort_by.sort_orders(); let sort_key_mapper = HitSortingMapper { order1, order2 }; IncrementalCollector { @@ -1225,6 +1255,7 @@ impl IncrementalCollector { num_attempted_splits: 0, num_successful_splits: 0, resource_stats: None, + merge_level, } } @@ -1298,7 +1329,8 @@ impl IncrementalCollector { /// Finalize the merge, creating a LeafSearchResponse. pub(crate) fn finalize(self) -> tantivy::Result { - let intermediate_aggregation_result = self.incremental_aggregation.finalize()?; + let intermediate_aggregation_result = + self.incremental_aggregation.finalize(self.merge_level)?; let mut partial_hits = self.top_k_hits.finalize(); if self.start_offset != 0 { partial_hits.drain(0..self.start_offset.min(partial_hits.len())); @@ -1331,7 +1363,7 @@ mod tests { use super::{IncrementalCollector, make_merge_collector}; use crate::QuickwitAggregations; - use crate::collector::{merge_intermediate_aggregation_result, top_k_partial_hits}; + use crate::collector::{MergeLevel, merge_intermediate_aggregation_result, top_k_partial_hits}; #[test] fn test_merge_partial_hits_no_tie() { @@ -1778,7 +1810,8 @@ mod tests { request: &SearchRequest, results: Vec, ) -> LeafSearchResponse { - let collector = make_merge_collector(request, Default::default()).unwrap(); + let collector = + make_merge_collector(request, Default::default(), MergeLevel::Leaf).unwrap(); let mut incremental_collector = IncrementalCollector::new(collector.clone()); let result = collector @@ -2064,7 +2097,9 @@ mod tests { #[test] fn test_merge_empty_intermediate_aggregation_result() { - let merged = merge_intermediate_aggregation_result(&None, std::iter::empty()).unwrap(); + let merged = + merge_intermediate_aggregation_result(&None, std::iter::empty(), MergeLevel::Leaf) + .unwrap(); assert!(merged.is_none()); let aggregations_json = r#"{ @@ -2072,16 +2107,34 @@ mod tests { }"#; let ttv_aggregations: Aggregations = serde_json::from_str(aggregations_json).unwrap(); let qw_aggregations = QuickwitAggregations::TantivyAggregations(ttv_aggregations); - let serialized = - merge_intermediate_aggregation_result(&Some(qw_aggregations), std::iter::empty()) - .unwrap() - .unwrap(); + let serialized = merge_intermediate_aggregation_result( + &Some(qw_aggregations), + std::iter::empty(), + MergeLevel::Leaf, + ) + .unwrap() + .unwrap(); let _merged: IntermediateAggregationResults = postcard::from_bytes(&serialized).unwrap(); // Hopefully `_merged` is empty but the API does not allow us to assert that. } #[test] fn test_merge_intermediate_terms_prunes_to_segment_size() { + // Nine distinct terms are merged, and leaf pruning cuts them down to `segment_size`, + // keeping more candidates than the final `size` for the root to choose between. + assert_eq!(merged_terms_bucket_count(MergeLevel::Leaf), 4); + } + + /// The root has no downstream consumer to bound: finalization prunes to the requested + /// `size` on its own. Pruning here as well would throw away counts that finalization + /// could still have used, which shows up as an inflated `doc_count_error_upper_bound`. + #[test] + fn test_merge_intermediate_terms_does_not_prune_at_root() { + assert_eq!(merged_terms_bucket_count(MergeLevel::Root), 9); + } + + /// Merges three disjoint three-term fruits and reports how many term buckets survive. + fn merged_terms_bucket_count(merge_level: MergeLevel) -> usize { use tantivy::Index; use tantivy::aggregation::DistributedAggregationCollector; use tantivy::aggregation::intermediate_agg_result::{ @@ -2133,6 +2186,7 @@ mod tests { let serialized = merge_intermediate_aggregation_result( &Some(quickwit_aggregations), serialized_fruits.iter().map(Vec::as_slice), + merge_level, ) .unwrap() .unwrap(); @@ -2143,8 +2197,6 @@ mod tests { else { panic!("expected terms aggregation result"); }; - // Intermediate mode preserves more candidates than the final size (2), while bounding the - // merged response to segment_size (4). - assert_eq!(buckets.entries().len(), 4); + buckets.entries().len() } } diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 245a18db75a..e5c92852f0d 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -58,7 +58,9 @@ use tokio::task::{JoinError, JoinSet}; use tokio_util::sync::CancellationToken; use tracing::*; -use crate::collector::{IncrementalCollector, make_collector_for_split, make_merge_collector}; +use crate::collector::{ + IncrementalCollector, MergeLevel, make_collector_for_split, make_merge_collector, +}; use crate::leaf_cache::LeafSearchCache; use crate::metrics::{ LEAF_SEARCH_SINGLE_SPLIT_WARMUP_NUM_BYTES, LEAF_SEARCH_SPLIT_DURATION_SECS, @@ -1620,8 +1622,11 @@ pub async fn multi_index_leaf_search( } // Creates a collector which merges responses into one - let merge_collector = - make_merge_collector(&search_request, searcher_context.get_aggregation_limits())?; + let merge_collector = make_merge_collector( + &search_request, + searcher_context.get_aggregation_limits(), + MergeLevel::Leaf, + )?; let mut incremental_merge_collector = IncrementalCollector::new(merge_collector); while let Some(leaf_response_join_result) = leaf_request_futures.join_next().await { @@ -1978,8 +1983,11 @@ pub async fn single_doc_mapping_leaf_search( } let split_filter_arc: Arc> = Arc::new(RwLock::new(split_filter)); - let merge_collector = - make_merge_collector(&request, searcher_context.get_aggregation_limits())?; + let merge_collector = make_merge_collector( + &request, + searcher_context.get_aggregation_limits(), + MergeLevel::Leaf, + )?; let mut incremental_merge_collector = IncrementalCollector::new(merge_collector); let split_outcome_counters = Arc::new(SplitSearchOutcomeCounters::default()); diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index c31a848f8ea..289421b48c6 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -51,7 +51,7 @@ use tracing::{Span, debug, error, info, info_span, instrument}; use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::cluster_client::ClusterClient; -use crate::collector::{QuickwitAggregations, make_merge_collector}; +use crate::collector::{MergeLevel, QuickwitAggregations, make_merge_collector}; use crate::metrics_trackers::{RootSearchMetricsFuture, RootSearchMetricsStep}; use crate::scroll_context::{ScrollContext, ScrollKeyAndStartOffset}; use crate::search_job_placer::{Job, group_by, group_jobs_by_index_id}; @@ -840,8 +840,11 @@ pub(crate) async fn search_partial_hits_phase( num_failed_splits, ); - let merge_collector = - make_merge_collector(search_request, searcher_context.get_aggregation_limits())?; + let merge_collector = make_merge_collector( + search_request, + searcher_context.get_aggregation_limits(), + MergeLevel::Root, + )?; // Merging is a cpu-bound task. Prioritize it over queued split searches to avoid delaying the // final response once all leaf responses are available. @@ -1424,7 +1427,8 @@ pub async fn search_plan( true, None, )?; - let merge_collector = make_merge_collector(&search_request, Default::default())?; + let merge_collector = + make_merge_collector(&search_request, Default::default(), MergeLevel::Root)?; warmup_info.merge(merge_collector.warmup_info()); warmup_info.simplify(); From b20d28949441ea0dbbc08f7dfe328155d9dab045 Mon Sep 17 00:00:00 2001 From: DeviousCardi <115358213+DeviousCardi@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:20:25 +0530 Subject: [PATCH 2/3] keep pruning at the root when finalization is skipped `skip_aggregation_finalization` makes `finalize_aggregation_if_any` hand the merged bytes straight back to the caller, so the root merge is then the only place that can bound them. Skipping intermediate pruning there let a high-cardinality terms aggregation return close to the union of every leaf's candidate set. Rename the flag to say what it decides rather than where it is used: the question is whether anything downstream will prune the result, not whether the merge happens at a leaf or at the root. --- quickwit/quickwit-search/src/collector.rs | 88 +++++++++++++---------- quickwit/quickwit-search/src/leaf.rs | 6 +- quickwit/quickwit-search/src/root.rs | 49 +++++++++++-- 3 files changed, 98 insertions(+), 45 deletions(-) diff --git a/quickwit/quickwit-search/src/collector.rs b/quickwit/quickwit-search/src/collector.rs index b12fc2a2f0d..158fd937992 100644 --- a/quickwit/quickwit-search/src/collector.rs +++ b/quickwit/quickwit-search/src/collector.rs @@ -682,7 +682,10 @@ impl QuickwitIncrementalAggregations { } } - fn finalize(self, merge_level: MergeLevel) -> tantivy::Result>> { + fn finalize( + self, + intermediate_pruning: IntermediatePruning, + ) -> tantivy::Result>> { match self { QuickwitIncrementalAggregations::FindTraceIdsAggregation(collector, mut state) => { let merged_fruit = if state.len() > 1 { @@ -697,7 +700,7 @@ impl QuickwitIncrementalAggregations { merge_intermediate_aggregation_result( &Some(QuickwitAggregations::TantivyAggregations(aggregation)), state.iter().map(|vec| vec.as_slice()), - merge_level, + intermediate_pruning, ) } QuickwitIncrementalAggregations::NoAggregation => Ok(None), @@ -705,16 +708,19 @@ impl QuickwitIncrementalAggregations { } } -/// Where in the search tree a collector is merging results. +/// Whether a merge prunes its intermediate results back to `segment_size`. +/// +/// Pruning bounds what the merged result costs to carry, so it is worth paying whenever +/// something downstream still has to receive or merge that result: a leaf sending to the +/// root, or a root answering a caller that asked to skip finalization. /// -/// Leaf merges are followed by more merging upstream, so their intermediate results are -/// pruned back to `segment_size` to bound what crosses the network. A root merge has no -/// such consumer: finalization prunes it to the requested `size` anyway, and pruning it -/// twice only widens `doc_count_error_upper_bound`. +/// It is pure loss when finalization runs next. Finalization prunes to the requested +/// `size` on its own, so cutting the candidate set beforehand only discards counts it +/// could have used, which shows up as a wider `doc_count_error_upper_bound`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum MergeLevel { - Leaf, - Root, +pub(crate) enum IntermediatePruning { + Apply, + Skip, } /// The quickwit collector is the tantivy Collector used in Quickwit. @@ -729,7 +735,7 @@ pub(crate) struct QuickwitCollector { pub sort_by: SortByPair, pub aggregation: Option, pub agg_context_params: AggContextParams, - pub merge_level: MergeLevel, + pub intermediate_pruning: IntermediatePruning, search_after: Option, } @@ -860,7 +866,7 @@ impl Collector for QuickwitCollector { sort_order1, sort_order2, num_hits, - self.merge_level, + self.intermediate_pruning, )?; // ... and drop the first [..start_offsets) hits. // note that self.start_offset is 0 when merging from leaf_search, and is only set when @@ -885,7 +891,7 @@ fn map_error(error: postcard::Error) -> TantivyError { fn merge_intermediate_aggregation_result<'a>( aggregations_opt: &Option, intermediate_aggregation_results: impl Iterator, - merge_level: MergeLevel, + intermediate_pruning: IntermediatePruning, ) -> tantivy::Result>> { let merged_intermediate_aggregation_result = match aggregations_opt { Some(QuickwitAggregations::FindTraceIdsAggregation(collector)) => { @@ -924,7 +930,7 @@ fn merge_intermediate_aggregation_result<'a>( // the requested `size`, and cutting the candidate set here as well only discards // counts that finalization could still have used, inflating // `doc_count_error_upper_bound`. - if merge_level == MergeLevel::Leaf { + if intermediate_pruning == IntermediatePruning::Apply { merged.prune_intermediate_results(aggregations, PruneMode::Intermediate)?; } let serialized = postcard::to_allocvec(&merged).map_err(map_error)?; @@ -943,7 +949,7 @@ fn merge_leaf_responses( sort_order1: SortOrder, sort_order2: SortOrder, max_hits: usize, - merge_level: MergeLevel, + intermediate_pruning: IntermediatePruning, ) -> tantivy::Result { // Optimization: No merging needed if there is only one result. if leaf_responses.len() == 1 { @@ -961,7 +967,7 @@ fn merge_leaf_responses( leaf_responses.iter().filter_map(|leaf_response| { leaf_response.intermediate_aggregation_result.as_deref() }), - merge_level, + intermediate_pruning, )?; let num_attempted_splits = leaf_responses .iter() @@ -1075,7 +1081,7 @@ pub(crate) fn make_collector_for_split( sort_by, aggregation, agg_context_params, - merge_level: MergeLevel::Leaf, + intermediate_pruning: IntermediatePruning::Apply, search_after: search_request.search_after.clone(), }) } @@ -1084,7 +1090,7 @@ pub(crate) fn make_collector_for_split( pub(crate) fn make_merge_collector( search_request: &SearchRequest, agg_limits: AggregationLimitsGuard, - merge_level: MergeLevel, + intermediate_pruning: IntermediatePruning, ) -> crate::Result { // Note: at this point the tokenizer manager is not used anymore by aggregations (filter query), // so we can create an empty one. So if it will ever be used, it would panic. @@ -1105,7 +1111,7 @@ pub(crate) fn make_merge_collector( sort_by, aggregation, agg_context_params, - merge_level, + intermediate_pruning, search_after: search_request.search_after.clone(), }) } @@ -1232,7 +1238,7 @@ pub(crate) struct IncrementalCollector { num_successful_splits: u64, start_offset: usize, resource_stats: Option, - merge_level: MergeLevel, + intermediate_pruning: IntermediatePruning, } impl IncrementalCollector { @@ -1243,7 +1249,7 @@ impl IncrementalCollector { .as_ref() .map(QuickwitAggregations::maybe_incremental_aggregator) .unwrap_or(QuickwitIncrementalAggregations::NoAggregation); - let merge_level = collector.merge_level; + let intermediate_pruning = collector.intermediate_pruning; let (order1, order2) = collector.sort_by.sort_orders(); let sort_key_mapper = HitSortingMapper { order1, order2 }; IncrementalCollector { @@ -1255,7 +1261,7 @@ impl IncrementalCollector { num_attempted_splits: 0, num_successful_splits: 0, resource_stats: None, - merge_level, + intermediate_pruning, } } @@ -1329,8 +1335,9 @@ impl IncrementalCollector { /// Finalize the merge, creating a LeafSearchResponse. pub(crate) fn finalize(self) -> tantivy::Result { - let intermediate_aggregation_result = - self.incremental_aggregation.finalize(self.merge_level)?; + let intermediate_aggregation_result = self + .incremental_aggregation + .finalize(self.intermediate_pruning)?; let mut partial_hits = self.top_k_hits.finalize(); if self.start_offset != 0 { partial_hits.drain(0..self.start_offset.min(partial_hits.len())); @@ -1363,7 +1370,9 @@ mod tests { use super::{IncrementalCollector, make_merge_collector}; use crate::QuickwitAggregations; - use crate::collector::{MergeLevel, merge_intermediate_aggregation_result, top_k_partial_hits}; + use crate::collector::{ + IntermediatePruning, merge_intermediate_aggregation_result, top_k_partial_hits, + }; #[test] fn test_merge_partial_hits_no_tie() { @@ -1811,7 +1820,7 @@ mod tests { results: Vec, ) -> LeafSearchResponse { let collector = - make_merge_collector(request, Default::default(), MergeLevel::Leaf).unwrap(); + make_merge_collector(request, Default::default(), IntermediatePruning::Apply).unwrap(); let mut incremental_collector = IncrementalCollector::new(collector.clone()); let result = collector @@ -2097,9 +2106,12 @@ mod tests { #[test] fn test_merge_empty_intermediate_aggregation_result() { - let merged = - merge_intermediate_aggregation_result(&None, std::iter::empty(), MergeLevel::Leaf) - .unwrap(); + let merged = merge_intermediate_aggregation_result( + &None, + std::iter::empty(), + IntermediatePruning::Apply, + ) + .unwrap(); assert!(merged.is_none()); let aggregations_json = r#"{ @@ -2110,7 +2122,7 @@ mod tests { let serialized = merge_intermediate_aggregation_result( &Some(qw_aggregations), std::iter::empty(), - MergeLevel::Leaf, + IntermediatePruning::Apply, ) .unwrap() .unwrap(); @@ -2122,19 +2134,19 @@ mod tests { fn test_merge_intermediate_terms_prunes_to_segment_size() { // Nine distinct terms are merged, and leaf pruning cuts them down to `segment_size`, // keeping more candidates than the final `size` for the root to choose between. - assert_eq!(merged_terms_bucket_count(MergeLevel::Leaf), 4); + assert_eq!(merged_terms_bucket_count(IntermediatePruning::Apply), 4); } - /// The root has no downstream consumer to bound: finalization prunes to the requested - /// `size` on its own. Pruning here as well would throw away counts that finalization - /// could still have used, which shows up as an inflated `doc_count_error_upper_bound`. + /// When finalization runs next it prunes to the requested `size` on its own, so the + /// merge keeps every candidate. Pruning here as well would throw away counts that + /// finalization could have used, inflating `doc_count_error_upper_bound`. #[test] - fn test_merge_intermediate_terms_does_not_prune_at_root() { - assert_eq!(merged_terms_bucket_count(MergeLevel::Root), 9); + fn test_merge_intermediate_terms_kept_when_finalization_follows() { + assert_eq!(merged_terms_bucket_count(IntermediatePruning::Skip), 9); } /// Merges three disjoint three-term fruits and reports how many term buckets survive. - fn merged_terms_bucket_count(merge_level: MergeLevel) -> usize { + fn merged_terms_bucket_count(intermediate_pruning: IntermediatePruning) -> usize { use tantivy::Index; use tantivy::aggregation::DistributedAggregationCollector; use tantivy::aggregation::intermediate_agg_result::{ @@ -2186,7 +2198,7 @@ mod tests { let serialized = merge_intermediate_aggregation_result( &Some(quickwit_aggregations), serialized_fruits.iter().map(Vec::as_slice), - merge_level, + intermediate_pruning, ) .unwrap() .unwrap(); diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index e5c92852f0d..f5297b96a46 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -59,7 +59,7 @@ use tokio_util::sync::CancellationToken; use tracing::*; use crate::collector::{ - IncrementalCollector, MergeLevel, make_collector_for_split, make_merge_collector, + IncrementalCollector, IntermediatePruning, make_collector_for_split, make_merge_collector, }; use crate::leaf_cache::LeafSearchCache; use crate::metrics::{ @@ -1625,7 +1625,7 @@ pub async fn multi_index_leaf_search( let merge_collector = make_merge_collector( &search_request, searcher_context.get_aggregation_limits(), - MergeLevel::Leaf, + IntermediatePruning::Apply, )?; let mut incremental_merge_collector = IncrementalCollector::new(merge_collector); @@ -1986,7 +1986,7 @@ pub async fn single_doc_mapping_leaf_search( let merge_collector = make_merge_collector( &request, searcher_context.get_aggregation_limits(), - MergeLevel::Leaf, + IntermediatePruning::Apply, )?; let mut incremental_merge_collector = IncrementalCollector::new(merge_collector); diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index 289421b48c6..72ad7b9ea53 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -51,7 +51,7 @@ use tracing::{Span, debug, error, info, info_span, instrument}; use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::cluster_client::ClusterClient; -use crate::collector::{MergeLevel, QuickwitAggregations, make_merge_collector}; +use crate::collector::{IntermediatePruning, QuickwitAggregations, make_merge_collector}; use crate::metrics_trackers::{RootSearchMetricsFuture, RootSearchMetricsStep}; use crate::scroll_context::{ScrollContext, ScrollKeyAndStartOffset}; use crate::search_job_placer::{Job, group_by, group_jobs_by_index_id}; @@ -843,7 +843,7 @@ pub(crate) async fn search_partial_hits_phase( let merge_collector = make_merge_collector( search_request, searcher_context.get_aggregation_limits(), - MergeLevel::Root, + root_intermediate_pruning(search_request), )?; // Merging is a cpu-bound task. Prioritize it over queued split searches to avoid delaying the @@ -1154,6 +1154,20 @@ fn finalize_aggregation( Ok(Some(merge_aggregation_result)) } +/// Whether the root merge has to bound its own intermediate results. +/// +/// [`finalize_aggregation_if_any`] hands the merged bytes back untouched when the caller +/// asked to skip finalization, which makes the root merge the last place able to bound +/// them. When finalization does run it prunes to the requested `size` on its own, and +/// pruning beforehand would only widen `doc_count_error_upper_bound`. +fn root_intermediate_pruning(search_request: &SearchRequest) -> IntermediatePruning { + if search_request.skip_aggregation_finalization { + IntermediatePruning::Apply + } else { + IntermediatePruning::Skip + } +} + fn finalize_aggregation_if_any( search_request: &SearchRequest, intermediate_aggregation_result_bytes_opt: Option>, @@ -1427,8 +1441,11 @@ pub async fn search_plan( true, None, )?; - let merge_collector = - make_merge_collector(&search_request, Default::default(), MergeLevel::Root)?; + let merge_collector = make_merge_collector( + &search_request, + Default::default(), + IntermediatePruning::Apply, + )?; warmup_info.merge(merge_collector.warmup_info()); warmup_info.simplify(); @@ -5742,6 +5759,30 @@ mod tests { Ok(()) } + #[test] + fn test_root_intermediate_pruning_follows_finalization() { + // Finalization prunes to the requested `size`, so the merge keeps every candidate. + let finalizing = SearchRequest { + skip_aggregation_finalization: false, + ..Default::default() + }; + assert_eq!( + root_intermediate_pruning(&finalizing), + IntermediatePruning::Skip + ); + + // Without finalization the merged bytes go back to the caller as they are, so this + // merge has to bound them itself. + let skipping = SearchRequest { + skip_aggregation_finalization: true, + ..Default::default() + }; + assert_eq!( + root_intermediate_pruning(&skipping), + IntermediatePruning::Apply + ); + } + #[tokio::test] async fn test_finalize_aggregation_if_any_no_aggregation_request() { let search_request = SearchRequest { From c3d7ae0561cba091a704ae96656e58f4c57e49e4 Mon Sep 17 00:00:00 2001 From: Aarav Date: Fri, 16 Oct 2026 10:09:15 +0530 Subject: [PATCH 3/3] run root aggregation finalization on the search thread pool With root pruning skipped when finalization follows, `into_final_result` now deserializes and sorts the candidates of every leaf rather than a set already cut to `segment_size`. That work grows with the number of leaves and was running directly on the Tokio worker. Move it to `search_thread_pool` at high priority, like the root merge just before it. The early returns (no aggregation, finalization skipped) stay on the async path since they do no work. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y1b55sEL1gA1c3FCMcQCbm --- quickwit/quickwit-search/src/root.rs | 39 ++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index 86e37c8e0d1..aef9603ad55 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -43,6 +43,7 @@ use quickwit_query::query_ast::{ }; use serde::{Deserialize, Serialize}; use tantivy::TantivyError; +use tantivy::aggregation::AggregationLimitsGuard; use tantivy::aggregation::agg_result::AggregationResults; use tantivy::aggregation::intermediate_agg_result::IntermediateAggregationResults; use tantivy::collector::Collector; @@ -1101,7 +1102,8 @@ async fn root_search_aux( &search_request, first_phase_result.intermediate_aggregation_result, searcher_context, - )?; + ) + .await?; // In case there is no index, we don't want the response to contain any aggregation structure if indexes_metas_for_leaf_search.is_empty() { aggregation_result_postcard_opt = None; @@ -1129,7 +1131,7 @@ async fn root_search_aux( fn finalize_aggregation( intermediate_aggregation_result_bytes_opt: Option>, aggregations: QuickwitAggregations, - searcher_context: &SearcherContext, + aggregation_limits: AggregationLimitsGuard, ) -> crate::Result>> { let merge_aggregation_result = match aggregations { QuickwitAggregations::FindTraceIdsAggregation(_) => { @@ -1149,7 +1151,7 @@ fn finalize_aggregation( Default::default() }; let final_aggregation_results: AggregationResults = intermediate_aggregation_results - .into_final_result(aggregations, searcher_context.get_aggregation_limits())?; + .into_final_result(aggregations, aggregation_limits)?; let final_aggregation_proxy: quickwit_query::aggregations::AggregationResults = final_aggregation_results.into(); postcard::to_stdvec(&final_aggregation_proxy)? @@ -1172,7 +1174,7 @@ fn root_intermediate_pruning(search_request: &SearchRequest) -> IntermediatePrun } } -fn finalize_aggregation_if_any( +async fn finalize_aggregation_if_any( search_request: &SearchRequest, intermediate_aggregation_result_bytes_opt: Option>, searcher_context: &SearcherContext, @@ -1184,12 +1186,21 @@ fn finalize_aggregation_if_any( return Ok(intermediate_aggregation_result_bytes_opt); } let aggregations: QuickwitAggregations = serde_json::from_str(aggregations_json)?; - let aggregation_result_postcard = finalize_aggregation( - intermediate_aggregation_result_bytes_opt, - aggregations, - searcher_context, - )?; - Ok(aggregation_result_postcard) + let aggregation_limits = searcher_context.get_aggregation_limits(); + // The root merge does not prune when finalization follows, so the finalizer deserializes + // and sorts the candidates of every leaf. That is cpu-bound, keep it off the runtime. + let span = info_span!("finalize_aggregation"); + crate::search_thread_pool() + .run_cpu_intensive_with_priority(Priority::High, move || { + let _span_guard = span.enter(); + finalize_aggregation( + intermediate_aggregation_result_bytes_opt, + aggregations, + aggregation_limits, + ) + }) + .await + .context("failed to finalize aggregation")? } /// Checks that all of the index researched as found. @@ -5789,6 +5800,7 @@ mod tests { let searcher_context = SearcherContext::for_test(); let result = finalize_aggregation_if_any(&search_request, Some(vec![1, 2, 3]), &searcher_context) + .await .unwrap(); assert!(result.is_none()); } @@ -5808,6 +5820,7 @@ mod tests { Some(intermediate_bytes.clone()), &searcher_context, ) + .await .unwrap(); assert_eq!(result, Some(intermediate_bytes)); } @@ -5821,7 +5834,9 @@ mod tests { ..Default::default() }; let searcher_context = SearcherContext::for_test(); - let result = finalize_aggregation_if_any(&search_request, None, &searcher_context).unwrap(); + let result = finalize_aggregation_if_any(&search_request, None, &searcher_context) + .await + .unwrap(); assert!(result.is_none()); } @@ -5841,6 +5856,7 @@ mod tests { Some(intermediate_bytes.clone()), &searcher_context, ) + .await .unwrap(); // Result should be Some (finalized), but different from intermediate bytes assert!(result.is_some()); @@ -5863,6 +5879,7 @@ mod tests { Some(intermediate_bytes.clone()), &searcher_context, ) + .await .unwrap(); assert!(result.is_some()); assert_ne!(result.unwrap(), intermediate_bytes);