feat: limit total rows copied in COPY TABLE FROM stmt

src/operator/src/statement/copy_table_from.rs

@@ -52,6 +52,7 @@ use crate::statement::StatementExecutor;
 
 const DEFAULT_BATCH_SIZE: usize = 8192;
 const DEFAULT_READ_BUFFER: usize = 256 * 1024;
+const MAX_INSERT_ROWS: usize = 1000;
 
 enum FileMetadata {
     Parquet {
@@ -377,6 +378,11 @@ impl StatementExecutor {
 
         let mut rows_inserted = 0;
         let mut insert_cost = 0;
+        let max_insert_rows = req
+            .with
+            .get("max_insert_rows")
+            .and_then(|val| val.parse::<usize>().ok())
+            .unwrap_or(MAX_INSERT_ROWS);
         for (compat_schema, file_schema_projection, projected_table_schema, file_metadata) in files
         {
             let mut stream = self
@@ -427,6 +433,10 @@ impl StatementExecutor {
                     rows_inserted += rows;
                     insert_cost += cost;
                 }
+
+                if rows_inserted >= max_insert_rows {
+                    return Ok(gen_insert_output(rows_inserted, insert_cost));
+                }
             }
 
             if !pending.is_empty() {
@@ -436,13 +446,17 @@ impl StatementExecutor {
             }
         }
 
-        Ok(Output::new(
-            OutputData::AffectedRows(rows_inserted),
-            OutputMeta::new_with_cost(insert_cost),
-        ))
+        Ok(gen_insert_output(rows_inserted, insert_cost))
     }
 }
 
+fn gen_insert_output(rows_inserted: usize, insert_cost: usize) -> Output {
+    Output::new(
+        OutputData::AffectedRows(rows_inserted),
+        OutputMeta::new_with_cost(insert_cost),
+    )
+}
+
 /// Executes all pending inserts all at once, drain pending requests and reset pending bytes.
 async fn batch_insert(
     pending: &mut Vec<impl Future<Output = Result<Output>>>,

tests/cases/standalone/common/copy/copy_from_fs_parquet.result

@@ -61,6 +61,23 @@ select * from with_pattern order by ts;
 | host2 | 88.8 | 333.3  | 2022-06-15T07:02:38 |
 +-------+------+--------+---------------------+
 
+CREATE TABLE with_limit_rows(host string, cpu double, memory double, ts timestamp time index);
+
+Affected Rows: 0
+
+Copy with_limit_rows FROM '/tmp/demo/export/parquet/' WITH (MAX_INSERT_ROWS = 10);
+
+Affected Rows: 2
+
+select * from with_limit_rows order by ts;
+
++-------+------+--------+---------------------+
+| host  | cpu  | memory | ts                  |
++-------+------+--------+---------------------+
+| host1 | 66.6 | 1024.0 | 2022-06-15T07:02:37 |
+| host2 | 88.8 | 333.3  | 2022-06-15T07:02:38 |
++-------+------+--------+---------------------+
+
 drop table demo;
 
 Affected Rows: 0
@@ -77,3 +94,7 @@ drop table with_pattern;
 
 Affected Rows: 0
 
+drop table with_limit_rows;
+
+Affected Rows: 0
+

tests/cases/standalone/common/copy/copy_from_fs_parquet.sql

@@ -22,10 +22,18 @@ Copy with_pattern FROM '/tmp/demo/export/parquet/' WITH (PATTERN = 'demo.*');
 
 select * from with_pattern order by ts;
 
+CREATE TABLE with_limit_rows(host string, cpu double, memory double, ts timestamp time index);
+
+Copy with_limit_rows FROM '/tmp/demo/export/parquet/' WITH (MAX_INSERT_ROWS = 10);
+
+select * from with_limit_rows order by ts;
+
 drop table demo;
 
 drop table with_filename;
 
 drop table with_path;
 
 drop table with_pattern;
+
+drop table with_limit_rows;