From 969a64d483a6593ebe3a95c4240e2136544f7aa1 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 3 Feb 2026 18:20:35 +0800 Subject: [PATCH] fix(copy): treat null inferred columns as Utf8 to preserve later values (#7658) Signed-off-by: WenyXu --- src/common/datasource/src/file_format/csv.rs | 36 ++++++++++++--- src/common/datasource/src/file_format/json.rs | 22 ++++++++- src/common/datasource/src/util.rs | 25 ++++++++++- src/common/datasource/tests/csv/max_infer.csv | 10 +++++ .../datasource/tests/json/max_infer.json | 6 +++ .../common/copy/copy_from_fs_csv.result | 45 +++++++++++++++++++ .../common/copy/copy_from_fs_csv.sql | 22 +++++++++ .../common/copy/copy_from_fs_json.result | 45 +++++++++++++++++++ .../common/copy/copy_from_fs_json.sql | 22 +++++++++ 9 files changed, 224 insertions(+), 9 deletions(-) create mode 100644 src/common/datasource/tests/csv/max_infer.csv create mode 100644 src/common/datasource/tests/json/max_infer.json diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index 7832c7eae96..77ea553f356 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -32,6 +32,7 @@ use crate::compression::CompressionType; use crate::error::{self, Result}; use crate::file_format::{self, FileFormat, stream_to_file}; use crate::share_buffer::SharedBuffer; +use crate::util::normalize_infer_schema; #[derive(Debug, Clone, PartialEq, Eq)] pub struct CsvFormat { @@ -142,7 +143,8 @@ impl FileFormat for CsvFormat { let (schema, _records_read) = format .infer_schema(reader, schema_infer_max_record) .context(error::InferSchemaSnafu)?; - Ok(schema) + + Ok(normalize_infer_schema(schema)) }) .await .context(error::JoinHandleSnafu)? @@ -241,13 +243,35 @@ mod tests { } #[tokio::test] - async fn infer_schema_with_limit() { - let json = CsvFormat { + async fn normalize_infer_schema() { + let csv = CsvFormat { schema_infer_max_record: Some(3), ..CsvFormat::default() }; let store = test_store(&test_data_root()); - let schema = json + let schema = csv.infer_schema(&store, "max_infer.csv").await.unwrap(); + let formatted: Vec<_> = format_schema(schema); + + assert_eq!( + vec![ + "num: Int64: NULL", + "str: Utf8: NULL", + "ts: Utf8: NULL", + "t: Utf8: NULL", + "date: Date32: NULL" + ], + formatted, + ); + } + + #[tokio::test] + async fn infer_schema_with_limit() { + let csv = CsvFormat { + schema_infer_max_record: Some(3), + ..CsvFormat::default() + }; + let store = test_store(&test_data_root()); + let schema = csv .infer_schema(&store, "schema_infer_limit.csv") .await .unwrap(); @@ -263,9 +287,9 @@ mod tests { formatted ); - let json = CsvFormat::default(); + let csv = CsvFormat::default(); let store = test_store(&test_data_root()); - let schema = json + let schema = csv .infer_schema(&store, "schema_infer_limit.csv") .await .unwrap(); diff --git a/src/common/datasource/src/file_format/json.rs b/src/common/datasource/src/file_format/json.rs index 86376933bd1..3dd62e8d95b 100644 --- a/src/common/datasource/src/file_format/json.rs +++ b/src/common/datasource/src/file_format/json.rs @@ -34,6 +34,7 @@ use crate::compression::CompressionType; use crate::error::{self, Result}; use crate::file_format::{self, FileFormat, stream_to_file}; use crate::share_buffer::SharedBuffer; +use crate::util::normalize_infer_schema; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct JsonFormat { @@ -102,7 +103,7 @@ impl FileFormat for JsonFormat { let schema = infer_json_schema_from_iterator(iter).context(error::InferSchemaSnafu)?; - Ok(schema) + Ok(normalize_infer_schema(schema)) }) .await .context(error::JoinHandleSnafu)? @@ -124,7 +125,11 @@ pub async fn stream_to_json( threshold, concurrency, format.compression_type, - json::LineDelimitedWriter::new, + |b| { + arrow::json::writer::WriterBuilder::new() + .with_explicit_nulls(true) + .build::<_, LineDelimited>(b) + }, ) .await } @@ -178,6 +183,19 @@ mod tests { ); } + #[tokio::test] + async fn normalize_infer_schema() { + let json = JsonFormat { + schema_infer_max_record: Some(3), + ..JsonFormat::default() + }; + let store = test_store(&test_data_root()); + let schema = json.infer_schema(&store, "max_infer.json").await.unwrap(); + let formatted: Vec<_> = format_schema(schema); + + assert_eq!(vec!["num: Int64: NULL", "str: Utf8: NULL"], formatted,); + } + #[tokio::test] async fn infer_schema_with_limit() { let json = JsonFormat { diff --git a/src/common/datasource/src/util.rs b/src/common/datasource/src/util.rs index fc0340b148d..63374e7cdba 100644 --- a/src/common/datasource/src/util.rs +++ b/src/common/datasource/src/util.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + pub fn find_dir_and_filename(path: &str) -> (String, Option) { if path.is_empty() { ("/".to_string(), None) @@ -27,9 +29,30 @@ pub fn find_dir_and_filename(path: &str) -> (String, Option) { } } +/// Normalize the schema inferred from the data. +/// If the data type is null, set the data type to Utf8. +pub fn normalize_infer_schema(schema: arrow_schema::Schema) -> arrow_schema::Schema { + let fields = schema + .fields + .iter() + .map(|f| { + if f.data_type().is_null() { + // Set the data type to Utf8 for null fields + Arc::new((**f).clone().with_data_type(arrow_schema::DataType::Utf8)) + } else { + f.clone() + } + }) + .collect::>(); + + arrow_schema::Schema { + fields: arrow_schema::Fields::from(fields), + metadata: schema.metadata, + } +} + #[cfg(test)] mod tests { - use url::Url; use super::*; diff --git a/src/common/datasource/tests/csv/max_infer.csv b/src/common/datasource/tests/csv/max_infer.csv new file mode 100644 index 00000000000..81e96225d7e --- /dev/null +++ b/src/common/datasource/tests/csv/max_infer.csv @@ -0,0 +1,10 @@ +num,str,ts,t,date +1,,,,2023-04-01 +1,,,,2023-04-02 +1,,,,2023-04-03 +1,,,,2023-04-04 +1,,,,2023-04-05 +1,,,,2023-04-06 +5,test,2023-04-01 00:00:00,10,2023-04-01 +2,hello,2023-04-01 00:00:00,20,2023-04-01 +4,foo,2023-04-01 00:00:00,30,2023-04-01 diff --git a/src/common/datasource/tests/json/max_infer.json b/src/common/datasource/tests/json/max_infer.json new file mode 100644 index 00000000000..e1f826d9f3d --- /dev/null +++ b/src/common/datasource/tests/json/max_infer.json @@ -0,0 +1,6 @@ +{"num":5,"str":null} +{"num":2,"str":null} +{"num":4,"str":null} +{"num":5,"str":null} +{"num":2,"str":null} +{"num":4,"str":"foo"} \ No newline at end of file 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 4ebd4f20000..2e3373af154 100644 --- a/tests/cases/standalone/common/copy/copy_from_fs_csv.result +++ b/tests/cases/standalone/common/copy/copy_from_fs_csv.result @@ -146,6 +146,43 @@ select * from demo_with_less_columns order by ts; | host3 | 444.4 | 2024-07-27T10:47:43 | +-------+--------+---------------------+ +CREATE TABLE csv_null_prefix(note string, ts timestamp time index); + +Affected Rows: 0 + +insert into + csv_null_prefix(note, ts) +values + (NULL, 1700000000000), + (NULL, 1700000001000), + (NULL, 1700000002000), + ('final', 1700000003000); + +Affected Rows: 4 + +Copy csv_null_prefix TO '${SQLNESS_HOME}/demo/export/csv_null_prefix.csv' with (format='csv'); + +Affected Rows: 4 + +CREATE TABLE csv_null_prefix_import(note string, ts timestamp time index); + +Affected Rows: 0 + +Copy csv_null_prefix_import FROM '${SQLNESS_HOME}/demo/export/csv_null_prefix.csv' with (format='csv', schema_infer_max_record=2); + +Affected Rows: 4 + +select * from csv_null_prefix_import; + ++-------+---------------------+ +| note | ts | ++-------+---------------------+ +| | 2023-11-14T22:13:20 | +| | 2023-11-14T22:13:21 | +| | 2023-11-14T22:13:22 | +| final | 2023-11-14T22:13:23 | ++-------+---------------------+ + drop table demo; Affected Rows: 0 @@ -174,3 +211,11 @@ drop table demo_with_less_columns; Affected Rows: 0 +drop table csv_null_prefix; + +Affected Rows: 0 + +drop table csv_null_prefix_import; + +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 49550d9ecb3..09010481770 100644 --- a/tests/cases/standalone/common/copy/copy_from_fs_csv.sql +++ b/tests/cases/standalone/common/copy/copy_from_fs_csv.sql @@ -55,6 +55,24 @@ Copy demo_with_less_columns FROM '${SQLNESS_HOME}/demo/export/csv/demo.csv' WITH select * from demo_with_less_columns order by ts; +CREATE TABLE csv_null_prefix(note string, ts timestamp time index); + +insert into + csv_null_prefix(note, ts) +values + (NULL, 1700000000000), + (NULL, 1700000001000), + (NULL, 1700000002000), + ('final', 1700000003000); + +Copy csv_null_prefix TO '${SQLNESS_HOME}/demo/export/csv_null_prefix.csv' with (format='csv'); + +CREATE TABLE csv_null_prefix_import(note string, ts timestamp time index); + +Copy csv_null_prefix_import FROM '${SQLNESS_HOME}/demo/export/csv_null_prefix.csv' with (format='csv', schema_infer_max_record=2); + +select * from csv_null_prefix_import; + drop table demo; drop table with_filename; @@ -68,3 +86,7 @@ drop table with_pattern; drop table demo_with_external_column; drop table demo_with_less_columns; + +drop table csv_null_prefix; + +drop table csv_null_prefix_import; diff --git a/tests/cases/standalone/common/copy/copy_from_fs_json.result b/tests/cases/standalone/common/copy/copy_from_fs_json.result index 87fae5b9129..f18044bd499 100644 --- a/tests/cases/standalone/common/copy/copy_from_fs_json.result +++ b/tests/cases/standalone/common/copy/copy_from_fs_json.result @@ -146,6 +146,43 @@ select * from demo_with_less_columns order by ts; | host3 | 444.4 | 2024-07-27T10:47:43 | +-------+--------+---------------------+ +CREATE TABLE json_null_prefix(note string, ts timestamp time index); + +Affected Rows: 0 + +insert into + json_null_prefix(note, ts) +values + (NULL, 1700000000000), + (NULL, 1700000001000), + (NULL, 1700000002000), + ('final', 1700000003000); + +Affected Rows: 4 + +Copy json_null_prefix TO '${SQLNESS_HOME}/demo/export/json_null_prefix.json' with (format='json'); + +Affected Rows: 4 + +CREATE TABLE json_null_prefix_import(note string, ts timestamp time index); + +Affected Rows: 0 + +Copy json_null_prefix_import FROM '${SQLNESS_HOME}/demo/export/json_null_prefix.json' with (format='json', schema_infer_max_record=2); + +Affected Rows: 4 + +select * from json_null_prefix_import; + ++-------+---------------------+ +| note | ts | ++-------+---------------------+ +| | 2023-11-14T22:13:20 | +| | 2023-11-14T22:13:21 | +| | 2023-11-14T22:13:22 | +| final | 2023-11-14T22:13:23 | ++-------+---------------------+ + drop table demo; Affected Rows: 0 @@ -174,3 +211,11 @@ drop table demo_with_less_columns; Affected Rows: 0 +drop table json_null_prefix; + +Affected Rows: 0 + +drop table json_null_prefix_import; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/copy/copy_from_fs_json.sql b/tests/cases/standalone/common/copy/copy_from_fs_json.sql index 215877fae6e..f109c10c98e 100644 --- a/tests/cases/standalone/common/copy/copy_from_fs_json.sql +++ b/tests/cases/standalone/common/copy/copy_from_fs_json.sql @@ -55,6 +55,24 @@ Copy demo_with_less_columns FROM '${SQLNESS_HOME}/demo/export/json/demo.json' WI select * from demo_with_less_columns order by ts; +CREATE TABLE json_null_prefix(note string, ts timestamp time index); + +insert into + json_null_prefix(note, ts) +values + (NULL, 1700000000000), + (NULL, 1700000001000), + (NULL, 1700000002000), + ('final', 1700000003000); + +Copy json_null_prefix TO '${SQLNESS_HOME}/demo/export/json_null_prefix.json' with (format='json'); + +CREATE TABLE json_null_prefix_import(note string, ts timestamp time index); + +Copy json_null_prefix_import FROM '${SQLNESS_HOME}/demo/export/json_null_prefix.json' with (format='json', schema_infer_max_record=2); + +select * from json_null_prefix_import; + drop table demo; drop table with_filename; @@ -68,3 +86,7 @@ drop table with_pattern; drop table demo_with_external_column; drop table demo_with_less_columns; + +drop table json_null_prefix; + +drop table json_null_prefix_import;