polars_core/chunked_array/iterator/
mod.rs

1use arrow::array::*;
2
3use crate::prelude::*;
4
5pub mod par;
6
7impl<T> ChunkedArray<T>
8where
9    T: PolarsDataType,
10{
11    #[inline]
12    pub fn iter(&self) -> impl PolarsIterator<Item = Option<T::Physical<'_>>> {
13        // SAFETY: we set the correct length of the iterator.
14        unsafe {
15            self.downcast_iter()
16                .flat_map(|arr| arr.iter())
17                .trust_my_length(self.len())
18        }
19    }
20}
21
22/// A [`PolarsIterator`] is an iterator over a [`ChunkedArray`] which contains polars types. A [`PolarsIterator`]
23/// must implement [`ExactSizeIterator`] and [`DoubleEndedIterator`].
24pub trait PolarsIterator:
25    ExactSizeIterator + DoubleEndedIterator + Send + Sync + TrustedLen
26{
27}
28unsafe impl<I> TrustedLen for Box<dyn PolarsIterator<Item = I> + '_> {}
29
30/// Implement [`PolarsIterator`] for every iterator that implements the needed traits.
31impl<T: ?Sized> PolarsIterator for T where
32    T: ExactSizeIterator + DoubleEndedIterator + Send + Sync + TrustedLen
33{
34}
35
36impl<'a, T> IntoIterator for &'a ChunkedArray<T>
37where
38    T: PolarsNumericType,
39{
40    type Item = Option<T::Native>;
41    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
42    fn into_iter(self) -> Self::IntoIter {
43        Box::new(
44            // we know that we only iterate over length == self.len()
45            unsafe {
46                self.downcast_iter()
47                    .flatten()
48                    .map(|x| x.copied())
49                    .trust_my_length(self.len())
50            },
51        )
52    }
53}
54
55impl<'a> IntoIterator for &'a BooleanChunked {
56    type Item = Option<bool>;
57    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
58    fn into_iter(self) -> Self::IntoIter {
59        // we know that we only iterate over length == self.len()
60        unsafe { Box::new(self.downcast_iter().flatten().trust_my_length(self.len())) }
61    }
62}
63
64/// The no null iterator for a [`BooleanArray`]
65pub struct BoolIterNoNull<'a> {
66    array: &'a BooleanArray,
67    current: usize,
68    current_end: usize,
69}
70
71impl<'a> BoolIterNoNull<'a> {
72    /// create a new iterator
73    pub fn new(array: &'a BooleanArray) -> Self {
74        BoolIterNoNull {
75            array,
76            current: 0,
77            current_end: array.len(),
78        }
79    }
80}
81
82impl Iterator for BoolIterNoNull<'_> {
83    type Item = bool;
84
85    fn next(&mut self) -> Option<Self::Item> {
86        if self.current == self.current_end {
87            None
88        } else {
89            let old = self.current;
90            self.current += 1;
91            unsafe { Some(self.array.value_unchecked(old)) }
92        }
93    }
94
95    fn size_hint(&self) -> (usize, Option<usize>) {
96        (
97            self.array.len() - self.current,
98            Some(self.array.len() - self.current),
99        )
100    }
101}
102
103impl DoubleEndedIterator for BoolIterNoNull<'_> {
104    fn next_back(&mut self) -> Option<Self::Item> {
105        if self.current_end == self.current {
106            None
107        } else {
108            self.current_end -= 1;
109            unsafe { Some(self.array.value_unchecked(self.current_end)) }
110        }
111    }
112}
113
114/// all arrays have known size.
115impl ExactSizeIterator for BoolIterNoNull<'_> {}
116
117impl BooleanChunked {
118    #[allow(clippy::wrong_self_convention)]
119    #[doc(hidden)]
120    pub fn into_no_null_iter(
121        &self,
122    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = bool> + DoubleEndedIterator + TrustedLen
123    {
124        // we know that we only iterate over length == self.len()
125        unsafe {
126            self.downcast_iter()
127                .flat_map(BoolIterNoNull::new)
128                .trust_my_length(self.len())
129        }
130    }
131}
132
133impl<'a> IntoIterator for &'a StringChunked {
134    type Item = Option<&'a str>;
135    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
136    fn into_iter(self) -> Self::IntoIter {
137        // we know that we only iterate over length == self.len()
138        unsafe { Box::new(self.downcast_iter().flatten().trust_my_length(self.len())) }
139    }
140}
141
142impl StringChunked {
143    #[allow(clippy::wrong_self_convention)]
144    #[doc(hidden)]
145    pub fn into_no_null_iter(
146        &self,
147    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = &str> + DoubleEndedIterator + TrustedLen
148    {
149        // we know that we only iterate over length == self.len()
150        unsafe {
151            self.downcast_iter()
152                .flat_map(|arr| arr.values_iter())
153                .trust_my_length(self.len())
154        }
155    }
156}
157
158impl<'a> IntoIterator for &'a BinaryChunked {
159    type Item = Option<&'a [u8]>;
160    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
161    fn into_iter(self) -> Self::IntoIter {
162        // we know that we only iterate over length == self.len()
163        unsafe { Box::new(self.downcast_iter().flatten().trust_my_length(self.len())) }
164    }
165}
166
167impl BinaryChunked {
168    #[allow(clippy::wrong_self_convention)]
169    #[doc(hidden)]
170    pub fn into_no_null_iter(
171        &self,
172    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = &[u8]> + DoubleEndedIterator + TrustedLen
173    {
174        // we know that we only iterate over length == self.len()
175        unsafe {
176            self.downcast_iter()
177                .flat_map(|arr| arr.values_iter())
178                .trust_my_length(self.len())
179        }
180    }
181}
182
183impl<'a> IntoIterator for &'a BinaryOffsetChunked {
184    type Item = Option<&'a [u8]>;
185    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
186    fn into_iter(self) -> Self::IntoIter {
187        // we know that we only iterate over length == self.len()
188        unsafe { Box::new(self.downcast_iter().flatten().trust_my_length(self.len())) }
189    }
190}
191
192impl BinaryOffsetChunked {
193    #[allow(clippy::wrong_self_convention)]
194    #[doc(hidden)]
195    pub fn into_no_null_iter(
196        &self,
197    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = &[u8]> + DoubleEndedIterator + TrustedLen
198    {
199        // we know that we only iterate over length == self.len()
200        unsafe {
201            self.downcast_iter()
202                .flat_map(|arr| arr.values_iter())
203                .trust_my_length(self.len())
204        }
205    }
206}
207
208impl<'a> IntoIterator for &'a ListChunked {
209    type Item = Option<Series>;
210    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
211    fn into_iter(self) -> Self::IntoIter {
212        let dtype = self.inner_dtype();
213
214        if self.null_count() == 0 {
215            // we know that we only iterate over length == self.len()
216            unsafe {
217                Box::new(
218                    self.downcast_iter()
219                        .flat_map(|arr| arr.iter().unwrap_required())
220                        .trust_my_length(self.len())
221                        .map(move |arr| {
222                            Some(Series::from_chunks_and_dtype_unchecked(
223                                PlSmallStr::EMPTY,
224                                vec![arr],
225                                dtype,
226                            ))
227                        }),
228                )
229            }
230        } else {
231            // we know that we only iterate over length == self.len()
232            unsafe {
233                Box::new(
234                    self.downcast_iter()
235                        .flat_map(|arr| arr.iter())
236                        .trust_my_length(self.len())
237                        .map(move |arr| {
238                            arr.map(|arr| {
239                                Series::from_chunks_and_dtype_unchecked(
240                                    PlSmallStr::EMPTY,
241                                    vec![arr],
242                                    dtype,
243                                )
244                            })
245                        }),
246                )
247            }
248        }
249    }
250}
251
252impl ListChunked {
253    #[allow(clippy::wrong_self_convention)]
254    #[doc(hidden)]
255    pub fn into_no_null_iter(
256        &self,
257    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = Series> + DoubleEndedIterator + TrustedLen
258    {
259        let inner_type = self.inner_dtype();
260        unsafe {
261            self.downcast_iter()
262                .flat_map(|arr| arr.values_iter())
263                .map(move |arr| {
264                    Series::from_chunks_and_dtype_unchecked(
265                        PlSmallStr::EMPTY,
266                        vec![arr],
267                        inner_type,
268                    )
269                })
270                .trust_my_length(self.len())
271        }
272    }
273}
274
275#[cfg(feature = "dtype-array")]
276impl<'a> IntoIterator for &'a ArrayChunked {
277    type Item = Option<Series>;
278    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
279    fn into_iter(self) -> Self::IntoIter {
280        let dtype = self.inner_dtype();
281
282        if self.null_count() == 0 {
283            // we know that we only iterate over length == self.len()
284            unsafe {
285                Box::new(
286                    self.downcast_iter()
287                        .flat_map(|arr| arr.iter().unwrap_required())
288                        .trust_my_length(self.len())
289                        .map(move |arr| {
290                            Some(Series::from_chunks_and_dtype_unchecked(
291                                PlSmallStr::EMPTY,
292                                vec![arr],
293                                dtype,
294                            ))
295                        }),
296                )
297            }
298        } else {
299            // we know that we only iterate over length == self.len()
300            unsafe {
301                Box::new(
302                    self.downcast_iter()
303                        .flat_map(|arr| arr.iter())
304                        .trust_my_length(self.len())
305                        .map(move |arr| {
306                            arr.map(|arr| {
307                                Series::from_chunks_and_dtype_unchecked(
308                                    PlSmallStr::EMPTY,
309                                    vec![arr],
310                                    dtype,
311                                )
312                            })
313                        }),
314                )
315            }
316        }
317    }
318}
319
320#[cfg(feature = "dtype-array")]
321pub struct FixedSizeListIterNoNull<'a> {
322    array: &'a FixedSizeListArray,
323    inner_type: DataType,
324    current: usize,
325    current_end: usize,
326}
327
328#[cfg(feature = "dtype-array")]
329impl<'a> FixedSizeListIterNoNull<'a> {
330    /// create a new iterator
331    pub fn new(array: &'a FixedSizeListArray, inner_type: DataType) -> Self {
332        FixedSizeListIterNoNull {
333            array,
334            inner_type,
335            current: 0,
336            current_end: array.len(),
337        }
338    }
339}
340
341#[cfg(feature = "dtype-array")]
342impl Iterator for FixedSizeListIterNoNull<'_> {
343    type Item = Series;
344
345    fn next(&mut self) -> Option<Self::Item> {
346        if self.current == self.current_end {
347            None
348        } else {
349            let old = self.current;
350            self.current += 1;
351            unsafe {
352                Some(Series::from_chunks_and_dtype_unchecked(
353                    PlSmallStr::EMPTY,
354                    vec![self.array.value_unchecked(old)],
355                    &self.inner_type,
356                ))
357            }
358        }
359    }
360
361    fn size_hint(&self) -> (usize, Option<usize>) {
362        (
363            self.array.len() - self.current,
364            Some(self.array.len() - self.current),
365        )
366    }
367}
368
369#[cfg(feature = "dtype-array")]
370impl DoubleEndedIterator for FixedSizeListIterNoNull<'_> {
371    fn next_back(&mut self) -> Option<Self::Item> {
372        if self.current_end == self.current {
373            None
374        } else {
375            self.current_end -= 1;
376            unsafe {
377                Some(
378                    Series::try_from((
379                        PlSmallStr::EMPTY,
380                        self.array.value_unchecked(self.current_end),
381                    ))
382                    .unwrap(),
383                )
384            }
385        }
386    }
387}
388
389/// all arrays have known size.
390#[cfg(feature = "dtype-array")]
391impl ExactSizeIterator for FixedSizeListIterNoNull<'_> {}
392
393#[cfg(feature = "dtype-array")]
394impl ArrayChunked {
395    #[allow(clippy::wrong_self_convention)]
396    #[doc(hidden)]
397    pub fn into_no_null_iter(
398        &self,
399    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = Series> + DoubleEndedIterator + TrustedLen
400    {
401        // we know that we only iterate over length == self.len()
402        let inner_type = self.inner_dtype();
403        unsafe {
404            self.downcast_iter()
405                .flat_map(move |arr| FixedSizeListIterNoNull::new(arr, inner_type.clone()))
406                .trust_my_length(self.len())
407        }
408    }
409}
410
411#[cfg(feature = "object")]
412impl<'a, T> IntoIterator for &'a ObjectChunked<T>
413where
414    T: PolarsObject,
415{
416    type Item = Option<&'a T>;
417    type IntoIter = Box<dyn PolarsIterator<Item = Self::Item> + 'a>;
418    fn into_iter(self) -> Self::IntoIter {
419        // we know that we only iterate over length == self.len()
420        unsafe { Box::new(self.downcast_iter().flatten().trust_my_length(self.len())) }
421    }
422}
423
424#[cfg(feature = "object")]
425impl<T: PolarsObject> ObjectChunked<T> {
426    #[allow(clippy::wrong_self_convention)]
427    #[doc(hidden)]
428    pub fn into_no_null_iter(
429        &self,
430    ) -> impl '_ + Send + Sync + ExactSizeIterator<Item = &T> + DoubleEndedIterator + TrustedLen
431    {
432        // we know that we only iterate over length == self.len()
433        unsafe {
434            self.downcast_iter()
435                .flat_map(|arr| arr.values_iter())
436                .trust_my_length(self.len())
437        }
438    }
439}
440
441/// Wrapper struct to convert an iterator of type `T` into one of type [`Option<T>`].  It is useful to make the
442/// [`IntoIterator`] trait, in which every iterator shall return an [`Option<T>`].
443pub struct SomeIterator<I>(I)
444where
445    I: Iterator;
446
447impl<I> Iterator for SomeIterator<I>
448where
449    I: Iterator,
450{
451    type Item = Option<I::Item>;
452
453    fn next(&mut self) -> Option<Self::Item> {
454        self.0.next().map(Some)
455    }
456
457    fn size_hint(&self) -> (usize, Option<usize>) {
458        self.0.size_hint()
459    }
460}
461
462impl<I> DoubleEndedIterator for SomeIterator<I>
463where
464    I: DoubleEndedIterator,
465{
466    fn next_back(&mut self) -> Option<Self::Item> {
467        self.0.next_back().map(Some)
468    }
469}
470
471impl<I> ExactSizeIterator for SomeIterator<I> where I: ExactSizeIterator {}
472
473#[cfg(test)]
474mod test {
475    use crate::prelude::*;
476
477    #[test]
478    fn out_of_bounds() {
479        let mut a = UInt32Chunked::from_slice(PlSmallStr::from_static("a"), &[1, 2, 3]);
480        let b = UInt32Chunked::from_slice(PlSmallStr::from_static("a"), &[1, 2, 3]);
481        a.append(&b).unwrap();
482
483        let v = a.into_iter().collect::<Vec<_>>();
484        assert_eq!(
485            vec![Some(1u32), Some(2), Some(3), Some(1), Some(2), Some(3)],
486            v
487        )
488    }
489
490    /// Generate test for [`IntoIterator`] trait for chunked arrays with just one chunk and no null values.
491    /// The expected return value of the iterator generated by [`IntoIterator`] trait is [`Option<T>`], where
492    /// `T` is the chunked array type.
493    ///
494    /// # Input
495    ///
496    /// test_name: The name of the test to generate.
497    /// ca_type: The chunked array to use for this test. Ex: [`StringChunked`], [`UInt32Chunked`] ...
498    /// first_val: The first value contained in the chunked array.
499    /// second_val: The second value contained in the chunked array.
500    /// third_val: The third value contained in the chunked array.
501    macro_rules! impl_test_iter_single_chunk {
502        ($test_name:ident, $ca_type:ty, $first_val:expr, $second_val:expr, $third_val:expr) => {
503            #[test]
504            fn $test_name() {
505                let a = <$ca_type>::from_slice(
506                    PlSmallStr::from_static("test"),
507                    &[$first_val, $second_val, $third_val],
508                );
509
510                // normal iterator
511                let mut it = a.into_iter();
512                assert_eq!(it.next(), Some(Some($first_val)));
513                assert_eq!(it.next(), Some(Some($second_val)));
514                assert_eq!(it.next(), Some(Some($third_val)));
515                assert_eq!(it.next(), None);
516                // ensure both sides are consumes.
517                assert_eq!(it.next_back(), None);
518
519                // reverse iterator
520                let mut it = a.into_iter();
521                assert_eq!(it.next_back(), Some(Some($third_val)));
522                assert_eq!(it.next_back(), Some(Some($second_val)));
523                assert_eq!(it.next_back(), Some(Some($first_val)));
524                assert_eq!(it.next_back(), None);
525                // ensure both sides are consumes.
526                assert_eq!(it.next(), None);
527
528                // iterators should not cross
529                let mut it = a.into_iter();
530                assert_eq!(it.next_back(), Some(Some($third_val)));
531                assert_eq!(it.next(), Some(Some($first_val)));
532                assert_eq!(it.next(), Some(Some($second_val)));
533                // should stop here as we took this one from the back
534                assert_eq!(it.next(), None);
535                // ensure both sides are consumes.
536                assert_eq!(it.next_back(), None);
537
538                // do the same from the right side
539                let mut it = a.into_iter();
540                assert_eq!(it.next(), Some(Some($first_val)));
541                assert_eq!(it.next_back(), Some(Some($third_val)));
542                assert_eq!(it.next_back(), Some(Some($second_val)));
543                assert_eq!(it.next_back(), None);
544                // ensure both sides are consumes.
545                assert_eq!(it.next(), None);
546            }
547        };
548    }
549
550    impl_test_iter_single_chunk!(num_iter_single_chunk, UInt32Chunked, 1, 2, 3);
551    impl_test_iter_single_chunk!(utf8_iter_single_chunk, StringChunked, "a", "b", "c");
552    impl_test_iter_single_chunk!(bool_iter_single_chunk, BooleanChunked, true, true, false);
553
554    /// Generate test for [`IntoIterator`] trait for chunked arrays with just one chunk and null values.
555    /// The expected return value of the iterator generated by [`IntoIterator`] trait is [`Option<T>`], where
556    /// `T` is the chunked array type.
557    ///
558    /// # Input
559    ///
560    /// test_name: The name of the test to generate.
561    /// ca_type: The chunked array to use for this test. Ex: [`StringChunked`], [`UInt32Chunked`] ...
562    /// first_val: The first value contained in the chunked array. Must be an [`Option<T>`].
563    /// second_val: The second value contained in the chunked array. Must be an [`Option<T>`].
564    /// third_val: The third value contained in the chunked array. Must be an [`Option<T>`].
565    macro_rules! impl_test_iter_single_chunk_null_check {
566        ($test_name:ident, $ca_type:ty, $first_val:expr, $second_val:expr, $third_val:expr) => {
567            #[test]
568            fn $test_name() {
569                let a = <$ca_type>::new(
570                    PlSmallStr::from_static("test"),
571                    &[$first_val, $second_val, $third_val],
572                );
573
574                // normal iterator
575                let mut it = a.into_iter();
576                assert_eq!(it.next(), Some($first_val));
577                assert_eq!(it.next(), Some($second_val));
578                assert_eq!(it.next(), Some($third_val));
579                assert_eq!(it.next(), None);
580                // ensure both sides are consumes.
581                assert_eq!(it.next_back(), None);
582
583                // reverse iterator
584                let mut it = a.into_iter();
585                assert_eq!(it.next_back(), Some($third_val));
586                assert_eq!(it.next_back(), Some($second_val));
587                assert_eq!(it.next_back(), Some($first_val));
588                assert_eq!(it.next_back(), None);
589                // ensure both sides are consumes.
590                assert_eq!(it.next(), None);
591
592                // iterators should not cross
593                let mut it = a.into_iter();
594                assert_eq!(it.next_back(), Some($third_val));
595                assert_eq!(it.next(), Some($first_val));
596                assert_eq!(it.next(), Some($second_val));
597                // should stop here as we took this one from the back
598                assert_eq!(it.next(), None);
599                // ensure both sides are consumes.
600                assert_eq!(it.next_back(), None);
601
602                // do the same from the right side
603                let mut it = a.into_iter();
604                assert_eq!(it.next(), Some($first_val));
605                assert_eq!(it.next_back(), Some($third_val));
606                assert_eq!(it.next_back(), Some($second_val));
607                assert_eq!(it.next_back(), None);
608                // ensure both sides are consumes.
609                assert_eq!(it.next(), None);
610            }
611        };
612    }
613
614    impl_test_iter_single_chunk_null_check!(
615        num_iter_single_chunk_null_check,
616        UInt32Chunked,
617        Some(1),
618        None,
619        Some(3)
620    );
621    impl_test_iter_single_chunk_null_check!(
622        utf8_iter_single_chunk_null_check,
623        StringChunked,
624        Some("a"),
625        None,
626        Some("c")
627    );
628    impl_test_iter_single_chunk_null_check!(
629        bool_iter_single_chunk_null_check,
630        BooleanChunked,
631        Some(true),
632        None,
633        Some(false)
634    );
635
636    /// Generate test for [`IntoIterator`] trait for chunked arrays with many chunks and no null values.
637    /// The expected return value of the iterator generated by [`IntoIterator`] trait is [`Option<T>`], where
638    /// `T` is the chunked array type.
639    ///
640    /// # Input
641    ///
642    /// test_name: The name of the test to generate.
643    /// ca_type: The chunked array to use for this test. Ex: [`StringChunked`], [`UInt32Chunked`] ...
644    /// first_val: The first value contained in the chunked array.
645    /// second_val: The second value contained in the chunked array.
646    /// third_val: The third value contained in the chunked array.
647    macro_rules! impl_test_iter_many_chunk {
648        ($test_name:ident, $ca_type:ty, $first_val:expr, $second_val:expr, $third_val:expr) => {
649            #[test]
650            fn $test_name() {
651                let mut a = <$ca_type>::from_slice(
652                    PlSmallStr::from_static("test"),
653                    &[$first_val, $second_val],
654                );
655                let a_b = <$ca_type>::from_slice(PlSmallStr::EMPTY, &[$third_val]);
656                a.append(&a_b).unwrap();
657
658                // normal iterator
659                let mut it = a.into_iter();
660                assert_eq!(it.next(), Some(Some($first_val)));
661                assert_eq!(it.next(), Some(Some($second_val)));
662                assert_eq!(it.next(), Some(Some($third_val)));
663                assert_eq!(it.next(), None);
664                // ensure both sides are consumes.
665                assert_eq!(it.next_back(), None);
666
667                // reverse iterator
668                let mut it = a.into_iter();
669                assert_eq!(it.next_back(), Some(Some($third_val)));
670                assert_eq!(it.next_back(), Some(Some($second_val)));
671                assert_eq!(it.next_back(), Some(Some($first_val)));
672                assert_eq!(it.next_back(), None);
673                // ensure both sides are consumes.
674                assert_eq!(it.next(), None);
675
676                // iterators should not cross
677                let mut it = a.into_iter();
678                assert_eq!(it.next_back(), Some(Some($third_val)));
679                assert_eq!(it.next(), Some(Some($first_val)));
680                assert_eq!(it.next(), Some(Some($second_val)));
681                // should stop here as we took this one from the back
682                assert_eq!(it.next(), None);
683                // ensure both sides are consumes.
684                assert_eq!(it.next_back(), None);
685
686                // do the same from the right side
687                let mut it = a.into_iter();
688                assert_eq!(it.next(), Some(Some($first_val)));
689                assert_eq!(it.next_back(), Some(Some($third_val)));
690                assert_eq!(it.next_back(), Some(Some($second_val)));
691                assert_eq!(it.next_back(), None);
692                // ensure both sides are consumes.
693                assert_eq!(it.next(), None);
694            }
695        };
696    }
697
698    impl_test_iter_many_chunk!(num_iter_many_chunk, UInt32Chunked, 1, 2, 3);
699    impl_test_iter_many_chunk!(utf8_iter_many_chunk, StringChunked, "a", "b", "c");
700    impl_test_iter_many_chunk!(bool_iter_many_chunk, BooleanChunked, true, true, false);
701
702    /// Generate test for [`IntoIterator`] trait for chunked arrays with many chunk and null values.
703    /// The expected return value of the iterator generated by [`IntoIterator`] trait is [`Option<T>`], where
704    /// `T` is the chunked array type.
705    ///
706    /// # Input
707    ///
708    /// test_name: The name of the test to generate.
709    /// ca_type: The chunked array to use for this test. Ex: [`StringChunked`], [`UInt32Chunked`] ...
710    /// first_val: The first value contained in the chunked array. Must be an [`Option<T>`].
711    /// second_val: The second value contained in the chunked array. Must be an [`Option<T>`].
712    /// third_val: The third value contained in the chunked array. Must be an [`Option<T>`].
713    macro_rules! impl_test_iter_many_chunk_null_check {
714        ($test_name:ident, $ca_type:ty, $first_val:expr, $second_val:expr, $third_val:expr) => {
715            #[test]
716            fn $test_name() {
717                let mut a =
718                    <$ca_type>::new(PlSmallStr::from_static("test"), &[$first_val, $second_val]);
719                let a_b = <$ca_type>::new(PlSmallStr::EMPTY, &[$third_val]);
720                a.append(&a_b).unwrap();
721
722                // normal iterator
723                let mut it = a.into_iter();
724                assert_eq!(it.next(), Some($first_val));
725                assert_eq!(it.next(), Some($second_val));
726                assert_eq!(it.next(), Some($third_val));
727                assert_eq!(it.next(), None);
728                // ensure both sides are consumes.
729                assert_eq!(it.next_back(), None);
730
731                // reverse iterator
732                let mut it = a.into_iter();
733                assert_eq!(it.next_back(), Some($third_val));
734                assert_eq!(it.next_back(), Some($second_val));
735                assert_eq!(it.next_back(), Some($first_val));
736                assert_eq!(it.next_back(), None);
737                // ensure both sides are consumes.
738                assert_eq!(it.next(), None);
739
740                // iterators should not cross
741                let mut it = a.into_iter();
742                assert_eq!(it.next_back(), Some($third_val));
743                assert_eq!(it.next(), Some($first_val));
744                assert_eq!(it.next(), Some($second_val));
745                // should stop here as we took this one from the back
746                assert_eq!(it.next(), None);
747                // ensure both sides are consumes.
748                assert_eq!(it.next_back(), None);
749
750                // do the same from the right side
751                let mut it = a.into_iter();
752                assert_eq!(it.next(), Some($first_val));
753                assert_eq!(it.next_back(), Some($third_val));
754                assert_eq!(it.next_back(), Some($second_val));
755                assert_eq!(it.next_back(), None);
756                // ensure both sides are consumes.
757                assert_eq!(it.next(), None);
758            }
759        };
760    }
761
762    impl_test_iter_many_chunk_null_check!(
763        num_iter_many_chunk_null_check,
764        UInt32Chunked,
765        Some(1),
766        None,
767        Some(3)
768    );
769    impl_test_iter_many_chunk_null_check!(
770        utf8_iter_many_chunk_null_check,
771        StringChunked,
772        Some("a"),
773        None,
774        Some("c")
775    );
776    impl_test_iter_many_chunk_null_check!(
777        bool_iter_many_chunk_null_check,
778        BooleanChunked,
779        Some(true),
780        None,
781        Some(false)
782    );
783
784    /// Generate test for [`IntoNoNullIterator`] trait for chunked arrays with just one chunk and no null values.
785    /// The expected return value of the iterator generated by [`IntoNoNullIterator`] trait is `T`, where
786    /// `T` is the chunked array type.
787    ///
788    /// # Input
789    ///
790    /// test_name: The name of the test to generate.
791    /// ca_type: The chunked array to use for this test. Ex: [`StringChunked`], [`UInt32Chunked`] ...
792    /// first_val: The first value contained in the chunked array.
793    /// second_val: The second value contained in the chunked array.
794    /// third_val: The third value contained in the chunked array.
795    macro_rules! impl_test_no_null_iter_single_chunk {
796        ($test_name:ident, $ca_type:ty, $first_val:expr, $second_val:expr, $third_val:expr) => {
797            #[test]
798            fn $test_name() {
799                let a = <$ca_type>::from_slice(
800                    PlSmallStr::from_static("test"),
801                    &[$first_val, $second_val, $third_val],
802                );
803
804                // normal iterator
805                let mut it = a.into_no_null_iter();
806                assert_eq!(it.next(), Some($first_val));
807                assert_eq!(it.next(), Some($second_val));
808                assert_eq!(it.next(), Some($third_val));
809                assert_eq!(it.next(), None);
810                // ensure both sides are consumes.
811                assert_eq!(it.next_back(), None);
812
813                // reverse iterator
814                let mut it = a.into_no_null_iter();
815                assert_eq!(it.next_back(), Some($third_val));
816                assert_eq!(it.next_back(), Some($second_val));
817                assert_eq!(it.next_back(), Some($first_val));
818                assert_eq!(it.next_back(), None);
819                // ensure both sides are consumes.
820                assert_eq!(it.next(), None);
821
822                // iterators should not cross
823                let mut it = a.into_no_null_iter();
824                assert_eq!(it.next_back(), Some($third_val));
825                assert_eq!(it.next(), Some($first_val));
826                assert_eq!(it.next(), Some($second_val));
827                // should stop here as we took this one from the back
828                assert_eq!(it.next(), None);
829                // ensure both sides are consumes.
830                assert_eq!(it.next_back(), None);
831
832                // do the same from the right side
833                let mut it = a.into_no_null_iter();
834                assert_eq!(it.next(), Some($first_val));
835                assert_eq!(it.next_back(), Some($third_val));
836                assert_eq!(it.next_back(), Some($second_val));
837                assert_eq!(it.next_back(), None);
838                // ensure both sides are consumes.
839                assert_eq!(it.next(), None);
840            }
841        };
842    }
843
844    impl_test_no_null_iter_single_chunk!(num_no_null_iter_single_chunk, UInt32Chunked, 1, 2, 3);
845    impl_test_no_null_iter_single_chunk!(
846        utf8_no_null_iter_single_chunk,
847        StringChunked,
848        "a",
849        "b",
850        "c"
851    );
852    impl_test_no_null_iter_single_chunk!(
853        bool_no_null_iter_single_chunk,
854        BooleanChunked,
855        true,
856        true,
857        false
858    );
859
860    /// Generate test for [`IntoNoNullIterator`] trait for chunked arrays with many chunks and no null values.
861    /// The expected return value of the iterator generated by [`IntoNoNullIterator`] trait is `T`, where
862    /// `T` is the chunked array type.
863    ///
864    /// # Input
865    ///
866    /// test_name: The name of the test to generate.
867    /// ca_type: The chunked array to use for this test. Ex: [`StringChunked`], [`UInt32Chunked`] ...
868    /// first_val: The first value contained in the chunked array.
869    /// second_val: The second value contained in the chunked array.
870    /// third_val: The third value contained in the chunked array.
871    macro_rules! impl_test_no_null_iter_many_chunk {
872        ($test_name:ident, $ca_type:ty, $first_val:expr, $second_val:expr, $third_val:expr) => {
873            #[test]
874            fn $test_name() {
875                let mut a = <$ca_type>::from_slice(
876                    PlSmallStr::from_static("test"),
877                    &[$first_val, $second_val],
878                );
879                let a_b = <$ca_type>::from_slice(PlSmallStr::EMPTY, &[$third_val]);
880                a.append(&a_b).unwrap();
881
882                // normal iterator
883                let mut it = a.into_no_null_iter();
884                assert_eq!(it.next(), Some($first_val));
885                assert_eq!(it.next(), Some($second_val));
886                assert_eq!(it.next(), Some($third_val));
887                assert_eq!(it.next(), None);
888                // ensure both sides are consumes.
889                assert_eq!(it.next_back(), None);
890
891                // reverse iterator
892                let mut it = a.into_no_null_iter();
893                assert_eq!(it.next_back(), Some($third_val));
894                assert_eq!(it.next_back(), Some($second_val));
895                assert_eq!(it.next_back(), Some($first_val));
896                assert_eq!(it.next_back(), None);
897                // ensure both sides are consumes.
898                assert_eq!(it.next(), None);
899
900                // iterators should not cross
901                let mut it = a.into_no_null_iter();
902                assert_eq!(it.next_back(), Some($third_val));
903                assert_eq!(it.next(), Some($first_val));
904                assert_eq!(it.next(), Some($second_val));
905                // should stop here as we took this one from the back
906                assert_eq!(it.next(), None);
907                // ensure both sides are consumes.
908                assert_eq!(it.next_back(), None);
909
910                // do the same from the right side
911                let mut it = a.into_no_null_iter();
912                assert_eq!(it.next(), Some($first_val));
913                assert_eq!(it.next_back(), Some($third_val));
914                assert_eq!(it.next_back(), Some($second_val));
915                assert_eq!(it.next_back(), None);
916                // ensure both sides are consumes.
917                assert_eq!(it.next(), None);
918            }
919        };
920    }
921
922    impl_test_no_null_iter_many_chunk!(num_no_null_iter_many_chunk, UInt32Chunked, 1, 2, 3);
923    impl_test_no_null_iter_many_chunk!(utf8_no_null_iter_many_chunk, StringChunked, "a", "b", "c");
924    impl_test_no_null_iter_many_chunk!(
925        bool_no_null_iter_many_chunk,
926        BooleanChunked,
927        true,
928        true,
929        false
930    );
931
932    /// The size of the skip iterator.
933    const SKIP_ITERATOR_SIZE: usize = 10;
934
935    /// Generates tests to verify the correctness of the `skip` method.
936    ///
937    /// # Input
938    ///
939    /// test_name: The name of the test to implement, it is a function name so it shall be unique.
940    /// skip_values: The number of values to skip. Keep in mind that is the number of values to skip
941    ///   after performing the first next, then, skip_values = 8, will skip until index 1 + skip_values = 9.
942    /// first_val: The value before skip.
943    /// second_val: The value after skip.
944    /// ca_init_block: The block which initialize the chunked array. It shall return the chunked array.
945    macro_rules! impl_test_iter_skip {
946        ($test_name:ident, $skip_values:expr, $first_val:expr, $second_val:expr, $ca_init_block:block) => {
947            #[test]
948            fn $test_name() {
949                let a = $ca_init_block;
950
951                // Consume first position of iterator.
952                let mut it = a.into_iter();
953                assert_eq!(it.next(), Some($first_val));
954
955                // Consume `$skip_values` and check the result.
956                let mut it = it.skip($skip_values);
957                assert_eq!(it.next(), Some($second_val));
958
959                // Consume more values than available and check result is None.
960                let mut it = it.skip(SKIP_ITERATOR_SIZE);
961                assert_eq!(it.next(), None);
962            }
963        };
964    }
965
966    /// Generates a `Vec` of `Strings`, where every position is the `String` representation of its index.
967    fn generate_utf8_vec(size: usize) -> Vec<String> {
968        (0..size).map(|n| n.to_string()).collect()
969    }
970
971    /// Generate a `Vec` of `Option<String>`, where even indexes are `Some("{idx}")` and odd indexes are `None`.
972    fn generate_opt_utf8_vec(size: usize) -> Vec<Option<String>> {
973        (0..size)
974            .map(|n| {
975                if n % 2 == 0 {
976                    Some(n.to_string())
977                } else {
978                    None
979                }
980            })
981            .collect()
982    }
983
984    impl_test_iter_skip!(utf8_iter_single_chunk_skip, 8, Some("0"), Some("9"), {
985        StringChunked::from_slice(
986            PlSmallStr::from_static("test"),
987            &generate_utf8_vec(SKIP_ITERATOR_SIZE),
988        )
989    });
990
991    impl_test_iter_skip!(
992        utf8_iter_single_chunk_null_check_skip,
993        8,
994        Some("0"),
995        None,
996        {
997            StringChunked::new(
998                PlSmallStr::from_static("test"),
999                &generate_opt_utf8_vec(SKIP_ITERATOR_SIZE),
1000            )
1001        }
1002    );
1003
1004    impl_test_iter_skip!(utf8_iter_many_chunk_skip, 18, Some("0"), Some("9"), {
1005        let mut a = StringChunked::from_slice(
1006            PlSmallStr::from_static("test"),
1007            &generate_utf8_vec(SKIP_ITERATOR_SIZE),
1008        );
1009        let a_b = StringChunked::from_slice(
1010            PlSmallStr::from_static("test"),
1011            &generate_utf8_vec(SKIP_ITERATOR_SIZE),
1012        );
1013        a.append(&a_b).unwrap();
1014        a
1015    });
1016
1017    impl_test_iter_skip!(utf8_iter_many_chunk_null_check_skip, 18, Some("0"), None, {
1018        let mut a = StringChunked::new(
1019            PlSmallStr::from_static("test"),
1020            &generate_opt_utf8_vec(SKIP_ITERATOR_SIZE),
1021        );
1022        let a_b = StringChunked::new(
1023            PlSmallStr::from_static("test"),
1024            &generate_opt_utf8_vec(SKIP_ITERATOR_SIZE),
1025        );
1026        a.append(&a_b).unwrap();
1027        a
1028    });
1029
1030    /// Generates a [`Vec`] of [`bool`], with even indexes are true, and odd indexes are false.
1031    fn generate_boolean_vec(size: usize) -> Vec<bool> {
1032        (0..size).map(|n| n % 2 == 0).collect()
1033    }
1034
1035    /// Generate a [`Vec`] of [`Option<bool>`], where:
1036    /// - If the index is divisible by 3, then, the value is `None`.
1037    /// - If the index is not divisible by 3 and it is even, then, the value is `Some(true)`.
1038    /// - Otherwise, the value is `Some(false)`.
1039    fn generate_opt_boolean_vec(size: usize) -> Vec<Option<bool>> {
1040        (0..size)
1041            .map(|n| if n % 3 == 0 { None } else { Some(n % 2 == 0) })
1042            .collect()
1043    }
1044
1045    impl_test_iter_skip!(bool_iter_single_chunk_skip, 8, Some(true), Some(false), {
1046        BooleanChunked::from_slice(
1047            PlSmallStr::from_static("test"),
1048            &generate_boolean_vec(SKIP_ITERATOR_SIZE),
1049        )
1050    });
1051
1052    impl_test_iter_skip!(bool_iter_single_chunk_null_check_skip, 8, None, None, {
1053        BooleanChunked::new(
1054            PlSmallStr::from_static("test"),
1055            &generate_opt_boolean_vec(SKIP_ITERATOR_SIZE),
1056        )
1057    });
1058
1059    impl_test_iter_skip!(bool_iter_many_chunk_skip, 18, Some(true), Some(false), {
1060        let mut a = BooleanChunked::from_slice(
1061            PlSmallStr::from_static("test"),
1062            &generate_boolean_vec(SKIP_ITERATOR_SIZE),
1063        );
1064        let a_b = BooleanChunked::from_slice(
1065            PlSmallStr::from_static("test"),
1066            &generate_boolean_vec(SKIP_ITERATOR_SIZE),
1067        );
1068        a.append(&a_b).unwrap();
1069        a
1070    });
1071
1072    impl_test_iter_skip!(bool_iter_many_chunk_null_check_skip, 18, None, None, {
1073        let mut a = BooleanChunked::new(
1074            PlSmallStr::from_static("test"),
1075            &generate_opt_boolean_vec(SKIP_ITERATOR_SIZE),
1076        );
1077        let a_b = BooleanChunked::new(
1078            PlSmallStr::from_static("test"),
1079            &generate_opt_boolean_vec(SKIP_ITERATOR_SIZE),
1080        );
1081        a.append(&a_b).unwrap();
1082        a
1083    });
1084}
polars_core/chunked_array/iterator/mod.rs

polars_core/chunked_array/iterator/
mod.rs
