diff --git a/arrow-array/src/iterator.rs b/arrow-array/src/iterator.rs index e72b259ef049..c281231a2e79 100644 --- a/arrow-array/src/iterator.rs +++ b/arrow-array/src/iterator.rs @@ -56,7 +56,7 @@ impl ArrayIter { /// create a new iterator pub fn new(array: T) -> Self { let len = array.len(); - let logical_nulls = array.logical_nulls(); + let logical_nulls = array.logical_nulls().filter(|x| x.null_count() > 0); ArrayIter { array, logical_nulls, @@ -102,6 +102,38 @@ impl Iterator for ArrayIter { Some(self.current_end - self.current), ) } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + // Check if we can advance to the desired offset + match self.current.checked_add(n) { + // Yes, and still within bounds + Some(new_current) if new_current < self.current_end => { + self.current = new_current; + } + + // Either overflow or would exceed current_end + _ => { + self.current = self.current_end; + return None; + } + } + + self.next() + } + + #[inline] + fn last(mut self) -> Option { + self.next_back() + } + + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } } impl DoubleEndedIterator for ArrayIter { @@ -122,6 +154,25 @@ impl DoubleEndedIterator for ArrayIter { }) } } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + // Check if we advance to the one before the desired offset + match self.current_end.checked_sub(n) { + // Yes, and still within bounds + Some(new_offset) if self.current < new_offset => { + self.current_end = new_offset; + } + + // Either underflow or would exceed current + _ => { + self.current = self.current_end; + return None; + } + } + + self.next_back() + } } /// all arrays have known size.