Improved inference and deserialization of CSV

Cargo.toml

@@ -102,7 +102,7 @@ full = [
 ]
 merge_sort = ["itertools"]
 io_csv = ["io_csv_read", "io_csv_write"]
-io_csv_read = ["csv", "lazy_static", "regex", "lexical-core"]
+io_csv_read = ["csv", "lexical-core"]
 io_csv_write = ["csv", "streaming-iterator", "lexical-core"]
 io_json = ["serde", "serde_json", "indexmap"]
 io_ipc = ["flatbuffers"]

src/io/csv/read/deserialize.rs

@@ -8,10 +8,12 @@ use crate::{
     datatypes::*,
     error::{ArrowError, Result},
     record_batch::RecordBatch,
-    temporal_conversions::EPOCH_DAYS_FROM_CE,
+    temporal_conversions,
     types::{NativeType, NaturalDataType},
 };
 
+use super::infer_schema::RFC3339;
+
 fn deserialize_primitive<T, F>(
     rows: &[ByteRecord],
     column: usize,
@@ -63,6 +65,22 @@ fn deserialize_binary<O: Offset>(rows: &[ByteRecord], column: usize) -> Arc<dyn
     Arc::new(BinaryArray::<O>::from_trusted_len_iter(iter))
 }
 
+#[inline]
+fn deserialize_datetime<T: chrono::TimeZone>(string: &str, tz: &T) -> Option<i64> {
+    let mut parsed = chrono::format::Parsed::new();
+    let fmt = chrono::format::StrftimeItems::new(RFC3339);
+    if chrono::format::parse(&mut parsed, string, fmt).is_ok() {
+        parsed
+            .to_datetime()
+            .map(|x| x.naive_utc())
+            .map(|x| tz.from_utc_datetime(&x))
+            .map(|x| x.timestamp_nanos())
+            .ok()
+    } else {
+        None
+    }
+}
+
 /// Deserializes `column` of `rows` into an [`Array`] of [`DataType`] `datatype`.
 pub fn deserialize_column(
     rows: &[ByteRecord],
@@ -115,7 +133,7 @@ pub fn deserialize_column(
             simdutf8::basic::from_utf8(bytes)
                 .ok()
                 .and_then(|x| x.parse::<chrono::NaiveDate>().ok())
-                .map(|x| x.num_days_from_ce() - EPOCH_DAYS_FROM_CE)
+                .map(|x| x.num_days_from_ce() - temporal_conversions::EPOCH_DAYS_FROM_CE)
         }),
         Date64 => deserialize_primitive(rows, column, datatype, |bytes| {
             simdutf8::basic::from_utf8(bytes)
@@ -139,20 +157,30 @@ pub fn deserialize_column(
                     .map(|x| x.timestamp_nanos() / 1000)
             })
         }
-        Timestamp(TimeUnit::Millisecond, None) => {
-            deserialize_primitive(rows, column, datatype, |bytes| {
-                simdutf8::basic::from_utf8(bytes)
-                    .ok()
-                    .and_then(|x| x.parse::<chrono::NaiveDateTime>().ok())
-                    .map(|x| x.timestamp_nanos() / 1_000_000)
-            })
-        }
-        Timestamp(TimeUnit::Second, None) => {
+        Timestamp(time_unit, None) => deserialize_primitive(rows, column, datatype, |bytes| {
+            simdutf8::basic::from_utf8(bytes)
+                .ok()
+                .and_then(|x| x.parse::<chrono::NaiveDateTime>().ok())
+                .map(|x| x.timestamp_nanos())
+                .map(|x| match time_unit {
+                    TimeUnit::Second => x / 1_000_000_000,
+                    TimeUnit::Millisecond => x / 1_000_000,
+                    TimeUnit::Microsecond => x / 1_000,
+                    TimeUnit::Nanosecond => x,
+                })
+        }),
+        Timestamp(time_unit, Some(ref tz)) => {
+            let tz = temporal_conversions::parse_offset(tz)?;
             deserialize_primitive(rows, column, datatype, |bytes| {
                 simdutf8::basic::from_utf8(bytes)
                     .ok()
-                    .and_then(|x| x.parse::<chrono::NaiveDateTime>().ok())
-                    .map(|x| x.timestamp_nanos() / 1_000_000_000)
+                    .and_then(|x| deserialize_datetime(x, &tz))
+                    .map(|x| match time_unit {
+                        TimeUnit::Second => x / 1_000_000_000,
+                        TimeUnit::Millisecond => x / 1_000_000,
+                        TimeUnit::Microsecond => x / 1_000,
+                        TimeUnit::Nanosecond => x,
+                    })
             })
         }
         Utf8 => deserialize_utf8::<i32>(rows, column),

src/io/csv/read/infer_schema.rs

@@ -3,20 +3,82 @@ use std::{
     io::{Read, Seek},
 };
 
-use super::Reader;
-use csv::StringRecord;
+use super::{ByteRecord, Reader};
 
-use crate::datatypes::DataType;
+use crate::datatypes::{DataType, TimeUnit};
 use crate::datatypes::{Field, Schema};
 use crate::error::Result;
 
-/// Infer the schema of a CSV file by reading through the first n records of the file,
-/// with `max_rows` controlling the maximum number of records to read.
-///
-/// If `max_rows` is not set, the whole file is read to infer its schema.
+pub(super) const RFC3339: &str = "%Y-%m-%dT%H:%M:%S%.f%:z";
+
+fn is_boolean(bytes: &[u8]) -> bool {
+    bytes.eq_ignore_ascii_case(b"true") | bytes.eq_ignore_ascii_case(b"false")
+}
+
+fn is_float(bytes: &[u8]) -> bool {
+    lexical_core::parse::<f64>(bytes).is_ok()
+}
+
+fn is_integer(bytes: &[u8]) -> bool {
+    lexical_core::parse::<i64>(bytes).is_ok()
+}
+
+fn is_date(string: &str) -> bool {
+    string.parse::<chrono::NaiveDate>().is_ok()
+}
+
+fn is_time(string: &str) -> bool {
+    string.parse::<chrono::NaiveTime>().is_ok()
+}
+
+fn is_naive_datetime(string: &str) -> bool {
+    string.parse::<chrono::NaiveDateTime>().is_ok()
+}
+
+fn is_datetime(string: &str) -> Option<String> {
+    let mut parsed = chrono::format::Parsed::new();
+    let fmt = chrono::format::StrftimeItems::new(RFC3339);
+    if chrono::format::parse(&mut parsed, string, fmt).is_ok() {
+        parsed.offset.map(|x| {
+            let hours = x / 60 / 60;
+            let minutes = x / 60 - hours * 60;
+            format!("{:03}:{:02}", hours, minutes)
+        })
+    } else {
+        None
+    }
+}
+
+/// Infers [`DataType`] from `bytes`
+pub fn infer(bytes: &[u8]) -> DataType {
+    if is_boolean(bytes) {
+        DataType::Boolean
+    } else if is_integer(bytes) {
+        DataType::Int64
+    } else if is_float(bytes) {
+        DataType::Float64
+    } else if let Ok(string) = simdutf8::basic::from_utf8(bytes) {
+        if is_date(string) {
+            DataType::Date32
+        } else if is_time(string) {
+            DataType::Time32(TimeUnit::Millisecond)
+        } else if is_naive_datetime(string) {
+            DataType::Timestamp(TimeUnit::Millisecond, None)
+        } else if let Some(offset) = is_datetime(string) {
+            DataType::Timestamp(TimeUnit::Millisecond, Some(offset))
+        } else {
+            DataType::Utf8
+        }
+    } else {
+        // invalid utf8
+        DataType::Binary
+    }
+}
+
+/// Infer the schema of a CSV file by reading through the first n records up to `max_rows`.
 ///
 /// Return infered schema and number of records used for inference.
-pub fn infer_schema<R: Read + Seek, F: Fn(&str) -> DataType>(
+pub fn infer_schema<R: Read + Seek, F: Fn(&[u8]) -> DataType>(
     reader: &mut Reader<R>,
     max_rows: Option<usize>,
     has_header: bool,
@@ -25,8 +87,7 @@ pub fn infer_schema<R: Read + Seek, F: Fn(&str) -> DataType>(
     // get or create header names
     // when has_header is false, creates default column names with column_ prefix
     let headers: Vec<String> = if has_header {
-        let headers = &reader.headers()?.clone();
-        headers.iter().map(|s| s.to_string()).collect()
+        reader.headers()?.iter().map(|s| s.to_string()).collect()
     } else {
         let first_record_count = &reader.headers()?.len();
         (0..*first_record_count)
@@ -42,12 +103,11 @@ pub fn infer_schema<R: Read + Seek, F: Fn(&str) -> DataType>(
     let mut column_types: Vec<HashSet<DataType>> = vec![HashSet::new(); header_length];
 
     let mut records_count = 0;
-    let mut fields = vec![];
 
-    let mut record = StringRecord::new();
+    let mut record = ByteRecord::new();
     let max_records = max_rows.unwrap_or(usize::MAX);
     while records_count < max_records {
-        if !reader.read_record(&mut record)? {
+        if !reader.read_byte_record(&mut record)? {
             break;
         }
         records_count += 1;
@@ -60,32 +120,30 @@ pub fn infer_schema<R: Read + Seek, F: Fn(&str) -> DataType>(
     }
 
     // build schema from inference results
-    for i in 0..header_length {
-        let possibilities = &column_types[i];
-        let field_name = &headers[i];
-
-        // determine data type based on possible types
-        // if there are incompatible types, use DataType::Utf8
-        match possibilities.len() {
-            1 => {
-                for dtype in possibilities.iter() {
-                    fields.push(Field::new(field_name, dtype.clone(), true));
-                }
-            }
-            2 => {
-                if possibilities.contains(&DataType::Int64)
-                    && possibilities.contains(&DataType::Float64)
-                {
-                    // we have an integer and double, fall down to double
-                    fields.push(Field::new(field_name, DataType::Float64, true));
-                } else {
-                    // default to Utf8 for conflicting datatypes (e.g bool and int)
-                    fields.push(Field::new(field_name, DataType::Utf8, true));
+    let fields = headers
+        .iter()
+        .zip(column_types.into_iter())
+        .map(|(field_name, mut possibilities)| {
+            // determine data type based on possible types
+            // if there are incompatible types, use DataType::Utf8
+            let data_type = match possibilities.len() {
+                1 => possibilities.drain().next().unwrap(),
+                2 => {
+                    if possibilities.contains(&DataType::Int64)
+                        && possibilities.contains(&DataType::Float64)
+                    {
+                        // we have an integer and double, fall down to double
+                        DataType::Float64
+                    } else {
+                        // default to Utf8 for conflicting datatypes (e.g bool and int)
+                        DataType::Utf8
+                    }
                 }
-            }
-            _ => fields.push(Field::new(field_name, DataType::Utf8, true)),
-        }
-    }
+                _ => DataType::Utf8,
+            };
+            Field::new(field_name, data_type, true)
+        })
+        .collect();
 
     // return the reader seek back to the start
     reader.seek(position)?;

src/io/csv/read/mod.rs

@@ -8,5 +8,5 @@ pub use csv::{ByteRecord, Reader, ReaderBuilder};
 mod infer_schema;
 
 pub use deserialize::{deserialize_batch, deserialize_column};
-pub use infer_schema::infer_schema;
+pub use infer_schema::{infer, infer_schema};
 pub use reader::*;

src/io/csv/read/reader.rs

@@ -1,8 +1,5 @@
 use std::io::Read;
 
-use lazy_static::lazy_static;
-use regex::{Regex, RegexBuilder};
-
 use super::{ByteRecord, Reader};
 
 use crate::{
@@ -52,37 +49,3 @@ pub fn read_rows<R: Read>(
     }
     Ok(row_number)
 }
-
-lazy_static! {
-    static ref DECIMAL_RE: Regex = Regex::new(r"^-?(\d+\.\d+)$").unwrap();
-    static ref INTEGER_RE: Regex = Regex::new(r"^-?(\d+)$").unwrap();
-    static ref BOOLEAN_RE: Regex = RegexBuilder::new(r"^(true)$|^(false)$")
-        .case_insensitive(true)
-        .build()
-        .unwrap();
-    static ref DATE_RE: Regex = Regex::new(r"^\d{4}-\d\d-\d\d$").unwrap();
-    static ref DATETIME_RE: Regex = Regex::new(r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d$").unwrap();
-}
-
-/// Infer the data type of a record
-pub fn infer(string: &str) -> DataType {
-    // when quoting is enabled in the reader, these quotes aren't escaped, we default to
-    // Utf8 for them
-    if string.starts_with('"') {
-        return DataType::Utf8;
-    }
-    // match regex in a particular order
-    if BOOLEAN_RE.is_match(string) {
-        DataType::Boolean
-    } else if DECIMAL_RE.is_match(string) {
-        DataType::Float64
-    } else if INTEGER_RE.is_match(string) {
-        DataType::Int64
-    } else if DATETIME_RE.is_match(string) {
-        DataType::Date64
-    } else if DATE_RE.is_match(string) {
-        DataType::Date32
-    } else {
-        DataType::Utf8
-    }
-}

tests/it/io/csv/read.rs

@@ -1,3 +1,5 @@
+use proptest::prelude::*;
+
 use std::io::Cursor;
 use std::sync::Arc;
 
@@ -175,7 +177,7 @@ fn float32() -> Result<()> {
 }
 
 #[test]
-fn binary() -> Result<()> {
+fn deserialize_binary() -> Result<()> {
     let input = vec!["aa", "bb"];
     let input = input.join("\n");
 
@@ -185,3 +187,41 @@ fn binary() -> Result<()> {
     assert_eq!(expected, result.as_ref());
     Ok(())
 }
+
+#[test]
+fn deserialize_timestamp() -> Result<()> {
+    let input = vec!["1996-12-19T16:34:57-02:00", "1996-12-19T16:34:58-02:00"];
+    let input = input.join("\n");
+
+    let data_type = DataType::Timestamp(TimeUnit::Millisecond, Some("-01:00".to_string()));
+
+    let expected = Int64Array::from([Some(851020497000), Some(851020498000)]).to(data_type.clone());
+
+    let result = test_deserialize(&input, data_type)?;
+    assert_eq!(expected, result.as_ref());
+    Ok(())
+}
+
+proptest! {
+    #[test]
+    #[cfg_attr(miri, ignore)] // miri and proptest do not work well :(
+    fn i64(v in any::<i64>()) {
+        assert_eq!(infer(v.to_string().as_bytes()), DataType::Int64);
+    }
+}
+
+proptest! {
+    #[test]
+    #[cfg_attr(miri, ignore)] // miri and proptest do not work well :(
+    fn utf8(v in "a.*") {
+        assert_eq!(infer(v.as_bytes()), DataType::Utf8);
+    }
+}
+
+proptest! {
+    #[test]
+    #[cfg_attr(miri, ignore)] // miri and proptest do not work well :(
+    fn dates(v in "1996-12-19T16:3[0-9]:57-02:00") {
+        assert_eq!(infer(v.as_bytes()), DataType::Timestamp(TimeUnit::Millisecond, Some("-02:00".to_string())));
+    }
+}