arrow_array/
iterator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Idiomatic iterators for [`Array`](crate::Array)
6
7use crate::array::{
8    ArrayAccessor, BooleanArray, FixedSizeBinaryArray, GenericBinaryArray, GenericListArray,
9    GenericStringArray, PrimitiveArray,
10};
11use crate::{FixedSizeListArray, GenericListViewArray, MapArray};
12use arrow_buffer::NullBuffer;
13
14/// An iterator that returns Some(T) or None, that can be used on any [`ArrayAccessor`]
15///
16/// # Performance
17///
18/// [`ArrayIter`] provides an idiomatic way to iterate over an array, however, this
19/// comes at the cost of performance. In particular the interleaved handling of
20/// the null mask is often sub-optimal.
21///
22/// If performing an infallible operation, it is typically faster to perform the operation
23/// on every index of the array, and handle the null mask separately. For [`PrimitiveArray`]
24/// this functionality is provided by [`compute::unary`]
25///
26/// If performing a fallible operation, it isn't possible to perform the operation independently
27/// of the null mask, as this might result in a spurious failure on a null index. However,
28/// there are more efficient ways to iterate over just the non-null indices, this functionality
29/// is provided by [`compute::try_unary`]
30///
31/// [`PrimitiveArray`]: crate::PrimitiveArray
32/// [`compute::unary`]: https://docs.rs/arrow/latest/arrow/compute/fn.unary.html
33/// [`compute::try_unary`]: https://docs.rs/arrow/latest/arrow/compute/fn.try_unary.html
34#[derive(Debug, Clone)]
35pub struct ArrayIter<T: ArrayAccessor> {
36    array: T,
37    logical_nulls: Option<NullBuffer>,
38    current: usize,
39    current_end: usize,
40}
41
42impl<T: ArrayAccessor> ArrayIter<T> {
43    /// create a new iterator
44    pub fn new(array: T) -> Self {
45        let len = array.len();
46        let logical_nulls = array.logical_nulls().filter(|x| x.null_count() > 0);
47        ArrayIter {
48            array,
49            logical_nulls,
50            current: 0,
51            current_end: len,
52        }
53    }
54
55    #[inline]
56    fn is_null(&self, idx: usize) -> bool {
57        self.logical_nulls.as_ref().is_some_and(|x| x.is_null(idx))
58    }
59}
60
61impl<T: ArrayAccessor> Iterator for ArrayIter<T> {
62    type Item = Option<T::Item>;
63
64    #[inline]
65    fn next(&mut self) -> Option<Self::Item> {
66        if self.current == self.current_end {
67            None
68        } else if self.is_null(self.current) {
69            self.current += 1;
70            Some(None)
71        } else {
72            let old = self.current;
73            self.current += 1;
74            // Safety:
75            // we just checked bounds in `self.current_end == self.current`
76            // this is safe on the premise that this struct is initialized with
77            // current = array.len()
78            // and that current_end is ever only decremented
79            unsafe { Some(Some(self.array.value_unchecked(old))) }
80        }
81    }
82
83    fn size_hint(&self) -> (usize, Option<usize>) {
84        (
85            self.current_end - self.current,
86            Some(self.current_end - self.current),
87        )
88    }
89
90    #[inline]
91    fn nth(&mut self, n: usize) -> Option<Self::Item> {
92        // Check if we can advance to the desired offset
93        match self.current.checked_add(n) {
94            // Yes, and still within bounds
95            Some(new_current) if new_current < self.current_end => {
96                self.current = new_current;
97            }
98
99            // Either overflow or would exceed current_end
100            _ => {
101                self.current = self.current_end;
102                return None;
103            }
104        }
105
106        self.next()
107    }
108
109    #[inline]
110    fn last(mut self) -> Option<Self::Item> {
111        self.next_back()
112    }
113
114    #[inline]
115    fn count(self) -> usize
116    where
117        Self: Sized,
118    {
119        self.len()
120    }
121}
122
123impl<T: ArrayAccessor> DoubleEndedIterator for ArrayIter<T> {
124    fn next_back(&mut self) -> Option<Self::Item> {
125        if self.current_end == self.current {
126            None
127        } else {
128            self.current_end -= 1;
129            Some(if self.is_null(self.current_end) {
130                None
131            } else {
132                // Safety:
133                // we just checked bounds in `self.current_end == self.current`
134                // this is safe on the premise that this struct is initialized with
135                // current = array.len()
136                // and that current_end is ever only decremented
137                unsafe { Some(self.array.value_unchecked(self.current_end)) }
138            })
139        }
140    }
141
142    #[inline]
143    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
144        // Check if we advance to the one before the desired offset
145        match self.current_end.checked_sub(n) {
146            // Yes, and still within bounds
147            Some(new_offset) if self.current < new_offset => {
148                self.current_end = new_offset;
149            }
150
151            // Either underflow or would exceed current
152            _ => {
153                self.current = self.current_end;
154                return None;
155            }
156        }
157
158        self.next_back()
159    }
160}
161
162/// all arrays have known size.
163impl<T: ArrayAccessor> ExactSizeIterator for ArrayIter<T> {}
164
165/// an iterator that returns Some(T) or None, that can be used on any PrimitiveArray
166pub type PrimitiveIter<'a, T> = ArrayIter<&'a PrimitiveArray<T>>;
167/// an iterator that returns Some(T) or None, that can be used on any BooleanArray
168pub type BooleanIter<'a> = ArrayIter<&'a BooleanArray>;
169/// an iterator that returns Some(T) or None, that can be used on any Utf8Array
170pub type GenericStringIter<'a, T> = ArrayIter<&'a GenericStringArray<T>>;
171/// an iterator that returns Some(T) or None, that can be used on any BinaryArray
172pub type GenericBinaryIter<'a, T> = ArrayIter<&'a GenericBinaryArray<T>>;
173/// an iterator that returns Some(T) or None, that can be used on any FixedSizeBinaryArray
174pub type FixedSizeBinaryIter<'a> = ArrayIter<&'a FixedSizeBinaryArray>;
175/// an iterator that returns Some(T) or None, that can be used on any FixedSizeListArray
176pub type FixedSizeListIter<'a> = ArrayIter<&'a FixedSizeListArray>;
177/// an iterator that returns Some(T) or None, that can be used on any ListArray
178pub type GenericListArrayIter<'a, O> = ArrayIter<&'a GenericListArray<O>>;
179/// an iterator that returns Some(T) or None, that can be used on any MapArray
180pub type MapArrayIter<'a> = ArrayIter<&'a MapArray>;
181/// an iterator that returns Some(T) or None, that can be used on any ListArray
182pub type GenericListViewArrayIter<'a, O> = ArrayIter<&'a GenericListViewArray<O>>;
183#[cfg(test)]
184mod tests {
185    use crate::array::{ArrayRef, BinaryArray, BooleanArray, Int32Array, StringArray};
186    use crate::iterator::ArrayIter;
187    use rand::rngs::StdRng;
188    use rand::{RngExt, SeedableRng};
189    use std::fmt::Debug;
190    use std::sync::Arc;
191
192    #[test]
193    fn test_primitive_array_iter_round_trip() {
194        let array = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
195        let array = Arc::new(array) as ArrayRef;
196
197        let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
198
199        // to and from iter, with a +1
200        let result: Int32Array = array.iter().map(|e| e.map(|e| e + 1)).collect();
201
202        let expected = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
203        assert_eq!(result, expected);
204
205        // check if DoubleEndedIterator is implemented
206        let result: Int32Array = array.iter().rev().collect();
207        let rev_array = Int32Array::from(vec![Some(4), None, Some(2), None, Some(0)]);
208        assert_eq!(result, rev_array);
209        // check if ExactSizeIterator is implemented
210        let _ = array.iter().rposition(|opt_b| opt_b == Some(1));
211    }
212
213    #[test]
214    fn test_double_ended() {
215        let array = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]);
216        let mut a = array.iter();
217        assert_eq!(a.next(), Some(Some(0)));
218        assert_eq!(a.next(), Some(None));
219        assert_eq!(a.next_back(), Some(Some(4)));
220        assert_eq!(a.next_back(), Some(None));
221        assert_eq!(a.next_back(), Some(Some(2)));
222        // the two sides have met: None is returned by both
223        assert_eq!(a.next_back(), None);
224        assert_eq!(a.next(), None);
225    }
226
227    #[test]
228    fn test_string_array_iter_round_trip() {
229        let array = StringArray::from(vec![Some("a"), None, Some("aaa"), None, Some("aaaaa")]);
230        let array = Arc::new(array) as ArrayRef;
231
232        let array = array.as_any().downcast_ref::<StringArray>().unwrap();
233
234        // to and from iter, with a +1
235        let result: StringArray = array
236            .iter()
237            .map(|e| {
238                e.map(|e| {
239                    let mut a = e.to_string();
240                    a.push('b');
241                    a
242                })
243            })
244            .collect();
245
246        let expected =
247            StringArray::from(vec![Some("ab"), None, Some("aaab"), None, Some("aaaaab")]);
248        assert_eq!(result, expected);
249
250        // check if DoubleEndedIterator is implemented
251        let result: StringArray = array.iter().rev().collect();
252        let rev_array = StringArray::from(vec![Some("aaaaa"), None, Some("aaa"), None, Some("a")]);
253        assert_eq!(result, rev_array);
254        // check if ExactSizeIterator is implemented
255        let _ = array.iter().rposition(|opt_b| opt_b == Some("a"));
256    }
257
258    #[test]
259    fn test_binary_array_iter_round_trip() {
260        let array = BinaryArray::from(vec![
261            Some(b"a" as &[u8]),
262            None,
263            Some(b"aaa"),
264            None,
265            Some(b"aaaaa"),
266        ]);
267
268        // to and from iter
269        let result: BinaryArray = array.iter().collect();
270
271        assert_eq!(result, array);
272
273        // check if DoubleEndedIterator is implemented
274        let result: BinaryArray = array.iter().rev().collect();
275        let rev_array = BinaryArray::from(vec![
276            Some(b"aaaaa" as &[u8]),
277            None,
278            Some(b"aaa"),
279            None,
280            Some(b"a"),
281        ]);
282        assert_eq!(result, rev_array);
283
284        // check if ExactSizeIterator is implemented
285        let _ = array.iter().rposition(|opt_b| opt_b == Some(&[9]));
286    }
287
288    #[test]
289    fn test_boolean_array_iter_round_trip() {
290        let array = BooleanArray::from(vec![Some(true), None, Some(false)]);
291
292        // to and from iter
293        let result: BooleanArray = array.iter().collect();
294
295        assert_eq!(result, array);
296
297        // check if DoubleEndedIterator is implemented
298        let result: BooleanArray = array.iter().rev().collect();
299        let rev_array = BooleanArray::from(vec![Some(false), None, Some(true)]);
300        assert_eq!(result, rev_array);
301
302        // check if ExactSizeIterator is implemented
303        let _ = array.iter().rposition(|opt_b| opt_b == Some(true));
304    }
305
306    trait SharedBetweenArrayIterAndSliceIter:
307        ExactSizeIterator<Item = Option<i32>> + DoubleEndedIterator<Item = Option<i32>> + Clone
308    {
309    }
310    impl<T: Clone + ExactSizeIterator<Item = Option<i32>> + DoubleEndedIterator<Item = Option<i32>>>
311        SharedBetweenArrayIterAndSliceIter for T
312    {
313    }
314
315    fn get_int32_iterator_cases() -> impl Iterator<Item = (Int32Array, Vec<Option<i32>>)> {
316        let mut rng = StdRng::seed_from_u64(42);
317
318        let no_nulls_and_no_duplicates = (0..10).map(Some).collect::<Vec<Option<i32>>>();
319        let no_nulls_random_values = (0..10)
320            .map(|_| rng.random::<i32>())
321            .map(Some)
322            .collect::<Vec<Option<i32>>>();
323
324        let all_nulls = (0..10).map(|_| None).collect::<Vec<Option<i32>>>();
325        let only_start_nulls = (0..10)
326            .map(|item| if item < 4 { None } else { Some(item) })
327            .collect::<Vec<Option<i32>>>();
328        let only_end_nulls = (0..10)
329            .map(|item| if item > 8 { None } else { Some(item) })
330            .collect::<Vec<Option<i32>>>();
331        let only_middle_nulls = (0..10)
332            .map(|item| {
333                if (4..=8).contains(&item) && rng.random_bool(0.9) {
334                    None
335                } else {
336                    Some(item)
337                }
338            })
339            .collect::<Vec<Option<i32>>>();
340        let random_values_with_random_nulls = (0..10)
341            .map(|_| {
342                if rng.random_bool(0.3) {
343                    None
344                } else {
345                    Some(rng.random::<i32>())
346                }
347            })
348            .collect::<Vec<Option<i32>>>();
349
350        let no_nulls_and_some_duplicates = (0..10)
351            .map(|item| item % 3)
352            .map(Some)
353            .collect::<Vec<Option<i32>>>();
354        let no_nulls_and_all_same_value =
355            (0..10).map(|_| 9).map(Some).collect::<Vec<Option<i32>>>();
356        let no_nulls_and_continues_duplicates = [0, 0, 0, 1, 1, 2, 2, 2, 2, 3]
357            .map(Some)
358            .into_iter()
359            .collect::<Vec<Option<i32>>>();
360
361        let single_null_and_no_duplicates = (0..10)
362            .map(|item| if item == 4 { None } else { Some(item) })
363            .collect::<Vec<Option<i32>>>();
364        let multiple_nulls_and_no_duplicates = (0..10)
365            .map(|item| if item % 3 == 2 { None } else { Some(item) })
366            .collect::<Vec<Option<i32>>>();
367        let continues_nulls_and_no_duplicates = [
368            Some(0),
369            Some(1),
370            None,
371            None,
372            Some(2),
373            Some(3),
374            None,
375            Some(4),
376            Some(5),
377            None,
378        ]
379        .into_iter()
380        .collect::<Vec<Option<i32>>>();
381
382        [
383            no_nulls_and_no_duplicates,
384            no_nulls_random_values,
385            no_nulls_and_some_duplicates,
386            no_nulls_and_all_same_value,
387            no_nulls_and_continues_duplicates,
388            all_nulls,
389            only_start_nulls,
390            only_end_nulls,
391            only_middle_nulls,
392            random_values_with_random_nulls,
393            single_null_and_no_duplicates,
394            multiple_nulls_and_no_duplicates,
395            continues_nulls_and_no_duplicates,
396        ]
397        .map(|case| (Int32Array::from(case.clone()), case))
398        .into_iter()
399    }
400
401    trait SetupIter {
402        fn description(&self) -> String;
403        fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I);
404    }
405
406    struct NoSetup;
407    impl SetupIter for NoSetup {
408        fn description(&self) -> String {
409            "no setup".to_string()
410        }
411        fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, _iter: &mut I) {
412            // none
413        }
414    }
415
416    fn setup_and_assert_cases_on_single_operation(
417        o: &impl ConsumingArrayIteratorOp,
418        setup_iterator: impl SetupIter,
419    ) {
420        for (array, source) in get_int32_iterator_cases() {
421            let mut actual = ArrayIter::new(&array);
422            let mut expected = source.iter().copied();
423
424            setup_iterator.setup(&mut actual);
425            setup_iterator.setup(&mut expected);
426
427            let current_iterator_values: Vec<Option<i32>> = expected.clone().collect();
428
429            assert_eq!(
430                o.get_value(actual),
431                o.get_value(expected),
432                "Failed on op {} for {} (left actual, right expected) ({current_iterator_values:?})",
433                o.name(),
434                setup_iterator.description(),
435            );
436        }
437    }
438
439    /// Trait representing an operation on a [`ArrayIter`]
440    /// that can be compared against a slice iterator
441    ///
442    /// this is for consuming operations (e.g. `count`, `last`, etc)
443    trait ConsumingArrayIteratorOp {
444        /// What the operation returns (e.g. Option<i32> for last, usize for count, etc)
445        type Output: PartialEq + Debug;
446
447        /// The name of the operation, used for error messages
448        fn name(&self) -> String;
449
450        /// Get the value of the operation for the provided iterator
451        /// This will be either a [`ArrayIter`] or a slice iterator to make sure they produce the same result
452        ///
453        /// Example implementation:
454        /// 1. for `last` it will be the last value
455        /// 2. for `count` it will be the returned length
456        fn get_value<T: SharedBetweenArrayIterAndSliceIter>(&self, iter: T) -> Self::Output;
457    }
458
459    /// Trait representing an operation on a [`ArrayIter`]
460    /// that can be compared against a slice iterator.
461    ///
462    /// This is for mutating operations (e.g. `position`, `any`, `find`, etc)
463    trait MutatingArrayIteratorOp {
464        /// What the operation returns (e.g. Option<i32> for last, usize for count, etc)
465        type Output: PartialEq + Debug;
466
467        /// The name of the operation, used for error messages
468        fn name(&self) -> String;
469
470        /// Get the value of the operation for the provided iterator
471        /// This will be either a [`ArrayIter`] or a slice iterator to make sure they produce the same result
472        ///
473        /// Example implementation:
474        /// 1. for `for_each` it will be the iterator element that the function was called with
475        /// 2. for `fold` it will be the accumulator and the iterator element from each call, as well as the final result
476        fn get_value<T: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut T) -> Self::Output;
477    }
478
479    /// Helper function that will assert that the provided operation
480    /// produces the same result for both [`ArrayIter`] and slice iterator
481    /// under various consumption patterns (e.g. some calls to next/next_back/consume_all/etc)
482    fn assert_array_iterator_cases<O: ConsumingArrayIteratorOp>(o: O) {
483        setup_and_assert_cases_on_single_operation(&o, NoSetup);
484
485        struct Next;
486        impl SetupIter for Next {
487            fn description(&self) -> String {
488                "new iter after consuming 1 element from the start".to_string()
489            }
490            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
491                iter.next();
492            }
493        }
494        setup_and_assert_cases_on_single_operation(&o, Next);
495
496        struct NextBack;
497        impl SetupIter for NextBack {
498            fn description(&self) -> String {
499                "new iter after consuming 1 element from the end".to_string()
500            }
501
502            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
503                iter.next_back();
504            }
505        }
506
507        setup_and_assert_cases_on_single_operation(&o, NextBack);
508
509        struct NextAndBack;
510        impl SetupIter for NextAndBack {
511            fn description(&self) -> String {
512                "new iter after consuming 1 element from start and end".to_string()
513            }
514
515            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
516                iter.next();
517                iter.next_back();
518            }
519        }
520
521        setup_and_assert_cases_on_single_operation(&o, NextAndBack);
522
523        struct NextUntilLast;
524        impl SetupIter for NextUntilLast {
525            fn description(&self) -> String {
526                "new iter after consuming all from the start but 1".to_string()
527            }
528            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
529                let len = iter.len();
530                if len > 1 {
531                    iter.nth(len - 2);
532                }
533            }
534        }
535        setup_and_assert_cases_on_single_operation(&o, NextUntilLast);
536
537        struct NextBackUntilFirst;
538        impl SetupIter for NextBackUntilFirst {
539            fn description(&self) -> String {
540                "new iter after consuming all from the end but 1".to_string()
541            }
542
543            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
544                let len = iter.len();
545                if len > 1 {
546                    iter.nth_back(len - 2);
547                }
548            }
549        }
550        setup_and_assert_cases_on_single_operation(&o, NextBackUntilFirst);
551
552        struct NextFinish;
553        impl SetupIter for NextFinish {
554            fn description(&self) -> String {
555                "new iter after consuming all from the start".to_string()
556            }
557            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
558                iter.nth(iter.len());
559            }
560        }
561        setup_and_assert_cases_on_single_operation(&o, NextFinish);
562
563        struct NextBackFinish;
564        impl SetupIter for NextBackFinish {
565            fn description(&self) -> String {
566                "new iter after consuming all from the end".to_string()
567            }
568            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
569                iter.nth_back(iter.len());
570            }
571        }
572        setup_and_assert_cases_on_single_operation(&o, NextBackFinish);
573
574        struct NextUntilLastNone;
575        impl SetupIter for NextUntilLastNone {
576            fn description(&self) -> String {
577                "new iter that have no nulls left".to_string()
578            }
579            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
580                let last_null_position = iter.clone().rposition(|item| item.is_none());
581
582                // move the iterator to the location where there are no nulls anymore
583                if let Some(last_null_position) = last_null_position {
584                    iter.nth(last_null_position);
585                }
586            }
587        }
588        setup_and_assert_cases_on_single_operation(&o, NextUntilLastNone);
589
590        struct NextUntilLastSome;
591        impl SetupIter for NextUntilLastSome {
592            fn description(&self) -> String {
593                "iter that only have nulls left".to_string()
594            }
595            fn setup<I: SharedBetweenArrayIterAndSliceIter>(&self, iter: &mut I) {
596                let last_some_position = iter.clone().rposition(|item| item.is_some());
597
598                // move the iterator to the location where there are only nulls
599                if let Some(last_some_position) = last_some_position {
600                    iter.nth(last_some_position);
601                }
602            }
603        }
604        setup_and_assert_cases_on_single_operation(&o, NextUntilLastSome);
605    }
606
607    /// Helper function that will assert that the provided operation
608    /// produces the same result for both [`ArrayIter`] and slice iterator
609    /// under various consumption patterns (e.g. some calls to next/next_back/consume_all/etc)
610    ///
611    /// this is different from [`assert_array_iterator_cases`] as this also check that the state after the call is correct
612    /// to make sure we don't leave the iterator in incorrect state
613    fn assert_array_iterator_cases_mutate<O: MutatingArrayIteratorOp>(o: O) {
614        struct Adapter<O: MutatingArrayIteratorOp> {
615            o: O,
616        }
617
618        #[derive(Debug, PartialEq)]
619        struct AdapterOutput<Value> {
620            value: Value,
621            /// collect on the iterator after running the operation
622            leftover: Vec<Option<i32>>,
623        }
624
625        impl<O: MutatingArrayIteratorOp> ConsumingArrayIteratorOp for Adapter<O> {
626            type Output = AdapterOutput<O::Output>;
627
628            fn name(&self) -> String {
629                self.o.name()
630            }
631
632            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(
633                &self,
634                mut iter: T,
635            ) -> Self::Output {
636                let value = self.o.get_value(&mut iter);
637
638                // Get the rest of the iterator to make sure we leave the iterator in a valid state
639                let leftover: Vec<_> = iter.collect();
640
641                AdapterOutput { value, leftover }
642            }
643        }
644
645        assert_array_iterator_cases(Adapter { o })
646    }
647
648    #[derive(Debug, PartialEq)]
649    struct CallTrackingAndResult<Result: Debug + PartialEq, CallArgs: Debug + PartialEq> {
650        result: Result,
651        calls: Vec<CallArgs>,
652    }
653    type CallTrackingWithInputType<Result> = CallTrackingAndResult<Result, Option<i32>>;
654    type CallTrackingOnly = CallTrackingWithInputType<()>;
655
656    #[test]
657    fn assert_position() {
658        struct PositionOp {
659            reverse: bool,
660            number_of_false: usize,
661        }
662
663        impl MutatingArrayIteratorOp for PositionOp {
664            type Output = CallTrackingWithInputType<Option<usize>>;
665            fn name(&self) -> String {
666                if self.reverse {
667                    format!("rposition with {} false returned", self.number_of_false)
668                } else {
669                    format!("position with {} false returned", self.number_of_false)
670                }
671            }
672
673            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(
674                &self,
675                iter: &mut T,
676            ) -> Self::Output {
677                let mut items = vec![];
678
679                let mut count = 0;
680
681                let cb = |item| {
682                    items.push(item);
683
684                    if count < self.number_of_false {
685                        count += 1;
686                        false
687                    } else {
688                        true
689                    }
690                };
691
692                let position_result = if self.reverse {
693                    iter.rposition(cb)
694                } else {
695                    iter.position(cb)
696                };
697
698                CallTrackingAndResult {
699                    result: position_result,
700                    calls: items,
701                }
702            }
703        }
704
705        for reverse in [false, true] {
706            for number_of_false in [0, 1, 2, usize::MAX] {
707                assert_array_iterator_cases_mutate(PositionOp {
708                    reverse,
709                    number_of_false,
710                });
711            }
712        }
713    }
714
715    #[test]
716    fn assert_nth() {
717        for (array, source) in get_int32_iterator_cases() {
718            let actual = ArrayIter::new(&array);
719            let expected = source.iter().copied();
720            {
721                let mut actual = actual.clone();
722                let mut expected = expected.clone();
723                for _ in 0..expected.len() {
724                    #[expect(clippy::iter_nth_zero)]
725                    let actual_val = actual.nth(0);
726                    #[expect(clippy::iter_nth_zero)]
727                    let expected_val = expected.nth(0);
728                    assert_eq!(actual_val, expected_val, "Failed on nth(0)");
729                }
730            }
731
732            {
733                let mut actual = actual.clone();
734                let mut expected = expected.clone();
735                for _ in 0..expected.len() {
736                    let actual_val = actual.nth(1);
737                    let expected_val = expected.nth(1);
738                    assert_eq!(actual_val, expected_val, "Failed on nth(1)");
739                }
740            }
741
742            {
743                let mut actual = actual.clone();
744                let mut expected = expected.clone();
745                for _ in 0..expected.len() {
746                    let actual_val = actual.nth(2);
747                    let expected_val = expected.nth(2);
748                    assert_eq!(actual_val, expected_val, "Failed on nth(2)");
749                }
750            }
751        }
752    }
753
754    #[test]
755    fn assert_nth_back() {
756        for (array, source) in get_int32_iterator_cases() {
757            let actual = ArrayIter::new(&array);
758            let expected = source.iter().copied();
759            {
760                let mut actual = actual.clone();
761                let mut expected = expected.clone();
762                for _ in 0..expected.len() {
763                    let actual_val = actual.nth_back(0);
764                    let expected_val = expected.nth_back(0);
765                    assert_eq!(actual_val, expected_val, "Failed on nth_back(0)");
766                }
767            }
768
769            {
770                let mut actual = actual.clone();
771                let mut expected = expected.clone();
772                for _ in 0..expected.len() {
773                    let actual_val = actual.nth_back(1);
774                    let expected_val = expected.nth_back(1);
775                    assert_eq!(actual_val, expected_val, "Failed on nth_back(1)");
776                }
777            }
778
779            {
780                let mut actual = actual.clone();
781                let mut expected = expected.clone();
782                for _ in 0..expected.len() {
783                    let actual_val = actual.nth_back(2);
784                    let expected_val = expected.nth_back(2);
785                    assert_eq!(actual_val, expected_val, "Failed on nth_back(2)");
786                }
787            }
788        }
789    }
790
791    #[test]
792    fn assert_last() {
793        for (array, source) in get_int32_iterator_cases() {
794            let mut actual_forward = ArrayIter::new(&array);
795            let mut expected_forward = source.iter().copied();
796
797            for _ in 0..source.len() + 1 {
798                {
799                    let actual_forward_clone = actual_forward.clone();
800                    let expected_forward_clone = expected_forward.clone();
801
802                    assert_eq!(actual_forward_clone.last(), expected_forward_clone.last());
803                }
804
805                actual_forward.next();
806                expected_forward.next();
807            }
808
809            let mut actual_backward = ArrayIter::new(&array);
810            let mut expected_backward = source.iter().copied();
811            for _ in 0..source.len() + 1 {
812                {
813                    assert_eq!(
814                        actual_backward.clone().last(),
815                        expected_backward.clone().last()
816                    );
817                }
818
819                actual_backward.next_back();
820                expected_backward.next_back();
821            }
822        }
823    }
824
825    #[test]
826    fn assert_for_each() {
827        struct ForEachOp;
828
829        impl ConsumingArrayIteratorOp for ForEachOp {
830            type Output = CallTrackingOnly;
831
832            fn name(&self) -> String {
833                "for_each".to_string()
834            }
835
836            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(&self, iter: T) -> Self::Output {
837                let mut items = Vec::with_capacity(iter.len());
838
839                iter.for_each(|item| {
840                    items.push(item);
841                });
842
843                CallTrackingAndResult {
844                    calls: items,
845                    result: (),
846                }
847            }
848        }
849
850        assert_array_iterator_cases(ForEachOp)
851    }
852
853    #[test]
854    fn assert_fold() {
855        struct FoldOp {
856            reverse: bool,
857        }
858
859        #[derive(Debug, PartialEq)]
860        struct CallArgs {
861            acc: Option<i32>,
862            item: Option<i32>,
863        }
864
865        impl ConsumingArrayIteratorOp for FoldOp {
866            type Output = CallTrackingAndResult<Option<i32>, CallArgs>;
867
868            fn name(&self) -> String {
869                if self.reverse {
870                    "rfold".to_string()
871                } else {
872                    "fold".to_string()
873                }
874            }
875
876            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(&self, iter: T) -> Self::Output {
877                let mut items = Vec::with_capacity(iter.len());
878
879                let cb = |acc, item| {
880                    items.push(CallArgs { acc, item });
881
882                    item.map(|val| val + 100)
883                };
884
885                let result = if self.reverse {
886                    iter.rfold(Some(1), cb)
887                } else {
888                    iter.fold(Some(1), cb)
889                };
890
891                CallTrackingAndResult {
892                    calls: items,
893                    result,
894                }
895            }
896        }
897
898        assert_array_iterator_cases(FoldOp { reverse: false });
899        assert_array_iterator_cases(FoldOp { reverse: true });
900    }
901
902    #[test]
903    fn assert_count() {
904        struct CountOp;
905
906        impl ConsumingArrayIteratorOp for CountOp {
907            type Output = usize;
908
909            fn name(&self) -> String {
910                "count".to_string()
911            }
912
913            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(&self, iter: T) -> Self::Output {
914                iter.count()
915            }
916        }
917
918        assert_array_iterator_cases(CountOp)
919    }
920
921    #[test]
922    fn assert_any() {
923        struct AnyOp {
924            false_count: usize,
925        }
926
927        impl MutatingArrayIteratorOp for AnyOp {
928            type Output = CallTrackingWithInputType<bool>;
929
930            fn name(&self) -> String {
931                format!("any with {} false returned", self.false_count)
932            }
933
934            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(
935                &self,
936                iter: &mut T,
937            ) -> Self::Output {
938                let mut items = Vec::with_capacity(iter.len());
939
940                let mut count = 0;
941                let res = iter.any(|item| {
942                    items.push(item);
943
944                    if count < self.false_count {
945                        count += 1;
946                        false
947                    } else {
948                        true
949                    }
950                });
951
952                CallTrackingWithInputType {
953                    calls: items,
954                    result: res,
955                }
956            }
957        }
958
959        for false_count in [0, 1, 2, usize::MAX] {
960            assert_array_iterator_cases_mutate(AnyOp { false_count });
961        }
962    }
963
964    #[test]
965    fn assert_all() {
966        struct AllOp {
967            true_count: usize,
968        }
969
970        impl MutatingArrayIteratorOp for AllOp {
971            type Output = CallTrackingWithInputType<bool>;
972
973            fn name(&self) -> String {
974                format!("all with {} false returned", self.true_count)
975            }
976
977            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(
978                &self,
979                iter: &mut T,
980            ) -> Self::Output {
981                let mut items = Vec::with_capacity(iter.len());
982
983                let mut count = 0;
984                let res = iter.all(|item| {
985                    items.push(item);
986
987                    if count < self.true_count {
988                        count += 1;
989                        true
990                    } else {
991                        false
992                    }
993                });
994
995                CallTrackingWithInputType {
996                    calls: items,
997                    result: res,
998                }
999            }
1000        }
1001
1002        for true_count in [0, 1, 2, usize::MAX] {
1003            assert_array_iterator_cases_mutate(AllOp { true_count });
1004        }
1005    }
1006
1007    #[test]
1008    fn assert_find() {
1009        struct FindOp {
1010            reverse: bool,
1011            false_count: usize,
1012        }
1013
1014        impl MutatingArrayIteratorOp for FindOp {
1015            type Output = CallTrackingWithInputType<Option<Option<i32>>>;
1016
1017            fn name(&self) -> String {
1018                if self.reverse {
1019                    format!("rfind with {} false returned", self.false_count)
1020                } else {
1021                    format!("find with {} false returned", self.false_count)
1022                }
1023            }
1024
1025            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(
1026                &self,
1027                iter: &mut T,
1028            ) -> Self::Output {
1029                let mut items = vec![];
1030
1031                let mut count = 0;
1032
1033                let cb = |item: &Option<i32>| {
1034                    items.push(*item);
1035
1036                    if count < self.false_count {
1037                        count += 1;
1038                        false
1039                    } else {
1040                        true
1041                    }
1042                };
1043
1044                let position_result = if self.reverse {
1045                    iter.rfind(cb)
1046                } else {
1047                    iter.find(cb)
1048                };
1049
1050                CallTrackingWithInputType {
1051                    calls: items,
1052                    result: position_result,
1053                }
1054            }
1055        }
1056
1057        for reverse in [false, true] {
1058            for false_count in [0, 1, 2, usize::MAX] {
1059                assert_array_iterator_cases_mutate(FindOp {
1060                    reverse,
1061                    false_count,
1062                });
1063            }
1064        }
1065    }
1066
1067    #[test]
1068    fn assert_find_map() {
1069        struct FindMapOp {
1070            number_of_nones: usize,
1071        }
1072
1073        impl MutatingArrayIteratorOp for FindMapOp {
1074            type Output = CallTrackingWithInputType<Option<&'static str>>;
1075
1076            fn name(&self) -> String {
1077                format!("find_map with {} None returned", self.number_of_nones)
1078            }
1079
1080            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(
1081                &self,
1082                iter: &mut T,
1083            ) -> Self::Output {
1084                let mut items = vec![];
1085
1086                let mut count = 0;
1087
1088                let result = iter.find_map(|item| {
1089                    items.push(item);
1090
1091                    if count < self.number_of_nones {
1092                        count += 1;
1093                        None
1094                    } else {
1095                        Some("found it")
1096                    }
1097                });
1098
1099                CallTrackingAndResult {
1100                    result,
1101                    calls: items,
1102                }
1103            }
1104        }
1105
1106        for number_of_nones in [0, 1, 2, usize::MAX] {
1107            assert_array_iterator_cases_mutate(FindMapOp { number_of_nones });
1108        }
1109    }
1110
1111    #[test]
1112    fn assert_partition() {
1113        struct PartitionOp<F: Fn(usize, &Option<i32>) -> bool> {
1114            description: &'static str,
1115            predicate: F,
1116        }
1117
1118        #[derive(Debug, PartialEq)]
1119        struct PartitionResult {
1120            left: Vec<Option<i32>>,
1121            right: Vec<Option<i32>>,
1122        }
1123
1124        impl<F: Fn(usize, &Option<i32>) -> bool> ConsumingArrayIteratorOp for PartitionOp<F> {
1125            type Output = CallTrackingWithInputType<PartitionResult>;
1126
1127            fn name(&self) -> String {
1128                format!("partition by {}", self.description)
1129            }
1130
1131            fn get_value<T: SharedBetweenArrayIterAndSliceIter>(&self, iter: T) -> Self::Output {
1132                let mut items = vec![];
1133
1134                let mut index = 0;
1135
1136                let (left, right) = iter.partition(|item| {
1137                    items.push(*item);
1138
1139                    let res = (self.predicate)(index, item);
1140
1141                    index += 1;
1142                    res
1143                });
1144
1145                CallTrackingAndResult {
1146                    result: PartitionResult { left, right },
1147                    calls: items,
1148                }
1149            }
1150        }
1151
1152        assert_array_iterator_cases(PartitionOp {
1153            description: "None on one side and Some(*) on the other",
1154            predicate: |_, item| item.is_none(),
1155        });
1156
1157        assert_array_iterator_cases(PartitionOp {
1158            description: "all true",
1159            predicate: |_, _| true,
1160        });
1161
1162        assert_array_iterator_cases(PartitionOp {
1163            description: "all false",
1164            predicate: |_, _| false,
1165        });
1166
1167        let random_values = (0..100).map(|_| rand::random_bool(0.5)).collect::<Vec<_>>();
1168        assert_array_iterator_cases(PartitionOp {
1169            description: "random",
1170            predicate: |index, _| random_values[index % random_values.len()],
1171        });
1172    }
1173}
arrow_array/iterator.rs

arrow_array/
iterator.rs
