You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ag...@apache.org on 2018/11/21 17:39:26 UTC
[arrow] branch master updated: ARROW-3726: [Rust] Add CSV reader
with example
This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new c04a62b ARROW-3726: [Rust] Add CSV reader with example
c04a62b is described below
commit c04a62b9579d420da32ed6d962d92266508e6abe
Author: Andy Grove <an...@gmail.com>
AuthorDate: Wed Nov 21 10:39:14 2018 -0700
ARROW-3726: [Rust] Add CSV reader with example
This adds a CSV reader and an example that accessed the loaded data through the use of downcasting arrays to specific types. The CSV reader supports all primitive types + string (`List<u8>`).
Author: Andy Grove <an...@gmail.com>
Closes #2992 from andygrove/ARROW-3726 and squashes the following commits:
4d1bf98 <Andy Grove> Exclude Rust test data csv files from rat
70140c6 <Andy Grove> re-export csv::Reader
4674651 <Andy Grove> add missing file after rename
3ab0a47 <Andy Grove> create module for csv::reader
26857a3 <Andy Grove> Update Windows CI to run new example
5d43b1f <Andy Grove> Remove Arc<>
ab1b20f <Andy Grove> use isEmpty() instead of len() == 0
2b9d9e1 <Andy Grove> Remove support for Float16
6167223 <Andy Grove> Remove support for Float16
80c44ca <Andy Grove> Use BinaryArray instead of List<u8>
f726559 <Andy Grove> Merge branch 'master' into ARROW-3726
3928d6d <Andy Grove> Add documentation, rename CsvFile to CsvReader
247092d <Andy Grove> cargo fmt
aae53aa <Andy Grove> add value_slice method, clean up example code
e539814 <Andy Grove> update example to print city names to demonstrate usage of List<u8>
8974c60 <Andy Grove> Example displays data
638159d <Andy Grove> fix test
9e88791 <Andy Grove> Update CI script
517da28 <Andy Grove> Implement csv reader
---
ci/rust-build-main.bat | 1 +
ci/travis_script_rust.sh | 1 +
dev/release/rat_exclude_files.txt | 1 +
rust/Cargo.toml | 1 +
rust/README.md | 1 +
rust/examples/read_csv.rs | 70 ++++++++++
rust/src/array.rs | 14 ++
rust/src/{lib.rs => csv/mod.rs} | 20 +--
rust/src/csv/reader.rs | 264 ++++++++++++++++++++++++++++++++++++++
rust/src/lib.rs | 2 +
rust/test/data/null_test.csv | 6 +
rust/test/data/uk_cities.csv | 37 ++++++
12 files changed, 400 insertions(+), 18 deletions(-)
diff --git a/ci/rust-build-main.bat b/ci/rust-build-main.bat
index ea040e0..463fdc0 100644
--- a/ci/rust-build-main.bat
+++ b/ci/rust-build-main.bat
@@ -55,5 +55,6 @@ cargo test --target %TARGET% --release
@echo ---------------------
cargo run --example builders --target %TARGET% --release
cargo run --example dynamic_types --target %TARGET% --release
+cargo run --example read_csv --target %TARGET% --release
popd
diff --git a/ci/travis_script_rust.sh b/ci/travis_script_rust.sh
index d0889e1..1058ec2 100755
--- a/ci/travis_script_rust.sh
+++ b/ci/travis_script_rust.sh
@@ -38,5 +38,6 @@ cargo test
cargo bench
cargo run --example builders
cargo run --example dynamic_types
+cargo run --example read_csv
popd
diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt
index 27c04ac..88608b3 100644
--- a/dev/release/rat_exclude_files.txt
+++ b/dev/release/rat_exclude_files.txt
@@ -160,3 +160,4 @@ r/README.md
r/README.Rmd
r/man/*.Rd
.gitattributes
+rust/test/data/*.csv
\ No newline at end of file
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index c74e7a9..d000445 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -38,6 +38,7 @@ bytes = "0.4"
libc = "0.2"
serde_json = "1.0.13"
rand = "0.5"
+csv = "1.0.0"
[dev-dependencies]
criterion = "0.2"
diff --git a/rust/README.md b/rust/README.md
index 131c7d9..51beb46 100644
--- a/rust/README.md
+++ b/rust/README.md
@@ -47,6 +47,7 @@ Examples can be run using the `cargo run --example` command. For example:
```bash
cargo run --example builders
cargo run --example dynamic_types
+cargo run --example read_csv
```
## Run Tests
diff --git a/rust/examples/read_csv.rs b/rust/examples/read_csv.rs
new file mode 100644
index 0000000..a12cafb
--- /dev/null
+++ b/rust/examples/read_csv.rs
@@ -0,0 +1,70 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+extern crate arrow;
+
+use arrow::array::{BinaryArray, PrimitiveArray};
+use arrow::csv;
+use arrow::datatypes::{DataType, Field, Schema};
+use std::fs::File;
+use std::sync::Arc;
+
+fn main() {
+ let schema = Schema::new(vec![
+ Field::new("city", DataType::Utf8, false),
+ Field::new("lat", DataType::Float64, false),
+ Field::new("lng", DataType::Float64, false),
+ ]);
+
+ let file = File::open("test/data/uk_cities.csv").unwrap();
+
+ let mut csv = csv::Reader::new(file, Arc::new(schema), false, 1024, None);
+ let batch = csv.next().unwrap().unwrap();
+
+ println!(
+ "Loaded {} rows containing {} columns",
+ batch.num_rows(),
+ batch.num_columns()
+ );
+
+ let city = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap();
+ let lat = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<f64>>()
+ .unwrap();
+ let lng = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<f64>>()
+ .unwrap();
+
+ for i in 0..batch.num_rows() {
+ let city_name: String = String::from_utf8(city.get_value(i).to_vec()).unwrap();
+
+ println!(
+ "City: {}, Latitude: {}, Longitude: {}",
+ city_name,
+ lat.value(i),
+ lng.value(i)
+ );
+ }
+}
diff --git a/rust/src/array.rs b/rust/src/array.rs
index ff3e785..ab44dc0 100644
--- a/rust/src/array.rs
+++ b/rust/src/array.rs
@@ -151,6 +151,11 @@ macro_rules! def_primitive_array {
self.data.buffers()[0].clone()
}
+ /// Returns the length of this array
+ pub fn len(&self) -> i64 {
+ self.data.len()
+ }
+
/// Returns a raw pointer to the values of this array.
pub fn raw_values(&self) -> *const $native_ty {
unsafe { mem::transmute(self.raw_values.get().offset(self.data.offset() as isize)) }
@@ -163,6 +168,15 @@ macro_rules! def_primitive_array {
unsafe { *(self.raw_values().offset(i as isize)) }
}
+ /// Returns a slice for the given offset and length
+ ///
+ /// Note this doesn't do any bound checking, for performance reason.
+ pub fn value_slice(&self, offset: i64, len: i64) -> &[$native_ty] {
+ let raw =
+ unsafe { std::slice::from_raw_parts(self.raw_values(), self.len() as usize) };
+ &raw[offset as usize..offset as usize + len as usize]
+ }
+
/// Returns the minimum value in the array, according to the natural order.
pub fn min(&self) -> Option<$native_ty> {
self.min_max_helper(|a, b| a < b)
diff --git a/rust/src/lib.rs b/rust/src/csv/mod.rs
similarity index 73%
copy from rust/src/lib.rs
copy to rust/src/csv/mod.rs
index 52f6763..9f2bd1d 100644
--- a/rust/src/lib.rs
+++ b/rust/src/csv/mod.rs
@@ -15,22 +15,6 @@
// specific language governing permissions and limitations
// under the License.
-extern crate bytes;
-extern crate libc;
+pub mod reader;
-#[macro_use]
-extern crate serde_json;
-
-extern crate rand;
-
-pub mod array;
-pub mod array_data;
-pub mod bitmap;
-pub mod buffer;
-pub mod builder;
-pub mod datatypes;
-pub mod error;
-pub mod memory;
-pub mod record_batch;
-pub mod tensor;
-pub mod util;
+pub use self::reader::Reader;
diff --git a/rust/src/csv/reader.rs b/rust/src/csv/reader.rs
new file mode 100644
index 0000000..0a8a6e4
--- /dev/null
+++ b/rust/src/csv/reader.rs
@@ -0,0 +1,264 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! CSV Reader
+//!
+//! This CSV reader allows CSV files to be read into the Arrow memory model. Records are loaded in
+//! batches and are then converted from row-based data to columnar data.
+//!
+//! Example:
+//!
+//! ```
+//! use arrow::csv;
+//! use arrow::datatypes::{DataType, Field, Schema};
+//! use std::fs::File;
+//! use std::sync::Arc;
+//!
+//! let schema = Schema::new(vec![
+//! Field::new("city", DataType::Utf8, false),
+//! Field::new("lat", DataType::Float64, false),
+//! Field::new("lng", DataType::Float64, false),
+//! ]);
+//!
+//! let file = File::open("test/data/uk_cities.csv").unwrap();
+//!
+//! let mut csv = csv::Reader::new(file, Arc::new(schema), false, 1024, None);
+//! let batch = csv.next().unwrap().unwrap();
+//!```
+
+use std::fs::File;
+use std::io::BufReader;
+use std::sync::Arc;
+
+use array::{ArrayRef, BinaryArray};
+use builder::{ArrayBuilder, ListArrayBuilder, PrimitiveArrayBuilder};
+use datatypes::{DataType, Schema};
+use error::{ArrowError, Result};
+use record_batch::RecordBatch;
+
+use csv_crate::{StringRecord, StringRecordsIntoIter};
+
+/// CSV file reader
+pub struct Reader {
+ /// Explicit schema for the CSV file
+ schema: Arc<Schema>,
+ /// Optional projection for which columns to load (zero-based column indices)
+ projection: Option<Vec<usize>>,
+ /// File reader
+ record_iter: StringRecordsIntoIter<BufReader<File>>,
+ /// Batch size (number of records to load each time)
+ batch_size: usize,
+}
+
+impl Reader {
+ /// Create a new CsvReader
+ pub fn new(
+ file: File,
+ schema: Arc<Schema>,
+ has_headers: bool,
+ batch_size: usize,
+ projection: Option<Vec<usize>>,
+ ) -> Self {
+ let csv_reader = csv::ReaderBuilder::new()
+ .has_headers(has_headers)
+ .from_reader(BufReader::new(file));
+
+ let record_iter = csv_reader.into_records();
+ Reader {
+ schema: schema.clone(),
+ projection,
+ record_iter,
+ batch_size,
+ }
+ }
+}
+
+macro_rules! build_primitive_array {
+ ($ROWS:expr, $COL_INDEX:expr, $TY:ty) => {{
+ let mut builder = PrimitiveArrayBuilder::<$TY>::new($ROWS.len() as i64);
+ for row_index in 0..$ROWS.len() {
+ match $ROWS[row_index].get(*$COL_INDEX) {
+ Some(s) if s.len() > 0 => builder.push(s.parse::<$TY>().unwrap()).unwrap(),
+ _ => builder.push_null().unwrap(),
+ }
+ }
+ Ok(Arc::new(builder.finish()) as ArrayRef)
+ }};
+}
+
+impl Reader {
+ /// Read the next batch of rows
+ pub fn next(&mut self) -> Option<Result<RecordBatch>> {
+ // read a batch of rows into memory
+ let mut rows: Vec<StringRecord> = Vec::with_capacity(self.batch_size);
+ for _ in 0..self.batch_size {
+ match self.record_iter.next() {
+ Some(Ok(r)) => {
+ rows.push(r);
+ }
+ Some(Err(_)) => {
+ return Some(Err(ArrowError::ParseError(
+ "Error reading CSV file".to_string(),
+ )));
+ }
+ None => break,
+ }
+ }
+
+ // return early if no data was loaded
+ if rows.is_empty() {
+ return None;
+ }
+
+ let projection: Vec<usize> = match self.projection {
+ Some(ref v) => v.clone(),
+ None => self
+ .schema
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(i, _)| i)
+ .collect(),
+ };
+
+ let arrays: Result<Vec<ArrayRef>> = projection
+ .iter()
+ .map(|i| {
+ let field = self.schema.field(*i);
+
+ match field.data_type() {
+ &DataType::Boolean => build_primitive_array!(rows, i, bool),
+ &DataType::Int8 => build_primitive_array!(rows, i, i8),
+ &DataType::Int16 => build_primitive_array!(rows, i, i16),
+ &DataType::Int32 => build_primitive_array!(rows, i, i32),
+ &DataType::Int64 => build_primitive_array!(rows, i, i64),
+ &DataType::UInt8 => build_primitive_array!(rows, i, u8),
+ &DataType::UInt16 => build_primitive_array!(rows, i, u16),
+ &DataType::UInt32 => build_primitive_array!(rows, i, u32),
+ &DataType::UInt64 => build_primitive_array!(rows, i, u64),
+ &DataType::Float32 => build_primitive_array!(rows, i, f32),
+ &DataType::Float64 => build_primitive_array!(rows, i, f64),
+ &DataType::Utf8 => {
+ let mut values_builder: PrimitiveArrayBuilder<u8> =
+ PrimitiveArrayBuilder::<u8>::new(rows.len() as i64);
+ let mut list_builder = ListArrayBuilder::new(values_builder);
+ for row_index in 0..rows.len() {
+ match rows[row_index].get(*i) {
+ Some(s) => {
+ list_builder.values().push_slice(s.as_bytes()).unwrap();
+ list_builder.append(true).unwrap();
+ }
+ _ => {
+ list_builder.append(false).unwrap();
+ }
+ }
+ }
+ Ok(Arc::new(BinaryArray::from(list_builder.finish())) as ArrayRef)
+ }
+ other => Err(ArrowError::ParseError(format!(
+ "Unsupported data type {:?}",
+ other
+ ))),
+ }
+ }).collect();
+
+ match arrays {
+ Ok(arr) => Some(Ok(RecordBatch::new(self.schema.clone(), arr))),
+ Err(e) => Some(Err(e)),
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+
+ use super::*;
+ use array::PrimitiveArray;
+ use datatypes::Field;
+
+ #[test]
+ fn test_csv() {
+ let schema = Schema::new(vec![
+ Field::new("city", DataType::Utf8, false),
+ Field::new("lat", DataType::Float64, false),
+ Field::new("lng", DataType::Float64, false),
+ ]);
+
+ let file = File::open("test/data/uk_cities.csv").unwrap();
+
+ let mut csv = Reader::new(file, Arc::new(schema), false, 1024, None);
+ let batch = csv.next().unwrap().unwrap();
+ assert_eq!(37, batch.num_rows());
+ assert_eq!(3, batch.num_columns());
+
+ // access data from a primitive array
+ let lat = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<f64>>()
+ .unwrap();
+ assert_eq!(57.653484, lat.value(0));
+
+ // access data from a string array (ListArray<u8>)
+ let city = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap();
+
+ let city_name: String = String::from_utf8(city.get_value(13).to_vec()).unwrap();
+
+ assert_eq!("Aberdeen, Aberdeen City, UK", city_name);
+ }
+
+ #[test]
+ fn test_csv_with_projection() {
+ let schema = Schema::new(vec![
+ Field::new("city", DataType::Utf8, false),
+ Field::new("lat", DataType::Float64, false),
+ Field::new("lng", DataType::Float64, false),
+ ]);
+
+ let file = File::open("test/data/uk_cities.csv").unwrap();
+
+ let mut csv = Reader::new(file, Arc::new(schema), false, 1024, Some(vec![0, 1]));
+ let batch = csv.next().unwrap().unwrap();
+ assert_eq!(37, batch.num_rows());
+ assert_eq!(2, batch.num_columns());
+ }
+
+ #[test]
+ fn test_nulls() {
+ let schema = Schema::new(vec![
+ Field::new("c_int", DataType::UInt64, false),
+ Field::new("c_float", DataType::Float32, false),
+ Field::new("c_string", DataType::Utf8, false),
+ ]);
+
+ let file = File::open("test/data/null_test.csv").unwrap();
+
+ let mut csv = Reader::new(file, Arc::new(schema), true, 1024, None);
+ let batch = csv.next().unwrap().unwrap();
+
+ assert_eq!(false, batch.column(1).is_null(0));
+ assert_eq!(false, batch.column(1).is_null(1));
+ assert_eq!(true, batch.column(1).is_null(2));
+ assert_eq!(false, batch.column(1).is_null(3));
+ assert_eq!(false, batch.column(1).is_null(4));
+ }
+
+}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 52f6763..859b0a1 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -16,6 +16,7 @@
// under the License.
extern crate bytes;
+extern crate csv as csv_crate;
extern crate libc;
#[macro_use]
@@ -28,6 +29,7 @@ pub mod array_data;
pub mod bitmap;
pub mod buffer;
pub mod builder;
+pub mod csv;
pub mod datatypes;
pub mod error;
pub mod memory;
diff --git a/rust/test/data/null_test.csv b/rust/test/data/null_test.csv
new file mode 100644
index 0000000..8083060
--- /dev/null
+++ b/rust/test/data/null_test.csv
@@ -0,0 +1,6 @@
+c_int,c_float,c_string,c_bool
+1,1.1,"1.11",true
+2,2.2,"2.22",true
+3,,"3.33",true
+4,4.4,,false
+5,6.6,"",false
\ No newline at end of file
diff --git a/rust/test/data/uk_cities.csv b/rust/test/data/uk_cities.csv
new file mode 100644
index 0000000..db9e6da
--- /dev/null
+++ b/rust/test/data/uk_cities.csv
@@ -0,0 +1,37 @@
+"Elgin, Scotland, the UK",57.653484,-3.335724
+"Stoke-on-Trent, Staffordshire, the UK",53.002666,-2.179404
+"Solihull, Birmingham, UK",52.412811,-1.778197
+"Cardiff, Cardiff county, UK",51.481583,-3.179090
+"Eastbourne, East Sussex, UK",50.768036,0.290472
+"Oxford, Oxfordshire, UK",51.752022,-1.257677
+"London, UK",51.509865,-0.118092
+"Swindon, Swindon, UK",51.568535,-1.772232
+"Gravesend, Kent, UK",51.441883,0.370759
+"Northampton, Northamptonshire, UK",52.240479,-0.902656
+"Rugby, Warwickshire, UK",52.370876,-1.265032
+"Sutton Coldfield, West Midlands, UK",52.570385,-1.824042
+"Harlow, Essex, UK",51.772938,0.102310
+"Aberdeen, Aberdeen City, UK",57.149651,-2.099075
+"Swansea, Swansea, UK",51.621441,-3.943646
+"Chesterfield, Derbyshire, UK",53.235046,-1.421629
+"Londonderry, Derry, UK",55.006763,-7.318268
+"Salisbury, Wiltshire, UK",51.068787,-1.794472
+"Weymouth, Dorset, UK",50.614429,-2.457621
+"Wolverhampton, West Midlands, UK",52.591370,-2.110748
+"Preston, Lancashire, UK",53.765762,-2.692337
+"Bournemouth, UK",50.720806,-1.904755
+"Doncaster, South Yorkshire, UK",53.522820,-1.128462
+"Ayr, South Ayrshire, UK",55.458565,-4.629179
+"Hastings, East Sussex, UK",50.854259,0.573453
+"Bedford, UK",52.136436,-0.460739
+"Basildon, Essex, UK",51.572376,0.470009
+"Chippenham, Wiltshire, UK",51.458057,-2.116074
+"Belfast, UK",54.607868,-5.926437
+"Uckfield, East Sussex, UK",50.967941,0.085831
+"Worthing, West Sussex, UK",50.825024,-0.383835
+"Leeds, West Yorkshire, UK",53.801277,-1.548567
+"Kendal, Cumbria, UK",54.328506,-2.743870
+"Plymouth, UK",50.376289,-4.143841
+"Haverhill, Suffolk, UK",52.080875,0.444517
+"Frankton, Warwickshire, UK",52.328415,-1.377561
+"Inverness, the UK",57.477772,-4.224721
\ No newline at end of file