diff --git a/quickwit/quickwit-search/src/collector.rs b/quickwit/quickwit-search/src/collector.rs index 08ab06e6a8f..e01e6a01801 100644 --- a/quickwit/quickwit-search/src/collector.rs +++ b/quickwit/quickwit-search/src/collector.rs @@ -682,7 +682,10 @@ impl QuickwitIncrementalAggregations { } } - fn finalize(self) -> 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,6 +700,7 @@ impl QuickwitIncrementalAggregations { merge_intermediate_aggregation_result( &Some(QuickwitAggregations::TantivyAggregations(aggregation)), state.iter().map(|vec| vec.as_slice()), + intermediate_pruning, ) } QuickwitIncrementalAggregations::NoAggregation => Ok(None), @@ -704,6 +708,21 @@ impl QuickwitIncrementalAggregations { } } +/// 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. +/// +/// 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 IntermediatePruning { + Apply, + Skip, +} + /// The quickwit collector is the tantivy Collector used in Quickwit. /// /// It defines the data that should be accumulated about the documents matching @@ -716,6 +735,7 @@ pub(crate) struct QuickwitCollector { pub sort_by: SortByPair, pub aggregation: Option, pub agg_context_params: AggContextParams, + pub intermediate_pruning: IntermediatePruning, search_after: Option, } @@ -846,6 +866,7 @@ impl Collector for QuickwitCollector { sort_order1, sort_order2, num_hits, + 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 @@ -870,6 +891,7 @@ fn map_error(error: postcard::Error) -> TantivyError { fn merge_intermediate_aggregation_result<'a>( aggregations_opt: &Option, intermediate_aggregation_results: impl Iterator, + intermediate_pruning: IntermediatePruning, ) -> tantivy::Result>> { let merged_intermediate_aggregation_result = match aggregations_opt { Some(QuickwitAggregations::FindTraceIdsAggregation(collector)) => { @@ -903,7 +925,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 intermediate_pruning == IntermediatePruning::Apply { + merged.prune_intermediate_results(aggregations, PruneMode::Intermediate)?; + } let serialized = postcard::to_allocvec(&merged).map_err(map_error)?; Some(serialized) } @@ -920,6 +949,7 @@ fn merge_leaf_responses( sort_order1: SortOrder, sort_order2: SortOrder, max_hits: usize, + intermediate_pruning: IntermediatePruning, ) -> tantivy::Result { // Optimization: No merging needed if there is only one result. if leaf_responses.len() == 1 { @@ -937,6 +967,7 @@ fn merge_leaf_responses( leaf_responses.iter().filter_map(|leaf_response| { leaf_response.intermediate_aggregation_result.as_deref() }), + intermediate_pruning, )?; let num_attempted_splits = leaf_responses .iter() @@ -1050,6 +1081,7 @@ pub(crate) fn make_collector_for_split( sort_by, aggregation, agg_context_params, + intermediate_pruning: IntermediatePruning::Apply, search_after: search_request.search_after.clone(), }) } @@ -1058,6 +1090,7 @@ pub(crate) fn make_collector_for_split( pub(crate) fn make_merge_collector( search_request: &SearchRequest, agg_limits: AggregationLimitsGuard, + 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. @@ -1075,6 +1108,7 @@ pub(crate) fn make_merge_collector( sort_by, aggregation, agg_context_params, + intermediate_pruning, search_after: search_request.search_after.clone(), }) } @@ -1201,6 +1235,7 @@ pub(crate) struct IncrementalCollector { num_successful_splits: u64, start_offset: usize, resource_stats: Option, + intermediate_pruning: IntermediatePruning, } impl IncrementalCollector { @@ -1211,6 +1246,7 @@ impl IncrementalCollector { .as_ref() .map(QuickwitAggregations::maybe_incremental_aggregator) .unwrap_or(QuickwitIncrementalAggregations::NoAggregation); + let intermediate_pruning = collector.intermediate_pruning; let (order1, order2) = collector.sort_by.sort_orders(); let sort_key_mapper = HitSortingMapper { order1, order2 }; IncrementalCollector { @@ -1222,6 +1258,7 @@ impl IncrementalCollector { num_attempted_splits: 0, num_successful_splits: 0, resource_stats: None, + intermediate_pruning, } } @@ -1295,7 +1332,9 @@ 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.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())); @@ -1328,7 +1367,9 @@ 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::{ + IntermediatePruning, merge_intermediate_aggregation_result, top_k_partial_hits, + }; #[test] fn test_merge_partial_hits_no_tie() { @@ -1775,7 +1816,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(), IntermediatePruning::Apply).unwrap(); let mut incremental_collector = IncrementalCollector::new(collector.clone()); let result = collector @@ -2061,7 +2103,12 @@ 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(), + IntermediatePruning::Apply, + ) + .unwrap(); assert!(merged.is_none()); let aggregations_json = r#"{ @@ -2069,16 +2116,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(), + IntermediatePruning::Apply, + ) + .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(IntermediatePruning::Apply), 4); + } + + /// 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_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(intermediate_pruning: IntermediatePruning) -> usize { use tantivy::Index; use tantivy::aggregation::DistributedAggregationCollector; use tantivy::aggregation::intermediate_agg_result::{ @@ -2130,6 +2195,7 @@ mod tests { let serialized = merge_intermediate_aggregation_result( &Some(quickwit_aggregations), serialized_fruits.iter().map(Vec::as_slice), + intermediate_pruning, ) .unwrap() .unwrap(); @@ -2140,8 +2206,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 e60dfad251a..98eccca3ed0 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, IntermediatePruning, make_collector_for_split, make_merge_collector, +}; use crate::cost::{compute_query_complexity_factor, compute_split_query_cost}; use crate::leaf_cache::LeafSearchCache; use crate::metrics::{ @@ -1684,8 +1686,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(), + IntermediatePruning::Apply, + )?; let mut incremental_merge_collector = IncrementalCollector::new(merge_collector); while let Some(leaf_response_join_result) = leaf_request_futures.join_next().await { @@ -2045,8 +2050,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(), + IntermediatePruning::Apply, + )?; 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 e316c60bf85..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; @@ -51,7 +52,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::{IntermediatePruning, QuickwitAggregations, make_merge_collector}; use crate::cost::{compute_query_complexity_factor, compute_split_query_cost}; use crate::metrics_trackers::{RootSearchMetricsFuture, RootSearchMetricsStep}; use crate::scroll_context::{ScrollContext, ScrollKeyAndStartOffset}; @@ -844,8 +845,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(), + root_intermediate_pruning(search_request), + )?; // Merging is a cpu-bound task. Prioritize it over queued split searches to avoid delaying the // final response once all leaf responses are available. @@ -1098,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; @@ -1126,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(_) => { @@ -1146,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)? @@ -1155,7 +1160,21 @@ fn finalize_aggregation( Ok(Some(merge_aggregation_result)) } -fn finalize_aggregation_if_any( +/// 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 + } +} + +async fn finalize_aggregation_if_any( search_request: &SearchRequest, intermediate_aggregation_result_bytes_opt: Option>, searcher_context: &SearcherContext, @@ -1167,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. @@ -1429,7 +1457,11 @@ 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(), + IntermediatePruning::Apply, + )?; warmup_info.merge(merge_collector.warmup_info()); warmup_info.simplify(); @@ -5734,6 +5766,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 { @@ -5744,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()); } @@ -5763,6 +5820,7 @@ mod tests { Some(intermediate_bytes.clone()), &searcher_context, ) + .await .unwrap(); assert_eq!(result, Some(intermediate_bytes)); } @@ -5776,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()); } @@ -5796,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()); @@ -5818,6 +5879,7 @@ mod tests { Some(intermediate_bytes.clone()), &searcher_context, ) + .await .unwrap(); assert!(result.is_some()); assert_ne!(result.unwrap(), intermediate_bytes);