CR comments

Author: fulmicoton-dd

quickwit/Cargo.lock

@@ -32,6 +32,7 @@ openssl = { workspace = true, optional = true }
 
 quickwit-common = { workspace = true }
 quickwit-config = { workspace = true }
+quickwit-doc-mapper = { workspace = true }
 quickwit-proto = { workspace = true }
 quickwit-search = { workspace = true }
 quickwit-storage = { workspace = true }

quickwit/quickwit-lambda-server/Cargo.toml

@@ -12,13 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::str::FromStr;
+use std::sync::Arc;
+
 use base64::prelude::*;
 use prost::Message;
+use quickwit_common::uri::Uri;
+use quickwit_doc_mapper::DocMapper;
 use quickwit_proto::search::lambda_single_split_result::Outcome;
 use quickwit_proto::search::{
-    LambdaSearchResponses, LambdaSingleSplitResult, LeafSearchRequest, SplitIdAndFooterOffsets,
+    LambdaSearchResponses, LambdaSingleSplitResult, LeafRequestRef, LeafSearchRequest,
+    SearchRequest,
 };
-use quickwit_search::leaf::multi_index_leaf_search;
+use quickwit_search::leaf::single_doc_mapping_leaf_search;
+use quickwit_storage::Storage;
 use serde::{Deserialize, Serialize};
 use tracing::{error, info, instrument, warn};
 
@@ -50,57 +57,106 @@ pub async fn handle_leaf_search(
     ctx: &LambdaSearcherContext,
 ) -> LambdaResult<LambdaSearchResponsePayload> {
     // Decode base64 payload
-    let request_bytes = BASE64_STANDARD
+    let request_bytes: Vec<u8> = BASE64_STANDARD
         .decode(&event.payload)
-        .map_err(|e| LambdaError::Serialization(format!("base64 decode error: {}", e)))?;
+        .map_err(|err| LambdaError::Serialization(format!("base64 decode error: {}", err)))?;
 
     // Deserialize LeafSearchRequest
     let leaf_search_request = LeafSearchRequest::decode(&request_bytes[..])?;
 
-    let all_splits: Vec<(usize, SplitIdAndFooterOffsets)> =
-        leaf_search_request
-            .leaf_requests
-            .iter()
-            .enumerate()
-            .flat_map(|(leaf_req_idx, leaf_request_ref)| {
-                leaf_request_ref.split_offsets.iter().cloned().map(
-                    move |split_id_and_footer_offsets| (leaf_req_idx, split_id_and_footer_offsets),
-                )
-            })
-            .collect();
+    // Unpack the shared fields once instead of cloning per split.
+    let search_request: Arc<SearchRequest> = leaf_search_request
+        .search_request
+        .ok_or_else(|| LambdaError::Internal("no search request".to_string()))?
+        .into();
 
-    let num_splits = all_splits.len();
-    info!(num_splits, "processing leaf search request (per-split)");
+    let doc_mappers: Vec<Arc<DocMapper>> = leaf_search_request
+        .doc_mappers
+        .iter()
+        .map(String::as_str)
+        .map(serde_json::from_str::<Arc<DocMapper>>)
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|err| {
+            LambdaError::Internal(format!("failed to deserialize doc mapper: `{err}`"))
+        })?;
+
+    // Resolve storage for every index URI upfront.
+    let mut storages: Vec<Arc<dyn quickwit_storage::Storage>> =
+        Vec::with_capacity(leaf_search_request.index_uris.len());
+    for uri_str in &leaf_search_request.index_uris {
+        let uri = Uri::from_str(uri_str)
+            .map_err(|err| LambdaError::Internal(format!("invalid index uri: {err}")))?;
+        let storage =
+            ctx.storage_resolver.resolve(&uri).await.map_err(|err| {
+                LambdaError::Internal(format!("failed to resolve storage: {err}"))
+            })?;
+        storages.push(storage);
+    }
+
+    let split_results: Vec<LambdaSingleSplitResult> = lambda_leaf_search(
+        search_request,
+        leaf_search_request.leaf_requests,
+        &doc_mappers[..],
+        &storages[..],
+        ctx,
+    )
+    .await?;
+    let wrapper = LambdaSearchResponses { split_results };
+    let response_bytes = wrapper.encode_to_vec();
+    let payload = BASE64_STANDARD.encode(&response_bytes);
+
+    Ok(LambdaSearchResponsePayload { payload })
+}
 
-    // Process each split in parallel using a JoinSet. The SearchPermitProvider
-    // inside SearcherContext gates concurrency based on memory budget.
+/// Lambda leaf search returns individual split results.
+async fn lambda_leaf_search(
+    search_request: Arc<SearchRequest>,
+    leaf_req_ref: Vec<LeafRequestRef>,
+    doc_mappers: &[Arc<DocMapper>],
+    storages: &[Arc<dyn Storage>],
+    ctx: &LambdaSearcherContext,
+) -> LambdaResult<Vec<LambdaSingleSplitResult>> {
+    // Flatten leaf_requests into per-split tasks using pre-resolved Arc references.
     let mut split_search_joinset: tokio::task::JoinSet<(String, Result<_, String>)> =
         tokio::task::JoinSet::new();
-    for (leaf_req_idx, split) in all_splits {
-        let split_id = split.split_id.clone();
-        let leaf_request_ref = &leaf_search_request.leaf_requests[leaf_req_idx];
-        let single_split_request = LeafSearchRequest {
-            search_request: leaf_search_request.search_request.clone(),
-            doc_mappers: leaf_search_request.doc_mappers.clone(),
-            index_uris: leaf_search_request.index_uris.clone(),
-            leaf_requests: vec![quickwit_proto::search::LeafRequestRef {
-                index_uri_ord: leaf_request_ref.index_uri_ord,
-                doc_mapper_ord: leaf_request_ref.doc_mapper_ord,
-                split_offsets: vec![split],
-            }],
-        };
-
-        let searcher_context = ctx.searcher_context.clone();
-        let storage_resolver = ctx.storage_resolver.clone();
-        split_search_joinset.spawn(async move {
-            let result =
-                multi_index_leaf_search(searcher_context, single_split_request, storage_resolver)
-                    .await
-                    .map_err(|err| format!("{err}"));
-            (split_id, result)
-        });
+
+    for leaf_req in leaf_req_ref {
+        let doc_mapper = doc_mappers
+            .get(leaf_req.doc_mapper_ord as usize)
+            .ok_or_else(|| {
+                LambdaError::Internal(format!(
+                    "doc_mapper_ord out of bounds: {}",
+                    leaf_req.doc_mapper_ord
+                ))
+            })?
+            .clone();
+        let storage = storages[leaf_req.index_uri_ord as usize].clone();
+
+        for split_id_and_footer_offsets in leaf_req.split_offsets {
+            let split_id = split_id_and_footer_offsets.split_id.clone();
+            let searcher_context = ctx.searcher_context.clone();
+            let search_request = search_request.clone();
+            let doc_mapper = doc_mapper.clone();
+            let storage = storage.clone();
+            let split = split_id_and_footer_offsets.clone();
+            split_search_joinset.spawn(async move {
+                let result = single_doc_mapping_leaf_search(
+                    searcher_context,
+                    search_request,
+                    storage,
+                    vec![split],
+                    doc_mapper,
+                )
+                .await
+                .map_err(|err| format!("{err}"));
+                (split_id, result)
+            });
+        }
     }
 
+    let num_splits = split_search_joinset.len();
+    info!(num_splits, "processing leaf search request (per-split)");
+
     // Collect results. Order is irrelevant: each result is tagged with its split_id.
     let mut split_results: Vec<LambdaSingleSplitResult> = Vec::with_capacity(num_splits);
     let mut num_successes: usize = 0;
@@ -140,9 +196,5 @@ pub async fn handle_leaf_search(
         num_failures, "leaf search completed (per-split)"
     );
 
-    let wrapper = LambdaSearchResponses { split_results };
-    let response_bytes = wrapper.encode_to_vec();
-    let payload = BASE64_STANDARD.encode(&response_bytes);
-
-    Ok(LambdaSearchResponsePayload { payload })
+    Ok(split_results)
 }

quickwit/quickwit-lambda-server/src/handler.rs

@@ -181,8 +181,8 @@ impl BuildTantivyAst for RangeQuery {
                     convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
                 let truncate_datetime =
                     |date: &DateTime| date.truncate(date_options.get_precision());
-                let lower_bound = map_bound(&lower_bound, truncate_datetime);
-                let upper_bound = map_bound(&upper_bound, truncate_datetime);
+                let lower_bound = lower_bound.as_ref().map(truncate_datetime);
+                let upper_bound = upper_bound.as_ref().map(truncate_datetime);
                 FastFieldRangeQuery::new(
                     lower_bound.map(|val| Term::from_field_date(field, val)),
                     upper_bound.map(|val| Term::from_field_date(field, val)),
@@ -277,14 +277,6 @@ impl BuildTantivyAst for RangeQuery {
     }
 }
 
-fn map_bound<TFrom, TTo>(bound: &Bound<TFrom>, transform: impl Fn(&TFrom) -> TTo) -> Bound<TTo> {
-    match bound {
-        Bound::Excluded(from_val) => Bound::Excluded(transform(from_val)),
-        Bound::Included(from_val) => Bound::Included(transform(from_val)),
-        Bound::Unbounded => Bound::Unbounded,
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::ops::Bound;

quickwit/quickwit-query/src/query_ast/range_query.rs

@@ -499,6 +499,17 @@ async fn leaf_search_single_split(
     let mut leaf_search_state_guard =
         SplitSearchStateGuard::new(ctx.split_outcome_counters.clone());
 
+    // We already checked if the result was already in the partial result cache,
+    // but it's not a bad idea to check again.
+    if let Some(cached_answer) = ctx
+        .searcher_context
+        .leaf_search_cache
+        .get(split.clone(), search_request.clone())
+    {
+        leaf_search_state_guard.set_state(SplitSearchState::CacheHit);
+        return Ok(Some(cached_answer));
+    }
+
     let query_ast: QueryAst = serde_json::from_str(search_request.query_ast.as_str())
         .map_err(|err| SearchError::InvalidQuery(err.to_string()))?;
 
@@ -721,17 +732,6 @@ fn visit_aggregation_mut(
     modified_something
 }
 
-/// Maps a `Bound<T>` to a `Bound<U>` by applying a function to the contained value.
-/// Equivalent to `Bound::map`, which is currently unstable.
-pub fn map_bound<T, U>(bound: Bound<T>, f: impl FnOnce(T) -> U) -> Bound<U> {
-    use Bound::*;
-    match bound {
-        Unbounded => Unbounded,
-        Included(x) => Included(f(x)),
-        Excluded(x) => Excluded(f(x)),
-    }
-}
-
 // returns the max of left and right, that isn't unbounded. Useful for making
 // the intersection of lower bound of ranges
 fn max_bound<T: Ord + Copy>(left: Bound<T>, right: Bound<T>) -> Bound<T> {
@@ -864,12 +864,8 @@ fn remove_redundant_timestamp_range(
     if final_start_timestamp != Bound::Unbounded || final_end_timestamp != Bound::Unbounded {
         let range = RangeQuery {
             field: timestamp_field.to_string(),
-            lower_bound: map_bound(final_start_timestamp, |bound| {
-                bound.into_timestamp_nanos().into()
-            }),
-            upper_bound: map_bound(final_end_timestamp, |bound| {
-                bound.into_timestamp_nanos().into()
-            }),
+            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()
@@ -1390,12 +1386,12 @@ fn disable_search_request_hits(search_request: &mut SearchRequest) {
 }
 
 /// Searches multiple splits for a specific index and a single doc mapping
-/// Offloads splits to Lambda invocations, distributing them accross batches
+/// Offloads splits to Lambda invocations, distributing them across batches
 /// balanced by document count. Each batch is invoked independently; a failure
 /// in one batch does not affect others.
 async fn run_offloaded_search_tasks(
     searcher_context: &SearcherContext,
-    request: &SearchRequest,
+    search_request: &SearchRequest,
     doc_mapper: &DocMapper,
     index_uri: Uri,
     splits_with_requests: Vec<(SplitIdAndFooterOffsets, SearchRequest)>,
@@ -1436,16 +1432,14 @@ async fn run_offloaded_search_tasks(
         lambda_config.max_splits_per_invocation,
     );
 
-    let mut search_request_for_leaf = request.clone();
-    search_request_for_leaf.start_offset = 0;
-    search_request_for_leaf.max_hits += request.start_offset;
-
     let mut lambda_tasks_joinset = JoinSet::new();
     for batch in batches {
         let batch_split_ids: Vec<String> =
             batch.iter().map(|split| split.split_id.clone()).collect();
         let leaf_request = LeafSearchRequest {
-            search_request: Some(search_request_for_leaf.clone()),
+            // Note this is not the split-specific rewritten request, we ship the main request,
+            // and the leaf will apply the split specific rewrite on its own.
+            search_request: Some(search_request.clone()),
             doc_mappers: vec![doc_mapper_str.clone()],
             index_uris: vec![index_uri.as_str().to_string()], //< careful here. Calling to_string() directly would return a redacted uri.
             leaf_requests: vec![quickwit_proto::search::LeafRequestRef {
@@ -1474,12 +1468,13 @@ async fn run_offloaded_search_tasks(
                 for split_result in split_results {
                     match split_result.outcome {
                         Some(Outcome::Response(response)) => {
-                            if let Some((split_info, search_req)) =
+                            if let Some((split_info, single_split_search_req)) =
                                 split_lookup.remove(&split_result.split_id)
                             {
+                                // We use the single_split_search_req to perform the search
                                 searcher_context.leaf_search_cache.put(
                                     split_info,
-                                    search_req,
+                                    single_split_search_req,
                                     response.clone(),
                                 );
                             }
@@ -1616,8 +1611,12 @@ pub async fn single_doc_mapping_leaf_search(
         CanSplitDoBetter::from_request(&request, doc_mapper.timestamp_field_name());
     let mut split_with_req: Vec<(SplitIdAndFooterOffsets, SearchRequest)> =
         split_filter.optimize(&request, splits)?;
-    for (split, search_request) in &mut split_with_req {
-        rewrite_request(search_request, split, doc_mapper.timestamp_field_name());
+    for (split, single_split_search_request) in &mut split_with_req {
+        rewrite_request(
+            single_split_search_request,
+            split,
+            doc_mapper.timestamp_field_name(),
+        );
     }
     let split_filter_arc: Arc<RwLock<CanSplitDoBetter>> = Arc::new(RwLock::new(split_filter));
 

quickwit/quickwit-search/src/leaf.rs

@@ -153,7 +153,8 @@ impl SearchServiceImpl {
     }
 }
 
-pub fn deserialize_doc_mapper(doc_mapper_str: &str) -> crate::Result<Arc<DocMapper>> {
+/// Deserializes a JSON-encoded doc mapper string into an `Arc<DocMapper>`.
+pub(crate) fn deserialize_doc_mapper(doc_mapper_str: &str) -> crate::Result<Arc<DocMapper>> {
     let doc_mapper = serde_json::from_str::<Arc<DocMapper>>(doc_mapper_str).map_err(|err| {
         SearchError::Internal(format!("failed to deserialize doc mapper: `{err}`"))
     })?;