mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-20 10:52:32 +00:00
feat: support CSV copy skip bad records (#8198)
* feat: support CSV copy headers and skip bad records Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com> * refactor: reuse CSV skippable error helper Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com> * fix: address CSV skip bad records review Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com> * docs: clarify CSV skip bad records performance tradeoff Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com> --------- Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
This commit is contained in:
QuakeWang
@@ -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";
|
||||
|
||||
@@ -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<usize>,
|
||||
pub compression_type: CompressionType,
|
||||
@@ -76,13 +88,11 @@ impl TryFrom<&HashMap<String, String>> 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<String, String>> for CsvFormat {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_bool(key: &'static str, value: &str) -> Result<bool> {
|
||||
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<SharedBuffer> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<usize>,
|
||||
format: &CsvFormat,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
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<Bytes>>,
|
||||
path: &str,
|
||||
schema: SchemaRef,
|
||||
projection: Vec<usize>,
|
||||
has_header: bool,
|
||||
delimiter: u8,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
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<RecordBatch, ArrowError>| 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::<Vec<_>>().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::<Vec<_>>().await.unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("incorrect number of fields"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String>) -> 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<RecordBatch>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
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::<ArrowError>()
|
||||
.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<impl Future<Output = Result<Output>>>,
|
||||
@@ -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::<Vec<_>>();
|
||||
|
||||
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()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
Reference in New Issue
Block a user