Add support for reading CSV files with comments (#10467)

This patch adds support for parsing CSV files containing comment lines.

Closes #10262.

Author: bbannier

datafusion-examples/examples/csv_opener.rs

@@ -48,6 +48,7 @@ async fn main() -> Result<()> {
         b',',
         b'"',
         object_store,
+        Some(b'#'),
     );
 
     let opener = CsvOpener::new(Arc::new(config), FileCompressionType::UNCOMPRESSED);

datafusion/common/src/config.rs

@@ -1567,6 +1567,7 @@ config_namespace! {
         pub timestamp_tz_format: Option<String>, default = None
         pub time_format: Option<String>, default = None
         pub null_value: Option<String>, default = None
+        pub comment: Option<u8>, default = None
     }
 }
 

datafusion/core/src/datasource/file_format/csv.rs

@@ -147,6 +147,12 @@ impl CsvFormat {
         self.options.has_header
     }
 
+    /// Lines beginning with this byte are ignored.
+    pub fn with_comment(mut self, comment: Option<u8>) -> Self {
+        self.options.comment = comment;
+        self
+    }
+
     /// The character separating values within a row.
     /// - default to ','
     pub fn with_delimiter(mut self, delimiter: u8) -> Self {
@@ -252,6 +258,7 @@ impl FileFormat for CsvFormat {
             self.options.delimiter,
             self.options.quote,
             self.options.escape,
+            self.options.comment,
             self.options.compression.into(),
         );
         Ok(Arc::new(exec))
@@ -300,7 +307,7 @@ impl CsvFormat {
         pin_mut!(stream);
 
         while let Some(chunk) = stream.next().await.transpose()? {
-            let format = arrow::csv::reader::Format::default()
+            let mut format = arrow::csv::reader::Format::default()
                 .with_header(
                     first_chunk
                         && self
@@ -310,6 +317,10 @@ impl CsvFormat {
                 )
                 .with_delimiter(self.options.delimiter);
 
+            if let Some(comment) = self.options.comment {
+                format = format.with_comment(comment);
+            }
+
             let (Schema { fields, .. }, records_read) =
                 format.infer_schema(chunk.reader(), Some(records_to_read))?;
 

datafusion/core/src/datasource/file_format/options.rs

@@ -61,6 +61,8 @@ pub struct CsvReadOptions<'a> {
     pub quote: u8,
     /// An optional escape character. Defaults to None.
     pub escape: Option<u8>,
+    /// If enabled, lines beginning with this byte are ignored.
+    pub comment: Option<u8>,
     /// An optional schema representing the CSV files. If None, CSV reader will try to infer it
     /// based on data in file.
     pub schema: Option<&'a Schema>,
@@ -97,6 +99,7 @@ impl<'a> CsvReadOptions<'a> {
             table_partition_cols: vec![],
             file_compression_type: FileCompressionType::UNCOMPRESSED,
             file_sort_order: vec![],
+            comment: None,
         }
     }
 
@@ -106,6 +109,12 @@ impl<'a> CsvReadOptions<'a> {
         self
     }
 
+    /// Specify comment char to use for CSV read
+    pub fn comment(mut self, comment: u8) -> Self {
+        self.comment = Some(comment);
+        self
+    }
+
     /// Specify delimiter to use for CSV read
     pub fn delimiter(mut self, delimiter: u8) -> Self {
         self.delimiter = delimiter;
@@ -477,6 +486,7 @@ impl ReadOptions<'_> for CsvReadOptions<'_> {
         let file_format = CsvFormat::default()
             .with_options(table_options.csv)
             .with_has_header(self.has_header)
+            .with_comment(self.comment)
             .with_delimiter(self.delimiter)
             .with_quote(self.quote)
             .with_escape(self.escape)

datafusion/core/src/datasource/physical_plan/csv.rs

@@ -58,6 +58,7 @@ pub struct CsvExec {
     delimiter: u8,
     quote: u8,
     escape: Option<u8>,
+    comment: Option<u8>,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
     /// Compression type of the file associated with CsvExec
@@ -73,6 +74,7 @@ impl CsvExec {
         delimiter: u8,
         quote: u8,
         escape: Option<u8>,
+        comment: Option<u8>,
         file_compression_type: FileCompressionType,
     ) -> Self {
         let (projected_schema, projected_statistics, projected_output_ordering) =
@@ -92,6 +94,7 @@ impl CsvExec {
             metrics: ExecutionPlanMetricsSet::new(),
             file_compression_type,
             cache,
+            comment,
         }
     }
 
@@ -113,6 +116,11 @@ impl CsvExec {
         self.quote
     }
 
+    /// Lines beginning with this byte are ignored.
+    pub fn comment(&self) -> Option<u8> {
+        self.comment
+    }
+
     /// The escape character
     pub fn escape(&self) -> Option<u8> {
         self.escape
@@ -234,6 +242,7 @@ impl ExecutionPlan for CsvExec {
             quote: self.quote,
             escape: self.escape,
             object_store,
+            comment: self.comment,
         });
 
         let opener = CsvOpener {
@@ -265,9 +274,11 @@ pub struct CsvConfig {
     quote: u8,
     escape: Option<u8>,
     object_store: Arc<dyn ObjectStore>,
+    comment: Option<u8>,
 }
 
 impl CsvConfig {
+    #[allow(clippy::too_many_arguments)]
     /// Returns a [`CsvConfig`]
     pub fn new(
         batch_size: usize,
@@ -277,6 +288,7 @@ impl CsvConfig {
         delimiter: u8,
         quote: u8,
         object_store: Arc<dyn ObjectStore>,
+        comment: Option<u8>,
     ) -> Self {
         Self {
             batch_size,
@@ -287,6 +299,7 @@ impl CsvConfig {
             quote,
             escape: None,
             object_store,
+            comment,
         }
     }
 }
@@ -309,6 +322,9 @@ impl CsvConfig {
         if let Some(escape) = self.escape {
             builder = builder.with_escape(escape)
         }
+        if let Some(comment) = self.comment {
+            builder = builder.with_comment(comment);
+        }
 
         builder
     }
@@ -570,6 +586,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -636,6 +653,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -702,6 +720,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -765,6 +784,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(14, csv.base_config.file_schema.fields().len());
@@ -827,6 +847,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -921,6 +942,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
 

datafusion/core/src/physical_optimizer/enforce_distribution.rs

@@ -1471,6 +1471,7 @@ pub(crate) mod tests {
             b',',
             b'"',
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }
@@ -1494,6 +1495,7 @@ pub(crate) mod tests {
             b',',
             b'"',
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }
@@ -3767,6 +3769,7 @@ pub(crate) mod tests {
                     b',',
                     b'"',
                     None,
+                    None,
                     compression_type,
                 )),
                 vec![("a".to_string(), "a".to_string())],

datafusion/core/src/physical_optimizer/projection_pushdown.rs

@@ -185,6 +185,7 @@ fn try_swapping_with_csv(
             csv.delimiter(),
             csv.quote(),
             csv.escape(),
+            csv.comment(),
             csv.file_compression_type,
         )) as _
     })
@@ -1686,6 +1687,7 @@ mod tests {
             0,
             0,
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }
@@ -1708,6 +1710,7 @@ mod tests {
             0,
             0,
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }

datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -1502,6 +1502,7 @@ mod tests {
             0,
             b'"',
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }

datafusion/core/src/test/mod.rs

@@ -98,6 +98,7 @@ pub fn scan_partitioned_csv(partitions: usize, work_dir: &Path) -> Result<Arc<Cs
         b',',
         b'"',
         None,
+        None,
         FileCompressionType::UNCOMPRESSED,
     )))
 }
@@ -282,6 +283,7 @@ pub fn csv_exec_sorted(
         0,
         0,
         None,
+        None,
         FileCompressionType::UNCOMPRESSED,
     ))
 }
@@ -337,6 +339,7 @@ pub fn csv_exec_ordered(
         0,
         b'"',
         None,
+        None,
         FileCompressionType::UNCOMPRESSED,
     ))
 }

datafusion/proto-common/proto/datafusion_common.proto

@@ -397,6 +397,7 @@ message CsvOptions {
   string timestamp_tz_format = 10; // Optional timestamp with timezone format
   string time_format = 11; // Optional time format
   string null_value = 12; // Optional representation of null value
+  bytes comment = 13; // Optional comment character as a byte
 }
 
 // Options controlling CSV format