Moving lambda execution to the leaf.

Author: fulmicoton-dd

quickwit/quickwit-config/src/node_config/mod.rs

@@ -320,6 +320,11 @@ pub struct LambdaConfig {
     /// Maximum number of splits per Lambda invocation.
     #[serde(default = "LambdaConfig::default_max_splits_per_invocation")]
     pub max_splits_per_invocation: usize,
+    /// When the number of pending split searches exceeds this threshold,
+    /// new splits are offloaded to Lambda instead of being queued locally.
+    /// A value of 0 disables offloading (all splits are processed locally).
+    #[serde(default = "LambdaConfig::default_offload_threshold")]
+    pub offload_threshold: usize,
     /// Auto-deploy configuration. If set, Quickwit will automatically deploy
     /// the Lambda function at startup.
     /// If deploying a lambda fails, Quickwit will log an error and fail.
@@ -356,6 +361,7 @@ impl Default for LambdaConfig {
         Self {
             function_name: Self::default_function_name(),
             max_splits_per_invocation: Self::default_max_splits_per_invocation(),
+            offload_threshold: Self::default_offload_threshold(),
             auto_deploy: None,
         }
     }
@@ -368,6 +374,9 @@ impl LambdaConfig {
     fn default_max_splits_per_invocation() -> usize {
         10
     }
+    fn default_offload_threshold() -> usize {
+        20
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]

quickwit/quickwit-lambda-client/src/invoker.rs

@@ -21,7 +21,7 @@ use aws_sdk_lambda::types::InvocationType;
 use base64::prelude::*;
 use prost::Message;
 use quickwit_lambda_server::{LeafSearchPayload, LeafSearchResponsePayload};
-use quickwit_proto::search::{LeafSearchRequest, LeafSearchResponse};
+use quickwit_proto::search::{LeafSearchRequest, LeafSearchResponse, LeafSearchResponses};
 use quickwit_search::{LambdaLeafSearchInvoker, SearchError};
 use tracing::{debug, info, instrument};
 
@@ -91,7 +91,7 @@ impl LambdaLeafSearchInvoker for AwsLambdaInvoker {
     async fn invoke_leaf_search(
         &self,
         request: LeafSearchRequest,
-    ) -> Result<LeafSearchResponse, SearchError> {
+    ) -> Result<Vec<LeafSearchResponse>, SearchError> {
         let start = std::time::Instant::now();
 
         let result = self.invoke_leaf_search_inner(request).await;
@@ -115,7 +115,7 @@ impl AwsLambdaInvoker {
     async fn invoke_leaf_search_inner(
         &self,
         request: LeafSearchRequest,
-    ) -> Result<LeafSearchResponse, SearchError> {
+    ) -> Result<Vec<LeafSearchResponse>, SearchError> {
         // Serialize request to protobuf bytes, then base64 encode
         let request_bytes = request.encode_to_vec();
         let payload = LeafSearchPayload {
@@ -178,14 +178,14 @@ impl AwsLambdaInvoker {
             .decode(&lambda_response.payload)
             .map_err(|e| SearchError::Internal(format!("Base64 decode error: {}", e)))?;
 
-        let leaf_response = LeafSearchResponse::decode(&response_bytes[..])
+        let leaf_responses = LeafSearchResponses::decode(&response_bytes[..])
             .map_err(|e| SearchError::Internal(format!("Protobuf decode error: {}", e)))?;
 
         debug!(
-            num_hits = leaf_response.num_hits,
+            num_responses = leaf_responses.responses.len(),
             "Lambda invocation completed"
         );
 
-        Ok(leaf_response)
+        Ok(leaf_responses.responses)
     }
 }

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

@@ -34,7 +34,7 @@ impl LambdaSearcherContext {
 
         let config = LambdaSearcherConfig::try_from_env()?;
         let searcher_config = create_searcher_config(&config);
-        let searcher_context = Arc::new(SearcherContext::new(searcher_config, None));
+        let searcher_context = Arc::new(SearcherContext::new(searcher_config, None, None));
         let storage_resolver = StorageResolver::configured(&Default::default());
 
         Ok(Self {

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

@@ -14,9 +14,12 @@
 
 use base64::prelude::*;
 use prost::Message;
-use quickwit_proto::search::{LeafSearchRequest, LeafSearchResponse};
+use quickwit_proto::search::{
+    LeafSearchRequest, LeafSearchResponse, LeafSearchResponses, SplitIdAndFooterOffsets,
+};
 use quickwit_search::leaf::multi_index_leaf_search;
 use serde::{Deserialize, Serialize};
+use tokio::task::JoinSet;
 use tracing::{info, instrument};
 
 use crate::context::LambdaSearcherContext;
@@ -32,11 +35,14 @@ pub struct LeafSearchPayload {
 /// Response from leaf search Lambda invocation.
 #[derive(Debug, Serialize, Deserialize)]
 pub struct LeafSearchResponsePayload {
-    /// Base64-encoded serialized LeafSearchResponse protobuf.
+    /// Base64-encoded serialized `LeafSearchResponses` protobuf (one per split).
     pub payload: String,
 }
 
 /// Handle a leaf search request in Lambda.
+///
+/// Returns one `LeafSearchResponse` per split. Each split is processed
+/// independently so that the caller can cache and merge results individually.
 #[instrument(skip(ctx), fields(request_id))]
 pub async fn handle_leaf_search(
     event: LeafSearchPayload,
@@ -50,30 +56,73 @@ pub async fn handle_leaf_search(
     // Deserialize LeafSearchRequest
     let leaf_search_request = LeafSearchRequest::decode(&request_bytes[..])?;
 
-    let num_splits: usize = leaf_search_request
+    let all_splits: Vec<(usize, SplitIdAndFooterOffsets)> = leaf_search_request
         .leaf_requests
         .iter()
-        .map(|leaf_request_ref| leaf_request_ref.split_offsets.len())
-        .sum();
+        .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();
 
-    info!(num_splits, "processing leaf search request");
+    let num_splits = all_splits.len();
+    info!(num_splits, "processing leaf search request (per-split)");
 
-    // Execute leaf search
-    let leaf_search_response = multi_index_leaf_search(
-        ctx.searcher_context.clone(),
-        leaf_search_request,
-        &ctx.storage_resolver,
-    )
-    .await?;
+    // Process each split in parallel. The SearchPermitProvider inside
+    // SearcherContext gates concurrency based on memory budget.
+    let mut join_set = JoinSet::new();
+    for (split_idx, (leaf_req_idx, split)) in all_splits.into_iter().enumerate() {
+        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();
+        join_set.spawn(async move {
+            let response = multi_index_leaf_search(
+                searcher_context,
+                single_split_request,
+                &storage_resolver,
+            )
+            .await;
+            (split_idx, response)
+        });
+    }
+
+    // Collect results, preserving split order.
+    let mut indexed_responses: Vec<(usize, LeafSearchResponse)> =
+        Vec::with_capacity(num_splits);
+    while let Some(join_result) = join_set.join_next().await {
+        let (split_idx, search_result) = join_result
+            .map_err(|e| LambdaError::Internal(format!("split search task failed: {e}")))?;
+        let response = search_result
+            .map_err(|e| LambdaError::Internal(format!("leaf search failed: {e}")))?;
+        indexed_responses.push((split_idx, response));
+    }
+    indexed_responses.sort_by_key(|(idx, _)| *idx);
+    let responses: Vec<LeafSearchResponse> =
+        indexed_responses.into_iter().map(|(_, r)| r).collect();
 
     info!(
-        num_hits = leaf_search_response.num_hits,
-        num_successful_splits = leaf_search_response.num_successful_splits,
-        "leaf search completed"
+        num_responses = responses.len(),
+        "leaf search completed (per-split)"
     );
 
-    // Serialize response
-    let response_bytes = leaf_search_response.encode_to_vec();
+    // Serialize as LeafSearchResponses wrapper
+    let wrapper = LeafSearchResponses { responses };
+    let response_bytes = wrapper.encode_to_vec();
     let payload = BASE64_STANDARD.encode(&response_bytes);
 
     Ok(LeafSearchResponsePayload { payload })

quickwit/quickwit-proto/protos/quickwit/search.proto

@@ -182,14 +182,6 @@ message ListFields {
 }
 // -- Search -------------------
 
-// Execution mode for leaf search operations.
-enum ExecutionMode {
-  // Default: use gRPC to searcher nodes (existing behavior).
-  EXECUTION_MODE_GRPC = 0;
-  // Execute leaf search via remote serverless functions (e.g., AWS Lambda).
-  EXECUTION_MODE_REMOTE_FUNCTION = 1;
-}
-
 message SearchRequest {
   // Index ID patterns
   repeated string index_id_patterns = 1;
@@ -255,11 +247,6 @@ message SearchRequest {
   // When an exact index ID is provided (not a pattern), the query fails only if
   // that index is not found and this parameter is set to `false`.
   bool ignore_missing_indexes = 18;
-
-  // Execution mode for leaf search operations.
-  // When set to EXECUTION_MODE_REMOTE_FUNCTION, leaf search is executed via
-  // serverless functions (e.g., AWS Lambda) instead of gRPC to searcher nodes.
-  ExecutionMode execution_mode = 19;
 }
 
 enum CountHits {
@@ -506,6 +493,12 @@ message LeafSearchResponse {
   ResourceStats resource_stats = 8;
 }
 
+// Wrapper for multiple LeafSearchResponse messages, used by Lambda to return
+// per-split results.
+message LeafSearchResponses {
+  repeated LeafSearchResponse responses = 1;
+}
+
 message SnippetRequest {
   repeated string snippet_fields = 1;
   string query_ast_resolved = 2;

quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs

@@ -26,8 +26,10 @@ use crate::SearchError;
 #[async_trait]
 pub trait LambdaLeafSearchInvoker: Send + Sync + 'static {
     /// Invoke the remote function with a LeafSearchRequest.
+    ///
+    /// Returns one `LeafSearchResponse` per split in the request.
     async fn invoke_leaf_search(
         &self,
         request: LeafSearchRequest,
-    ) -> Result<LeafSearchResponse, SearchError>;
+    ) -> Result<Vec<LeafSearchResponse>, SearchError>;
 }