Skip to content

Commit

Permalink
Implement the Iterator trait for the json Reader. (#451)
Browse files Browse the repository at this point in the history
* Implement the Iterator trait for the json Reader.

* Use transpose.
  • Loading branch information
LaurentMazare authored and alamb committed Jun 22, 2021
1 parent b5e50ef commit c45a11a
Showing 1 changed file with 39 additions and 0 deletions.
39 changes: 39 additions & 0 deletions arrow/src/json/reader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1568,6 +1568,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::{
Expand Down Expand Up @@ -2945,4 +2953,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);
}
}

0 comments on commit c45a11a

Please sign in to comment.