Merge pull request #10695 from youngsofun/fix

refactor(copy): calc suitable parallelism for reading parquet splits.

Author: mergify[bot]

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

@@ -34,7 +34,6 @@ use common_arrow::parquet::metadata::ColumnChunkMetaData;
 use common_arrow::parquet::metadata::FileMetaData;
 use common_arrow::parquet::metadata::RowGroupMetaData;
 use common_arrow::parquet::read::read_metadata;
-use common_arrow::read_columns_async;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::DataBlock;
@@ -48,7 +47,9 @@ use common_settings::Settings;
 use common_storage::read_parquet_metas_in_parallel;
 use common_storage::StageFileInfo;
 use futures::AsyncRead;
+use futures::AsyncReadExt;
 use futures::AsyncSeek;
+use futures::AsyncSeekExt;
 use opendal::Operator;
 use serde::Deserializer;
 use serde::Serializer;
@@ -174,16 +175,9 @@ impl InputFormatPipe for ParquetFormatPipe {
     ) -> Result<Self::RowBatch> {
         let meta = Self::get_split_meta(&split_info).expect("must success");
         let op = ctx.source.get_operator()?;
-        let mut reader = op.reader(&split_info.file.path).await?;
         let input_fields = Arc::new(get_used_fields(&meta.file.fields, &ctx.schema)?);
 
-        RowGroupInMemory::read_async(
-            split_info.to_string(),
-            &mut reader,
-            meta.meta.clone(),
-            input_fields,
-        )
-        .await
+        RowGroupInMemory::read_async(split_info.clone(), op, meta.meta.clone(), input_fields).await
     }
 }
 
@@ -269,25 +263,48 @@ impl RowGroupInMemory {
         })
     }
 
-    async fn read_async<R: AsyncRead + AsyncSeek + Send + Unpin>(
-        split_info: String,
-        reader: &mut R,
+    async fn read_field_async(
+        op: Operator,
+        path: String,
+        col_metas: Vec<ColumnChunkMetaData>,
+        index: usize,
+    ) -> Result<(usize, Vec<Vec<u8>>)> {
+        let mut cols = Vec::with_capacity(col_metas.len());
+        let mut reader = op.reader(&path).await?;
+        for meta in &col_metas {
+            cols.push(read_single_column_async(&mut reader, meta).await?)
+        }
+        Ok((index, cols))
+    }
+
+    async fn read_async(
+        split_info: Arc<SplitInfo>,
+        operator: Operator,
         meta: RowGroupMetaData,
         fields: Arc<Vec<Field>>,
     ) -> Result<Self> {
         let field_names = fields.iter().map(|x| x.name.as_str()).collect::<Vec<_>>();
         let field_meta_indexes = split_column_metas_by_field(meta.columns(), &field_names);
-        let mut filed_arrays = vec![];
-        for field_name in field_names {
-            let meta_data = read_columns_async(reader, meta.columns(), field_name).await?;
-            let data = meta_data.into_iter().map(|t| t.1).collect::<Vec<_>>();
-            filed_arrays.push(data)
+        let mut join_handlers = Vec::with_capacity(field_names.len());
+        for (i, field_name) in field_names.iter().enumerate() {
+            let col_metas = get_field_columns(meta.columns(), field_name)
+                .into_iter()
+                .cloned()
+                .collect::<Vec<_>>();
+            let op = operator.clone();
+            let path = split_info.file.path.clone();
+            join_handlers.push(async move { Self::read_field_async(op, path, col_metas, i).await });
         }
+
+        let mut field_arrays = futures::future::try_join_all(join_handlers).await?;
+        field_arrays.sort();
+        let field_arrays = field_arrays.into_iter().map(|t| t.1).collect::<Vec<_>>();
+
         Ok(Self {
-            split_info,
+            split_info: split_info.to_string(),
             meta,
             field_meta_indexes,
-            field_arrays: filed_arrays,
+            field_arrays,
             fields_to_read: fields,
         })
     }
@@ -489,3 +506,27 @@ pub fn split_column_metas_by_field(
     });
     r
 }
+
+fn get_field_columns<'a>(
+    columns: &'a [ColumnChunkMetaData],
+    field_name: &str,
+) -> Vec<&'a ColumnChunkMetaData> {
+    columns
+        .iter()
+        .filter(|x| x.descriptor().path_in_schema[0] == field_name)
+        .collect()
+}
+
+async fn read_single_column_async<R>(
+    reader: &mut R,
+    meta: &ColumnChunkMetaData,
+) -> Result<Vec<u8>>
+where
+    R: AsyncRead + AsyncSeek + Send + Unpin,
+{
+    let (start, len) = meta.byte_range();
+    reader.seek(std::io::SeekFrom::Start(start)).await?;
+    let mut chunk = vec![0; len as usize];
+    reader.read_exact(&mut chunk).await?;
+    Ok(chunk)
+}

src/query/pipeline/sources/src/input_formats/input_pipeline.rs

@@ -18,7 +18,6 @@ use std::sync::Arc;
 use common_base::base::tokio;
 use common_base::base::tokio::sync::mpsc::Receiver;
 use common_base::base::tokio::sync::mpsc::Sender;
-use common_base::runtime::CatchUnwindFuture;
 use common_base::runtime::GlobalIORuntime;
 use common_base::runtime::TrySpawn;
 use common_compress::CompressAlgorithm;
@@ -27,9 +26,8 @@ use common_exception::Result;
 use common_expression::DataBlock;
 use common_pipeline_core::Pipeline;
 use futures::AsyncRead;
-use futures_util::stream::FuturesUnordered;
 use futures_util::AsyncReadExt;
-use futures_util::StreamExt;
+use parking_lot::Mutex;
 
 use crate::input_formats::Aligner;
 use crate::input_formats::BeyondEndReader;
@@ -191,39 +189,59 @@ pub trait InputFormatPipe: Sized + Send + 'static {
     }
 
     fn execute_copy_aligned(ctx: Arc<InputContext>, pipeline: &mut Pipeline) -> Result<()> {
-        let (data_tx, data_rx) = async_channel::bounded(ctx.num_prefetch_splits()?);
+        let (data_tx, data_rx) = async_channel::bounded(1);
         Self::build_pipeline_aligned(&ctx, data_rx, pipeline)?;
-
-        let ctx_clone = ctx.clone();
-        let p = 3;
-        GlobalIORuntime::instance().spawn(async move {
-            for splits in ctx_clone.splits.chunks(p) {
-                let ctx_clone2 = ctx_clone.clone();
-                let row_batch_tx = data_tx.clone();
-                let splits = splits.to_owned().clone();
-                tokio::spawn(async move {
-                    let mut futs = FuturesUnordered::new();
-                    for s in splits.into_iter() {
-                        let fut =
-                            CatchUnwindFuture::create(Self::read_split(ctx_clone2.clone(), s));
-                        futs.push(fut);
-                    }
-                    while let Some(row_batch) = futs.next().await {
-                        match row_batch {
-                            Ok(row_batch) => {
-                                if row_batch_tx.send(row_batch).await.is_err() {
-                                    break;
-                                }
-                            }
-                            Err(cause) => {
-                                row_batch_tx.send(Err(cause)).await.ok();
+        let max_storage_io_requests = ctx.settings.get_max_storage_io_requests()?;
+        let per_split_io = ctx.schema.fields().len();
+        let max_splits = max_storage_io_requests as usize / per_split_io;
+        let mut max_splits = std::cmp::max(max_splits, 1);
+        let mut sizes = ctx.splits.iter().map(|s| s.size).collect::<Vec<usize>>();
+        sizes.sort_by(|a, b| b.cmp(a));
+        let max_memory = ctx.settings.get_max_memory_usage()? as usize;
+        let mut mem = 0;
+        for (i, s) in sizes.iter().enumerate() {
+            let m = mem + s;
+            if m > max_memory {
+                max_splits = std::cmp::min(max_splits, std::cmp::max(i, 1));
+                break;
+            } else {
+                mem = m
+            }
+        }
+        tracing::info!(
+            "copy read {max_splits} splits in parallel, according to max_memory={max_memory}, num_fields={per_split_io}, max_storage_io_requests={max_storage_io_requests}, max_split_size={}",
+            sizes[0]
+        );
+        let splits = ctx.splits.to_vec();
+        let splits = Arc::new(Mutex::new(splits));
+        for _ in 0..max_splits {
+            let splits = splits.clone();
+            let ctx_clone = ctx.clone();
+            let data_tx = data_tx.clone();
+            GlobalIORuntime::instance().spawn(async move {
+                loop {
+                    let split = {
+                        let mut splits = splits.lock();
+                        if let Some(split) = splits.pop() {
+                            split
+                        } else {
+                            break;
+                        }
+                    };
+                    match Self::read_split(ctx_clone.clone(), split.clone()).await {
+                        Ok(row_batch) => {
+                            if data_tx.send(Ok(row_batch)).await.is_err() {
                                 break;
                             }
                         }
+                        Err(cause) => {
+                            data_tx.send(Err(cause)).await.ok();
+                            break;
+                        }
                     }
-                });
-            }
-        });
+                }
+            });
+        }
         Ok(())
     }