1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
use crate::datatypes::ArrowPrimitiveType;
use super::{Array, PrimitiveArray, PrimitiveArrayOps};
#[derive(Debug)]
pub struct PrimitiveIter<'a, T: ArrowPrimitiveType> {
array: &'a PrimitiveArray<T>,
i: usize,
len: usize,
}
impl<'a, T: ArrowPrimitiveType> PrimitiveIter<'a, T> {
pub fn new(array: &'a PrimitiveArray<T>) -> Self {
PrimitiveIter::<T> {
array,
i: 0,
len: array.len(),
}
}
}
impl<'a, T: ArrowPrimitiveType> std::iter::Iterator for PrimitiveIter<'a, T> {
type Item = Option<T::Native>;
fn next(&mut self) -> Option<Self::Item> {
let i = self.i;
if i >= self.len {
None
} else if self.array.is_null(i) {
self.i += 1;
Some(None)
} else {
self.i += 1;
Some(Some(self.array.value(i)))
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
}
impl<'a, T: ArrowPrimitiveType> std::iter::ExactSizeIterator for PrimitiveIter<'a, T> {}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use crate::array::{ArrayRef, Int32Array};
#[test]
fn test_primitive_array_iter_round_trip() {
let array = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
let array = Arc::new(array) as ArrayRef;
let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
let result: Int32Array =
array.iter().map(|e| e.and_then(|e| Some(e + 1))).collect();
let expected = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
assert_eq!(result, expected);
}
}