feat(query): copy with transform support max_files (#10958)

* feat(query): copy with transform support max_files

* modify debug test

Author: TCeason

src/common/storage/src/stage.rs

@@ -159,7 +159,7 @@ impl StageFilesInfo {
     }
 
     pub fn blocking_first_file(&self, operator: &Operator) -> Result<StageFileInfo> {
-        let mut files = self.blocking_list(operator, true)?;
+        let mut files = self.blocking_list(operator, true, None)?;
         match files.pop() {
             None => Err(ErrorCode::BadArguments("no file found")),
             Some(f) => Ok(f),
@@ -170,7 +170,10 @@ impl StageFilesInfo {
         &self,
         operator: &Operator,
         first_only: bool,
+        max_files: Option<usize>,
     ) -> Result<Vec<StageFileInfo>> {
+        let max_files = max_files.unwrap_or(usize::MAX);
+        let mut limit = 0;
         if let Some(files) = &self.files {
             let mut res = Vec::new();
             for file in files {
@@ -189,11 +192,15 @@ impl StageFilesInfo {
                 if first_only {
                     break;
                 }
+                limit += 1;
+                if limit == max_files {
+                    return Ok(res);
+                }
             }
             Ok(res)
         } else {
             let pattern = self.get_pattern()?;
-            blocking_list_files_with_pattern(operator, &self.path, pattern, first_only)
+            blocking_list_files_with_pattern(operator, &self.path, pattern, first_only, max_files)
         }
     }
 
@@ -260,6 +267,7 @@ fn blocking_list_files_with_pattern(
     path: &str,
     pattern: Option<Regex>,
     first_only: bool,
+    max_files: usize,
 ) -> Result<Vec<StageFileInfo>> {
     let operator = operator.blocking();
 
@@ -282,6 +290,7 @@ fn blocking_list_files_with_pattern(
     // path is a dir
     let mut files = Vec::new();
     let list = operator.list(path)?;
+    let mut limit = 0;
     for obj in list {
         let obj = obj?;
         let meta = operator.metadata(&obj, StageFileInfo::meta_query())?;
@@ -290,6 +299,10 @@ fn blocking_list_files_with_pattern(
             if first_only {
                 return Ok(files);
             }
+            limit += 1;
+            if limit == max_files {
+                return Ok(files);
+            }
         }
     }
     Ok(files)

src/query/sql/src/planner/binder/copy.rs

@@ -41,7 +41,6 @@ use common_exception::Result;
 use common_meta_app::principal::OnErrorMode;
 use common_meta_app::principal::StageInfo;
 use common_storage::init_stage_operator;
-use common_storage::StageFileStatus;
 use common_storage::StageFilesInfo;
 use common_users::UserApiProvider;
 use tracing::debug;
@@ -547,15 +546,30 @@ impl<'a> Binder {
         }
 
         let operator = init_stage_operator(&stage_info)?;
+        let max_files = stage_info.copy_options.max_files;
+        let max_files = if max_files == 0 {
+            None
+        } else {
+            Some(max_files)
+        };
         let mut files = if operator.info().can_blocking() {
-            files_info.blocking_list(&operator, false)
+            if stmt.force {
+                files_info.blocking_list(&operator, false, max_files)
+            } else {
+                files_info.blocking_list(&operator, false, None)
+            }
+        } else if stmt.force {
+            files_info.list(&operator, false, max_files).await
         } else {
             files_info.list(&operator, false, None).await
         }?;
 
-        info!("end to list files: {}", files.len());
+        let num_all_files = files.len();
+        info!("end to list files: {}", num_all_files);
 
-        if !stmt.force {
+        let need_copy_file_infos = if stmt.force {
+            files
+        } else {
             // Status.
             {
                 let status = "begin to color copied files";
@@ -570,24 +584,19 @@ impl<'a> Binder {
                     dst_database_name,
                     dst_table_name,
                     &files,
-                    None,
+                    max_files,
                 )
                 .await?;
 
             info!("end to color copied files: {}", files.len());
-        }
-
-        let mut need_copy_file_infos = vec![];
-        for file in &files {
-            if file.status == StageFileStatus::NeedCopy {
-                need_copy_file_infos.push(file.clone());
-            }
-        }
+            files
+        };
 
         info!(
-            "copy: read all files finished, all:{}, need copy:{}, elapsed:{}",
-            files.len(),
+            "copy: read all files finished, all:{}, need copy:{}, max_files:{:?}, elapsed:{}",
+            num_all_files,
             need_copy_file_infos.len(),
+            max_files,
             start.elapsed().as_secs()
         );
 

src/query/storages/parquet/src/parquet_table/partition.rs

@@ -114,7 +114,7 @@ impl ParquetTable {
                 .map(|f| (f.path.clone(), f.size))
                 .collect::<Vec<_>>(),
             None => if self.operator.info().can_blocking() {
-                self.files_info.blocking_list(&self.operator, false)
+                self.files_info.blocking_list(&self.operator, false, None)
             } else {
                 self.files_info.list(&self.operator, false, None).await
             }?

tests/data/00_0005/data/f1.parquet

@@ -0,0 +1,35 @@
+--- force = false, purge = false
+6
+4
+remain 3 files
+6
+remain 3 files
+ERROR 1105 (HY000) at line 1: Code: 1016, Text = no file need to copy.
+6
+remain 3 files
+--- force = false, purge = true
+6
+4
+remain 1 files
+6
+remain 0 files
+ERROR 1105 (HY000) at line 1: Code: 1016, Text = no file need to copy.
+6
+remain 0 files
+--- force = true, purge = false
+6
+4
+remain 3 files
+8
+remain 3 files
+12
+remain 3 files
+--- force = true, purge = true
+6
+4
+remain 1 files
+6
+remain 0 files
+ERROR 1105 (HY000) at line 1: Code: 1016, Text = no file need to copy.
+6
+remain 0 files

tests/data/00_0005/data/f2.parquet

@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+
+# Should be <root>/tests/data/
+
+for force in 'false'  'true'
+do
+	for purge in 'false'  'true'
+	do
+		table="test_max_files_force_${force}_purge_${purge}"
+		echo "drop table if exists ${table}" | $MYSQL_CLIENT_CONNECT
+		echo "CREATE TABLE ${table} (
+      id INT,
+      c1 INT
+    ) ENGINE=FUSE;" | $MYSQL_CLIENT_CONNECT
+	done
+done
+
+gen_files() {
+  rm -rf /tmp/00_0005
+  cp -r "$CURDIR"/../../../data/00_0005 /tmp
+}
+
+echo "drop stage if exists s5;" | $MYSQL_CLIENT_CONNECT
+echo "create stage s5 url = 'fs:///tmp/00_0005/data/' FILE_FORMAT = (type = PARQUET)" | $MYSQL_CLIENT_CONNECT
+
+
+for force in 'false'  'true'
+do
+	for purge in 'false'  'true'
+	do
+	  gen_files
+		echo "--- force = ${force}, purge = ${purge}"
+	  echo "select count(*) from @s5" | $MYSQL_CLIENT_CONNECT
+		for i in {1..3}
+		do
+			table="test_max_files_force_${force}_purge_${purge}"
+			echo "copy into ${table} from (select * from @s5 t) max_files=2 force=${force} purge=${purge}" | $MYSQL_CLIENT_CONNECT
+			echo "select count(*) from ${table}" | $MYSQL_CLIENT_CONNECT
+		  remain=$(ls -1 /tmp/00_0005/data/ | wc -l |  sed 's/ //g')
+			echo "remain ${remain} files"
+		done
+	done
+done
+
+for force in 'false'  'true'
+do
+	for purge in 'false'  'true'
+	do
+		echo "drop table if exists test_max_files_force_${force}_purge_${purge}" | $MYSQL_CLIENT_CONNECT
+	done
+done