Merge pull request #116 from Embucket/df55-csv-null-regex

Propagate CSV null regex through scans and physical plans

Author: osipovartem

Cargo.lock

@@ -89,30 +89,30 @@ version = "55.1.0"
 #
 # See for more details: https://github.com/rust-lang/cargo/issues/11329
 apache-avro = { version = "0.21", default-features = false }
-arrow = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", features = [
+arrow = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", features = [
     "prettyprint",
     "chrono-tz",
 ] }
-arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false, features = [
+arrow-avro = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false, features = [
     "deflate",
     "snappy",
     "zstd",
     "bzip2",
     "xz",
 ] }
-arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false }
-arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false }
-arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", features = [
+arrow-buffer = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false }
+arrow-data = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false }
+arrow-flight = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", features = [
     "flight-sql-experimental",
 ] }
 # Both codecs are required here to make sure that code paths like
 # file-spilling have access to all compression codecs.
-arrow-ipc = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false, features = [
+arrow-ipc = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false, features = [
     "lz4",
     "zstd",
 ] }
-arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false }
-arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false }
+arrow-ord = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false }
+arrow-schema = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false }
 async-trait = "0.1.89"
 bigdecimal = "0.4.8"
 bytes = "1.11"
@@ -178,7 +178,7 @@ memchr = "2.8.1"
 num-traits = { version = "0.2" }
 object_store = { version = "0.13.2", default-features = false }
 parking_lot = "0.12"
-parquet = { git = "https://github.com/Embucket/arrow-rs.git", rev = "539206bde", default-features = false, features = [
+parquet = { git = "https://github.com/Embucket/arrow-rs.git", rev = "821c2dd1e", default-features = false, features = [
     "arrow",
     "async",
     "object_store",

Cargo.toml

@@ -3689,6 +3689,8 @@ config_namespace! {
         pub null_value: Option<String>, default = None
         // The input regex for Nulls when loading CSVs.
         pub null_regex: Option<String>, default = None
+        /// Keep quoted empty CSV fields distinct from unquoted empty fields when matching nulls.
+        pub preserve_quoted_empty: Option<bool>, default = None
         pub comment: Option<u8>, default = None
         /// Whether to allow truncated rows when parsing, both within a single file and across files.
         ///

datafusion/common/src/config.rs

@@ -54,6 +54,7 @@ datafusion-session = { workspace = true }
 futures = { workspace = true }
 object_store = { workspace = true }
 regex = { workspace = true }
+regex-syntax = "0.8"
 tokio = { workspace = true }
 
 # Note: add additional linter rules in lib.rs.

datafusion/datasource-csv/Cargo.toml

@@ -249,6 +249,15 @@ impl CsvFormat {
         self
     }
 
+    /// Keep quoted empty fields distinct from unquoted empty fields during CSV scans.
+    ///
+    /// If the null matcher includes an empty field, this requires an explicit
+    /// table schema: CSV schema inference cannot recover quoting information.
+    pub fn with_preserve_quoted_empty(mut self, preserve: bool) -> Self {
+        self.options.preserve_quoted_empty = Some(preserve);
+        self
+    }
+
     /// Returns `Some(true)` if the first line is a header, `Some(false)` if
     /// it is not, and `None` if it is not specified.
     pub fn has_header(&self) -> Option<bool> {
@@ -559,12 +568,19 @@ impl CsvFormat {
                 )
                 .with_delimiter(self.options.delimiter)
                 .with_quote(self.options.quote)
+                .with_preserve_quoted_empty(
+                    self.options.preserve_quoted_empty.unwrap_or(false),
+                )
                 .with_truncated_rows(self.options.truncated_rows.unwrap_or(false));
 
             if let Some(null_regex) = &self.options.null_regex {
-                let regex = Regex::new(null_regex.as_str())
-                    .expect("Unable to parse CSV null regex.");
-                format = format.with_null_regex(regex);
+                if let Some(values) = crate::exact_null_values(null_regex) {
+                    format = format.with_null_values(values);
+                } else {
+                    let regex = Regex::new(null_regex.as_str())
+                        .expect("Unable to parse CSV null regex.");
+                    format = format.with_null_regex(regex);
+                }
             }
 
             if let Some(escape) = self.options.escape {
@@ -971,6 +987,7 @@ impl From<&CsvFormatFactory> for datafusion_proto_models::protobuf::CsvOptions {
                 time_format: options.time_format.clone().unwrap_or_default(),
                 null_value: options.null_value.clone().unwrap_or_default(),
                 null_regex: options.null_regex.clone().unwrap_or_default(),
+                preserve_quoted_empty: options.preserve_quoted_empty,
                 comment: options.comment.map_or(vec![], |v| vec![v]),
                 newlines_in_values: options
                     .newlines_in_values

datafusion/datasource-csv/src/file_format.rs

@@ -31,6 +31,42 @@ use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_datasource::{file::FileSource, file_scan_config::FileScanConfig};
 use datafusion_execution::object_store::ObjectStoreUrl;
 pub use file_format::*;
+use regex_syntax::hir::{Hir, HirKind, Look};
+
+/// Extract exact values only when a regex matches the entire field without flags or classes.
+fn exact_null_values(pattern: &str) -> Option<Vec<String>> {
+    if !pattern.starts_with(r"\A") || !pattern.ends_with(r"\z") {
+        return None;
+    }
+    let hir = regex_syntax::Parser::new().parse(pattern).ok()?;
+    let HirKind::Concat(parts) = hir.kind() else {
+        return None;
+    };
+    let [start, values, end] = parts.as_slice() else {
+        return None;
+    };
+    if !matches!(start.kind(), HirKind::Look(Look::Start))
+        || !matches!(end.kind(), HirKind::Look(Look::End))
+    {
+        return None;
+    }
+    let alternatives = match values.kind() {
+        HirKind::Alternation(parts) => parts.as_slice(),
+        _ => std::slice::from_ref(values),
+    };
+    if alternatives.len() > 64 {
+        return None;
+    }
+    alternatives.iter().map(exact_literal).collect()
+}
+
+fn exact_literal(hir: &Hir) -> Option<String> {
+    match hir.kind() {
+        HirKind::Empty => Some(String::new()),
+        HirKind::Literal(value) => std::str::from_utf8(&value.0).ok().map(str::to_owned),
+        _ => None,
+    }
+}
 
 /// Returns a [`FileScanConfig`] for given `file_groups`
 pub fn partitioned_csv_config(
@@ -43,3 +79,24 @@ pub fn partitioned_csv_config(
             .build(),
     )
 }
+
+#[cfg(test)]
+mod null_value_tests {
+    use super::exact_null_values;
+
+    #[test]
+    fn extracts_only_exact_null_values() {
+        assert_eq!(
+            exact_null_values(r"\A(?:\\N|NA|)\z"),
+            Some(vec![r"\N".to_owned(), "NA".to_owned(), String::new()])
+        );
+        assert_eq!(
+            exact_null_values(r"\A(?:NA|)\z"),
+            Some(vec!["NA".to_owned(), String::new()])
+        );
+        assert_eq!(exact_null_values(r"\Afoo\z"), Some(vec!["foo".to_owned()]));
+        assert_eq!(exact_null_values("^foo$"), None);
+        assert_eq!(exact_null_values(r"\A(?i:foo)\z"), None);
+        assert_eq!(exact_null_values(r"\Afo.*\z"), None);
+    }
+}