Rename "encoding" to "charset" for clarity, add in missing functions,

make schema inference charset-aware

Author: Rafferty97

datafusion/common/src/config.rs

@@ -2921,7 +2921,7 @@ config_namespace! {
         ///
         /// The default behaviour depends on the `datafusion.catalog.newlines_in_values` setting.
         pub newlines_in_values: Option<bool>, default = None
-        pub encoding: Option<String>, default = None
+        pub charset: Option<String>, default = None
         pub compression: CompressionTypeVariant, default = CompressionTypeVariant::UNCOMPRESSED
         /// Compression level for the output file. The valid range depends on the
         /// compression algorithm:
@@ -3034,6 +3034,13 @@ impl CsvOptions {
         self
     }
 
+    /// Specifies the character encoding the file is encoded with.
+    /// - defaults to UTF-8
+    pub fn with_charset(mut self, charset: impl Into<String>) -> Self {
+        self.charset = Some(charset.into());
+        self
+    }
+
     /// Set a `CompressionTypeVariant` of CSV
     /// - defaults to `CompressionTypeVariant::UNCOMPRESSED`
     pub fn with_file_compression_type(

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

@@ -1646,7 +1646,7 @@ mod tests {
 
         // Read the file
         let ctx = SessionContext::new();
-        let opts = CsvReadOptions::new().has_header(true);
+        let opts = CsvReadOptions::new().has_header(true).charset("SHIFT-JIS");
         let batches = ctx.read_csv(path, opts).await?.collect().await?;
 
         // Check

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

@@ -85,6 +85,8 @@ pub struct CsvReadOptions<'a> {
     pub file_extension: &'a str,
     /// Partition Columns
     pub table_partition_cols: Vec<(String, DataType)>,
+    /// Character encoding
+    pub charset: Option<&'a str>,
     /// File compression type
     pub file_compression_type: FileCompressionType,
     /// Indicates how the file is sorted
@@ -118,6 +120,7 @@ impl<'a> CsvReadOptions<'a> {
             newlines_in_values: false,
             file_extension: DEFAULT_CSV_EXTENSION,
             table_partition_cols: vec![],
+            charset: None,
             file_compression_type: FileCompressionType::UNCOMPRESSED,
             file_sort_order: vec![],
             comment: None,
@@ -209,6 +212,12 @@ impl<'a> CsvReadOptions<'a> {
         self
     }
 
+    /// Configure the character set encoding
+    pub fn charset(mut self, charset: &'a str) -> Self {
+        self.charset = Some(charset);
+        self
+    }
+
     /// Configure file compression type
     pub fn file_compression_type(
         mut self,
@@ -633,6 +642,7 @@ impl ReadOptions<'_> for CsvReadOptions<'_> {
             .with_terminator(self.terminator)
             .with_newlines_in_values(self.newlines_in_values)
             .with_schema_infer_max_rec(self.schema_infer_max_records)
+            .with_charset(self.charset.map(ToOwned::to_owned))
             .with_file_compression_type(self.file_compression_type.to_owned())
             .with_null_regex(self.null_regex.clone())
             .with_truncated_rows(self.truncated_rows);

datafusion/datasource-csv/src/encoding.rs datafusion/datasource-csv/src/charset.rsdatafusion/datasource-csv/src/encoding.rs renamed to datafusion/datasource-csv/src/charset.rs

@@ -16,18 +16,31 @@
 // under the License.
 
 use std::fmt::Debug;
+use std::io::{BufRead, Read};
 
 use arrow::array::RecordBatch;
 use arrow::error::ArrowError;
-use datafusion_common::Result;
+use datafusion_common::{DataFusionError, Result};
 use datafusion_datasource::decoder::Decoder;
-use encoding_rs::{CoderResult, Encoding};
+use encoding_rs::{CoderResult, Encoding, UTF_8};
 
 use self::buffer::Buffer;
 
 /// Default capacity of the buffer used to decode non-UTF-8 charset streams
 static DECODE_BUFFER_CAP: usize = 8 * 1024;
 
+pub fn lookup_charset(enc: Option<&str>) -> Result<Option<&'static Encoding>> {
+    match enc {
+        Some(enc) => match Encoding::for_label(enc.as_bytes()) {
+            Some(enc) => Ok(Some(enc).filter(|enc| *enc != UTF_8)),
+            None => Err(DataFusionError::Configuration(format!(
+                "Unknown character set '{enc}'"
+            )))?,
+        },
+        None => Ok(None),
+    }
+}
+
 /// A `Decoder` that decodes input bytes from the specified character encoding
 /// to UTF-8 before passing them onto the inner `Decoder`.
 pub struct CharsetDecoder<T> {
@@ -100,6 +113,54 @@ impl<T: Debug> Debug for CharsetDecoder<T> {
     }
 }
 
+pub struct CharsetReader<R> {
+    inner: R,
+    charset_decoder: encoding_rs::Decoder,
+    buffer: Buffer,
+}
+
+impl<R: BufRead> CharsetReader<R> {
+    pub fn new(inner: R, encoding: &'static Encoding) -> Self {
+        Self {
+            inner,
+            charset_decoder: encoding.new_decoder(),
+            buffer: Buffer::with_capacity(DECODE_BUFFER_CAP),
+        }
+    }
+}
+
+impl<R: BufRead> Read for CharsetReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let src = self.fill_buf()?;
+        let len = src.len().min(buf.len());
+        buf[..len].copy_from_slice(&src[..len]);
+        Ok(len)
+    }
+}
+
+impl<R: BufRead> BufRead for CharsetReader<R> {
+    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
+        if self.buffer.is_empty() {
+            self.buffer.backshift();
+
+            let buf = self.inner.fill_buf()?;
+            let (_, read, written, _) = self.charset_decoder.decode_to_utf8(
+                buf,
+                self.buffer.write_buf(),
+                buf.is_empty(),
+            );
+            self.inner.consume(read);
+            self.buffer.advance(written);
+        }
+
+        Ok(self.buffer.read_buf())
+    }
+
+    fn consume(&mut self, amount: usize) {
+        self.buffer.consume(amount);
+    }
+}
+
 mod buffer {
     /// A fixed-sized buffer that maintains both
     /// a read position and a write position

datafusion/datasource-csv/src/file_format.rs

@@ -22,6 +22,7 @@ use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Debug};
 use std::sync::Arc;
 
+use crate::charset::lookup_charset;
 use crate::source::CsvSource;
 
 use arrow::array::RecordBatch;
@@ -294,6 +295,13 @@ impl CsvFormat {
         self
     }
 
+    /// Sets the character encoding of the CSV.
+    /// Defaults to UTF-8 if unspecified.
+    pub fn with_charset(mut self, charset: Option<String>) -> Self {
+        self.options.charset = charset;
+        self
+    }
+
     /// Set a `FileCompressionType` of CSV
     /// - defaults to `FileCompressionType::UNCOMPRESSED`
     pub fn with_file_compression_type(
@@ -540,6 +548,8 @@ impl CsvFormat {
 
         pin_mut!(stream);
 
+        let charset = lookup_charset(self.options.charset.as_deref())?;
+
         while let Some(chunk) = stream.next().await.transpose()? {
             record_number += 1;
             let first_chunk = record_number == 0;
@@ -569,8 +579,15 @@ impl CsvFormat {
                 format = format.with_comment(comment);
             }
 
-            let (Schema { fields, .. }, records_read) =
-                format.infer_schema(chunk.reader(), Some(records_to_read))?;
+            let (Schema { fields, .. }, records_read) = match charset {
+                #[cfg(feature = "encoding_rs")]
+                Some(enc) => {
+                    use crate::charset::CharsetReader;
+                    let reader = CharsetReader::new(chunk.reader(), enc);
+                    format.infer_schema(reader, Some(records_to_read))?
+                }
+                None => format.infer_schema(chunk.reader(), Some(records_to_read))?,
+            };
 
             records_to_read -= records_read;
             total_records_read += records_read;

datafusion/datasource-csv/src/mod.rs

@@ -21,7 +21,7 @@
 #![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))]
 
 #[cfg(feature = "encoding_rs")]
-mod encoding;
+mod charset;
 pub mod file_format;
 pub mod source;
 
@@ -32,6 +32,7 @@ use datafusion_datasource::file_groups::FileGroup;
 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::*;
 
 /// Returns a [`FileScanConfig`] for given `file_groups`
@@ -45,3 +46,17 @@ pub fn partitioned_csv_config(
             .build(),
     )
 }
+
+#[cfg(not(feature = "encoding_rs"))]
+mod encoding {
+    use datafusion_common::{DataFusionError, Result};
+
+    pub fn find_encoding(enc: Option<&str>) -> Result<Option<core::convert::Infallible>> {
+        match enc {
+            Some(_) => Err(DataFusionError::NotImplemented(format!(
+                "The 'encoding_rs' feature must be enabled to decode non-UTF-8 encodings"
+            )))?,
+            None => Ok(None),
+        }
+    }
+}

datafusion/datasource-csv/src/source.rs

@@ -46,8 +46,7 @@ use datafusion_physical_plan::{
     DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
 };
 
-#[cfg(feature = "encoding_rs")]
-use crate::encoding::CharsetDecoder;
+use crate::charset::lookup_charset;
 use crate::file_format::CsvDecoder;
 use futures::{StreamExt, TryStreamExt};
 use object_store::buffered::BufWriter;
@@ -232,30 +231,6 @@ impl CsvOpener {
             partition_index: 0,
         }
     }
-
-    #[cfg(feature = "encoding_rs")]
-    fn encoding(&self) -> Result<Option<&'static encoding_rs::Encoding>> {
-        match self.config.options.encoding.as_ref() {
-            Some(enc) => match encoding_rs::Encoding::for_label(enc.as_bytes()) {
-                Some(enc) => Ok(Some(enc)),
-                None => Err(DataFusionError::Configuration(format!(
-                    "Unknown character set '{enc}'"
-                )))?,
-            },
-            None => Ok(None),
-        }
-    }
-
-    #[cfg(not(feature = "encoding_rs"))]
-    fn encoding(&self) -> Result<Option<core::convert::Infallible>> {
-        match &self.config.options.encoding {
-            Some(_) => Err(DataFusionError::NotImplemented(
-                "The 'encoding_rs' feature must be enabled to decode non-UTF-8 encodings"
-                    .to_owned(),
-            ))?,
-            None => Ok(None),
-        }
-    }
 }
 
 impl From<CsvSource> for Arc<dyn FileSource> {
@@ -379,7 +354,7 @@ impl FileOpener for CsvOpener {
         config.options.truncated_rows = Some(config.truncate_rows());
 
         let file_compression_type = self.file_compression_type.to_owned();
-        let encoding = self.encoding()?;
+        let charset = lookup_charset(self.config.options.charset.as_deref())?;
 
         if partitioned_file.range.is_some() {
             assert!(
@@ -437,9 +412,10 @@ impl FileOpener for CsvOpener {
 
                     let reader = BufReader::new(reader);
 
-                    let mut reader = match encoding {
+                    let mut reader = match charset {
                         #[cfg(feature = "encoding_rs")]
                         Some(enc) => {
+                            use crate::charset::CharsetDecoder;
                             let decoder = CharsetDecoder::new(decoder, enc);
                             deserialize_reader(reader, decoder)
                         }
@@ -461,9 +437,10 @@ impl FileOpener for CsvOpener {
 
                     let stream = file_compression_type.convert_stream(stream)?.fuse();
 
-                    let stream = match encoding {
+                    let stream = match charset {
                         #[cfg(feature = "encoding_rs")]
                         Some(enc) => {
+                            use crate::charset::CharsetDecoder;
                             let decoder = CharsetDecoder::new(decoder, enc);
                             deserialize_stream(stream, DecoderDeserializer::new(decoder))
                         }

datafusion/proto-common/proto/datafusion_common.proto

@@ -476,7 +476,7 @@ message CsvOptions {
   bytes terminator = 17; // Optional terminator character as a byte
   bytes truncated_rows = 18; // Indicates if truncated rows are allowed
   optional uint32 compression_level = 19; // Optional compression level
-  string encoding = 20; // Optional character encoding
+  string charset = 20; // Optional character encoding
 }
 
 // Options controlling CSV format

datafusion/proto-common/src/from_proto/mod.rs

@@ -984,7 +984,7 @@ impl TryFrom<&protobuf::CsvOptions> for CsvOptions {
             escape: proto_opts.escape.first().copied(),
             double_quote: proto_opts.double_quote.first().map(|h| *h != 0),
             newlines_in_values: proto_opts.newlines_in_values.first().map(|h| *h != 0),
-            encoding: (!proto_opts.encoding.is_empty())
+            charset: (!proto_opts.encoding.is_empty())
                 .then(|| proto_opts.encoding.clone()),
             compression: proto_opts.compression().into(),
             compression_level: proto_opts.compression_level,

datafusion/proto-common/src/to_proto/mod.rs

@@ -986,7 +986,7 @@ impl TryFrom<&CsvOptions> for protobuf::CsvOptions {
             newlines_in_values: opts
                 .newlines_in_values
                 .map_or_else(Vec::new, |h| vec![h as u8]),
-            encoding: opts.encoding.clone().unwrap_or_default(),
+            encoding: opts.charset.clone().unwrap_or_default(),
             compression: compression.into(),
             schema_infer_max_rec: opts.schema_infer_max_rec.map(|h| h as u64),
             date_format: opts.date_format.clone().unwrap_or_default(),