Merge pull request #10514 from BohuTANG/dev-fast-parquet-splits

feat: improve the parquet get splits to parallel

Author: BohuTANG

src/common/base/src/runtime/mod.rs

@@ -23,6 +23,7 @@ mod thread_pool;
 pub use catch_unwind::catch_unwind;
 pub use catch_unwind::CatchUnwindFuture;
 pub use global_runtime::GlobalIORuntime;
+pub use runtime::execute_futures_in_parallel;
 pub use runtime::match_join_handle;
 pub use runtime::Dropper;
 pub use runtime::Runtime;

src/common/base/src/runtime/runtime.rs

@@ -21,6 +21,7 @@ use std::time::Instant;
 
 use common_exception::ErrorCode;
 use common_exception::Result;
+use futures::future;
 use tokio::runtime::Builder;
 use tokio::runtime::Handle;
 use tokio::sync::oneshot;
@@ -325,3 +326,32 @@ pub async fn match_join_handle<T>(handle: JoinHandle<Result<T>>) -> Result<T> {
         },
     }
 }
+
+/// Run multiple futures parallel
+/// using a semaphore to limit the parallelism number, and a specified thread pool to run the futures.
+/// It waits for all futures to complete and returns their results.
+pub async fn execute_futures_in_parallel<Fut>(
+    futures: impl IntoIterator<Item = Fut>,
+    thread_nums: usize,
+    permit_nums: usize,
+    thread_name: String,
+) -> Result<Vec<Fut::Output>>
+where
+    Fut: Future + Send + 'static,
+    Fut::Output: Send + 'static,
+{
+    // 1. build the runtime.
+    let semaphore = Semaphore::new(permit_nums);
+    let runtime = Arc::new(Runtime::with_worker_threads(
+        thread_nums,
+        Some(thread_name),
+    )?);
+
+    // 2. spawn all the tasks to the runtime with semaphore.
+    let join_handlers = runtime.try_spawn_batch(semaphore, futures).await?;
+
+    // 3. get all the result.
+    future::try_join_all(join_handlers)
+        .await
+        .map_err(|e| ErrorCode::Internal(format!("try join all futures failure, {}", e)))
+}

src/query/pipeline/sources/src/input_formats/impls/input_format_parquet.rs

@@ -34,6 +34,7 @@ use common_arrow::parquet::metadata::ColumnChunkMetaData;
 use common_arrow::parquet::metadata::RowGroupMetaData;
 use common_arrow::parquet::read::read_metadata;
 use common_arrow::read_columns_async;
+use common_base::runtime::execute_futures_in_parallel;
 use common_catalog::plan::StageFileInfo;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -63,18 +64,10 @@ use crate::input_formats::SplitInfo;
 
 pub struct InputFormatParquet;
 
-fn col_offset(meta: &ColumnChunkMetaData) -> i64 {
-    meta.data_page_offset()
-}
-
-#[async_trait::async_trait]
-impl InputFormat for InputFormatParquet {
-    async fn get_splits(
-        &self,
+impl InputFormatParquet {
+    async fn get_split_batch(
         file_infos: Vec<StageFileInfo>,
-        _stage_info: &StageInfo,
-        op: &Operator,
-        _settings: &Arc<Settings>,
+        op: Operator,
     ) -> Result<Vec<Arc<SplitInfo>>> {
         let mut infos = vec![];
         let mut schema = None;
@@ -124,8 +117,56 @@ impl InputFormat for InputFormatParquet {
                 }
             }
         }
+
         Ok(infos)
     }
+}
+
+fn col_offset(meta: &ColumnChunkMetaData) -> i64 {
+    meta.data_page_offset()
+}
+
+#[async_trait::async_trait]
+impl InputFormat for InputFormatParquet {
+    async fn get_splits(
+        &self,
+        file_infos: Vec<StageFileInfo>,
+        _stage_info: &StageInfo,
+        op: &Operator,
+        _settings: &Arc<Settings>,
+    ) -> Result<Vec<Arc<SplitInfo>>> {
+        let batch_size = 1000;
+
+        if file_infos.len() <= batch_size {
+            Self::get_split_batch(file_infos, op.clone()).await
+        } else {
+            let mut chunks = file_infos.chunks(batch_size);
+
+            let tasks = std::iter::from_fn(move || {
+                chunks
+                    .next()
+                    .map(|location| Self::get_split_batch(location.to_vec(), op.clone()))
+            });
+
+            // TODO: Get from ctx.
+            let thread_nums = 16;
+            let permit_nums = 64;
+            let result = execute_futures_in_parallel(
+                tasks,
+                thread_nums,
+                permit_nums,
+                "get-parquet-splits-worker".to_owned(),
+            )
+            .await?
+            .into_iter()
+            .collect::<Result<Vec<Vec<_>>>>()?
+            .into_iter()
+            .flatten()
+            .collect();
+
+            Ok(result)
+        }
+    }
 
     async fn infer_schema(&self, path: &str, op: &Operator) -> Result<TableSchemaRef> {
         let mut reader = op.reader(path).await?;

src/query/storages/fuse/src/io/files.rs

@@ -14,12 +14,11 @@
 
 use std::sync::Arc;
 
+use common_base::runtime::execute_futures_in_parallel;
 use common_catalog::table_context::TableContext;
 use common_exception::Result;
 use opendal::Operator;
 
-use crate::io::execute_futures_in_parallel;
-
 // File related operations.
 pub struct Files {
     ctx: Arc<dyn TableContext>,
@@ -41,7 +40,7 @@ impl Files {
         let batch_size = 1000;
         let locations = Vec::from_iter(file_locations.into_iter().map(|v| v.as_ref().to_string()));
 
-        if locations.len() < batch_size {
+        if locations.len() <= batch_size {
             Self::delete_files(self.operator.clone(), locations).await?;
         } else {
             let mut chunks = locations.chunks(batch_size);
@@ -52,9 +51,12 @@ impl Files {
                     .map(|location| Self::delete_files(self.operator.clone(), location.to_vec()))
             });
 
+            let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
+            let permit_nums = self.ctx.get_settings().get_max_storage_io_requests()? as usize;
             execute_futures_in_parallel(
-                self.ctx.clone(),
                 tasks,
+                threads_nums,
+                permit_nums,
                 "batch-remove-files-worker".to_owned(),
             )
             .await?;

src/query/storages/fuse/src/io/mod.rs

@@ -14,15 +14,13 @@
 
 mod files;
 mod locations;
-mod parallel;
 mod read;
 mod segments;
 mod snapshots;
 mod write;
 
 pub use files::Files;
 pub use locations::TableMetaLocationGenerator;
-pub use parallel::execute_futures_in_parallel;
 pub use read::BlockReader;
 pub use read::BloomBlockFilterReader;
 pub use read::MergeIOReadResult;

src/query/storages/fuse/src/io/parallel.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use common_base::runtime::execute_futures_in_parallel;
 use common_catalog::table_context::TableContext;
 use common_exception::Result;
 use common_expression::TableSchemaRef;
@@ -23,7 +24,6 @@ use storages_common_table_meta::meta::Location;
 use storages_common_table_meta::meta::SegmentInfo;
 use tracing::Instrument;
 
-use crate::io::execute_futures_in_parallel;
 use crate::io::MetaReaders;
 
 // Read segment related operations.
@@ -88,9 +88,12 @@ impl SegmentsIO {
             }
         });
 
+        let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
+        let permit_nums = self.ctx.get_settings().get_max_storage_io_requests()? as usize;
         execute_futures_in_parallel(
-            self.ctx.clone(),
             tasks,
+            threads_nums,
+            permit_nums,
             "fuse-req-segments-worker".to_owned(),
         )
         .await
@@ -143,9 +146,12 @@ impl SegmentsIO {
             })
         });
 
+        let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
+        let permit_nums = self.ctx.get_settings().get_max_storage_io_requests()? as usize;
         execute_futures_in_parallel(
-            self.ctx.clone(),
             tasks,
+            threads_nums,
+            permit_nums,
             "fuse-req-segments-worker".to_owned(),
         )
         .await

src/query/storages/fuse/src/io/segments.rs

@@ -20,6 +20,7 @@ use std::time::Instant;
 
 use chrono::DateTime;
 use chrono::Utc;
+use common_base::runtime::execute_futures_in_parallel;
 use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -37,7 +38,6 @@ use tracing::info;
 use tracing::warn;
 use tracing::Instrument;
 
-use crate::io::execute_futures_in_parallel;
 use crate::io::MetaReaders;
 use crate::io::SnapshotHistoryReader;
 use crate::io::TableMetaLocationGenerator;
@@ -168,9 +168,12 @@ impl SnapshotsIO {
             })
         });
 
+        let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
+        let permit_nums = self.ctx.get_settings().get_max_storage_io_requests()? as usize;
         execute_futures_in_parallel(
-            self.ctx.clone(),
             tasks,
+            threads_nums,
+            permit_nums,
             "fuse-req-snapshots-worker".to_owned(),
         )
         .await

src/query/storages/fuse/src/io/snapshots.rs

@@ -15,6 +15,7 @@
 
 use std::sync::Arc;
 
+use common_base::runtime::execute_futures_in_parallel;
 use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -30,7 +31,6 @@ use storages_common_table_meta::meta::Location;
 use storages_common_table_meta::meta::SegmentInfo;
 use tracing::debug;
 
-use crate::io::execute_futures_in_parallel;
 use crate::io::SegmentsIO;
 use crate::io::TableMetaLocationGenerator;
 use crate::operations::merge_into::mutation_meta::mutation_log::CommitMeta;
@@ -148,10 +148,10 @@ impl TableMutationAggregator {
 
     // TODO use batch_meta_writer
     async fn write_segments(&self, segments: Vec<SerializedSegment>) -> Result<()> {
-        let mut handles = Vec::with_capacity(segments.len());
+        let mut tasks = Vec::with_capacity(segments.len());
         for segment in segments {
             let op = self.dal.clone();
-            handles.push(async move {
+            tasks.push(async move {
                 op.write(&segment.path, segment.raw_data).await?;
                 if let Some(segment_cache) = CacheManager::instance().get_table_segment_cache() {
                     segment_cache.put(segment.path.clone(), segment.segment.clone());
@@ -160,9 +160,12 @@ impl TableMutationAggregator {
             });
         }
 
+        let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
+        let permit_nums = self.ctx.get_settings().get_max_storage_io_requests()? as usize;
         execute_futures_in_parallel(
-            self.ctx.clone(),
-            handles,
+            tasks,
+            threads_nums,
+            permit_nums,
             "mutation-write-segments-worker".to_owned(),
         )
         .await?