limit the concurrency of row fetch

Author: SkyFan2002

src/query/service/src/pipelines/builders/builder_row_fetch.rs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
+
+use databend_common_base::runtime::Runtime;
 use databend_common_exception::Result;
 use databend_common_pipeline_core::processors::InputPort;
 use databend_common_pipeline_core::processors::OutputPort;
@@ -21,18 +24,28 @@ use databend_common_pipeline_transforms::processors::create_dummy_item;
 use databend_common_sql::executor::physical_plans::RowFetch;
 use databend_common_sql::executor::PhysicalPlan;
 use databend_common_storages_fuse::operations::row_fetch_processor;
+use databend_common_storages_fuse::TableContext;
+use tokio::sync::Semaphore;
 
 use crate::pipelines::PipelineBuilder;
-
 impl PipelineBuilder {
     pub(crate) fn build_row_fetch(&mut self, row_fetch: &RowFetch) -> Result<()> {
         self.build_pipeline(&row_fetch.input)?;
+        let max_threads = self.ctx.get_settings().get_max_threads()? as usize;
+        let max_io_requests = self.ctx.get_settings().get_max_storage_io_requests()? as usize;
+        let row_fetch_runtime = Arc::new(Runtime::with_worker_threads(
+            max_threads,
+            Some("row-fetch-worker".to_owned()),
+        )?);
+        let row_fetch_semaphore = Arc::new(Semaphore::new(max_io_requests));
         let processor = row_fetch_processor(
             self.ctx.clone(),
             row_fetch.row_id_col_offset,
             &row_fetch.source,
             row_fetch.cols_to_fetch.clone(),
             row_fetch.need_wrap_nullable,
+            row_fetch_semaphore,
+            row_fetch_runtime,
         )?;
         if !matches!(&*row_fetch.input, PhysicalPlan::MutationSplit(_)) {
             self.main_pipeline.add_transform(processor)?;

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

@@ -15,6 +15,8 @@
 use std::collections::HashSet;
 use std::sync::Arc;
 
+use databend_common_base::base::tokio::sync::Semaphore;
+use databend_common_base::runtime::Runtime;
 use databend_common_catalog::plan::split_row_id;
 use databend_common_catalog::plan::DataSourcePlan;
 use databend_common_catalog::plan::Projection;
@@ -49,6 +51,8 @@ pub fn row_fetch_processor(
     source: &DataSourcePlan,
     projection: Projection,
     need_wrap_nullable: bool,
+    semaphore: Arc<Semaphore>,
+    runtime: Arc<Runtime>,
 ) -> Result<RowFetcher> {
     let table = ctx.build_table_from_source_plan(source)?;
     let fuse_table = table
@@ -108,6 +112,8 @@ pub fn row_fetch_processor(
                             block_reader.clone(),
                             read_settings,
                             max_threads,
+                            semaphore.clone(),
+                            runtime.clone(),
                         ),
                         need_wrap_nullable,
                     )
@@ -123,6 +129,8 @@ pub fn row_fetch_processor(
                             block_reader.clone(),
                             read_settings,
                             max_threads,
+                            semaphore.clone(),
+                            runtime.clone(),
                         ),
                         need_wrap_nullable,
                     )

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

@@ -16,7 +16,8 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
 
-use databend_common_base::runtime::execute_futures_in_parallel;
+use databend_common_base::base::tokio::sync::Semaphore;
+use databend_common_base::runtime::Runtime;
 use databend_common_catalog::plan::block_idx_in_segment;
 use databend_common_catalog::plan::split_prefix;
 use databend_common_catalog::plan::split_row_id;
@@ -32,6 +33,7 @@ use databend_storages_common_cache::LoadParams;
 use databend_storages_common_io::ReadSettings;
 use databend_storages_common_table_meta::meta::BlockMeta;
 use databend_storages_common_table_meta::meta::TableSnapshot;
+use futures_util::future;
 use itertools::Itertools;
 
 use super::fuse_rows_fetcher::RowsFetcher;
@@ -56,6 +58,9 @@ pub(super) struct ParquetRowsFetcher<const BLOCKING_IO: bool> {
 
     // To control the parallelism of fetching blocks.
     max_threads: usize,
+
+    semaphore: Arc<Semaphore>,
+    runtime: Arc<Runtime>,
 }
 
 #[async_trait::async_trait]
@@ -125,19 +130,25 @@ impl<const BLOCKING_IO: bool> RowsFetcher for ParquetRowsFetcher<BLOCKING_IO> {
             begin = end;
         }
 
-        let num_task = tasks.len();
-        let blocks = execute_futures_in_parallel(
-            tasks,
-            num_task,
-            num_task * 2,
-            "parqeut rows fetch".to_string(),
-        )
-        .await?
-        .into_iter()
-        .collect::<Result<Vec<_>>>()?
-        .into_iter()
-        .flatten()
-        .collect::<Vec<_>>();
+        let tasks = tasks.into_iter().map(|v| {
+            |permit| async {
+                let r = v.await;
+                drop(permit);
+                r
+            }
+        });
+        let join_handlers = self
+            .runtime
+            .try_spawn_batch_with_owned_semaphore(self.semaphore.clone(), tasks)
+            .await?;
+
+        let joint = future::try_join_all(join_handlers).await?;
+        let blocks = joint
+            .into_iter()
+            .collect::<Result<Vec<_>>>()?
+            .into_iter()
+            .flatten()
+            .collect::<Vec<_>>();
         // Take result rows from blocks.
         let indices = row_set
             .iter()
@@ -171,6 +182,8 @@ impl<const BLOCKING_IO: bool> ParquetRowsFetcher<BLOCKING_IO> {
         reader: Arc<BlockReader>,
         settings: ReadSettings,
         max_threads: usize,
+        semaphore: Arc<Semaphore>,
+        runtime: Arc<Runtime>,
     ) -> Self {
         let schema = table.schema();
         let segment_reader =
@@ -186,6 +199,8 @@ impl<const BLOCKING_IO: bool> ParquetRowsFetcher<BLOCKING_IO> {
             part_map: HashMap::new(),
             segment_blocks_cache: HashMap::new(),
             max_threads,
+            semaphore,
+            runtime,
         }
     }