Merge pull request #10536 from youngsofun/trans

refactor(stage): unify list files and read parquet metas

Author: mergify[bot]

Cargo.lock

@@ -20,6 +20,7 @@ common-meta-app = { path = "../../meta/app" }
 anyhow = { workspace = true }
 async-trait = "0.1"
 bytes = "1"
+chrono = { workspace = true }
 futures = "0.3"
 opendal = { workspace = true }
 regex = "1.6.0"

src/common/storage/Cargo.toml

@@ -50,8 +50,12 @@ pub use column_node::ColumnNode;
 pub use column_node::ColumnNodes;
 
 mod parquet;
+pub use parquet::read_parquet_metas_in_parallel;
 pub use parquet::read_parquet_schema_async;
 
 mod stage;
 pub use stage::init_stage_operator;
+pub use stage::FileWithMeta;
+pub use stage::StageFileInfo;
+pub use stage::StageFileStatus;
 pub use stage::StageFilesInfo;

src/common/storage/src/lib.rs

@@ -14,6 +14,8 @@
 
 use common_arrow::arrow::datatypes::Schema as ArrowSchema;
 use common_arrow::arrow::io::parquet::read as pread;
+use common_arrow::parquet::metadata::FileMetaData;
+use common_base::runtime::execute_futures_in_parallel;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use opendal::Operator;
@@ -27,3 +29,51 @@ pub async fn read_parquet_schema_async(operator: &Operator, path: &str) -> Resul
     let arrow_schema = pread::infer_schema(&meta)?;
     Ok(arrow_schema)
 }
+
+async fn read_parquet_metas_batch(
+    file_infos: Vec<(String, u64)>,
+    op: Operator,
+) -> Result<Vec<FileMetaData>> {
+    // todo(youngsofun): we should use size in StageFileInfo, but parquet2 do not have the interface for now
+    let mut metas = vec![];
+    for (path, _size) in file_infos {
+        let mut reader = op.reader(&path).await?;
+        metas.push(pread::read_metadata_async(&mut reader).await?)
+    }
+    Ok(metas)
+}
+
+pub async fn read_parquet_metas_in_parallel(
+    op: Operator,
+    file_infos: Vec<(String, u64)>,
+    thread_nums: usize,
+    permit_nums: usize,
+) -> Result<Vec<FileMetaData>> {
+    let batch_size = 1000;
+    if file_infos.len() <= batch_size {
+        read_parquet_metas_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| read_parquet_metas_batch(location.to_vec(), op.clone()))
+        });
+
+        let result = execute_futures_in_parallel(
+            tasks,
+            thread_nums,
+            permit_nums,
+            "read-parquet-metas-worker".to_owned(),
+        )
+        .await?
+        .into_iter()
+        .collect::<Result<Vec<Vec<_>>>>()?
+        .into_iter()
+        .flatten()
+        .collect();
+
+        Ok(result)
+    }
+}

src/common/storage/src/parquet.rs

@@ -14,12 +14,18 @@
 
 use std::path::Path;
 
+use chrono::DateTime;
+use chrono::TimeZone;
+use chrono::Utc;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_app::principal::StageInfo;
 use common_meta_app::principal::StageType;
+use common_meta_app::principal::UserIdentity;
 use futures::TryStreamExt;
+use opendal::Entry;
 use opendal::EntryMode;
+use opendal::Metadata;
 use opendal::Metakey;
 use opendal::Operator;
 use regex::Regex;
@@ -29,16 +35,57 @@ use crate::DataOperator;
 
 pub struct FileWithMeta {
     pub path: String,
+    pub metadata: Metadata,
 }
 
 impl FileWithMeta {
-    fn new(path: &str) -> Self {
+    fn new(path: &str, meta: Metadata) -> Self {
         Self {
             path: path.to_string(),
+            metadata: meta,
         }
     }
 }
 
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
+pub enum StageFileStatus {
+    NeedCopy,
+    AlreadyCopied,
+}
+
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
+pub struct StageFileInfo {
+    pub path: String,
+    pub size: u64,
+    pub md5: Option<String>,
+    pub last_modified: DateTime<Utc>,
+    pub etag: Option<String>,
+    pub status: StageFileStatus,
+    pub creator: Option<UserIdentity>,
+}
+
+impl StageFileInfo {
+    pub fn new(path: String, meta: &Metadata) -> StageFileInfo {
+        StageFileInfo {
+            path,
+            size: meta.content_length(),
+            md5: meta.content_md5().map(str::to_string),
+            last_modified: meta
+                .last_modified()
+                .map_or(Utc::now(), |t| Utc.timestamp(t.unix_timestamp(), 0)),
+            etag: meta.etag().map(str::to_string),
+            status: StageFileStatus::NeedCopy,
+            creator: None,
+        }
+    }
+}
+
+impl From<FileWithMeta> for StageFileInfo {
+    fn from(value: FileWithMeta) -> Self {
+        StageFileInfo::new(value.path, &value.metadata)
+    }
+}
+
 pub fn init_stage_operator(stage_info: &StageInfo) -> Result<Operator> {
     if stage_info.stage_type == StageType::External {
         Ok(init_operator(&stage_info.stage_params.storage)?)
@@ -84,7 +131,7 @@ impl StageFilesInfo {
                     .to_string();
                 let meta = operator.stat(&full_path).await?;
                 if meta.mode().is_file() {
-                    res.push(FileWithMeta::new(&full_path))
+                    res.push(FileWithMeta::new(&full_path, meta))
                 } else {
                     return Err(ErrorCode::BadArguments(format!(
                         "{full_path} is not a file"
@@ -97,7 +144,7 @@ impl StageFilesInfo {
             Ok(res)
         } else {
             let pattern = self.get_pattern()?;
-            list_files_with_pattern(operator, &self.path, pattern, first_only).await
+            StageFilesInfo::list_files_with_pattern(operator, &self.path, pattern, first_only).await
         }
     }
 
@@ -131,7 +178,7 @@ impl StageFilesInfo {
                     .to_string();
                 let meta = operator.blocking().stat(&full_path)?;
                 if meta.mode().is_file() {
-                    res.push(FileWithMeta::new(&full_path))
+                    res.push(FileWithMeta::new(&full_path, meta))
                 } else {
                     return Err(ErrorCode::BadArguments(format!(
                         "{full_path} is not a file"
@@ -147,43 +194,46 @@ impl StageFilesInfo {
             blocking_list_files_with_pattern(operator, &self.path, pattern, first_only)
         }
     }
-}
 
-async fn list_files_with_pattern(
-    operator: &Operator,
-    path: &str,
-    pattern: Option<Regex>,
-    first_only: bool,
-) -> Result<Vec<FileWithMeta>> {
-    let root_meta = operator.stat(path).await;
-    match root_meta {
-        Ok(meta) => match meta.mode() {
-            EntryMode::FILE => return Ok(vec![FileWithMeta::new(path)]),
-            EntryMode::DIR => {}
-            EntryMode::Unknown => return Err(ErrorCode::BadArguments("object mode is unknown")),
-        },
-        Err(e) => {
-            if e.kind() == opendal::ErrorKind::NotFound {
-                return Ok(vec![]);
-            } else {
-                return Err(e.into());
+    pub async fn list_files_with_pattern(
+        operator: &Operator,
+        path: &str,
+        pattern: Option<Regex>,
+        first_only: bool,
+    ) -> Result<Vec<FileWithMeta>> {
+        let root_meta = operator.stat(path).await;
+        match root_meta {
+            Ok(meta) => match meta.mode() {
+                EntryMode::FILE => return Ok(vec![FileWithMeta::new(path, meta)]),
+                EntryMode::DIR => {}
+                EntryMode::Unknown => {
+                    return Err(ErrorCode::BadArguments("object mode is unknown"));
+                }
+            },
+            Err(e) => {
+                if e.kind() == opendal::ErrorKind::NotFound {
+                    return Ok(vec![]);
+                } else {
+                    return Err(e.into());
+                }
             }
-        }
-    };
+        };
 
-    // path is a dir
-    let mut files = Vec::new();
-    let mut list = operator.scan(path).await?;
-    while let Some(obj) = list.try_next().await? {
-        let meta = operator.metadata(&obj, Metakey::Mode).await?;
-        if check_file(obj.path(), meta.mode(), &pattern) {
-            files.push(FileWithMeta::new(obj.path()));
-            if first_only {
-                return Ok(files);
+        // path is a dir
+        let mut files = Vec::new();
+        let mut list = operator.scan(path).await?;
+        while let Some(obj) = list.try_next().await? {
+            // todo(youngsofun): not always need Metakey::Complete
+            let meta = operator.metadata(&obj, Metakey::Complete).await?;
+            if check_file(obj.path(), meta.mode(), &pattern) {
+                files.push(FileWithMeta::new(obj.path(), meta));
+                if first_only {
+                    return Ok(files);
+                }
             }
         }
+        Ok(files)
     }
-    Ok(files)
 }
 
 fn check_file(path: &str, mode: EntryMode, pattern: &Option<Regex>) -> bool {
@@ -208,7 +258,7 @@ fn blocking_list_files_with_pattern(
     let root_meta = operator.stat(path);
     match root_meta {
         Ok(meta) => match meta.mode() {
-            EntryMode::FILE => return Ok(vec![FileWithMeta::new(path)]),
+            EntryMode::FILE => return Ok(vec![FileWithMeta::new(path, meta)]),
             EntryMode::DIR => {}
             EntryMode::Unknown => return Err(ErrorCode::BadArguments("object mode is unknown")),
         },
@@ -226,13 +276,38 @@ fn blocking_list_files_with_pattern(
     let list = operator.list(path)?;
     for obj in list {
         let obj = obj?;
-        let meta = operator.metadata(&obj, Metakey::Mode)?;
+        let meta = operator.metadata(&obj, Metakey::Complete)?;
         if check_file(obj.path(), meta.mode(), &pattern) {
-            files.push(FileWithMeta::new(obj.path()));
+            files.push(FileWithMeta::new(obj.path(), meta));
             if first_only {
                 return Ok(files);
             }
         }
     }
     Ok(files)
 }
+
+/// # Behavior
+///
+///
+/// - `Ok(Some(v))` if given object is a file and no error happened.
+/// - `Ok(None)` if given object is not a file.
+/// - `Err(err)` if there is an error happened.
+#[allow(unused)]
+pub async fn stat_file(op: Operator, de: Entry) -> Result<Option<StageFileInfo>> {
+    let meta = op
+        .metadata(&de, {
+            Metakey::Mode
+                | Metakey::ContentLength
+                | Metakey::ContentMd5
+                | Metakey::LastModified
+                | Metakey::Etag
+        })
+        .await?;
+
+    if !meta.mode().is_file() {
+        return Ok(None);
+    }
+
+    Ok(Some(StageFileInfo::new(de.path().to_string(), &meta)))
+}

src/common/storage/src/stage.rs

@@ -705,9 +705,9 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         children.push(self.children.pop().unwrap());
         self.visit_copy_unit(&copy.dst);
         children.push(self.children.pop().unwrap());
-        if !copy.files.is_empty() {
-            let mut files_children = Vec::with_capacity(copy.files.len());
-            for file in copy.files.iter() {
+        if let Some(files) = &copy.files {
+            let mut files_children = Vec::with_capacity(files.len());
+            for file in files.iter() {
                 let file_name = format!("File {}", file);
                 let file_format_ctx = AstFormatContext::new(file_name);
                 let file_node = FormatTreeNode::new(file_format_ctx);
@@ -719,8 +719,8 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
             let files_node = FormatTreeNode::with_children(files_format_ctx, files_children);
             children.push(files_node);
         }
-        if !copy.pattern.is_empty() {
-            let pattern_name = format!("Pattern {}", copy.pattern);
+        if let Some(pattern) = &copy.pattern {
+            let pattern_name = format!("Pattern {}", pattern);
             let pattern_format_ctx = AstFormatContext::new(pattern_name);
             let pattern_node = FormatTreeNode::new(pattern_format_ctx);
             children.push(pattern_node);