Add support for reading CSV files with comments

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

Closes apache#10262.

datafusion-examples/examples/csv_opener.rs

@@ -49,6 +49,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

@@ -1570,12 +1570,13 @@ config_namespace! {
         pub escape: Option<u8>, default = None
         pub compression: CompressionTypeVariant, default = CompressionTypeVariant::UNCOMPRESSED
         pub schema_infer_max_rec: usize, default = 100
-        pub date_format: Option<String>,  default = None
-        pub datetime_format: Option<String>,  default = None
-        pub timestamp_format: Option<String>,  default = None
-        pub timestamp_tz_format: Option<String>,  default = None
-        pub time_format: Option<String>,  default = None
-        pub null_value: Option<String>,  default = None
+        pub date_format: Option<String>, default = None
+        pub datetime_format: Option<String>, default = None
+        pub timestamp_format: Option<String>, default = None
+        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

@@ -140,6 +140,12 @@ impl CsvFormat {
         self
     }
 
+    /// Lines beginning with this byte are ignored.
+    pub fn with_comment(mut self, comment: Option<u8>) -> Self {
+        self.options.comment = comment;
+        self
+    }
+
     /// True if the first line is a header.
     pub fn has_header(&self) -> bool {
         self.options.has_header
@@ -246,6 +252,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))
@@ -297,10 +304,14 @@ 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(self.options.has_header && first_chunk)
                 .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: Option<u8>) -> Self {
+        self.comment = 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());
@@ -635,6 +652,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -700,6 +718,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -763,6 +782,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(14, csv.base_config.file_schema.fields().len());
@@ -825,6 +845,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
         assert_eq!(13, csv.base_config.file_schema.fields().len());
@@ -919,6 +940,7 @@ mod tests {
             b',',
             b'"',
             None,
+            None,
             file_compression_type.to_owned(),
         );
 

datafusion/core/src/physical_optimizer/enforce_distribution.rs

@@ -1496,6 +1496,7 @@ pub(crate) mod tests {
             b',',
             b'"',
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }
@@ -1526,6 +1527,7 @@ pub(crate) mod tests {
             b',',
             b'"',
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }
@@ -3803,6 +3805,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 _
     })
@@ -1694,6 +1695,7 @@ mod tests {
             0,
             0,
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }
@@ -1720,6 +1722,7 @@ mod tests {
             0,
             0,
             None,
+            None,
             FileCompressionType::UNCOMPRESSED,
         ))
     }

datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -1508,6 +1508,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,
     )))
 }
@@ -297,6 +298,7 @@ pub fn csv_exec_sorted(
         0,
         0,
         None,
+        None,
         FileCompressionType::UNCOMPRESSED,
     ))
 }
@@ -359,6 +361,7 @@ pub fn csv_exec_ordered(
         0,
         b'"',
         None,
+        None,
         FileCompressionType::UNCOMPRESSED,
     ))
 }

datafusion/proto/proto/datafusion.proto

@@ -1118,6 +1118,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
@@ -1478,6 +1479,9 @@ message CsvScanExecNode {
   oneof optional_escape {
     string escape = 5;
   }
+  oneof optional_comment {
+    string comment = 6;
+  }
 }
 
 message AvroScanExecNode {