fix(max_file_size): fix up max_file_size may oom

src/query/storages/stage/src/stage_table_sink.rs

@@ -79,12 +79,7 @@ impl StageTableSink {
             &ctx.get_settings(),
         )?;
 
-        let mut max_file_size = table_info.user_stage_info.copy_options.max_file_size;
-        if max_file_size == 0 {
-            // 64M per file by default
-            max_file_size = 64 * 1024 * 1024;
-        }
-
+        let max_file_size = Self::adjust_max_file_size(&ctx, &table_info)?;
         let single = table_info.user_stage_info.copy_options.single;
 
         Ok(ProcessorPtr::create(Box::new(StageTableSink {
@@ -95,9 +90,7 @@ impl StageTableSink {
             output,
             single,
             output_format,
-            working_buffer: Vec::with_capacity(
-                (max_file_size.min(256 * 1024) as f64 * 1.2) as usize,
-            ),
+            working_buffer: Vec::with_capacity((max_file_size as f64 * 1.2) as usize),
             working_datablocks: vec![],
             write_header: false,
 
@@ -108,6 +101,25 @@ impl StageTableSink {
         })))
     }
 
+    fn adjust_max_file_size(
+        ctx: &Arc<dyn TableContext>,
+        stage_info: &StageTableInfo,
+    ) -> Result<usize> {
+        // 256M per file by default.
+        const DEFAULT_SIZE: usize = 256 * 1024 * 1024;
+        // max is half of the max memory usage.
+        let max_size = (ctx.get_settings().get_max_memory_usage()? / 2) as usize;
+
+        let mut max_file_size = stage_info.user_stage_info.copy_options.max_file_size;
+        if max_file_size == 0 {
+            max_file_size = DEFAULT_SIZE;
+        } else if max_file_size > max_size {
+            max_file_size = max_size
+        }
+
+        Ok(max_file_size)
+    }
+
     pub fn unload_path(&self) -> String {
         let format_name = format!(
             "{:?}",

tests/sqllogictests/suites/base/05_ddl/05_0027_ddl_copy

@@ -2,7 +2,7 @@ statement ok
 DROP TABLE IF EXISTS t
 
 statement ok
-DROP STAGE IF EXISTS t_copy
+DROP TABLE IF EXISTS t_copy
 
 statement ok
 DROP STAGE IF EXISTS st
@@ -23,30 +23,37 @@ INSERT INTO t SELECT * FROM numbers(3)
 statement ok
 INSERT INTO t SELECT * FROM numbers(4)
 
+### Type tests ###
+
+###### TSV ######
 statement ok
 COPY INTO @st FROM t FILE_FORMAT=(type=TSV)
 
 statement ok
-COPY INTO t_copy from @st PATTERN='.*[.]tsv' file_format = (type = TSV);
+COPY INTO t_copy from @st PATTERN='.*[.]tsv' file_format = (type = TSV) PURGE = true
 
+###### CSV ######
 statement ok
 COPY INTO @st FROM t FILE_FORMAT=(type=CSV);
 
 statement ok
-COPY INTO t_copy from @st PATTERN='.*[.]csv' file_format = (type = CSV);
+COPY INTO t_copy from @st PATTERN='.*[.]csv' file_format = (type = CSV) PURGE = true
 
+###### NDJSON ######
 statement ok
 COPY INTO @st FROM t FILE_FORMAT=(type=NDJSON)
 
 statement ok
-COPY INTO t_copy from @st PATTERN='.*[.]ndjson' file_format = (type = NDJSON);
+COPY INTO t_copy from @st PATTERN='.*[.]ndjson' file_format = (type = NDJSON) PURGE = true
 
+###### PARQUET ######
 statement ok
 COPY INTO @st FROM t FILE_FORMAT=(type=PARQUET)
 
 statement ok
-COPY INTO t_copy from @st PATTERN='.*[.]parquet' file_format = (type = PARQUET);
+COPY INTO t_copy from @st PATTERN='.*[.]parquet' file_format = (type = PARQUET) PURGE = true
 
+###### JSON ######
 statement ok
 COPY INTO @st FROM t FILE_FORMAT=(type=JSON)
 
@@ -56,6 +63,11 @@ SELECT COUNT(*) FROM t_copy;
 ----
 28
 
+###### max_file_size check ######
+## 107374182400 = 100G
+statement ok
+COPY INTO @st FROM t FILE_FORMAT=(type=PARQUET) MAX_FILE_SIZE = 107374182400
+
 statement ok
 DROP TABLE t