feat: impl max_file_count for infer_schema

Author: KKould

src/common/storage/src/stage.rs

@@ -98,7 +98,7 @@ pub fn init_stage_operator(stage_info: &StageInfo) -> Result<Operator> {
 }
 /// select * from @s1/<path> (FILES => <files> PATTERN => <pattern>)
 /// copy from @s1/<path> FILES = <files> PATTERN => <pattern>
-#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq, Debug, Hash)]
+#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq, Debug)]
 pub struct StageFilesInfo {
     pub path: String,
     pub files: Option<Vec<String>>,

src/meta/app/src/principal/file_format.rs

@@ -52,7 +52,7 @@ const OPT_BINARY_FORMAT: &str = "binary_format";
 const OPT_USE_LOGIC_TYPE: &str = "use_logic_type";
 
 /// File format parameters after checking and parsing.
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(tag = "type")]
 pub enum FileFormatParams {
     Csv(CsvFileFormatParams),
@@ -446,7 +446,7 @@ impl FileFormatOptionsReader {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct CsvFileFormatParams {
     pub compression: StageFileCompression,
 
@@ -498,7 +498,7 @@ impl CsvFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct TsvFileFormatParams {
     pub compression: StageFileCompression,
     pub headers: u64,
@@ -532,7 +532,7 @@ impl TsvFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct XmlFileFormatParams {
     pub compression: StageFileCompression,
     pub row_tag: String,
@@ -558,7 +558,7 @@ impl Default for XmlFileFormatParams {
 
 /// used for both `missing_field_as` and `null_field_as`
 /// for extensibility, it is stored as PB string in meta
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
 pub enum NullAs {
     /// for `missing_field_as` only, and is default for it for safety,
     /// in case of wrong field names when creating table.
@@ -570,7 +570,7 @@ pub enum NullAs {
     FieldDefault,
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
 pub enum EmptyFieldAs {
     #[default]
     Null,
@@ -638,7 +638,7 @@ impl Display for NullAs {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
 pub enum BinaryFormat {
     #[default]
     Hex,
@@ -668,7 +668,7 @@ impl Display for BinaryFormat {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct JsonFileFormatParams {
     pub compression: StageFileCompression,
 }
@@ -690,7 +690,7 @@ impl Default for JsonFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct NdJsonFileFormatParams {
     pub compression: StageFileCompression,
     pub missing_field_as: NullAs,
@@ -741,7 +741,7 @@ impl NdJsonFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct AvroFileFormatParams {
     pub compression: StageFileCompression,
     pub missing_field_as: NullAs,
@@ -791,7 +791,7 @@ impl AvroFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ParquetFileFormatParams {
     // used only for unload
     pub compression: StageFileCompression,
@@ -828,7 +828,7 @@ impl ParquetFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 pub struct OrcFileFormatParams {
     pub missing_field_as: NullAs,
 }

src/meta/app/src/principal/user_stage.rs

@@ -60,7 +60,7 @@ pub const COPY_MAX_FILES_PER_COMMIT: usize = 15000;
 /// Instruction for exceeding 'copy into table' file limit.
 pub const COPY_MAX_FILES_COMMIT_MSG: &str = "Commit limit reached: 15,000 files for 'copy into table'. To handle more files, adjust 'CopyOption' with 'max_files=<num>'(e.g., 'max_files=10000') and perform several operations until all files are processed.";
 
-#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq)]
 pub enum StageType {
     /// LegacyInternal will be deprecated.
     ///
@@ -96,7 +96,7 @@ impl Default for StageType {
     }
 }
 
-#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, Eq, PartialEq, Hash)]
+#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, Eq, PartialEq)]
 pub enum StageFileCompression {
     Auto,
     Gzip,
@@ -396,13 +396,13 @@ impl Display for FileFormatOptions {
     }
 }
 
-#[derive(serde::Serialize, serde::Deserialize, Default, Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(serde::Serialize, serde::Deserialize, Default, Clone, Debug, Eq, PartialEq)]
 #[serde(default)]
 pub struct StageParams {
     pub storage: StorageParams,
 }
 
-#[derive(serde::Serialize, serde::Deserialize, Clone, Default, Debug, Eq, PartialEq, Hash)]
+#[derive(serde::Serialize, serde::Deserialize, Clone, Default, Debug, Eq, PartialEq)]
 #[serde(default)]
 pub struct CopyOptions {
     pub on_error: OnErrorMode,
@@ -419,7 +419,7 @@ pub struct CopyOptions {
     pub detailed_output: bool,
 }
 
-#[derive(serde::Serialize, serde::Deserialize, Default, Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(serde::Serialize, serde::Deserialize, Default, Clone, Debug, Eq, PartialEq)]
 #[serde(default)]
 pub struct StageInfo {
     pub stage_name: String,

src/query/ast/src/ast/statements/copy.rs

@@ -648,9 +648,7 @@ impl Display for FileFormatValue {
     }
 }
 
-#[derive(
-    serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Drive, DriveMut, Eq, Hash,
-)]
+#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Drive, DriveMut, Eq)]
 pub enum OnErrorMode {
     Continue,
     SkipFileNum(u64),

src/query/service/src/table_functions/infer_schema/infer_schema_table.rs

@@ -102,6 +102,7 @@ impl InferSchemaTable {
             TableField::new("column_name", TableDataType::String),
             TableField::new("type", TableDataType::String),
             TableField::new("nullable", TableDataType::Boolean),
+            TableField::new("filenames", TableDataType::String),
             TableField::new("order_id", TableDataType::Number(NumberDataType::UInt64)),
         ])
     }

src/query/service/src/table_functions/infer_schema/parquet.rs

@@ -32,6 +32,7 @@ use databend_common_storage::init_stage_operator;
 use databend_common_storage::read_parquet_schema_async_rs;
 use databend_common_storage::StageFileInfo;
 use futures_util::future::try_join_all;
+use itertools::Itertools;
 
 use crate::table_functions::infer_schema::infer_schema_table::INFER_SCHEMA;
 
@@ -79,13 +80,20 @@ impl AsyncSource for ParquetInferSchemaSource {
         let mut names: Vec<String> = vec![];
         let mut types: Vec<String> = vec![];
         let mut nulls: Vec<bool> = vec![];
+        let mut filenames: Vec<String> = vec![];
+        let filenames_str = self
+            .stage_file_infos
+            .iter()
+            .map(|info| &info.path)
+            .join(", ");
 
         for field in table_schema.fields().iter() {
             names.push(field.name().to_string());
 
             let non_null_type = field.data_type().remove_recursive_nullable();
             types.push(non_null_type.sql_name());
             nulls.push(field.is_nullable());
+            filenames.push(filenames_str.clone());
         }
 
         let order_ids = (0..table_schema.fields().len() as u64).collect::<Vec<_>>();
@@ -94,6 +102,7 @@ impl AsyncSource for ParquetInferSchemaSource {
             StringType::from_data(names),
             StringType::from_data(types),
             BooleanType::from_data(nulls),
+            StringType::from_data(filenames),
             UInt64Type::from_data(order_ids),
         ]);
         Ok(Some(block))

src/query/service/src/table_functions/infer_schema/separator.rs

@@ -32,15 +32,19 @@ use databend_common_expression::TableSchema;
 use databend_common_meta_app::principal::FileFormatParams;
 use databend_common_pipeline_transforms::AccumulatingTransform;
 use databend_common_storages_stage::BytesBatch;
+use itertools::Itertools;
 
 use crate::table_functions::infer_schema::merge::merge_schema;
 
+const MAX_SINGLE_FILE_BYTES: usize = 100 * 1024 * 1024;
+
 pub struct InferSchemaSeparator {
     pub file_format_params: FileFormatParams,
     files: HashMap<String, Vec<u8>>,
     pub max_records: Option<usize>,
     schemas: Option<TableSchema>,
-    remaining_files_len: usize,
+    files_len: usize,
+    filenames: Vec<String>,
     is_finished: bool,
 }
 
@@ -55,7 +59,8 @@ impl InferSchemaSeparator {
             files: HashMap::new(),
             max_records,
             schemas: None,
-            remaining_files_len: files_len,
+            files_len,
+            filenames: Vec::with_capacity(files_len),
             is_finished: false,
         }
     }
@@ -76,6 +81,14 @@ impl AccumulatingTransform for InferSchemaSeparator {
         let bytes = self.files.entry(batch.path.clone()).or_default();
         bytes.extend(batch.data);
 
+        if bytes.len() > MAX_SINGLE_FILE_BYTES {
+            return Err(ErrorCode::InvalidArgument(format!(
+                "The file '{}' is too large(maximum allowed: {})",
+                batch.path,
+                human_readable_size(MAX_SINGLE_FILE_BYTES),
+            )));
+        }
+
         // When max_records exists, it will try to use the current bytes to read, otherwise it will buffer all bytes
         if self.max_records.is_none() && !batch.is_eof {
             return Ok(vec![DataBlock::empty()]);
@@ -138,15 +151,15 @@ impl AccumulatingTransform for InferSchemaSeparator {
             }
         };
         self.files.remove(&batch.path);
+        self.filenames.push(batch.path);
 
         let merge_schema = match self.schemas.take() {
             None => TableSchema::try_from(&arrow_schema)?,
             Some(schema) => merge_schema(schema, TableSchema::try_from(&arrow_schema)?),
         };
         self.schemas = Some(merge_schema);
 
-        self.remaining_files_len = self.remaining_files_len.saturating_sub(1);
-        if self.remaining_files_len > 0 {
+        if self.files_len > self.filenames.len() {
             return Ok(vec![DataBlock::empty()]);
         }
         self.is_finished = true;
@@ -157,13 +170,16 @@ impl AccumulatingTransform for InferSchemaSeparator {
         let mut names: Vec<String> = vec![];
         let mut types: Vec<String> = vec![];
         let mut nulls: Vec<bool> = vec![];
+        let mut filenames: Vec<String> = vec![];
+        let filenames_str = self.filenames.iter().join(", ");
 
         for field in table_schema.fields().iter() {
             names.push(field.name().to_string());
 
             let non_null_type = field.data_type().remove_recursive_nullable();
             types.push(non_null_type.sql_name());
             nulls.push(field.is_nullable());
+            filenames.push(filenames_str.clone());
         }
 
         let order_ids = (0..table_schema.fields().len() as u64).collect::<Vec<_>>();
@@ -172,8 +188,26 @@ impl AccumulatingTransform for InferSchemaSeparator {
             StringType::from_data(names),
             StringType::from_data(types),
             BooleanType::from_data(nulls),
+            StringType::from_data(filenames),
             UInt64Type::from_data(order_ids),
         ]);
         Ok(vec![block])
     }
 }
+
+fn human_readable_size(bytes: usize) -> String {
+    const KB: f64 = 1024.0;
+    const MB: f64 = KB * 1024.0;
+    const GB: f64 = MB * 1024.0;
+
+    let b = bytes as f64;
+    if b >= GB {
+        format!("{:.2} GB", b / GB)
+    } else if b >= MB {
+        format!("{:.2} MB", b / MB)
+    } else if b >= KB {
+        format!("{:.2} KB", b / KB)
+    } else {
+        format!("{} B", bytes)
+    }
+}

src/query/storages/stage/src/infer_schema.rs

@@ -13,9 +13,6 @@
 // limitations under the License.
 
 use std::any::Any;
-use std::hash::DefaultHasher;
-use std::hash::Hash;
-use std::hash::Hasher;
 use std::sync::Arc;
 
 use databend_common_catalog::plan::PartInfo;
@@ -48,11 +45,7 @@ impl PartInfo for InferSchemaPartInfo {
     }
 
     fn hash(&self) -> u64 {
-        let mut s = DefaultHasher::new();
-        self.files_info.hash(&mut s);
-        self.file_format_params.hash(&mut s);
-        self.stage_info.hash(&mut s);
-        s.finish()
+        0
     }
 
     fn part_type(&self) -> PartInfoType {

tests/data/csv/max_file_count/numbers0.csv

@@ -1,5 +1,4 @@
 col1,col2,col3,col4,col5
 0,1,2,3,4
 5,6,7,8,9
-10,11,12,13,14
-a,b,c,d,e
+10,11,12,13,14

tests/data/ndjson/max_file_count/numbers0.ndjson

@@ -0,0 +1,3 @@
+{"id": 1, "value": 100}
+{"id": 2, "value": 200}
+{"id": 3, "value": 300}