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

use crate::chunked_array::object::{ObjectArray, PolarsObject};
use arrow::array::Array;

/// An iterator that returns Some(T) or None, that can be used on any ObjectArray
// Note: This implementation is based on std's [Vec]s' [IntoIter].
#[derive(Debug)]
pub struct ObjectIter<'a, T: PolarsObject> {
    array: &'a ObjectArray<T>,
    current: usize,
    current_end: usize,
}

impl<'a, T: PolarsObject> ObjectIter<'a, T> {
    /// create a new iterator
    pub fn new(array: &'a ObjectArray<T>) -> Self {
        ObjectIter::<T> {
            array,
            current: 0,
            current_end: array.len(),
        }
    }
}

impl<'a, T: PolarsObject> std::iter::Iterator for ObjectIter<'a, T> {
    type Item = Option<&'a T>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.current == self.current_end {
            None
        } else if self.array.is_null(self.current) {
            self.current += 1;
            Some(None)
        } else {
            let old = self.current;
            self.current += 1;
            // Safety:
            // we just checked bounds in `self.current_end == self.current`
            // this is safe on the premise that this struct is initialized with
            // current = array.len()
            // and that current_end is ever only decremented
            unsafe { Some(Some(self.array.value_unchecked(old))) }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (
            self.array.len() - self.current,
            Some(self.array.len() - self.current),
        )
    }
}

impl<'a, T: PolarsObject> std::iter::DoubleEndedIterator for ObjectIter<'a, T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.current_end == self.current {
            None
        } else {
            self.current_end -= 1;
            Some(if self.array.is_null(self.current_end) {
                None
            } else {
                // Safety:
                // we just checked bounds in `self.current_end == self.current`
                // this is safe on the premise that this struct is initialized with
                // current = array.len()
                // and that current_end is ever only decremented
                unsafe { Some(self.array.value_unchecked(self.current_end)) }
            })
        }
    }
}

/// all arrays have known size.
impl<'a, T: PolarsObject> std::iter::ExactSizeIterator for ObjectIter<'a, T> {}

impl<'a, T: PolarsObject> IntoIterator for &'a ObjectArray<T> {
    type Item = Option<&'a T>;
    type IntoIter = ObjectIter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        ObjectIter::<'a, T>::new(self)
    }
}