arrow_array/
iterator.rs

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

arrow_array/
iterator.rs
