fix(copy): treat null inferred columns as Utf8 to preserve later values (#7658)

Signed-off-by: WenyXu <wenymedia@gmail.com>

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();

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 {

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<String>) {
     if path.is_empty() {
         ("/".to_string(), None)
@@ -27,9 +29,30 @@ pub fn find_dir_and_filename(path: &str) -> (String, Option<String>) {
     }
 }
 
+/// 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::<Vec<_>>();
+
+    arrow_schema::Schema {
+        fields: arrow_schema::Fields::from(fields),
+        metadata: schema.metadata,
+    }
+}
+
 #[cfg(test)]
 mod tests {
-
     use url::Url;
 
     use super::*;

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

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"}

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
+

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;

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
+

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;