schedule leaf searches by (priority, remaining_splits) lexicographically

quickwit/quickwit-proto/build.rs

@@ -222,6 +222,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     prost_config
         .file_descriptor_set_path("src/codegen/quickwit/search_descriptor.bin")
         .protoc_arg("--experimental_allow_proto3_optional")
+        .field_attribute("SearchRequest.priority", "#[serde(default)]")
         // Box the large `LeafSearchResponse` variant so the oneof stays small
         // (the `Error` variant only carries a `String`).
         .boxed("LambdaSingleSplitResult.outcome.response");

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

@@ -273,6 +273,13 @@ message SearchRequest {
   // When true, skip finalization of aggregation results and return
   // the raw IntermediateAggregationResults bytes instead.
   bool skip_aggregation_finalization = 19;
+
+  // Field 20 is intentionally reserved for wire compatibility with a downstream fork.
+  reserved 20;
+
+  // Scheduling priority for leaf search execution. Negative values are allowed,
+  // and lower values have higher priority. Callers that omit it get priority 0.
+  int32 priority = 21;
 }
 
 enum CountHits {

quickwit/quickwit-search/src/leaf.rs

@@ -1891,7 +1891,11 @@ async fn schedule_search_tasks(
     mut splits: Vec<(SplitIdAndFooterOffsets, SearchRequest)>,
     searcher_context: &SearcherContext,
 ) -> ScheduleSearchTaskResult {
-    let task_metadata: Vec<crate::search_permit_provider::SplitSearchTaskMetadata> = splits
+    let priority = splits
+        .first()
+        .map(|(_, search_request)| search_request.priority)
+        .unwrap_or_default();
+    let split_metadatas = splits
         .iter()
         .map(|(split, _)| {
             let memory_allocation = compute_initial_memory_allocation(
@@ -1907,6 +1911,10 @@ async fn schedule_search_tasks(
             }
         })
         .collect();
+    let task_metadata = crate::search_permit_provider::LeafSearchTaskMetadata {
+        priority,
+        splits: split_metadatas,
+    };
 
     let offload_threshold: usize = if searcher_context.lambda_invoker.is_some()
         && let Some(lambda_config) = &searcher_context.searcher_config.lambda

quickwit/quickwit-search/src/leaf_cache.rs

@@ -107,6 +107,8 @@ impl CacheKey {
         // it doesn't matter whether or not we count all hits at the scale of a
         // single split: either we did process it and got everything, or we didn't.
         search_request.count_hits = CountHits::CountAll.into();
+        // Priority only affects scheduling, not search results.
+        search_request.priority = 0;
 
         CacheKey {
             split_id: split_info.split_id,
@@ -436,4 +438,30 @@ mod tests {
         assert!(cache.get(split_3.clone(), query_2).is_none());
         assert!(cache.get(split_3, query_2bis).is_some());
     }
+
+    #[test]
+    fn test_leaf_search_cache_ignores_priority() {
+        let cache = LeafSearchCache::new(&ByteSize::mb(64).into());
+        let split = SplitIdAndFooterOffsets {
+            split_id: "split".to_string(),
+            ..Default::default()
+        };
+        let high_priority_request = SearchRequest {
+            priority: -10,
+            ..Default::default()
+        };
+        let low_priority_request = SearchRequest {
+            priority: 10,
+            ..Default::default()
+        };
+
+        cache.put(
+            split.clone(),
+            high_priority_request,
+            LeafSearchResponse::default(),
+        );
+
+        // the two requests differ only in priority, so it should be a cache-hit
+        assert!(cache.get(split, low_priority_request).is_some());
+    }
 }

quickwit/quickwit-search/src/list_terms.rs

@@ -326,7 +326,7 @@ pub async fn leaf_list_terms(
     splits: &[SplitIdAndFooterOffsets],
 ) -> Result<LeafListTermsResponse, SearchError> {
     info!(split_offsets = ?PrettySample::new(splits, 5));
-    let task_metadata: Vec<crate::search_permit_provider::SplitSearchTaskMetadata> = splits
+    let split_metadatas = splits
         .iter()
         .map(|split| {
             let memory_allocation = compute_initial_memory_allocation(
@@ -342,6 +342,10 @@ pub async fn leaf_list_terms(
             }
         })
         .collect();
+    let task_metadata = crate::search_permit_provider::LeafSearchTaskMetadata {
+        priority: 0,
+        splits: split_metadatas,
+    };
     // We have added offloading leaf search to lambdas, but not for list_terms yet.
     // TODO (Add it)
     // https://github.com/quickwit-oss/quickwit/issues/6150

quickwit/quickwit-search/src/root.rs

@@ -376,6 +376,7 @@ fn simplify_search_request_for_scroll_api(req: &SearchRequest) -> crate::Result<
         count_hits: quickwit_proto::search::CountHits::Underestimate as i32,
         ignore_missing_indexes: req.ignore_missing_indexes,
         skip_aggregation_finalization: false,
+        priority: req.priority,
     })
 }
 
@@ -1320,6 +1321,7 @@ pub async fn root_search(
         count_required = search_request.count_hits().as_str_name(),
         num_docs = num_docs,
         num_splits = num_splits,
+        priority = search_request.priority,
         "root_search"
     );