From ca07a53debc3e64c03a03f2521c500c237453528 Mon Sep 17 00:00:00 2001 From: QuakeWang <45645138+QuakeWang@users.noreply.github.com> Date: Wed, 3 Jun 2026 20:11:02 +0800 Subject: [PATCH] feat: support CSV copy skip bad records (#8198) * feat: support CSV copy headers and skip bad records Signed-off-by: QuakeWang * refactor: reuse CSV skippable error helper Signed-off-by: QuakeWang * fix: address CSV skip bad records review Signed-off-by: QuakeWang * docs: clarify CSV skip bad records performance tradeoff Signed-off-by: QuakeWang --------- Signed-off-by: QuakeWang --- src/common/datasource/src/file_format.rs | 1 + src/common/datasource/src/file_format/csv.rs | 238 +++++++++++++++++- src/operator/src/statement/copy_table_from.rs | 182 +++++++++++++- src/sql/src/parsers/copy_parser.rs | 22 ++ .../common/copy/copy_from_fs_csv.result | 22 ++ .../common/copy/copy_from_fs_csv.sql | 9 + tests/data/csv/skip_bad_records.csv | 4 + 7 files changed, 457 insertions(+), 21 deletions(-) create mode 100644 tests/data/csv/skip_bad_records.csv diff --git a/src/common/datasource/src/file_format.rs b/src/common/datasource/src/file_format.rs index e36f94c0d2..d9d7b8b648 100644 --- a/src/common/datasource/src/file_format.rs +++ b/src/common/datasource/src/file_format.rs @@ -61,6 +61,7 @@ pub const FORMAT_COMPRESSION_TYPE: &str = "compression_type"; pub const FORMAT_DELIMITER: &str = "delimiter"; pub const FORMAT_SCHEMA_INFER_MAX_RECORD: &str = "schema_infer_max_record"; pub const FORMAT_HAS_HEADER: &str = "has_header"; +pub const FORMAT_SKIP_BAD_RECORDS: &str = "skip_bad_records"; pub const FORMAT_TYPE: &str = "format"; pub const FILE_PATTERN: &str = "pattern"; pub const TIMESTAMP_FORMAT: &str = "timestamp_format"; diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index 77ea553f35..2b39051b48 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -13,15 +13,24 @@ // limitations under the License. use std::collections::HashMap; +use std::io; use std::str::FromStr; +use std::sync::Arc; +use std::task::Poll; use arrow::csv::reader::Format; use arrow::csv::{self, WriterBuilder}; +use arrow::error::ArrowError; use arrow::record_batch::RecordBatch; -use arrow_schema::Schema; +use arrow_schema::{Schema, SchemaRef}; use async_trait::async_trait; +use bytes::{Buf, Bytes}; use common_runtime; +use common_telemetry::warn; use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::StreamExt; +use futures::stream::BoxStream; use object_store::ObjectStore; use snafu::ResultExt; use tokio_util::compat::FuturesAsyncReadCompatExt; @@ -34,9 +43,12 @@ use crate::file_format::{self, FileFormat, stream_to_file}; use crate::share_buffer::SharedBuffer; use crate::util::normalize_infer_schema; +const SKIP_BAD_RECORDS_BATCH_SIZE: usize = 1; + #[derive(Debug, Clone, PartialEq, Eq)] pub struct CsvFormat { pub has_header: bool, + pub skip_bad_records: bool, pub delimiter: u8, pub schema_infer_max_record: Option, pub compression_type: CompressionType, @@ -76,13 +88,11 @@ impl TryFrom<&HashMap> for CsvFormat { })?); }; if let Some(has_header) = value.get(file_format::FORMAT_HAS_HEADER) { - format.has_header = has_header.parse().map_err(|_| { - error::ParseFormatSnafu { - key: file_format::FORMAT_HAS_HEADER, - value: has_header, - } - .build() - })?; + format.has_header = parse_bool(file_format::FORMAT_HAS_HEADER, has_header)?; + }; + if let Some(skip_bad_records) = value.get(file_format::FORMAT_SKIP_BAD_RECORDS) { + format.skip_bad_records = + parse_bool(file_format::FORMAT_SKIP_BAD_RECORDS, skip_bad_records)?; }; if let Some(timestamp_format) = value.get(file_format::TIMESTAMP_FORMAT) { format.timestamp_format = Some(timestamp_format.clone()); @@ -97,10 +107,17 @@ impl TryFrom<&HashMap> for CsvFormat { } } +fn parse_bool(key: &'static str, value: &str) -> Result { + value + .parse() + .map_err(|_| error::ParseFormatSnafu { key, value }.build()) +} + impl Default for CsvFormat { fn default() -> Self { Self { has_header: true, + skip_bad_records: false, delimiter: b',', schema_infer_max_record: Some(file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD), compression_type: CompressionType::Uncompressed, @@ -189,10 +206,136 @@ impl DfRecordBatchEncoder for csv::Writer { } } +/// Builds a CSV stream that can skip selected record-level parse/cast errors. +/// +/// This recovery path intentionally uses one-record batches. It is slower than +/// normal CSV scanning, but keeps each parse/cast failure isolated to a single +/// record. Arrow's CSV decoder clears buffered rows before type parsing, so a +/// failed multi-row flush cannot be safely retried row by row without replaying +/// input bytes. +pub async fn tolerant_csv_stream( + store: &ObjectStore, + path: &str, + schema: SchemaRef, + projection: Vec, + format: &CsvFormat, +) -> Result { + let meta = store + .stat(path) + .await + .context(error::ReadObjectSnafu { path })?; + + let reader = store + .reader(path) + .await + .context(error::ReadObjectSnafu { path })? + .into_bytes_stream(0..meta.content_length()) + .await + .context(error::ReadObjectSnafu { path })?; + + let reader = format.compression_type.convert_stream(reader).boxed(); + tolerant_csv_stream_from_reader( + reader, + path, + schema, + projection, + format.has_header, + format.delimiter, + ) +} + +fn tolerant_csv_stream_from_reader( + reader: BoxStream<'static, io::Result>, + path: &str, + schema: SchemaRef, + projection: Vec, + has_header: bool, + delimiter: u8, +) -> Result { + let projected_schema = Arc::new( + schema + .project(&projection) + .context(error::InferSchemaSnafu)?, + ); + let mut decoder = csv::ReaderBuilder::new(schema) + .with_header(has_header) + .with_delimiter(delimiter) + .with_batch_size(SKIP_BAD_RECORDS_BATCH_SIZE) + .with_projection(projection) + .build_decoder(); + + let path = path.to_string(); + let mut upstream = reader.fuse(); + let mut buffered = Bytes::new(); + let mut input_finished = false; + let stream = futures::stream::poll_fn(move |cx| { + loop { + while !input_finished { + if buffered.is_empty() { + match futures::ready!(upstream.poll_next_unpin(cx)) { + Some(Ok(bytes)) if bytes.is_empty() => continue, + Some(Ok(bytes)) => buffered = bytes, + Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), + None => input_finished = true, + } + } + + let decoded = decoder.decode(buffered.as_ref())?; + if decoded > 0 { + buffered.advance(decoded); + continue; + } + + if decoder.capacity() == 0 || input_finished { + break; + } + + if buffered.is_empty() { + continue; + } + + return Poll::Ready(Some(Err(ArrowError::ParseError( + "CSV decoder made no progress while input bytes remain".to_string(), + )))); + } + + match decoder.flush() { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) if input_finished => return Poll::Ready(None), + Ok(None) => continue, + Err(error) if is_skippable_arrow_error(&error) => { + warn!( + "Skipping bad CSV record while copying from {}: {}", + path, error + ); + } + Err(error) => return Poll::Ready(Some(Err(error))), + } + } + }) + .map(|result: std::result::Result| result.map_err(Into::into)); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + projected_schema, + stream, + ))) +} + +pub fn is_skippable_arrow_error(error: &ArrowError) -> bool { + matches!( + error, + ArrowError::ParseError(_) + | ArrowError::CastError(_) + | ArrowError::ComputeError(_) + | ArrowError::InvalidArgumentError(_) + ) +} + #[cfg(test)] mod tests { use std::sync::Arc; + use arrow_schema::{DataType, Field}; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; use common_recordbatch::{RecordBatch, RecordBatches}; use common_test_util::find_workspace_path; @@ -205,7 +348,7 @@ mod tests { use super::*; use crate::file_format::{ FORMAT_COMPRESSION_TYPE, FORMAT_DELIMITER, FORMAT_HAS_HEADER, - FORMAT_SCHEMA_INFER_MAX_RECORD, FileFormat, file_to_stream, + FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FileFormat, file_to_stream, }; use crate::test_util::{format_schema, test_store}; @@ -331,11 +474,29 @@ mod tests { schema_infer_max_record: Some(2000), delimiter: b'\t', has_header: false, + skip_bad_records: false, timestamp_format: None, time_format: None, date_format: None } ); + + let map = HashMap::from([(FORMAT_SKIP_BAD_RECORDS.to_string(), "true".to_string())]); + let format = CsvFormat::try_from(&map).unwrap(); + + assert_eq!( + format, + CsvFormat { + skip_bad_records: true, + ..CsvFormat::default() + } + ); + } + + #[test] + fn test_try_from_rejects_invalid_bool_options() { + let map = HashMap::from([(FORMAT_SKIP_BAD_RECORDS.to_string(), "yes".to_string())]); + assert!(CsvFormat::try_from(&map).is_err()); } #[tokio::test] @@ -496,4 +657,63 @@ mod tests { assert_eq!(expected, pretty_print); } } + + #[tokio::test] + async fn test_tolerant_csv_stream_continues_after_parse_error() { + let temp_dir = common_test_util::temp_dir::create_temp_dir("test_tolerant_csv_stream"); + let csv_file_path = temp_dir.path().join("input.csv"); + std::fs::write( + &csv_file_path, + "id,name,value\n1,Alice,10.5\nbad,Bad,20.0\nworse,Bad,21.0\n2,Bob,30.5", + ) + .unwrap(); + + let store = test_store("/"); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("name", DataType::Utf8, false), + Field::new("value", DataType::Float64, false), + ])); + let path = csv_file_path.to_str().unwrap(); + + let stream = + tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) + .await + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let pretty_print = arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + let expected = r#"+----+-------+-------+ +| id | name | value | ++----+-------+-------+ +| 1 | Alice | 10.5 | +| 2 | Bob | 30.5 | ++----+-------+-------+"#; + assert_eq!(expected, pretty_print); + } + + #[tokio::test] + async fn test_tolerant_csv_stream_fails_on_structural_csv_error() { + let temp_dir = + common_test_util::temp_dir::create_temp_dir("test_tolerant_csv_stream_csv_error"); + let csv_file_path = temp_dir.path().join("input.csv"); + std::fs::write(&csv_file_path, "id,name,value\n1,Alice,10.5\n2,Bob\n").unwrap(); + + let store = test_store("/"); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("name", DataType::Utf8, false), + Field::new("value", DataType::Float64, false), + ])); + let path = csv_file_path.to_str().unwrap(); + + let stream = + tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) + .await + .unwrap(); + let error = stream.try_collect::>().await.unwrap_err(); + + assert!(error.to_string().contains("incorrect number of fields")); + } } diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index 6f58603247..cae2835242 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -15,11 +15,15 @@ use std::collections::HashMap; use std::future::Future; use std::path::Path; +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use client::{Output, OutputData, OutputMeta}; use common_base::readable_size::ReadableSize; -use common_datasource::file_format::csv::CsvFormat; +use common_datasource::file_format::csv::{ + CsvFormat, is_skippable_arrow_error, tolerant_csv_stream, +}; use common_datasource::file_format::json::JsonFormat; use common_datasource::file_format::orc::{ReaderAdapter, infer_orc_schema, new_orc_stream_reader}; use common_datasource::file_format::{FileFormat, Format, file_to_stream}; @@ -33,10 +37,13 @@ use common_telemetry::{debug, tracing}; use datafusion::datasource::physical_plan::{CsvSource, FileSource, JsonSource}; use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder; use datafusion::parquet::arrow::arrow_reader::ArrowReaderMetadata; +use datafusion_common::DataFusionError; +use datafusion_common::arrow::error::ArrowError; use datafusion_common::config::CsvOptions; use datafusion_expr::Expr; use datatypes::arrow::compute::can_cast_types; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef}; +use datatypes::arrow::record_batch::RecordBatch; use datatypes::vectors::Helper; use futures_util::StreamExt; use object_store::{Entry, EntryMode, ObjectStore}; @@ -221,23 +228,42 @@ impl StatementExecutor { let csv_source = CsvSource::new(schema.clone()) .with_csv_options(options) .with_batch_size(DEFAULT_BATCH_SIZE); - let stream = file_to_stream( - object_store, - path, - csv_source, - Some(projection), - format.compression_type, - ) - .await - .context(error::BuildFileStreamSnafu)?; + let stream = if format.skip_bad_records { + let reader_schema = + csv_reader_schema_for_skip_bad_records(schema, &compat_schema); + tolerant_csv_stream( + object_store, + path, + Arc::new(reader_schema), + projection.clone(), + format, + ) + .await + .context(error::BuildFileStreamSnafu)? + } else { + file_to_stream( + object_store, + path, + csv_source, + Some(projection), + format.compression_type, + ) + .await + .context(error::BuildFileStreamSnafu)? + }; - Ok(Box::pin( + let stream = Box::pin( // The projection is already applied in the CSV reader when we created the stream, // so we pass None here to avoid double projection which would cause schema mismatch errors. RecordBatchStreamTypeAdapter::new(output_schema, stream, None) .with_filter(filters) .context(error::PhysicalExprSnafu)?, - )) + ); + if format.skip_bad_records { + Ok(Box::pin(SkipBadRecordsStream::new(stream, path))) + } else { + Ok(stream) + } } FileMetadata::Json { path, @@ -469,6 +495,58 @@ fn gen_insert_output(rows_inserted: usize, insert_cost: usize) -> Output { ) } +struct SkipBadRecordsStream { + inner: DfSendableRecordBatchStream, + path: String, +} + +impl SkipBadRecordsStream { + fn new(inner: DfSendableRecordBatchStream, path: impl Into) -> Self { + Self { + inner, + path: path.into(), + } + } +} + +impl datafusion::physical_plan::RecordBatchStream for SkipBadRecordsStream { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } +} + +impl futures::Stream for SkipBadRecordsStream { + type Item = datafusion_common::Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + match this.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Err(error))) if is_skippable_record_error(&error) => { + common_telemetry::warn!( + "Skipping bad record while copying from {}: {}", + this.path, + error + ); + continue; + } + other => return other, + } + } + } +} + +fn is_skippable_record_error(error: &DataFusionError) -> bool { + match error { + DataFusionError::ArrowError(error, _) => is_skippable_arrow_error(error), + DataFusionError::External(error) => error + .downcast_ref::() + .is_some_and(is_skippable_arrow_error), + DataFusionError::Context(_, error) => is_skippable_record_error(error), + _ => false, + } +} + /// Executes all pending inserts all at once, drain pending requests and reset pending bytes. async fn batch_insert( pending: &mut Vec>>, @@ -498,6 +576,59 @@ fn can_cast_types_for_greptime(from: &ArrowDataType, to: &ArrowDataType) -> bool can_cast_types(from, to) } +fn csv_reader_schema_for_skip_bad_records(file: &SchemaRef, compat: &SchemaRef) -> Schema { + let fields = file + .fields() + .iter() + .enumerate() + .map(|(idx, file_field)| { + let compat_field = compat + .fields() + .find(file_field.name()) + .map(|(_, field)| field); + + match compat_field { + Some(compat_field) if can_csv_reader_parse_type(compat_field.data_type()) => { + compat_field.clone() + } + _ => file.fields()[idx].clone(), + } + }) + .collect::>(); + + Schema::new_with_metadata(fields, file.metadata().clone()) +} + +fn can_csv_reader_parse_type(data_type: &ArrowDataType) -> bool { + match data_type { + ArrowDataType::Boolean + | ArrowDataType::Decimal32(_, _) + | ArrowDataType::Decimal64(_, _) + | ArrowDataType::Decimal128(_, _) + | ArrowDataType::Decimal256(_, _) + | ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + | ArrowDataType::Date32 + | ArrowDataType::Date64 + | ArrowDataType::Time32(_) + | ArrowDataType::Time64(_) + | ArrowDataType::Timestamp(_, _) + | ArrowDataType::Null + | ArrowDataType::Utf8 + | ArrowDataType::Utf8View => true, + ArrowDataType::Dictionary(_, value_type) => value_type.as_ref() == &ArrowDataType::Utf8, + _ => false, + } +} + fn ensure_schema_compatible(from: &SchemaRef, to: &SchemaRef) -> Result<()> { let not_match = from .fields @@ -780,4 +911,31 @@ mod tests { assert_eq!(test.0.project(&fp).unwrap(), test.1.project(&tp).unwrap()); } } + + #[test] + fn test_csv_reader_schema_for_skip_bad_records() { + let file_schema = make_test_schema(&[ + Field::new("id", DataType::Utf8, true), + Field::new("jsons", DataType::Utf8, true), + Field::new("ts", DataType::Utf8, true), + ]); + let compat_schema = make_test_schema(&[ + Field::new("id", DataType::UInt32, true), + Field::new("jsons", DataType::Binary, true), + Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + true, + ), + ]); + + let reader_schema = csv_reader_schema_for_skip_bad_records(&file_schema, &compat_schema); + + assert_eq!(reader_schema.field(0).data_type(), &DataType::UInt32); + assert_eq!(reader_schema.field(1).data_type(), &DataType::Utf8); + assert_eq!( + reader_schema.field(2).data_type(), + compat_schema.field(2).data_type() + ); + } } diff --git a/src/sql/src/parsers/copy_parser.rs b/src/sql/src/parsers/copy_parser.rs index 9a2eddcc78..491912c82e 100644 --- a/src/sql/src/parsers/copy_parser.rs +++ b/src/sql/src/parsers/copy_parser.rs @@ -401,6 +401,28 @@ mod tests { } } + #[test] + fn test_parse_copy_table_from_csv_options() { + let sql = + "COPY my_table FROM '/tmp/test.csv' WITH (FORMAT = 'CSV', SKIP_BAD_RECORDS = 'false')"; + let mut result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert_eq!(1, result.len()); + + let statement = result.remove(0); + assert_matches!(statement, Statement::Copy { .. }); + match statement { + Statement::Copy(crate::statements::copy::Copy::CopyTable(CopyTable::From( + copy_table, + ))) => { + assert_eq!(copy_table.with.get("format"), Some("CSV")); + assert_eq!(copy_table.with.get("skip_bad_records"), Some("false")); + } + _ => unreachable!(), + } + } + #[test] fn test_parse_copy_table_to() { struct Test<'a> { diff --git a/tests/cases/standalone/common/copy/copy_from_fs_csv.result b/tests/cases/standalone/common/copy/copy_from_fs_csv.result index 2e3373af15..ad81cf1f77 100644 --- a/tests/cases/standalone/common/copy/copy_from_fs_csv.result +++ b/tests/cases/standalone/common/copy/copy_from_fs_csv.result @@ -183,6 +183,24 @@ select * from csv_null_prefix_import; | final | 2023-11-14T22:13:23 | +-------+---------------------+ +CREATE TABLE csv_skip_bad_records(host_id int, host_name string, reading_value double, ts timestamp time index); + +Affected Rows: 0 + +-- SQLNESS ENV PWD +Copy csv_skip_bad_records FROM '$PWD/tests/data/csv/skip_bad_records.csv' WITH (format='csv', skip_bad_records='true'); + +Affected Rows: 2 + +select * from csv_skip_bad_records order by ts; + ++---------+-----------+---------------+---------------------+ +| host_id | host_name | reading_value | ts | ++---------+-----------+---------------+---------------------+ +| 1 | Alice | 10.5 | 2024-01-01T00:00:00 | +| 2 | Bob | 30.5 | 2024-01-01T00:00:02 | ++---------+-----------+---------------+---------------------+ + drop table demo; Affected Rows: 0 @@ -219,3 +237,7 @@ drop table csv_null_prefix_import; Affected Rows: 0 +drop table csv_skip_bad_records; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/copy/copy_from_fs_csv.sql b/tests/cases/standalone/common/copy/copy_from_fs_csv.sql index 0901048177..cd6cfcbd23 100644 --- a/tests/cases/standalone/common/copy/copy_from_fs_csv.sql +++ b/tests/cases/standalone/common/copy/copy_from_fs_csv.sql @@ -73,6 +73,13 @@ Copy csv_null_prefix_import FROM '${SQLNESS_HOME}/demo/export/csv_null_prefix.cs select * from csv_null_prefix_import; +CREATE TABLE csv_skip_bad_records(host_id int, host_name string, reading_value double, ts timestamp time index); + +-- SQLNESS ENV PWD +Copy csv_skip_bad_records FROM '$PWD/tests/data/csv/skip_bad_records.csv' WITH (format='csv', skip_bad_records='true'); + +select * from csv_skip_bad_records order by ts; + drop table demo; drop table with_filename; @@ -90,3 +97,5 @@ drop table demo_with_less_columns; drop table csv_null_prefix; drop table csv_null_prefix_import; + +drop table csv_skip_bad_records; diff --git a/tests/data/csv/skip_bad_records.csv b/tests/data/csv/skip_bad_records.csv new file mode 100644 index 0000000000..f4c40d5d6e --- /dev/null +++ b/tests/data/csv/skip_bad_records.csv @@ -0,0 +1,4 @@ +host_id,host_name,reading_value,ts +1,Alice,10.5,2024-01-01T00:00:00 +bad,Bad,20.0,2024-01-01T00:00:01 +2,Bob,30.5,2024-01-01T00:00:02