Changes in the lambda offloading

Lambda now details (but does not retry) which split failed or was successful.
Leaf cache lambda individual split result, (keep track of the rewritten)
request to do so).
Lambda handler returned split-id named results in any order.

Author: fulmicoton-dd

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

@@ -21,8 +21,8 @@ use aws_sdk_lambda::primitives::Blob;
 use aws_sdk_lambda::types::InvocationType;
 use base64::prelude::*;
 use prost::Message;
-use quickwit_lambda_server::{LeafSearchRequestPayload, LeafSearchResponsePayload};
-use quickwit_proto::search::{LeafSearchRequest, LeafSearchResponse, LeafSearchResponses};
+use quickwit_lambda_server::{LambdaSearchRequestPayload, LambdaSearchResponsePayload};
+use quickwit_proto::search::{LambdaSearchResponses, LambdaSingleSplitResult, LeafSearchRequest};
 use quickwit_search::{LambdaLeafSearchInvoker, SearchError};
 use tracing::{debug, info, instrument};
 
@@ -85,7 +85,7 @@ impl LambdaLeafSearchInvoker for AwsLambdaInvoker {
     async fn invoke_leaf_search(
         &self,
         request: LeafSearchRequest,
-    ) -> Result<Vec<LeafSearchResponse>, SearchError> {
+    ) -> Result<Vec<LambdaSingleSplitResult>, SearchError> {
         let start = std::time::Instant::now();
 
         let result = self.invoke_leaf_search_inner(request).await;
@@ -109,10 +109,10 @@ impl AwsLambdaInvoker {
     async fn invoke_leaf_search_inner(
         &self,
         request: LeafSearchRequest,
-    ) -> Result<Vec<LeafSearchResponse>, SearchError> {
+    ) -> Result<Vec<LambdaSingleSplitResult>, SearchError> {
         // Serialize request to protobuf bytes, then base64 encode
         let request_bytes = request.encode_to_vec();
-        let payload = LeafSearchRequestPayload {
+        let payload = LambdaSearchRequestPayload {
             payload: BASE64_STANDARD.encode(&request_bytes),
         };
 
@@ -164,22 +164,22 @@ impl AwsLambdaInvoker {
             .leaf_search_response_payload_size_bytes
             .observe(response_payload.as_ref().len() as f64);
 
-        let lambda_response: LeafSearchResponsePayload =
+        let lambda_response: LambdaSearchResponsePayload =
             serde_json::from_slice(response_payload.as_ref())
                 .map_err(|e| SearchError::Internal(format!("json deserialization error: {}", e)))?;
 
         let response_bytes = BASE64_STANDARD
             .decode(&lambda_response.payload)
             .map_err(|e| SearchError::Internal(format!("base64 decode error: {}", e)))?;
 
-        let leaf_responses = LeafSearchResponses::decode(&response_bytes[..])
+        let leaf_responses = LambdaSearchResponses::decode(&response_bytes[..])
             .map_err(|e| SearchError::Internal(format!("protobuf decode error: {}", e)))?;
 
         debug!(
-            num_responses = leaf_responses.responses.len(),
+            num_results = leaf_responses.split_results.len(),
             "lambda invocation completed"
         );
 
-        Ok(leaf_responses.responses)
+        Ok(leaf_responses.split_results)
     }
 }

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

@@ -34,4 +34,4 @@ mod metrics;
 pub use deploy::try_get_or_deploy_invoker;
 pub use metrics::LAMBDA_METRICS;
 // Re-export payload types from server crate for convenience
-pub use quickwit_lambda_server::{LeafSearchRequestPayload, LeafSearchResponsePayload};
+pub use quickwit_lambda_server::{LambdaSearchRequestPayload, LambdaSearchResponsePayload};

quickwit/quickwit-lambda-server/src/bin/leaf_search.rs

@@ -17,7 +17,9 @@
 use std::sync::Arc;
 
 use lambda_runtime::{Error, LambdaEvent, service_fn};
-use quickwit_lambda_server::{LambdaSearcherContext, LeafSearchRequestPayload, handle_leaf_search};
+use quickwit_lambda_server::{
+    LambdaSearchRequestPayload, LambdaSearcherContext, handle_leaf_search,
+};
 use tracing::info;
 use tracing_subscriber::EnvFilter;
 
@@ -36,7 +38,7 @@ async fn main() -> Result<(), Error> {
 
     // Run the Lambda handler
     lambda_runtime::run(service_fn(
-        |event: LambdaEvent<LeafSearchRequestPayload>| {
+        |event: LambdaEvent<LambdaSearchRequestPayload>| {
             let ctx = Arc::clone(&context);
             async move {
                 let (payload, _event_ctx) = event.into_parts();

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

@@ -14,8 +14,9 @@
 
 use base64::prelude::*;
 use prost::Message;
+use quickwit_proto::search::lambda_single_split_result::Outcome;
 use quickwit_proto::search::{
-    LeafSearchRequest, LeafSearchResponse, LeafSearchResponses, SplitIdAndFooterOffsets,
+    LambdaSearchResponses, LambdaSingleSplitResult, LeafSearchRequest, SplitIdAndFooterOffsets,
 };
 use quickwit_search::leaf::multi_index_leaf_search;
 use serde::{Deserialize, Serialize};
@@ -26,27 +27,28 @@ use crate::error::{LambdaError, LambdaResult};
 
 /// Payload for leaf search Lambda invocation.
 #[derive(Debug, Serialize, Deserialize)]
-pub struct LeafSearchRequestPayload {
+pub struct LambdaSearchRequestPayload {
     /// Base64-encoded serialized LeafSearchRequest protobuf.
     pub payload: String,
 }
 
 /// Response from leaf search Lambda invocation.
 #[derive(Debug, Serialize, Deserialize)]
-pub struct LeafSearchResponsePayload {
-    /// Base64-encoded serialized `LeafSearchResponses` protobuf (one per split).
+pub struct LambdaSearchResponsePayload {
+    /// Base64-encoded serialized `LambdaSearchResponses` 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.
+/// Returns one `LambdaSingleSplitResult` per split, each tagged with its
+/// split_id. Individual split failures are reported per-split rather than
+/// failing the entire invocation, so the caller can retry only failed splits.
 #[instrument(skip(ctx), fields(request_id))]
 pub async fn handle_leaf_search(
-    event: LeafSearchRequestPayload,
+    event: LambdaSearchRequestPayload,
     ctx: &LambdaSearcherContext,
-) -> LambdaResult<LeafSearchResponsePayload> {
+) -> LambdaResult<LambdaSearchResponsePayload> {
     // Decode base64 payload
     let request_bytes = BASE64_STANDARD
         .decode(&event.payload)
@@ -70,11 +72,12 @@ pub async fn handle_leaf_search(
     let num_splits = all_splits.len();
     info!(num_splits, "processing leaf search request (per-split)");
 
-    // Process each split in parallel. The SearchPermitProvider inside
-    // SearcherContext gates concurrency based on memory budget.
-    let mut split_search_futures: Vec<tokio::task::JoinHandle<_>> =
-        Vec::with_capacity(all_splits.len());
+    // Process each split in parallel using a JoinSet. The SearchPermitProvider
+    // inside SearcherContext gates concurrency based on memory budget.
+    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(),
@@ -89,43 +92,57 @@ pub async fn handle_leaf_search(
 
         let searcher_context = ctx.searcher_context.clone();
         let storage_resolver = ctx.storage_resolver.clone();
-        split_search_futures.push(tokio::task::spawn(multi_index_leaf_search(
-            searcher_context,
-            single_split_request,
-            storage_resolver,
-        )));
+        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)
+        });
     }
 
-    // Collect results, preserving split order.
-    let mut responses: Vec<LeafSearchResponse> = Vec::with_capacity(num_splits);
-    for split_search_fut in split_search_futures {
-        match split_search_fut.await {
-            Ok(Ok(response)) => responses.push(response),
-            Ok(Err(e)) => {
-                return Err(LambdaError::Internal(format!("leaf search failed: {e}")));
+    // 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;
+    let mut num_failures: usize = 0;
+    while let Some(join_result) = split_search_joinset.join_next().await {
+        match join_result {
+            Ok((split_id, Ok(response))) => {
+                num_successes += 1;
+                split_results.push(LambdaSingleSplitResult {
+                    split_id,
+                    outcome: Some(Outcome::Response(response)),
+                });
+            }
+            Ok((split_id, Err(error_msg))) => {
+                num_failures += 1;
+                warn!(split_id = %split_id, error = %error_msg, "split search failed");
+                split_results.push(LambdaSingleSplitResult {
+                    split_id,
+                    outcome: Some(Outcome::Error(error_msg)),
+                });
             }
             Err(join_error) if join_error.is_cancelled() => {
                 warn!("search task was cancelled");
                 return Err(LambdaError::Cancelled);
             }
             Err(join_error) => {
+                // Panics lose the captured split_id, so we fail the entire invocation.
                 error!(error = %join_error, "search task panicked");
                 return Err(LambdaError::Internal(format!(
                     "search task panicked: {join_error}"
                 )));
             }
         }
     }
-
     info!(
-        num_responses = responses.len(),
-        "leaf search completed (per-split)"
+        num_successes,
+        num_failures, "leaf search completed (per-split)"
     );
 
-    // Serialize as LeafSearchResponses wrapper
-    let wrapper = LeafSearchResponses { responses };
+    let wrapper = LambdaSearchResponses { split_results };
     let response_bytes = wrapper.encode_to_vec();
     let payload = BASE64_STANDARD.encode(&response_bytes);
 
-    Ok(LeafSearchResponsePayload { payload })
+    Ok(LambdaSearchResponsePayload { payload })
 }

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

@@ -23,4 +23,4 @@ mod handler;
 
 pub use context::LambdaSearcherContext;
 pub use error::{LambdaError, LambdaResult};
-pub use handler::{LeafSearchRequestPayload, LeafSearchResponsePayload, handle_leaf_search};
+pub use handler::{LambdaSearchRequestPayload, LambdaSearchResponsePayload, handle_leaf_search};

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

@@ -333,8 +333,7 @@ message SplitSearchError {
   bool retryable_error = 3;
 }
 
-/// A LeafSearchRequest can span multiple indices.
-///
+// A LeafSearchRequest can span multiple indices.
 message LeafSearchRequest {
   // Search request. This is a perfect copy of the original search request
   // that was sent to root apart from the start_offset, max_hits params and index_id_patterns.
@@ -361,7 +360,7 @@ message ResourceStats {
     uint64 cpu_microsecs = 5;
 }
 
-/// LeafRequestRef references data in LeafSearchRequest to deduplicate data.
+// LeafRequestRef references data in LeafSearchRequest to deduplicate data.
 message LeafRequestRef {
   // The ordinal of the doc_mapper in `LeafSearchRequest.doc_mappers`
   uint32 doc_mapper_ord = 1;
@@ -497,10 +496,23 @@ 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;
+// The result of searching a single split in a Lambda invocation.
+// Each result is tagged with its split_id so that ordering is irrelevant.
+message LambdaSingleSplitResult {
+  // The split that was searched.
+  string split_id = 1;
+  oneof outcome {
+    // On success, the leaf search response for this split.
+    LeafSearchResponse response = 2;
+    // On failure, the error message.
+    string error = 3;
+  }
+}
+
+// Wrapper for per-split results from a Lambda invocation.
+message LambdaSearchResponses {
+  reserved 1; // was: repeated LeafSearchResponse responses
+  repeated LambdaSingleSplitResult split_results = 2;
 }
 
 message SnippetRequest {

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

@@ -15,7 +15,7 @@
 //! Trait for invoking remote serverless functions for leaf search.
 
 use async_trait::async_trait;
-use quickwit_proto::search::{LeafSearchRequest, LeafSearchResponse};
+use quickwit_proto::search::{LambdaSingleSplitResult, LeafSearchRequest};
 
 use crate::SearchError;
 
@@ -27,9 +27,12 @@ use crate::SearchError;
 pub trait LambdaLeafSearchInvoker: Send + Sync + 'static {
     /// Invoke the remote function with a LeafSearchRequest.
     ///
-    /// Returns one `LeafSearchResponse` per split in the request.
+    /// Returns one `LambdaSingleSplitResult` per split in the request.
+    /// Each result is tagged with its split_id so ordering is irrelevant.
+    /// Individual split failures are reported per-split; the outer `Result`
+    /// only represents transport-level errors.
     async fn invoke_leaf_search(
         &self,
         request: LeafSearchRequest,
-    ) -> Result<Vec<LeafSearchResponse>, SearchError>;
+    ) -> Result<Vec<LambdaSingleSplitResult>, SearchError>;
 }