diff --git a/quickwit/quickwit-query/src/query_ast/cache_node.rs b/quickwit/quickwit-query/src/query_ast/cache_node.rs index 48f0fabf419..b17201c64e9 100644 --- a/quickwit/quickwit-query/src/query_ast/cache_node.rs +++ b/quickwit/quickwit-query/src/query_ast/cache_node.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::sync::Arc; use bitpacking::{BitPacker, BitPacker1x}; @@ -105,14 +106,18 @@ impl BuildTantivyAst for CacheNode { } .into()), CacheState::CacheMiss(cache_filler) => { - let tantivy_query: Box = self - .inner - .build_tantivy_ast_call(context)? - .simplify() - .into(); + let tantivy_query_ast = self.inner.build_tantivy_ast_call(context)?.simplify(); + let mut required_terms = HashSet::new(); + super::required_terms::collect_required_terms( + &tantivy_query_ast, + context.schema, + &mut required_terms, + ); + let tantivy_query: Box = tantivy_query_ast.into(); Ok(CacheFillerQuery { inner_query: Box::new(tantivy_query), cache_filler: cache_filler.clone(), + required_terms, } .into()) } @@ -123,7 +128,7 @@ impl BuildTantivyAst for CacheNode { use tantivy::directory::OwnedBytes; use tantivy::index::SegmentId; use tantivy::query::{EnableScoring, Explanation, Query, Scorer, Weight}; -use tantivy::{DocId, DocSet, Score, SegmentReader, TantivyError}; +use tantivy::{DocId, DocSet, Score, SegmentReader, TantivyError, Term}; #[derive(Clone, Debug)] pub struct CacheHitQuery { @@ -397,6 +402,7 @@ impl CacheFiller { pub struct CacheFillerQuery { inner_query: Box, cache_filler: CacheFiller, + required_terms: HashSet, } impl Clone for CacheFillerQuery { @@ -404,10 +410,17 @@ impl Clone for CacheFillerQuery { Self { inner_query: self.inner_query.box_clone(), cache_filler: self.cache_filler.clone(), + required_terms: self.required_terms.clone(), } } } +impl CacheFillerQuery { + pub(crate) fn required_terms(&self) -> &HashSet { + &self.required_terms + } +} + impl Query for CacheFillerQuery { fn weight(&self, enable_scoring: EnableScoring<'_>) -> tantivy::Result> { if enable_scoring.is_scoring_enabled() { @@ -716,6 +729,56 @@ mod tests { } } + #[test] + fn test_cache_filler_preserves_required_terms() { + let mut schema_builder = Schema::builder(); + let body_field = schema_builder.add_text_field("body", TEXT); + let schema = schema_builder.build(); + let context = BuildTantivyAstContext::for_test(&schema); + let term_query: QueryAst = TermQuery { + field: "body".to_string(), + value: "val".to_string(), + } + .into(); + let cache_node = CacheNode::new(term_query); + let query_json = serde_json::to_string(&cache_node.inner).unwrap(); + let ast: QueryAst = cache_node.into(); + let cache = Arc::new(Mutex::new(HashMap::new())); + + let cache_miss_ast = PredicateCacheInjector { + cache: cache.clone(), + split_id: "split".to_string(), + } + .transform(ast.clone()) + .unwrap() + .unwrap(); + let (_query, required_terms) = cache_miss_ast + .build_tantivy_query_and_required_terms(&context) + .unwrap(); + assert_eq!( + required_terms, + HashSet::from([Term::from_field_text(body_field, "val")]) + ); + + cache.put( + "split".to_string(), + query_json, + SegmentId::from_uuid_string("1686a000d4f7a91939d0e71df1646d7a").unwrap(), + HitSet::empty(), + ); + let cache_hit_ast = PredicateCacheInjector { + cache, + split_id: "split".to_string(), + } + .transform(ast) + .unwrap() + .unwrap(); + let (_query, required_terms) = cache_hit_ast + .build_tantivy_query_and_required_terms(&context) + .unwrap(); + assert!(required_terms.is_empty()); + } + #[test] fn test_cache_preigniter_fills_cache() { let term_query: QueryAst = TermQuery { diff --git a/quickwit/quickwit-query/src/query_ast/required_terms.rs b/quickwit/quickwit-query/src/query_ast/required_terms.rs index 45e78b719a6..003b3f3d0cf 100644 --- a/quickwit/quickwit-query/src/query_ast/required_terms.rs +++ b/quickwit/quickwit-query/src/query_ast/required_terms.rs @@ -37,6 +37,7 @@ use tantivy::Term; use tantivy::query::TermQuery as TantivyTermQuery; use tantivy::schema::Schema; +use crate::query_ast::cache_node::CacheFillerQuery; use crate::query_ast::tantivy_query_ast::{TantivyBoolQuery, TantivyQueryAst}; /// Collects the terms that must be present for the (already simplified) query to @@ -48,12 +49,15 @@ pub(crate) fn collect_required_terms( ) { match ast { TantivyQueryAst::Leaf(query) => { - let Some(term_query) = query.downcast_ref::() else { + if let Some(term_query) = query.downcast_ref::() { + let term = term_query.term(); + if schema.get_field_entry(term.field()).is_indexed() { + required_terms.insert(term.clone()); + } return; - }; - let term = term_query.term(); - if schema.get_field_entry(term.field()).is_indexed() { - required_terms.insert(term.clone()); + } + if let Some(cache_filler_query) = query.downcast_ref::() { + required_terms.extend(cache_filler_query.required_terms().iter().cloned()); } } TantivyQueryAst::Bool(bool_query) => { diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 245a18db75a..69c73d2ca23 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -729,22 +729,46 @@ async fn leaf_search_single_split( agg_context_params, )?; - let predicate_cache = if collector.requires_scoring() { - // at the moment the predicate cache doesn't support scoring - None - } else { - Some(( - ctx.searcher_context.predicate_cache.clone() as _, - split.split_id.clone(), - )) - }; + let predicate_cache = + if collector.requires_scoring() || !ctx.searcher_context.predicate_cache.is_enabled() { + None + } else { + Some(( + ctx.searcher_context.predicate_cache.clone() as Arc, + split.split_id.clone(), + )) + }; + let predicate_cache_miss_ast = predicate_cache + .as_ref() + .and_then(|(cache, cache_split_id)| { + let timestamp_field = ctx.doc_mapper.timestamp_field_name()?; + let predicate_ast = time_bounded_cached_predicate(&query_ast, timestamp_field)?; + let predicate_key = serde_json::to_string(&predicate_ast).ok()?; + if cache.get(cache_split_id.clone(), predicate_key).is_some() { + return None; + } + Some(predicate_ast) + }); let split_schema = index.schema(); let (query, mut warmup_info) = ctx.doc_mapper.query( split_schema.clone(), query_ast.clone(), false, - predicate_cache, + predicate_cache.clone(), )?; + let predicate_cache_warmup_query = if let Some(predicate_ast) = predicate_cache_miss_ast { + let warmup_ast = QueryAst::from(CacheNode::new(predicate_ast)); + let (warmup_query, predicate_warmup_info) = ctx.doc_mapper.query( + split_schema.clone(), + warmup_ast, + false, + predicate_cache.clone(), + )?; + warmup_info.merge(predicate_warmup_info); + Some(warmup_query) + } else { + None + }; let collector_warmup_info = collector.warmup_info(); warmup_info.merge(collector_warmup_info); @@ -915,6 +939,22 @@ async fn leaf_search_single_split( return Ok(None); }; collector.update_search_param(&simplified_search_request); + let query = if let Some(predicate_warmup_query) = predicate_cache_warmup_query { + let predicate_warmup_span = info_span!("predicate_cache_warmup"); + let _predicate_warmup_span_guard = predicate_warmup_span.enter(); + predicate_warmup_query.count(&searcher).inspect_err(|_| { + leaf_search_state_guard + .set_state(SplitSearchState::Error(SplitSearchErrorKind::TantivySearch)) + })?; + drop(_predicate_warmup_span_guard); + ctx_clone + .doc_mapper + .query(split_schema, query_ast.clone(), false, predicate_cache) + .map_err(|error| TantivyError::InvalidArgument(error.to_string()))? + .0 + } else { + query + }; let mut leaf_search_response: LeafSearchResponse = if is_metadata_count_request_with_ast(&query_ast, &simplified_search_request) { get_leaf_resp_from_count(searcher.num_docs()) @@ -974,7 +1014,7 @@ async fn leaf_search_single_split( /// /// This include things such as sorting result by a field or _score when no document is requested, /// or applying date range when the range covers the entire split. -fn rewrite_request( +pub(crate) fn rewrite_request( search_request: &mut SearchRequest, split: &SplitIdAndFooterOffsets, timestamp_field: Option<&str>, @@ -983,7 +1023,7 @@ fn rewrite_request( search_request.sort_fields = Vec::new(); } if let Some(timestamp_field) = timestamp_field { - remove_redundant_timestamp_range(search_request, split, timestamp_field); + normalize_timestamp_range(search_request, split, timestamp_field, true); } rewrite_aggregation(search_request); // we add a top level cache node when search_after is set, this won't help for this query (which @@ -993,6 +1033,83 @@ fn rewrite_request( } } +fn is_trivial_time_bounded_predicate(query_ast: &QueryAst) -> bool { + match query_ast { + QueryAst::MatchAll | QueryAst::MatchNone => true, + QueryAst::Bool(bool_query) => { + let has_required_match_none = bool_query + .must + .iter() + .chain(&bool_query.filter) + .any(is_trivial_match_none); + let is_match_all = bool_query.must_not.is_empty() + && bool_query.should.is_empty() + && bool_query.must.iter().all(is_trivial_match_all) + && bool_query.filter.iter().all(is_trivial_match_all); + has_required_match_none || is_match_all + } + _ => false, + } +} + +fn is_trivial_match_none(query_ast: &QueryAst) -> bool { + match query_ast { + QueryAst::MatchNone => true, + QueryAst::Bool(bool_query) => bool_query + .must + .iter() + .chain(&bool_query.filter) + .any(is_trivial_match_none), + _ => false, + } +} + +fn is_trivial_match_all(query_ast: &QueryAst) -> bool { + match query_ast { + QueryAst::MatchAll => true, + QueryAst::Bool(bool_query) => { + bool_query.must_not.is_empty() + && bool_query.should.is_empty() + && bool_query.must.iter().all(is_trivial_match_all) + && bool_query.filter.iter().all(is_trivial_match_all) + } + _ => false, + } +} + +/// Returns the predicate cached after timestamp normalization. +/// An outer cache node can be present for search-after requests. +pub(crate) fn time_bounded_cached_predicate( + query_ast: &QueryAst, + timestamp_field: &str, +) -> Option { + match query_ast { + QueryAst::Cache(cache_node) => { + time_bounded_cached_predicate(&cache_node.inner, timestamp_field) + } + QueryAst::Bool(bool_query) + if bool_query.must.len() == 1 + && bool_query.filter.len() <= 1 + && bool_query.must_not.is_empty() + && bool_query.should.is_empty() => + { + let QueryAst::Cache(cache_node) = &bool_query.must[0] else { + return None; + }; + if let Some(time_filter) = bool_query.filter.first() { + let QueryAst::Range(time_range) = time_filter else { + return None; + }; + if time_range.field != timestamp_field { + return None; + } + } + Some((*cache_node.inner).clone()) + } + _ => None, + } +} + fn add_top_cache_node(search_request: &mut SearchRequest) { let Ok(query_ast) = serde_json::from_str(search_request.query_ast.as_str()) else { // an error will get raised a bit after anyway @@ -1099,14 +1216,11 @@ fn min_bound(left: Bound, right: Bound) -> Bound { } } -/// remove timestamp range that would be present both in QueryAst and SearchRequest -/// -/// this can save us from doing double the work in some cases, and help with the partial request -/// cache. -fn remove_redundant_timestamp_range( +fn normalize_timestamp_range( search_request: &mut SearchRequest, split: &SplitIdAndFooterOffsets, timestamp_field: &str, + install_predicate_cache_node: bool, ) { let Ok(query_ast) = serde_json::from_str(search_request.query_ast.as_str()) else { // an error will get raised a bit after anyway @@ -1133,6 +1247,8 @@ fn remove_redundant_timestamp_range( .transform(query_ast) .expect("can't fail unwrapping Infallible") .unwrap_or(QueryAst::MatchAll); + let is_time_bounded = + visitor.start_timestamp != Bound::Unbounded || visitor.end_timestamp != Bound::Unbounded; let final_start_timestamp = match ( visitor.start_timestamp, @@ -1172,12 +1288,23 @@ fn remove_redundant_timestamp_range( (Bound::Unbounded, Some(_)) => Bound::Unbounded, (query_bound, None) => query_bound, }; + if install_predicate_cache_node + && is_time_bounded + && !is_trivial_time_bounded_predicate(&new_ast) + { + new_ast = BoolQuery { + must: vec![QueryAst::from(CacheNode::new(new_ast))], + ..Default::default() + } + .into(); + } + if final_start_timestamp != Bound::Unbounded || final_end_timestamp != Bound::Unbounded { - let range = RangeQuery { + let time_range = QueryAst::from(RangeQuery { field: timestamp_field.to_string(), lower_bound: final_start_timestamp.map(|bound| bound.into_timestamp_nanos().into()), upper_bound: final_end_timestamp.map(|bound| bound.into_timestamp_nanos().into()), - }; + }); new_ast = if let QueryAst::Bool(mut bool_query) = new_ast { if bool_query.must.is_empty() && bool_query.filter.is_empty() @@ -1187,22 +1314,22 @@ fn remove_redundant_timestamp_range( // add a new layer of bool query BoolQuery { must: vec![bool_query.into()], - filter: vec![range.into()], + filter: vec![time_range], ..Default::default() } .into() } else { - bool_query.filter.push(range.into()); + bool_query.filter.push(time_range); QueryAst::Bool(bool_query) } } else { BoolQuery { must: vec![new_ast], - filter: vec![range.into()], + filter: vec![time_range], ..Default::default() } .into() - } + }; } search_request.query_ast = serde_json::to_string(&new_ast).unwrap(); @@ -2327,8 +2454,6 @@ async fn leaf_search_single_split_wrapper( #[cfg(test)] mod tests { - use std::ops::Bound; - use async_trait::async_trait; use bytes::BufMut; use quickwit_config::{LambdaConfig, SearcherConfig}; @@ -2380,302 +2505,6 @@ mod tests { assert_eq!(counters.cancel_warmup.get(), 1); } - fn bool_filter(ast: impl Into) -> QueryAst { - BoolQuery { - must: vec![QueryAst::MatchAll], - filter: vec![ast.into()], - ..Default::default() - } - .into() - } - - #[track_caller] - fn assert_ast_eq(got: &SearchRequest, expected: &QueryAst) { - let got_ast: QueryAst = serde_json::from_str(&got.query_ast).unwrap(); - assert_eq!(&got_ast, expected); - assert!(got.start_timestamp.is_none()); - assert!(got.end_timestamp.is_none()); - } - - #[track_caller] - fn remove_timestamp_test_case( - request: &SearchRequest, - split: &SplitIdAndFooterOffsets, - expected: Option, - ) { - let timestamp_field = "timestamp"; - - // test the query directly - let mut request_direct = request.clone(); - remove_redundant_timestamp_range(&mut request_direct, split, timestamp_field); - let expected_direct = expected - .clone() - .map(bool_filter) - .unwrap_or(QueryAst::MatchAll); - assert_ast_eq(&request_direct, &expected_direct); - } - - #[test] - fn test_remove_timestamp_range() { - const S_TO_NS: i64 = 1_000_000_000; - let time1 = 1700001000; - let time2 = 1700002000; - let time3 = 1700003000; - let time4 = 1700004000; - - let timestamp_field = "timestamp".to_string(); - - // cases where the bounds are larger than the split: no bound is emitted - let split = SplitIdAndFooterOffsets { - timestamp_start: Some(time2), - timestamp_end: Some(time3), - ..SplitIdAndFooterOffsets::default() - }; - - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time1.into()), - // *1000 has no impact, we detect timestamp in ms instead of s - upper_bound: Bound::Included((time4 * 1000).into()), - })) - .unwrap(), - ..SearchRequest::default() - }; - remove_timestamp_test_case(&search_request, &split, None); - - let expected_upper_inclusive = RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Unbounded, - upper_bound: Bound::Included((time3 * S_TO_NS).into()), - }; - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time1.into()), - upper_bound: Bound::Included(time3.into()), - })) - .unwrap(), - ..SearchRequest::default() - }; - remove_timestamp_test_case(&search_request, &split, Some(expected_upper_inclusive)); - - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::MatchAll).unwrap(), - start_timestamp: Some(time1), - end_timestamp: Some(time4), - ..SearchRequest::default() - }; - remove_timestamp_test_case(&search_request, &split, None); - - // request bound that are exclusive are treated properly - let expected_upper_exclusive = RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Unbounded, - upper_bound: Bound::Excluded((time3 * S_TO_NS).into()), - }; - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time1.into()), - upper_bound: Bound::Excluded(time3.into()), - })) - .unwrap(), - ..SearchRequest::default() - }; - remove_timestamp_test_case( - &search_request, - &split, - Some(expected_upper_exclusive.clone()), - ); - - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::MatchAll).unwrap(), - start_timestamp: Some(time1), - end_timestamp: Some(time3), - ..SearchRequest::default() - }; - remove_timestamp_test_case( - &search_request, - &split, - Some(expected_upper_exclusive.clone()), - ); - - let expected_lower_excl_upper_incl = RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Excluded((time2 * S_TO_NS).into()), - upper_bound: Bound::Included((time3 * S_TO_NS).into()), - }; - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Excluded(time2.into()), - upper_bound: Bound::Included(time3.into()), - })) - .unwrap(), - ..SearchRequest::default() - }; - remove_timestamp_test_case( - &search_request, - &split, - Some(expected_lower_excl_upper_incl.clone()), - ); - } - - #[test] - fn test_remove_timestamp_range_multiple_bounds() { - // When bounds are defined both in the AST and in the search request, - // make sure we take the most restrictive ones. - const S_TO_NS: i64 = 1_000_000_000; - let time1 = 1700001000; - let time2 = 1700002000; - let time3 = 1700003000; - let time4 = 1700004000; - - let timestamp_field = "timestamp".to_string(); - - let split = SplitIdAndFooterOffsets { - timestamp_start: Some(time1), - timestamp_end: Some(time4), - ..SplitIdAndFooterOffsets::default() - }; - - let expected_upper_2_ex = RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Unbounded, - upper_bound: Bound::Excluded((time2 * S_TO_NS).into()), - }; - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time1.into()), - upper_bound: Bound::Included(time3.into()), - })) - .unwrap(), - start_timestamp: Some(time1), - end_timestamp: Some(time2), - ..SearchRequest::default() - }; - remove_timestamp_test_case(&search_request, &split, Some(expected_upper_2_ex)); - - let expected_upper_2_inc = RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Unbounded, - upper_bound: Bound::Included((time2 * S_TO_NS).into()), - }; - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time1.into()), - upper_bound: Bound::Included(time2.into()), - })) - .unwrap(), - start_timestamp: Some(time1), - end_timestamp: Some(time3), - ..SearchRequest::default() - }; - remove_timestamp_test_case(&search_request, &split, Some(expected_upper_2_inc)); - - let expected_lower_3_upper_4 = RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included((time3 * S_TO_NS).into()), - upper_bound: Bound::Included((time4 * S_TO_NS).into()), - }; - - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time2.into()), - upper_bound: Bound::Included(time4.into()), - })) - .unwrap(), - start_timestamp: Some(time3), - end_timestamp: Some(time4 + 1), - ..SearchRequest::default() - }; - remove_timestamp_test_case( - &search_request, - &split, - Some(expected_lower_3_upper_4.clone()), - ); - - let search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Range(RangeQuery { - field: timestamp_field.to_string(), - lower_bound: Bound::Included(time3.into()), - upper_bound: Bound::Included(time4.into()), - })) - .unwrap(), - start_timestamp: Some(time2), - end_timestamp: Some(time4 + 1), - ..SearchRequest::default() - }; - remove_timestamp_test_case(&search_request, &split, Some(expected_lower_3_upper_4)); - - let mut search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::MatchAll).unwrap(), - start_timestamp: Some(time1), - end_timestamp: Some(time4), - ..SearchRequest::default() - }; - let split = SplitIdAndFooterOffsets { - timestamp_start: Some(time2), - timestamp_end: Some(time3), - ..SplitIdAndFooterOffsets::default() - }; - remove_redundant_timestamp_range(&mut search_request, &split, ×tamp_field); - assert_ast_eq(&search_request, &QueryAst::MatchAll); - } - - // regression test for #4935 - #[test] - fn test_remove_timestamp_range_keep_should() { - let time1 = 1700001000; - let time2 = 1700002000; - let time3 = 1700003000; - - let timestamp_field = "timestamp".to_string(); - - // cases where the bounds are larger than the split: no bound is emitted - let split = SplitIdAndFooterOffsets { - timestamp_start: Some(time1), - timestamp_end: Some(time3), - ..SplitIdAndFooterOffsets::default() - }; - - let mut search_request = SearchRequest { - query_ast: serde_json::to_string(&QueryAst::Bool(BoolQuery { - should: vec![QueryAst::MatchAll], - ..BoolQuery::default() - })) - .unwrap(), - start_timestamp: Some(time2), - end_timestamp: None, - ..SearchRequest::default() - }; - remove_redundant_timestamp_range(&mut search_request, &split, ×tamp_field); - assert_ast_eq( - &search_request, - &QueryAst::Bool(BoolQuery { - // original request - must: vec![QueryAst::Bool(BoolQuery { - should: vec![QueryAst::MatchAll], - ..BoolQuery::default() - })], - // time bound - filter: vec![ - RangeQuery { - field: "timestamp".to_string(), - lower_bound: Bound::Included(1_700_002_000_000_000_000u64.into()), - upper_bound: Bound::Unbounded, - } - .into(), - ], - ..BoolQuery::default() - }), - ); - } - #[test] fn test_remove_extended_bounds_from_histogram() { let histo_at_root = r#" diff --git a/quickwit/quickwit-search/src/leaf_cache.rs b/quickwit/quickwit-search/src/leaf_cache.rs index c48b9ef2bd9..e50ca469e4f 100644 --- a/quickwit/quickwit-search/src/leaf_cache.rs +++ b/quickwit/quickwit-search/src/leaf_cache.rs @@ -237,6 +237,7 @@ impl RangeBounds for HalfOpenRange { pub struct PredicateCacheImpl { content: MemorySizedCache, key_hasher: CacheKeyHasher, + enabled: bool, } impl PredicateCacheImpl { @@ -247,8 +248,13 @@ impl PredicateCacheImpl { &quickwit_storage::metrics::PREDICATE_CACHE, ), key_hasher: CacheKeyHasher::random(), + enabled: config.capacity().as_u64() > 0, } } + + pub fn is_enabled(&self) -> bool { + self.enabled + } } impl quickwit_query::query_ast::PredicateCache for PredicateCacheImpl { @@ -257,6 +263,9 @@ impl quickwit_query::query_ast::PredicateCache for PredicateCacheImpl { split_id: String, query_ast_json: String, ) -> Option<(SegmentId, quickwit_query::query_ast::HitSet)> { + if !self.enabled { + return None; + } let key = self .key_hasher .hash(&(split_id.as_str(), query_ast_json.as_str())); @@ -275,6 +284,9 @@ impl quickwit_query::query_ast::PredicateCache for PredicateCacheImpl { segment: SegmentId, hits: quickwit_query::query_ast::HitSet, ) { + if !self.enabled { + return; + } let hits_buffer = hits.into_buffer(); let mut buffer = Vec::with_capacity(32 + hits_buffer.len()); buffer.extend_from_slice(segment.uuid_string().as_bytes()); diff --git a/quickwit/quickwit-search/src/tests.rs b/quickwit/quickwit-search/src/tests.rs index 94183d16c2a..6e7a52236e4 100644 --- a/quickwit/quickwit-search/src/tests.rs +++ b/quickwit/quickwit-search/src/tests.rs @@ -14,6 +14,7 @@ use std::cmp::Ordering; use std::collections::{BTreeMap, BTreeSet}; +use std::ops::Bound; use assert_json_diff::{assert_json_eq, assert_json_include}; use quickwit_config::SearcherConfig; @@ -25,12 +26,13 @@ use quickwit_proto::search::{ SortValue, TraceId, }; use quickwit_query::query_ast::{ - HitSet, PredicateCache, QueryAst, qast_helper, qast_json_helper, query_ast_from_user_text, + BoolQuery, HitSet, PredicateCache, QueryAst, RangeQuery, qast_helper, qast_json_helper, + query_ast_from_user_text, }; use serde_json::{Value as JsonValue, json}; -use tantivy::Term; use tantivy::schema::OwnedValue as TantivyValue; use tantivy::time::OffsetDateTime; +use tantivy::{DocSet, Term}; use self::leaf::single_doc_mapping_leaf_search; use super::*; @@ -2137,6 +2139,278 @@ async fn negative_cache_ts_test_setup() -> ( ) } +#[test] +fn test_time_bounded_predicate_cache_eligibility_after_split_normalization() { + let split = SplitIdAndFooterOffsets { + timestamp_start: Some(100), + timestamp_end: Some(199), + ..Default::default() + }; + let predicate_ast = qast_helper("info", &["body"]); + + let mut partial_request = SearchRequest { + query_ast: serde_json::to_string(&predicate_ast).unwrap(), + start_timestamp: Some(120), + end_timestamp: Some(180), + ..Default::default() + }; + leaf::rewrite_request(&mut partial_request, &split, Some("ts")); + let partial_ast: QueryAst = serde_json::from_str(&partial_request.query_ast).unwrap(); + let partial_predicate = leaf::time_bounded_cached_predicate(&partial_ast, "ts") + .expect("a partial range should install a cached predicate"); + + let mut full_split_request = SearchRequest { + query_ast: serde_json::to_string(&predicate_ast).unwrap(), + start_timestamp: Some(50), + end_timestamp: Some(250), + ..Default::default() + }; + leaf::rewrite_request(&mut full_split_request, &split, Some("ts")); + let full_split_ast: QueryAst = serde_json::from_str(&full_split_request.query_ast).unwrap(); + let full_split_predicate = leaf::time_bounded_cached_predicate(&full_split_ast, "ts") + .expect("a full-split range should install a cached predicate"); + assert_eq!(full_split_predicate, partial_predicate); + + let mut full_split_search_after_request = SearchRequest { + query_ast: serde_json::to_string(&predicate_ast).unwrap(), + start_timestamp: Some(50), + end_timestamp: Some(250), + search_after: Some(Default::default()), + ..Default::default() + }; + leaf::rewrite_request(&mut full_split_search_after_request, &split, Some("ts")); + let full_split_search_after_ast: QueryAst = + serde_json::from_str(&full_split_search_after_request.query_ast).unwrap(); + assert_eq!( + leaf::time_bounded_cached_predicate(&full_split_search_after_ast, "ts"), + Some(partial_predicate.clone()) + ); + + let mut timeless_search_after_request = SearchRequest { + query_ast: serde_json::to_string(&predicate_ast).unwrap(), + search_after: Some(Default::default()), + ..Default::default() + }; + leaf::rewrite_request(&mut timeless_search_after_request, &split, Some("ts")); + let timeless_search_after_ast: QueryAst = + serde_json::from_str(&timeless_search_after_request.query_ast).unwrap(); + assert!( + leaf::time_bounded_cached_predicate(&timeless_search_after_ast, "ts").is_none(), + "the search-after cache node should not be treated as a predicate cache node" + ); + + for trivial_predicate in [QueryAst::MatchAll, QueryAst::MatchNone] { + let mut time_only_request = SearchRequest { + query_ast: serde_json::to_string(&trivial_predicate).unwrap(), + start_timestamp: Some(120), + end_timestamp: Some(180), + ..Default::default() + }; + leaf::rewrite_request(&mut time_only_request, &split, Some("ts")); + let time_only_ast: QueryAst = serde_json::from_str(&time_only_request.query_ast).unwrap(); + assert!(leaf::time_bounded_cached_predicate(&time_only_ast, "ts").is_none()); + } + + let user_semantic_time_range = QueryAst::from(RangeQuery { + field: "ts".to_string(), + lower_bound: Bound::Included(130i64.into()), + upper_bound: Bound::Unbounded, + }); + let nested_semantic_ast = QueryAst::from(BoolQuery { + must: vec![predicate_ast], + should: vec![user_semantic_time_range.clone()], + ..Default::default() + }); + let mut nested_semantic_request = SearchRequest { + query_ast: serde_json::to_string(&nested_semantic_ast).unwrap(), + start_timestamp: Some(120), + end_timestamp: Some(180), + ..Default::default() + }; + leaf::rewrite_request(&mut nested_semantic_request, &split, Some("ts")); + let nested_rewritten_ast: QueryAst = + serde_json::from_str(&nested_semantic_request.query_ast).unwrap(); + let cached_predicate = leaf::time_bounded_cached_predicate(&nested_rewritten_ast, "ts") + .expect("the non-time predicate should be cached"); + let QueryAst::Bool(cached_bool) = cached_predicate else { + panic!("expected normalized bool predicate"); + }; + assert_eq!(cached_bool.should, vec![user_semantic_time_range]); +} + +#[tokio::test] +async fn test_time_bounded_query_populates_and_reuses_complete_predicate_cache() { + let (test_sandbox, searcher_context, storage, splits, doc_mapper, start_timestamp) = + negative_cache_ts_test_setup().await; + let split_id = splits[0].split_id.clone(); + + let first_window_request = SearchRequest { + index_id_patterns: vec!["negative-cache-ts-index".to_string()], + query_ast: qast_json_helper("info", &["body"]), + start_timestamp: Some(start_timestamp), + end_timestamp: Some(start_timestamp + 3), + max_hits: 10, + ..Default::default() + }; + let mut rewritten_request = first_window_request.clone(); + leaf::rewrite_request( + &mut rewritten_request, + &splits[0], + doc_mapper.timestamp_field_name(), + ); + let rewritten_ast: QueryAst = serde_json::from_str(&rewritten_request.query_ast).unwrap(); + let predicate_ast = leaf::time_bounded_cached_predicate(&rewritten_ast, "ts") + .expect("a partial time range should install a cached predicate"); + let predicate_key = serde_json::to_string(&predicate_ast).unwrap(); + + let first_response = single_doc_mapping_leaf_search( + searcher_context.clone(), + std::sync::Arc::new(first_window_request), + storage.clone(), + splits.clone(), + doc_mapper.clone(), + ) + .await + .unwrap(); + assert_eq!(first_response.num_hits, 3); + let first_input_memory_bytes = first_response + .resource_stats + .as_ref() + .and_then(|stats| stats.split_resources_sum) + .expect("the split should report resource stats") + .input_memory_bytes; + + let (_segment_id, complete_hits) = searcher_context + .predicate_cache + .get(split_id.clone(), predicate_key.clone()) + .expect("the first window should populate the predicate cache"); + assert_eq!(complete_hits.size_hint(), 10); + + let full_split_window_request = SearchRequest { + index_id_patterns: vec!["negative-cache-ts-index".to_string()], + query_ast: qast_json_helper("info", &["body"]), + start_timestamp: Some(start_timestamp - 10), + end_timestamp: Some(start_timestamp + 20), + max_hits: 10, + ..Default::default() + }; + let full_split_response = single_doc_mapping_leaf_search( + searcher_context.clone(), + std::sync::Arc::new(full_split_window_request), + storage.clone(), + splits.clone(), + doc_mapper.clone(), + ) + .await + .unwrap(); + assert_eq!(full_split_response.num_hits, 10); + let full_split_input_memory_bytes = full_split_response + .resource_stats + .as_ref() + .and_then(|stats| stats.split_resources_sum) + .expect("the split should report resource stats") + .input_memory_bytes; + assert!( + full_split_input_memory_bytes < first_input_memory_bytes, + "a full-split predicate-cache hit should not warm the predicate posting lists" + ); + + let (_segment_id, reused_hits) = searcher_context + .predicate_cache + .get(split_id.clone(), predicate_key.clone()) + .expect("the full-split window should reuse the same predicate entry"); + assert_eq!(reused_hits.size_hint(), 10); + + let different_predicate_request = SearchRequest { + index_id_patterns: vec!["negative-cache-ts-index".to_string()], + query_ast: qast_json_helper("0", &["body"]), + start_timestamp: Some(start_timestamp - 10), + end_timestamp: Some(start_timestamp + 20), + max_hits: 10, + ..Default::default() + }; + let mut different_rewritten_request = different_predicate_request.clone(); + leaf::rewrite_request( + &mut different_rewritten_request, + &splits[0], + doc_mapper.timestamp_field_name(), + ); + let different_rewritten_ast: QueryAst = + serde_json::from_str(&different_rewritten_request.query_ast).unwrap(); + let different_predicate_ast = + leaf::time_bounded_cached_predicate(&different_rewritten_ast, "ts") + .expect("a full-split time range should install a cached predicate"); + let different_predicate_key = serde_json::to_string(&different_predicate_ast).unwrap(); + assert_ne!(different_predicate_key, predicate_key); + + let different_predicate_response = single_doc_mapping_leaf_search( + searcher_context.clone(), + std::sync::Arc::new(different_predicate_request), + test_sandbox.storage(), + splits.clone(), + doc_mapper.clone(), + ) + .await + .unwrap(); + assert_eq!(different_predicate_response.num_hits, 1); + let (_segment_id, different_hits) = searcher_context + .predicate_cache + .get(split_id.clone(), different_predicate_key.clone()) + .expect("a different predicate should populate a separate entry"); + assert_eq!(different_hits.size_hint(), 1); + + let different_partial_request = SearchRequest { + index_id_patterns: vec!["negative-cache-ts-index".to_string()], + query_ast: qast_json_helper("0", &["body"]), + start_timestamp: Some(start_timestamp), + end_timestamp: Some(start_timestamp + 2), + max_hits: 10, + ..Default::default() + }; + let mut different_partial_rewritten_request = different_partial_request.clone(); + leaf::rewrite_request( + &mut different_partial_rewritten_request, + &splits[0], + doc_mapper.timestamp_field_name(), + ); + let different_partial_ast: QueryAst = + serde_json::from_str(&different_partial_rewritten_request.query_ast).unwrap(); + let different_partial_predicate = + leaf::time_bounded_cached_predicate(&different_partial_ast, "ts") + .expect("a partial time range should install a cached predicate"); + assert_eq!( + serde_json::to_string(&different_partial_predicate).unwrap(), + different_predicate_key, + "full-split and partial windows should use the same predicate key" + ); + let different_partial_response = single_doc_mapping_leaf_search( + searcher_context.clone(), + std::sync::Arc::new(different_partial_request), + storage, + splits, + doc_mapper, + ) + .await + .unwrap(); + assert_eq!(different_partial_response.num_hits, 1); + assert!( + searcher_context + .predicate_cache + .get(split_id, different_predicate_key) + .is_some(), + "the partial window should reuse the full-split predicate entry" + ); + assert!( + searcher_context + .predicate_cache + .get("another-split".to_string(), predicate_key) + .is_none(), + "predicate entries must remain split-local" + ); + + test_sandbox.assert_quit().await; +} + #[tokio::test] async fn test_negative_cache_short_circuits_across_time_windows() { // A per-term absence does not depend on the time window: the required term