Refactor to clarify leaf search levels

Author: rdettai

quickwit/quickwit-search/src/leaf.rs

@@ -1174,9 +1174,10 @@ impl CanSplitDoBetter {
     }
 }
 
-/// `multi_leaf_search` searches multiple indices and multiple splits.
+/// Searches multiple splits, potentially in multiple indices, sitting on different storages and
+/// having different doc mappings.
 #[instrument(skip_all, fields(index = ?leaf_search_request.search_request.as_ref().unwrap().index_id_patterns))]
-pub async fn multi_leaf_search(
+pub async fn multi_index_leaf_search(
     searcher_context: Arc<SearcherContext>,
     leaf_search_request: LeafSearchRequest,
     storage_resolver: &StorageResolver,
@@ -1224,18 +1225,25 @@ pub async fn multi_leaf_search(
             })?
             .clone();
 
-        let leaf_request_future = tokio::spawn(
-            resolve_storage_and_leaf_search(
-                searcher_context.clone(),
-                search_request.clone(),
-                index_uri,
-                storage_resolver.clone(),
-                leaf_search_request_ref.split_offsets,
-                doc_mapper,
-                aggregation_limits.clone(),
-            )
-            .in_current_span(),
-        );
+        let leaf_request_future = tokio::spawn({
+            let storage_resolver = storage_resolver.clone();
+            let searcher_context = searcher_context.clone();
+            let search_request = search_request.clone();
+            let aggregation_limits = aggregation_limits.clone();
+            async move {
+                let storage = storage_resolver.resolve(&index_uri).await?;
+                single_doc_mapping_leaf_search(
+                    searcher_context,
+                    search_request,
+                    storage,
+                    leaf_search_request_ref.split_offsets,
+                    doc_mapper,
+                    aggregation_limits,
+                )
+                .await
+            }
+            .in_current_span()
+        });
         leaf_request_tasks.push(leaf_request_future);
     }
 
@@ -1268,29 +1276,6 @@ pub async fn multi_leaf_search(
         .context("failed to merge split search responses")?
 }
 
-/// Resolves storage and calls leaf_search
-#[allow(clippy::too_many_arguments)]
-async fn resolve_storage_and_leaf_search(
-    searcher_context: Arc<SearcherContext>,
-    search_request: Arc<SearchRequest>,
-    index_uri: quickwit_common::uri::Uri,
-    storage_resolver: StorageResolver,
-    splits: Vec<SplitIdAndFooterOffsets>,
-    doc_mapper: Arc<DocMapper>,
-    aggregations_limits: AggregationLimitsGuard,
-) -> crate::Result<LeafSearchResponse> {
-    let storage = storage_resolver.resolve(&index_uri).await?;
-    leaf_search(
-        searcher_context.clone(),
-        search_request.clone(),
-        storage.clone(),
-        splits,
-        doc_mapper,
-        aggregations_limits,
-    )
-    .await
-}
-
 /// Optimizes the search_request based on CanSplitDoBetter
 /// Returns true if the split can return better results
 fn check_optimize_search_request(
@@ -1314,14 +1299,14 @@ fn disable_search_request_hits(search_request: &mut SearchRequest) {
     search_request.sort_fields.clear();
 }
 
-/// `leaf` step of search.
+/// Searches multiple splits from a specific index and a single doc mapping
 ///
 /// The leaf search collects all kind of information, and returns a set of
 /// [PartialHit](quickwit_proto::search::PartialHit) candidates. The root will be in
 /// charge to consolidate, identify the actual final top hits to display, and
 /// fetch the actual documents to convert the partial hits into actual Hits.
 #[instrument(skip_all, fields(index = ?request.index_id_patterns))]
-pub async fn leaf_search(
+pub async fn single_doc_mapping_leaf_search(
     searcher_context: Arc<SearcherContext>,
     request: Arc<SearchRequest>,
     index_storage: Arc<dyn Storage>,

quickwit/quickwit-search/src/service.rs

@@ -38,7 +38,7 @@ use tantivy::aggregation::AggregationLimitsGuard;
 use tokio::sync::{Semaphore, oneshot};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
-use crate::leaf::multi_leaf_search;
+use crate::leaf::multi_index_leaf_search;
 use crate::leaf_cache::LeafSearchCache;
 use crate::list_fields::{leaf_list_fields, root_list_fields};
 use crate::list_fields_cache::ListFieldsCache;
@@ -204,7 +204,7 @@ impl SearchService for SearchServiceImpl {
             .map(|req| req.split_offsets.len())
             .sum::<usize>();
         let completion_tx = start_leaf_search_metric_recording(num_splits).await;
-        let leaf_search_response_result = multi_leaf_search(
+        let leaf_search_response_result = multi_index_leaf_search(
             self.searcher_context.clone(),
             leaf_search_request,
             &self.storage_resolver,

quickwit/quickwit-search/src/tests.rs

@@ -32,7 +32,7 @@ use tantivy::Term;
 use tantivy::schema::OwnedValue as TantivyValue;
 use tantivy::time::OffsetDateTime;
 
-use self::leaf::leaf_search;
+use self::leaf::single_doc_mapping_leaf_search;
 use super::*;
 use crate::find_trace_ids_collector::Span;
 use crate::list_terms::leaf_list_terms;
@@ -1051,7 +1051,7 @@ async fn test_search_util(test_sandbox: &TestSandbox, query: &str) -> Vec<u32> {
 
     let agg_limits = searcher_context.get_aggregation_limits();
 
-    let search_response = leaf_search(
+    let search_response = single_doc_mapping_leaf_search(
         searcher_context,
         request,
         test_sandbox.storage(),