perf: optimizer target partition config (#12051)

- close https://github.com/openobserve/o2-enterprise/issues/1691

src/config/src/config.rs

@@ -1555,8 +1555,6 @@ pub struct Limit {
     pub usage_reporting_thread_num: usize,
     #[env_config(name = "ZO_QUERY_THREAD_NUM", default = 0)]
     pub query_thread_num: usize,
-    #[env_config(name = "ZO_QUERY_INDEX_THREAD_NUM", default = 0)]
-    pub query_index_thread_num: usize,
     #[env_config(name = "ZO_FILE_DOWNLOAD_THREAD_NUM", default = 0)]
     pub file_download_thread_num: usize,
     #[env_config(name = "ZO_FILE_DOWNLOAD_PRIORITY_QUEUE_THREAD_NUM", default = 0)]
@@ -2602,13 +2600,6 @@ fn check_limit_config(cfg: &mut Config) -> Result<(), anyhow::Error> {
             cfg.limit.query_thread_num = cpu_num * 4;
         }
     }
-    if cfg.limit.query_index_thread_num == 0 {
-        if cfg.common.local_mode {
-            cfg.limit.query_index_thread_num = cpu_num;
-        } else {
-            cfg.limit.query_index_thread_num = cpu_num * 4;
-        }
-    }
 
     if cfg.limit.file_download_thread_num == 0 {
         cfg.limit.file_download_thread_num = std::cmp::max(1, cpu_num / 2);

src/service/promql/search/grpc/storage.rs

@@ -45,7 +45,7 @@ use crate::service::{
         datafusion::exec::register_metrics_table,
         grpc::{
             QueryParams,
-            storage::{cache_files, tantivy_search},
+            storage::{cache_files, calc_target_partitions, tantivy_search},
         },
         index::{Condition, IndexCondition},
         match_source,
@@ -198,13 +198,13 @@ pub(crate) async fn create_context(
             .observe(cached_ratio);
     }
 
-    // set target partitions based on cache type
     let cfg = get_config();
-    let target_partitions = if cache_type == file_data::CacheType::None {
-        cfg.limit.query_thread_num
-    } else {
-        cfg.limit.cpu_num
-    };
+    let target_partitions =
+        calc_target_partitions(cfg.limit.cpu_num, cfg.limit.query_thread_num, cached_ratio);
+
+    log::info!(
+        "[trace_id {trace_id}] promql->search->storage: session target_partitions: {target_partitions}"
+    );
 
     let schema = Arc::new(schema.to_owned().with_metadata(Default::default()));
 

src/service/search/grpc/storage.rs

@@ -239,14 +239,12 @@ pub async fn search(
             .observe(cached_ratio);
     }
 
-    // set target partitions based on cache type
-    let target_partitions = if cache_type == file_data::CacheType::None {
-        cfg.limit.query_thread_num
-    } else {
-        cfg.limit.cpu_num
-    };
+    let target_partitions =
+        calc_target_partitions(cfg.limit.cpu_num, cfg.limit.query_thread_num, cached_ratio);
 
-    log::debug!("search->storage: session target_partitions: {target_partitions}");
+    log::info!(
+        "[trace_id {trace_id}] search->storage: session target_partitions: {target_partitions}"
+    );
 
     let session = config::meta::search::Session {
         id: format!("{trace_id}-storage"),
@@ -426,6 +424,7 @@ pub async fn tantivy_search(
 ) -> Result<(usize, bool, TantivyMultiResult), Error> {
     let start = std::time::Instant::now();
     let cfg = get_config();
+    let trace_id = &query.trace_id;
 
     // Cache the corresponding Index files
     let mut scan_stats = ScanStats::new();
@@ -478,8 +477,7 @@ pub async fn tantivy_search(
         "{}",
         search_inspector_fields(
             format!(
-                "[trace_id {}] search->tantivy: stream {}/{}/{}, load tantivy index files {}, index size: {}, memory cached {}, disk cached {}, cached ratio {}%,{download_msg} took: {} ms",
-                query.trace_id,
+                "[trace_id {trace_id}] search->tantivy: stream {}/{}/{}, load tantivy index files {}, index size: {}, memory cached {}, disk cached {}, cached ratio {}%,{download_msg} took: {} ms",
                 query.org_id,
                 query.stream_type,
                 query.stream_name,
@@ -512,12 +510,12 @@ pub async fn tantivy_search(
             .observe(cached_ratio);
     }
 
-    // set target partitions based on cache type
-    let target_partitions = if cache_type == file_data::CacheType::None {
-        cfg.limit.query_thread_num
-    } else {
-        cfg.limit.query_index_thread_num
-    };
+    let target_partitions =
+        calc_target_partitions(cfg.limit.cpu_num, cfg.limit.query_thread_num, cached_ratio);
+
+    log::info!(
+        "[trace_id {trace_id}] search->tantivy: session target_partitions: {target_partitions}",
+    );
 
     let search_start = std::time::Instant::now();
     let mut is_add_filter_back = file_list_map.len() != index_file_names.len();
@@ -532,8 +530,7 @@ pub async fn tantivy_search(
     let max_group_len = index_parquet_files.len();
 
     log::info!(
-        "[trace_id {}] search->tantivy: target_partitions: {target_partitions}, group_num: {group_num}, max_group_len: {max_group_len}",
-        query.trace_id,
+        "[trace_id {trace_id}] search->tantivy: target_partitions: {target_partitions}, group_num: {group_num}, max_group_len: {max_group_len}",
     );
 
     for file_group in index_parquet_files {
@@ -587,8 +584,7 @@ pub async fn tantivy_search(
             Err(e) => {
                 let took = start.elapsed().as_millis() as usize;
                 log::error!(
-                    "[trace_id {}] search->tantivy: error filtering via index, error: {e:?}, took: {took} ms",
-                    query.trace_id,
+                    "[trace_id {trace_id}] search->tantivy: error filtering via index, error: {e:?}, took: {took} ms",
                 );
                 // search error, need add filter back
                 return Ok((took, true, TantivyMultiResult::RowNums(0)));
@@ -610,8 +606,7 @@ pub async fn tantivy_search(
                         total_row_ids_percent += result.percent();
                         if threshold_num == 0 {
                             log::warn!(
-                                "[trace_id {}] search->tantivy: skip tantivy search, too many row_ids returned from tantivy index, avg percent: {}, took: {took} ms",
-                                query.trace_id,
+                                "[trace_id {trace_id}] search->tantivy: skip tantivy search, too many row_ids returned from tantivy index, avg percent: {}, took: {took} ms",
                                 total_row_ids_percent as f64 / cfg.limit.cpu_num as f64,
                             );
                             file_list.extend(file_list_map.into_values());
@@ -659,8 +654,7 @@ pub async fn tantivy_search(
                 }
                 Err(e) => {
                     log::error!(
-                        "[trace_id {}] search->tantivy: error filtering via index. Keep file to search, error: {e}",
-                        query.trace_id,
+                        "[trace_id {trace_id}] search->tantivy: error filtering via index. Keep file to search, error: {e}"
                     );
                     is_add_filter_back = true;
                     continue;
@@ -680,8 +674,7 @@ pub async fn tantivy_search(
         "{}",
         search_inspector_fields(
             format!(
-                "[trace_id {}] search->tantivy: total hits for index_condition: {index_condition:?} found {tantivy_result}, is_add_filter_back: {is_add_filter_back}, file_num: {}, took: {} ms",
-                query.trace_id,
+                "[trace_id {trace_id}] search->tantivy: total hits for index_condition: {index_condition:?} found {tantivy_result}, is_add_filter_back: {is_add_filter_back}, file_num: {}, took: {} ms",
                 file_list_map.len(),
                 search_start.elapsed().as_millis()
             ),
@@ -999,6 +992,17 @@ async fn search_tantivy_index(
     Ok((key, result, has_skipped_conditions))
 }
 
+/// Linear interpolation: cached_ratio=0 -> query_thread_num, cached_ratio=1 -> cpu_num.
+pub(crate) fn calc_target_partitions(
+    cpu_num: usize,
+    query_thread_num: usize,
+    cached_ratio: f64,
+) -> usize {
+    (cpu_num as i64
+        + ((query_thread_num as i64 - cpu_num as i64) as f64 * (1.0 - cached_ratio)) as i64)
+        as usize
+}
+
 /// if simple distinct without filter, we need to warm up the field
 fn get_simple_distinct_field(idx_optimize_rule: &Option<IndexOptimizeMode>) -> Vec<String> {
     if let Some(IndexOptimizeMode::SimpleDistinct(field, ..)) = idx_optimize_rule {