Get count from split metadata on simple query with timerange

Author: tontinton

quickwit/quickwit-search/src/leaf.rs

@@ -463,7 +463,12 @@ async fn leaf_search_single_split(
     // This may be the case for AllQuery with a sort by date and time filter, where the current
     // split can't have better results.
     //
-    if is_metadata_count_request_with_ast(&query_ast, &search_request) {
+    if is_metadata_count_request_with_ast(
+        &query_ast,
+        &search_request,
+        split.timestamp_start,
+        split.timestamp_end,
+    ) {
         return Ok(get_leaf_resp_from_count(split.num_docs));
     }
 
@@ -535,7 +540,12 @@ async fn leaf_search_single_split(
                 check_optimize_search_request(&mut search_request, &split, &split_filter);
                 collector.update_search_param(&search_request);
                 let mut leaf_search_response: LeafSearchResponse =
-                    if is_metadata_count_request_with_ast(&query_ast, &search_request) {
+                    if is_metadata_count_request_with_ast(
+                        &query_ast,
+                        &search_request,
+                        split.timestamp_start,
+                        split.timestamp_end,
+                    ) {
                         get_leaf_resp_from_count(searcher.num_docs())
                     } else if collector.is_count_only() {
                         let count = query.count(&searcher)? as u64;

quickwit/quickwit-search/src/root.rs

@@ -636,9 +636,13 @@ async fn search_partial_hits_phase_with_scroll(
 /// metadata count.
 ///
 /// This is done by exclusion, so we will need to keep it up to date if fields are added.
-pub fn is_metadata_count_request(request: &SearchRequest) -> bool {
+pub fn is_metadata_count_request(request: &SearchRequest, split: &SplitMetadata) -> bool {
     let query_ast: QueryAst = serde_json::from_str(&request.query_ast).unwrap();
-    is_metadata_count_request_with_ast(&query_ast, request)
+
+    let start_time = split.time_range.as_ref().map(|x| x.start()).copied();
+    let end_time = split.time_range.as_ref().map(|x| x.end()).copied();
+
+    is_metadata_count_request_with_ast(&query_ast, request, start_time, end_time)
 }
 
 /// Check if the request is a count request without any filters, so we can just return the split
@@ -647,42 +651,53 @@ pub fn is_metadata_count_request(request: &SearchRequest) -> bool {
 /// This is done by exclusion, so we will need to keep it up to date if fields are added.
 ///
 /// The passed query_ast should match the serialized on in request.
-pub fn is_metadata_count_request_with_ast(query_ast: &QueryAst, request: &SearchRequest) -> bool {
+pub fn is_metadata_count_request_with_ast(
+    query_ast: &QueryAst,
+    request: &SearchRequest,
+    split_start_timestamp: Option<i64>,
+    split_end_timestamp: Option<i64>,
+) -> bool {
     if query_ast != &QueryAst::MatchAll {
         return false;
     }
     if request.max_hits != 0 {
         return false;
     }
 
-    // If the start and end timestamp encompass the whole split, it is still a count query.
-    // We remove this currently on the leaf level, but not yet on the root level.
-    // There's a small advantage when we would do this on the root level, since we have the
-    // counts available on the split. On the leaf it is currently required to open the split
-    // to get the count.
-    if request.start_timestamp.is_some() || request.end_timestamp.is_some() {
-        return false;
+    if let Some(request_start_timestamp) = request.start_timestamp {
+        let Some(split_start_timestamp) = split_start_timestamp else {
+            return false;
+        };
+        if split_start_timestamp < request_start_timestamp {
+            return false;
+        }
     }
+    if let Some(request_end_timestamp) = request.end_timestamp {
+        let Some(split_end_timestamp) = split_end_timestamp else {
+            return false;
+        };
+        if split_end_timestamp >= request_end_timestamp {
+            return false;
+        }
+    }
+
     if request.aggregation_request.is_some() || !request.snippet_fields.is_empty() {
         return false;
     }
     true
 }
 
 /// Get a leaf search response that returns the num_docs of the split
-pub fn get_count_from_metadata(split_metadatas: &[SplitMetadata]) -> Vec<LeafSearchResponse> {
-    split_metadatas
-        .iter()
-        .map(|metadata| LeafSearchResponse {
-            num_hits: metadata.num_docs as u64,
-            partial_hits: Vec::new(),
-            failed_splits: Vec::new(),
-            num_attempted_splits: 1,
-            num_successful_splits: 1,
-            intermediate_aggregation_result: None,
-            resource_stats: None,
-        })
-        .collect()
+pub fn get_count_from_metadata(metadata: &SplitMetadata) -> LeafSearchResponse {
+    LeafSearchResponse {
+        num_hits: metadata.num_docs as u64,
+        partial_hits: Vec::new(),
+        failed_splits: Vec::new(),
+        num_attempted_splits: 1,
+        num_successful_splits: 1,
+        intermediate_aggregation_result: None,
+        resource_stats: None,
+    }
 }
 
 /// Returns true if the query is particularly memory intensive.
@@ -730,26 +745,31 @@ pub(crate) async fn search_partial_hits_phase(
     split_metadatas: &[SplitMetadata],
     cluster_client: &ClusterClient,
 ) -> crate::Result<LeafSearchResponse> {
-    let leaf_search_responses: Vec<LeafSearchResponse> =
-        if is_metadata_count_request(search_request) {
-            get_count_from_metadata(split_metadatas)
+    let mut leaf_search_responses: Vec<LeafSearchResponse> =
+        Vec::with_capacity(split_metadatas.len());
+    let mut leaf_search_jobs = Vec::new();
+    for split in split_metadatas {
+        if is_metadata_count_request(search_request, split) {
+            leaf_search_responses.push(get_count_from_metadata(split));
         } else {
-            let jobs: Vec<SearchJob> = split_metadatas.iter().map(SearchJob::from).collect();
-            let assigned_leaf_search_jobs = cluster_client
-                .search_job_placer
-                .assign_jobs(jobs, &HashSet::default())
-                .await?;
-            let mut leaf_request_tasks = Vec::new();
-            for (client, client_jobs) in assigned_leaf_search_jobs {
-                let leaf_request = jobs_to_leaf_request(
-                    search_request,
-                    indexes_metas_for_leaf_search,
-                    client_jobs,
-                )?;
-                leaf_request_tasks.push(cluster_client.leaf_search(leaf_request, client.clone()));
-            }
-            try_join_all(leaf_request_tasks).await?
-        };
+            leaf_search_jobs.push(SearchJob::from(split));
+        }
+    }
+
+    if !leaf_search_jobs.is_empty() {
+        let assigned_leaf_search_jobs = cluster_client
+            .search_job_placer
+            .assign_jobs(leaf_search_jobs, &HashSet::default())
+            .await?;
+        let mut leaf_request_tasks = Vec::new();
+        for (client, client_jobs) in assigned_leaf_search_jobs {
+            let leaf_request =
+                jobs_to_leaf_request(search_request, indexes_metas_for_leaf_search, client_jobs)?;
+            leaf_request_tasks.push(cluster_client.leaf_search(leaf_request, client.clone()));
+        }
+        let executed_leaf_search_responses = try_join_all(leaf_request_tasks).await?;
+        leaf_search_responses.extend(executed_leaf_search_responses);
+    }
 
     // Creates a collector which merges responses into one
     let merge_collector =