[Rust] CSV reader performance & redesign [WIP]

dev/archery/archery/benchmark/compare.py

@@ -25,11 +25,11 @@ def items_per_seconds_fmt(value):
     if value < 1000:
         return "{} items/sec".format(value)
     if value < 1000**2:
-        return "{:.3f}k items/sec".format(value / 1000)
+        return "{:.3f}K items/sec".format(value / 1000)
     if value < 1000**3:
-        return "{:.3f}m items/sec".format(value / 1000**2)
+        return "{:.3f}M items/sec".format(value / 1000**2)
     else:
-        return "{:.3f}b items/sec".format(value / 1000**3)
+        return "{:.3f}G items/sec".format(value / 1000**3)
 
 
 def bytes_per_seconds_fmt(value):

rust/arrow/Cargo.toml

@@ -42,6 +42,7 @@ serde_json = { version = "1.0", features = ["preserve_order"] }
 indexmap = "1.6"
 rand = "0.7"
 csv = "1.1"
+csv-core = "0.1"
 num = "0.3"
 regex = "1.3"
 lazy_static = "1.4"

rust/arrow/examples/read_csv_infer_schema.rs

@@ -23,14 +23,14 @@ use arrow::util::pretty::print_batches;
 use std::fs::File;
 
 fn main() {
-    let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
-    let builder = csv::ReaderBuilder::new()
-        .has_header(true)
-        .infer_schema(Some(100));
-    let mut csv = builder.build(file).unwrap();
-    let _batch = csv.next().unwrap().unwrap();
-    #[cfg(feature = "prettyprint")]
-    {
-        print_batches(&[_batch]).unwrap();
-    }
+    // let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
+    // let builder = csv::ReaderBuilder::new()
+    //     .has_header(true)
+    //     .infer_schema(Some(100));
+    // let mut csv = builder.build(file).unwrap();
+    // let _batch = csv.next().unwrap().unwrap();
+    // #[cfg(feature = "prettyprint")]
+    // {
+    //     print_batches(&[_batch]).unwrap();
+    // }
 }

rust/arrow/src/compute/kernels/cast.rs

@@ -40,6 +40,7 @@ use std::sync::Arc;
 
 use crate::compute::kernels::arithmetic::{divide, multiply};
 use crate::compute::kernels::arity::unary;
+use crate::compute::kernels::cast_utils::string_to_timestamp_nanos;
 use crate::datatypes::*;
 use crate::error::{ArrowError, Result};
 use crate::{array::*, compute::take};
@@ -74,6 +75,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
 
         (Utf8, Date32) => true,
         (Utf8, Date64) => true,
+        (Utf8, Timestamp(TimeUnit::Nanosecond, None)) => true,
         (Utf8, _) => DataType::is_numeric(to_type),
         (_, Utf8) => DataType::is_numeric(from_type) || from_type == &Binary,
 
@@ -411,6 +413,22 @@ pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
                 }
                 Ok(Arc::new(builder.finish()) as ArrayRef)
             }
+            Timestamp(TimeUnit::Nanosecond, None) => {
+                let string_array = array.as_any().downcast_ref::<StringArray>().unwrap();
+                let mut builder =
+                    PrimitiveBuilder::<TimestampNanosecondType>::new(string_array.len());
+                for i in 0..string_array.len() {
+                    if string_array.is_null(i) {
+                        builder.append_null()?;
+                    } else {
+                        match string_to_timestamp_nanos(string_array.value(i)) {
+                            Ok(nanos) => builder.append_value(nanos)?,
+                            Err(_) => builder.append_null()?, // not a valid date
+                        };
+                    }
+                }
+                Ok(Arc::new(builder.finish()) as ArrayRef)
+            }
             _ => Err(ArrowError::ComputeError(format!(
                 "Casting from {:?} to {:?} not supported",
                 from_type, to_type,
@@ -1485,6 +1503,24 @@ mod tests {
         assert!(c.is_null(2));
     }
 
+    #[test]
+    fn test_cast_string_to_timestamp() {
+        let a = StringArray::from(vec![
+            Some("2020-09-08T12:00:00+00:00"),
+            Some("Not a valid date"),
+            None,
+        ]);
+        let array = Arc::new(a) as ArrayRef;
+        let b = cast(&array, &DataType::Timestamp(TimeUnit::Nanosecond, None)).unwrap();
+        let c = b
+            .as_any()
+            .downcast_ref::<TimestampNanosecondArray>()
+            .unwrap();
+        assert_eq!(1599566400000000000, c.value(0));
+        assert!(c.is_null(1));
+        assert!(c.is_null(2));
+    }
+
     #[test]
     fn test_cast_date32_to_int32() {
         let a = Date32Array::from(vec![10000, 17890]);