Merge pull request #10226 from sundy-li/limit-threads

Author: sundy-li

src/query/storages/common/table-meta/src/meta/v2/segment.rs

@@ -13,6 +13,7 @@
 //  limitations under the License.
 
 use std::collections::HashMap;
+use std::ops::Range;
 use std::sync::Arc;
 
 use common_arrow::native::ColumnMeta as NativeColumnMeta;
@@ -114,6 +115,38 @@ impl ColumnMeta {
             ColumnMeta::Native(v) => (v.offset, v.pages.iter().map(|page| page.length).sum()),
         }
     }
+
+    pub fn read_rows(&self, range: &Option<Range<usize>>) -> u64 {
+        match self {
+            ColumnMeta::Parquet(v) => v.num_values,
+            ColumnMeta::Native(v) => match range {
+                Some(range) => v
+                    .pages
+                    .iter()
+                    .skip(range.start)
+                    .take(range.end - range.start)
+                    .map(|page| page.num_values)
+                    .sum(),
+                None => v.pages.iter().map(|page| page.num_values).sum(),
+            },
+        }
+    }
+
+    pub fn read_bytes(&self, range: &Option<Range<usize>>) -> u64 {
+        match self {
+            ColumnMeta::Parquet(v) => v.len,
+            ColumnMeta::Native(v) => match range {
+                Some(range) => v
+                    .pages
+                    .iter()
+                    .skip(range.start)
+                    .take(range.end - range.start)
+                    .map(|page| page.length)
+                    .sum(),
+                None => v.pages.iter().map(|page| page.length).sum(),
+            },
+        }
+    }
 }
 
 impl SegmentInfo {

src/query/storages/fuse/src/operations/read/fuse_source.rs

@@ -26,6 +26,7 @@ use common_pipeline_core::Pipeline;
 use common_pipeline_core::SourcePipeBuilder;
 use tracing::info;
 
+use crate::fuse_part::FusePartInfo;
 use crate::io::BlockReader;
 use crate::operations::read::native_data_source_deserializer::NativeDeserializeDataTransform;
 use crate::operations::read::native_data_source_reader::ReadNativeDataSource;
@@ -41,7 +42,8 @@ pub fn build_fuse_native_source_pipeline(
     topk: Option<TopK>,
     mut max_io_requests: usize,
 ) -> Result<()> {
-    (max_threads, max_io_requests) = adjust_threads_and_request(max_threads, max_io_requests, plan);
+    (max_threads, max_io_requests) =
+        adjust_threads_and_request(true, max_threads, max_io_requests, plan);
 
     if topk.is_some() {
         max_threads = max_threads.min(16);
@@ -122,7 +124,8 @@ pub fn build_fuse_parquet_source_pipeline(
     mut max_threads: usize,
     mut max_io_requests: usize,
 ) -> Result<()> {
-    (max_threads, max_io_requests) = adjust_threads_and_request(max_threads, max_io_requests, plan);
+    (max_threads, max_io_requests) =
+        adjust_threads_and_request(false, max_threads, max_io_requests, plan);
 
     let mut source_builder = SourcePipeBuilder::create();
 
@@ -220,12 +223,38 @@ pub fn dispatch_partitions(
 }
 
 pub fn adjust_threads_and_request(
+    is_native: bool,
     mut max_threads: usize,
     mut max_io_requests: usize,
     plan: &DataSourcePlan,
 ) -> (usize, usize) {
     if !plan.parts.is_lazy {
-        let block_nums = plan.parts.partitions.len().max(1);
+        let mut block_nums = plan.parts.partitions.len();
+
+        // If the read bytes of a partition is small enough, less than 16k rows
+        // we will not use an extra heavy thread to process it.
+        // now only works for native reader
+        static MIN_ROWS_READ_PER_THREAD: u64 = 16 * 1024;
+        if is_native {
+            plan.parts.partitions.iter().for_each(|part| {
+                if let Some(part) = part.as_any().downcast_ref::<FusePartInfo>() {
+                    let to_read_rows = part
+                        .columns_meta
+                        .values()
+                        .map(|meta| meta.read_rows(&part.range))
+                        .find(|rows| *rows > 0)
+                        .unwrap_or(part.nums_rows as u64);
+
+                    if to_read_rows < MIN_ROWS_READ_PER_THREAD {
+                        block_nums -= 1;
+                    }
+                }
+            });
+        }
+
+        // At least max(1/8 of the original parts, 1), in case of too many small partitions but io threads is just one.
+        block_nums = std::cmp::max(block_nums, plan.parts.partitions.len() / 8);
+        block_nums = std::cmp::max(block_nums, 1);
 
         max_threads = std::cmp::min(max_threads, block_nums);
         max_io_requests = std::cmp::min(max_io_requests, block_nums);