feat: add continue_on_error option to skip files with invalid headers

Signed-off-by: Alan Tang <[email protected]>

Author: Standing-Man

src/common/datasource/src/file_format.rs

@@ -54,6 +54,7 @@ pub const FORMAT_SCHEMA_INFER_MAX_RECORD: &str = "schema_infer_max_record";
 pub const FORMAT_HAS_HEADER: &str = "has_header";
 pub const FORMAT_HEADER: &str = "header";
 pub const FORMAT_SKIP_BAD_RECORDS: &str = "skip_bad_records";
+pub const FORMAT_CONTINUE_ON_ERROR: &str = "continue_on_error";
 pub const FORMAT_TYPE: &str = "format";
 pub const FILE_PATTERN: &str = "pattern";
 pub const TIMESTAMP_FORMAT: &str = "timestamp_format";

src/common/datasource/src/file_format/csv.rs

@@ -38,6 +38,7 @@ pub struct CsvFormat {
     pub has_header: bool,
     pub header: bool,
     pub skip_bad_records: bool,
+    pub continue_on_error: bool,
     pub delimiter: u8,
     pub schema_infer_max_record: Option<usize>,
     pub compression_type: CompressionType,
@@ -103,6 +104,15 @@ impl TryFrom<&HashMap<String, String>> for CsvFormat {
                 .build()
             })?;
         }
+        if let Some(continue_on_error) = value.get(file_format::FORMAT_CONTINUE_ON_ERROR) {
+            format.continue_on_error = continue_on_error.parse().map_err(|_| {
+                error::ParseFormatSnafu {
+                    key: file_format::FORMAT_CONTINUE_ON_ERROR,
+                    value: continue_on_error,
+                }
+                .build()
+            })?;
+        }
         if let Some(timestamp_format) = value.get(file_format::TIMESTAMP_FORMAT) {
             format.timestamp_format = Some(timestamp_format.clone());
         }
@@ -122,6 +132,7 @@ impl Default for CsvFormat {
             has_header: true,
             header: true,
             skip_bad_records: false,
+            continue_on_error: true,
             delimiter: b',',
             schema_infer_max_record: Some(file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD),
             compression_type: CompressionType::Uncompressed,
@@ -313,6 +324,7 @@ mod tests {
                 schema_infer_max_record: Some(2000),
                 delimiter: b'\t',
                 has_header: false,
+                continue_on_error: true,
                 header: true,
                 skip_bad_records: false,
                 timestamp_format: None,

src/operator/src/error.rs

@@ -595,6 +595,18 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display(
+        "Header from CSV file not match, expected table schema: {}, actual file schema: {}",
+        table_schema,
+        file_schema
+    ))]
+    InvalidHeader {
+        table_schema: String,
+        file_schema: String,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Failed to project schema"))]
     ProjectSchema {
         #[snafu(source)]
@@ -929,6 +941,7 @@ impl ErrorExt for Error {
             | Error::InvalidTableName { .. }
             | Error::InvalidViewName { .. }
             | Error::InvalidFlowName { .. }
+            | Error::InvalidHeader { .. }
             | Error::InvalidView { .. }
             | Error::InvalidExpr { .. }
             | Error::AdminFunctionNotFound { .. }

src/operator/src/statement/copy_table_from.rs

@@ -22,7 +22,8 @@ use common_base::readable_size::ReadableSize;
 use common_datasource::file_format::csv::CsvFormat;
 use common_datasource::file_format::orc::{ReaderAdapter, infer_orc_schema, new_orc_stream_reader};
 use common_datasource::file_format::{
-    FORMAT_HAS_HEADER, FORMAT_SKIP_BAD_RECORDS, FORMAT_TYPE, FileFormat, Format,
+    FORMAT_CONTINUE_ON_ERROR, FORMAT_HAS_HEADER, FORMAT_SKIP_BAD_RECORDS, FORMAT_TYPE, FileFormat,
+    Format,
 };
 use common_datasource::lister::{Lister, Source};
 use common_datasource::object_store::{FS_SCHEMA, build_backend, parse_url};
@@ -384,6 +385,11 @@ impl StatementExecutor {
             .get(FORMAT_SKIP_BAD_RECORDS)
             .and_then(|v| v.parse::<bool>().ok())
             .unwrap_or(false);
+        let continue_on_error = req
+            .with
+            .get(FORMAT_CONTINUE_ON_ERROR)
+            .and_then(|v| v.parse::<bool>().ok())
+            .unwrap_or(false);
         let table = self.get_table(&table_ref).await?;
         let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?;
         let (object_store, entries) = self.list_copy_from_entries(&req).await?;
@@ -412,7 +418,15 @@ impl StatementExecutor {
                 && has_header
                 && !file_schema.equivalent_names_and_types(&table_schema)
             {
-                continue;
+                if continue_on_error {
+                    continue;
+                } else {
+                    return error::InvalidHeaderSnafu {
+                        table_schema: table_schema.to_string(),
+                        file_schema: file_schema.to_string(),
+                    }
+                    .fail();
+                }
             }
             let (file_schema_projection, table_schema_projection, compat_schema) =
                 generated_schema_projection_and_compatible_file_schema(file_schema, &table_schema);