feat: ndjson support option null_if. (#14938)

youngsofun · web-flow · commit 3c9853582bfe · 2024-03-14T15:59:11.000+08:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/meta/app/src/principal/file_format.rs b/src/meta/app/src/principal/file_format.rs
@@ -41,6 +41,7 @@ const OPT_ROW_TAG: &str = "row_tag";
 const OPT_ERROR_ON_COLUMN_COUNT_MISMATCH: &str = "error_on_column_count_mismatch";
 const MISSING_FIELD_AS: &str = "missing_field_as";
 const NULL_FIELD_AS: &str = "null_field_as";
+const NULL_IF: &str = "null_if";
 const OPT_EMPTY_FIELD_AS: &str = "empty_field_as";
 const OPT_BINARY_FORMAT: &str = "binary_format";
 
@@ -184,10 +185,24 @@ impl FileFormatParams {
                 let compression = ast.take_compression()?;
                 let missing_field_as = ast.options.remove(MISSING_FIELD_AS);
                 let null_field_as = ast.options.remove(NULL_FIELD_AS);
+                let null_if = ast.options.remove(NULL_IF);
+                let null_if = match null_if {
+                    None => {
+                        vec![]
+                    }
+                    Some(s) => {
+                        let values: Vec<String> = serde_json::from_str(&s).map_err(|_|
+                            ErrorCode::InvalidArgument(format!(
+                            "Invalid option value: NULL_IF is currently set to {s} (in JSON). The valid values are a list of strings."
+                            )))?;
+                        values
+                    }
+                };
                 FileFormatParams::NdJson(NdJsonFileFormatParams::try_create(
                     compression,
                     missing_field_as.as_deref(),
                     null_field_as.as_deref(),
+                    null_if,
                 )?)
             }
             StageFileFormatType::Parquet => {
@@ -586,13 +601,15 @@ pub struct NdJsonFileFormatParams {
     pub compression: StageFileCompression,
     pub missing_field_as: NullAs,
     pub null_field_as: NullAs,
+    pub null_if: Vec<String>,
 }
 
 impl NdJsonFileFormatParams {
     pub fn try_create(
         compression: StageFileCompression,
         missing_field_as: Option<&str>,
         null_field_as: Option<&str>,
+        null_if: Vec<String>,
     ) -> Result<Self> {
         let missing_field_as = NullAs::parse(missing_field_as, MISSING_FIELD_AS, NullAs::Error)?;
         let null_field_as = NullAs::parse(null_field_as, MISSING_FIELD_AS, NullAs::Null)?;
@@ -605,6 +622,7 @@ impl NdJsonFileFormatParams {
             compression,
             missing_field_as,
             null_field_as,
+            null_if,
         })
     }
 }
@@ -615,6 +633,7 @@ impl Default for NdJsonFileFormatParams {
             compression: StageFileCompression::None,
             missing_field_as: NullAs::Error,
             null_field_as: NullAs::FieldDefault,
+            null_if: vec![],
         }
     }
 }
diff --git a/src/meta/proto-conv/src/file_format_from_to_protobuf_impl.rs b/src/meta/proto-conv/src/file_format_from_to_protobuf_impl.rs
@@ -327,6 +327,7 @@ impl FromToProto for mt::principal::NdJsonFileFormatParams {
             compression,
             p.missing_field_as.as_deref(),
             p.null_field_as.as_deref(),
+            p.null_if,
         )
         .map_err(|e| Incompatible {
             reason: format!("{e}"),
@@ -342,6 +343,7 @@ impl FromToProto for mt::principal::NdJsonFileFormatParams {
             compression,
             missing_field_as: Some(self.missing_field_as.to_string()),
             null_field_as: Some(self.null_field_as.to_string()),
+            null_if: self.null_if.clone(),
         })
     }
 }
diff --git a/src/meta/proto-conv/src/util.rs b/src/meta/proto-conv/src/util.rs
@@ -111,7 +111,8 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (79, "2024-01-31: Add: udf.proto/UserDefinedFunction add created_on field", ),
     (80, "2024-02-01: Add: datatype.proto/DataType Geometry type"),
     (81, "2024-03-04: Add: udf.udf_script"),
-    (82, "2024-03-08: Add: table.inverted_index")
+    (82, "2024-03-08: Add: table.inverted_index"),
+    (83, "2024-03-14: Add: null_if in user.proto/NDJSONFileFormatParams")
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)
diff --git a/src/meta/proto-conv/tests/it/main.rs b/src/meta/proto-conv/tests/it/main.rs
@@ -85,3 +85,4 @@ mod v078_grantentry;
 mod v079_udf_created_on;
 mod v081_udf_script;
 mod v082_table_index;
+mod v083_ndjson_format_params;
diff --git a/src/meta/proto-conv/tests/it/v032_file_format_params.rs b/src/meta/proto-conv/tests/it/v032_file_format_params.rs
@@ -93,6 +93,7 @@ fn test_decode_v32_ndjson_file_format_params() -> anyhow::Result<()> {
             compression: StageFileCompression::Gzip,
             missing_field_as: NullAs::Error,
             null_field_as: NullAs::Null,
+            null_if: vec![],
         })
     };
     common::test_pb_from_to(func_name!(), want())?;
diff --git a/src/meta/proto-conv/tests/it/v064_ndjson_format_params.rs b/src/meta/proto-conv/tests/it/v064_ndjson_format_params.rs
@@ -41,6 +41,7 @@ fn test_decode_v64_ndjson_file_format_params() -> anyhow::Result<()> {
             compression: StageFileCompression::Gzip,
             missing_field_as: NullAs::FieldDefault,
             null_field_as: NullAs::Null,
+            null_if: vec![],
         })
     };
     common::test_pb_from_to(func_name!(), want())?;
diff --git a/src/meta/proto-conv/tests/it/v083_ndjson_format_params.rs b/src/meta/proto-conv/tests/it/v083_ndjson_format_params.rs
@@ -0,0 +1,52 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_meta_app::principal::NdJsonFileFormatParams;
+use databend_common_meta_app::principal::NullAs;
+use databend_common_meta_app::principal::StageFileCompression;
+use minitrace::func_name;
+
+use crate::common;
+
+// These bytes are built when a new version in introduced,
+
+// and are kept for backward compatibility test.
+//
+// *************************************************************
+// * These messages should never be updated,                   *
+// * only be added when a new version is added,                *
+// * or be removed when an old version is no longer supported. *
+// *************************************************************
+//
+#[test]
+fn test_decode_v83_ndjson_file_format_params() -> anyhow::Result<()> {
+    let nd_json_file_format_params_v83 = vec![
+        8, 1, 18, 13, 70, 73, 69, 76, 68, 95, 68, 69, 70, 65, 85, 76, 84, 26, 13, 70, 73, 69, 76,
+        68, 95, 68, 69, 70, 65, 85, 76, 84, 34, 0, 160, 6, 83, 168, 6, 24,
+    ];
+    let want = || NdJsonFileFormatParams {
+        compression: StageFileCompression::Gzip,
+        missing_field_as: NullAs::FieldDefault,
+        null_field_as: NullAs::FieldDefault,
+        null_if: vec!["".to_string()],
+    };
+    common::test_load_old(
+        func_name!(),
+        nd_json_file_format_params_v83.as_slice(),
+        83,
+        want(),
+    )?;
+    common::test_pb_from_to(func_name!(), want())?;
+    Ok(())
+}
diff --git a/src/meta/protos/proto/file_format.proto b/src/meta/protos/proto/file_format.proto
@@ -138,6 +138,7 @@ message NdJsonFileFormatParams {
   StageFileCompression compression = 1;
   optional string missing_field_as = 2;
   optional string null_field_as = 3;
+  repeated string null_if = 4;
 }
 
 message JsonFileFormatParams {
diff --git a/src/query/ast/Cargo.toml b/src/query/ast/Cargo.toml
@@ -30,6 +30,7 @@ nom-rule = "0.3.0"
 ordered-float = { workspace = true }
 pratt = "0.4.0"
 pretty = "0.11.3"
+serde_json = { workspace = true }
 strsim = "0.10"
 strum = "0.24"
 strum_macros = "0.24"
diff --git a/src/query/ast/src/parser/stage.rs b/src/query/ast/src/parser/stage.rs
@@ -133,6 +133,16 @@ pub fn format_options(i: Input) -> IResult<BTreeMap<String, String>> {
         |(k, _, v)| (k.text().to_string(), v.text().to_string()),
     );
 
+    let null_if = map(
+        rule! { NULL_IF ~ ^"=" ~ ^"(" ~ ^#comma_separated_list0(literal_string) ~ ^")" },
+        |(_, _, _, values, _)| {
+            (
+                "null_if".to_string(),
+                serde_json::to_string(&values).unwrap(),
+            )
+        },
+    );
+
     map(
         rule! { ((
         #option_type
@@ -141,7 +151,9 @@ pub fn format_options(i: Input) -> IResult<BTreeMap<String, String>> {
         | #string_options
         | #int_options
         | #bool_options
-        | #none_options) ~ ","?)* },
+        | #none_options
+        | #null_if
+        ) ~ ","?)* },
         |opts| BTreeMap::from_iter(opts.iter().map(|((k, v), _)| (k.to_lowercase(), v.clone()))),
     )(i)
 }
diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
@@ -566,6 +566,8 @@ pub enum TokenKind {
     NAN_DISPLAY,
     #[token("NULL_DISPLAY", ignore(ascii_case))]
     NULL_DISPLAY,
+    #[token("NULL_IF", ignore(ascii_case))]
+    NULL_IF,
     #[token("FILE_FORMAT", ignore(ascii_case))]
     FILE_FORMAT,
     #[token("FILE", ignore(ascii_case))]
diff --git a/src/query/pipeline/sources/src/input_formats/impls/input_format_ndjson.rs b/src/query/pipeline/sources/src/input_formats/impls/input_format_ndjson.rs
@@ -49,6 +49,7 @@ impl InputFormatNDJson {
         default_values: &Option<Vec<Scalar>>,
         null_field_as: &NullAs,
         missing_field_as: &NullAs,
+        null_if: &[&str],
     ) -> std::result::Result<(), FileParseError> {
         let mut json: serde_json::Value =
             serde_json::from_reader(buf).map_err(|e| FileParseError::InvalidNDJsonRow {
@@ -131,15 +132,23 @@ impl InputFormatNDJson {
                         }
                     },
                     Some(value) => {
-                        field_decoder.read_field(column, value).map_err(|e| {
-                            FileParseError::ColumnDecodeError {
-                                column_index,
-                                column_name: field.name().to_owned(),
-                                column_type: field.data_type.to_string(),
-                                decode_error: e.to_string(),
-                                column_data: truncate_column_data(value.to_string()),
-                            }
-                        })?;
+                        if !null_if.is_empty()
+                            && matches!(column, ColumnBuilder::Nullable(_))
+                            && value.is_string()
+                            && null_if.contains(&value.as_str().unwrap())
+                        {
+                            column.push_default();
+                        } else {
+                            field_decoder.read_field(column, value).map_err(|e| {
+                                FileParseError::ColumnDecodeError {
+                                    column_index,
+                                    column_name: field.name().to_owned(),
+                                    column_type: field.data_type.to_string(),
+                                    decode_error: e.to_string(),
+                                    column_data: truncate_column_data(value.to_string()),
+                                }
+                            })?;
+                        }
                     }
                 }
             }
@@ -190,6 +199,11 @@ impl InputFormatTextBase for InputFormatNDJson {
             _ => unreachable!(),
         };
 
+        let null_if = format_params
+            .null_if
+            .iter()
+            .map(|x| x.as_str())
+            .collect::<Vec<_>>();
         for (i, end) in batch.row_ends.iter().enumerate() {
             let buf = &batch.data[start..*end];
             let buf = buf.trim();
@@ -202,6 +216,7 @@ impl InputFormatTextBase for InputFormatNDJson {
                     &builder.ctx.default_values,
                     &format_params.null_field_as,
                     &format_params.missing_field_as,
+                    &null_if,
                 ) {
                     builder.ctx.on_error(
                         e,
diff --git a/tests/data/ndjson/null_if.ndjson b/tests/data/ndjson/null_if.ndjson
@@ -0,0 +1,2 @@
+{"a": "", "b": "Null", "c": "Null"}
+{"a": 1, "b": "null", "c": null}
diff --git a/tests/sqllogictests/suites/stage/formats/ndjson/ndjson_null_if.test b/tests/sqllogictests/suites/stage/formats/ndjson/ndjson_null_if.test
@@ -0,0 +1,42 @@
+statement ok
+drop table if exists v
+
+statement ok
+create table v (a variant, b string, c int)
+
+query 
+copy into v from @data/ndjson/null_if.ndjson file_format = (type = 'ndjson') on_error = continue
+----
+ndjson/null_if.ndjson 1 1 Invalid value '"Null"' for column 2 (c Int32 NULL): BadBytes. Code: 1046, Text = Incorrect json value, must be number. 1
+
+query 
+select * from v order by b
+----
+1 null NULL
+
+statement ok
+truncate table v
+
+query 
+copy into v from @data/ndjson/null_if.ndjson file_format = (type = 'ndjson' null_if = ()) force=true  on_error = continue
+----
+ndjson/null_if.ndjson 1 1 Invalid value '"Null"' for column 2 (c Int32 NULL): BadBytes. Code: 1046, Text = Incorrect json value, must be number. 1
+
+query 
+select * from v order by b
+----
+1 null NULL
+
+statement ok
+truncate table v
+
+query 
+copy into v from @data/ndjson/null_if.ndjson file_format = (type = 'ndjson' null_if = ('', 'Null')) force=true
+----
+ndjson/null_if.ndjson 2 0 NULL NULL
+
+query 
+select * from v order by b
+----
+1 null NULL
+NULL NULL NULL

Original file line number	Diff line number	Diff line change
`@@ -327,6 +327,7 @@ impl FromToProto for mt::principal::NdJsonFileFormatParams {`
`327`	`327`	`compression,`
`328`	`328`	`p.missing_field_as.as_deref(),`
`329`	`329`	`p.null_field_as.as_deref(),`
	`330`	`+ p.null_if,`
`330`	`331`	`)`
`331`	`332`	`.map_err(\|e\| Incompatible {`
`332`	`333`	`reason: format!("{e}"),`
`@@ -342,6 +343,7 @@ impl FromToProto for mt::principal::NdJsonFileFormatParams {`
`342`	`343`	`compression,`
`343`	`344`	`missing_field_as: Some(self.missing_field_as.to_string()),`
`344`	`345`	`null_field_as: Some(self.null_field_as.to_string()),`
	`346`	`+ null_if: self.null_if.clone(),`
`345`	`347`	`})`
`346`	`348`	`}`
`347`	`349`	`}`
Original file line number	Diff line number	Diff line change
`@@ -138,6 +138,7 @@ message NdJsonFileFormatParams {`
`138`	`138`	`StageFileCompression compression = 1;`
`139`	`139`	`optional string missing_field_as = 2;`
`140`	`140`	`optional string null_field_as = 3;`
	`141`	`+ repeated string null_if = 4;`
`141`	`142`	`}`
`142`	`143`
`143`	`144`	`message JsonFileFormatParams {`