[arrow-rs] branch master updated: Implement the Iterator trait for the json Reader. (#451)

[nevime]

This is an automated email from the ASF dual-hosted git repository.

nevime pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new e5cda31  Implement the Iterator trait for the json Reader. (#451)
e5cda31 is described below

commit e5cda312b697c3d610637b28c58b6f1b104b41cc
Author: Laurent Mazare <laurent.maz...@gmail.com>
AuthorDate: Sun Jun 13 08:22:38 2021 +0800

    Implement the Iterator trait for the json Reader. (#451)
    
    * Implement the Iterator trait for the json Reader.
    
    * Use transpose.
---
 arrow/src/json/reader.rs | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/arrow/src/json/reader.rs b/arrow/src/json/reader.rs
index d0b9c19..9235142 100644
--- a/arrow/src/json/reader.rs
+++ b/arrow/src/json/reader.rs
@@ -1569,6 +1569,14 @@ impl ReaderBuilder {
     }
 }
 
+impl<R: Read> Iterator for Reader<R> {
+    type Item = Result<RecordBatch>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next().transpose()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{
@@ -2946,4 +2954,35 @@ mod tests {
         assert_eq!(batch.num_columns(), 1);
         assert_eq!(batch.num_rows(), 3);
     }
+
+    #[test]
+    fn test_json_iterator() {
+        let builder = 
ReaderBuilder::new().infer_schema(None).with_batch_size(5);
+        let reader: Reader<File> = builder
+            .build::<File>(File::open("test/data/basic.json").unwrap())
+            .unwrap();
+        let schema = reader.schema();
+        let (col_a_index, _) = schema.column_with_name("a").unwrap();
+
+        let mut sum_num_rows = 0;
+        let mut num_batches = 0;
+        let mut sum_a = 0;
+        for batch in reader {
+            let batch = batch.unwrap();
+            assert_eq!(4, batch.num_columns());
+            sum_num_rows += batch.num_rows();
+            num_batches += 1;
+            let batch_schema = batch.schema();
+            assert_eq!(schema, batch_schema);
+            let a_array = batch
+                .column(col_a_index)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap();
+            sum_a += (0..a_array.len()).map(|i| a_array.value(i)).sum::<i64>();
+        }
+        assert_eq!(12, sum_num_rows);
+        assert_eq!(3, num_batches);
+        assert_eq!(100000000000011, sum_a);
+    }
 }