polars_core/chunked_array/
trusted_len.rs

1use std::borrow::Borrow;
2
3use arrow::legacy::trusted_len::{FromIteratorReversed, TrustedLenPush};
4
5use crate::chunked_array::from_iterator::PolarsAsRef;
6use crate::prelude::*;
7use crate::utils::{FromTrustedLenIterator, NoNull};
8
9impl<T> FromTrustedLenIterator<Option<T::Native>> for ChunkedArray<T>
10where
11    T: PolarsNumericType,
12{
13    fn from_iter_trusted_length<I: IntoIterator<Item = Option<T::Native>>>(iter: I) -> Self
14    where
15        I::IntoIter: TrustedLen,
16    {
17        // SAFETY: iter is TrustedLen.
18        let iter = iter.into_iter();
19        let arr = unsafe {
20            PrimitiveArray::from_trusted_len_iter_unchecked(iter)
21                .to(T::get_static_dtype().to_arrow(CompatLevel::newest()))
22        };
23        arr.into()
24    }
25}
26
27// NoNull is only a wrapper needed for specialization.
28impl<T> FromTrustedLenIterator<T::Native> for NoNull<ChunkedArray<T>>
29where
30    T: PolarsNumericType,
31{
32    // We use Vec because it is way faster than Arrows builder. We can do this
33    // because we know we don't have null values.
34    fn from_iter_trusted_length<I: IntoIterator<Item = T::Native>>(iter: I) -> Self
35    where
36        I::IntoIter: TrustedLen,
37    {
38        // SAFETY: iter is TrustedLen.
39        let iter = iter.into_iter();
40        let values = unsafe { Vec::from_trusted_len_iter_unchecked(iter) }.into();
41        let arr = PrimitiveArray::new(
42            T::get_static_dtype().to_arrow(CompatLevel::newest()),
43            values,
44            None,
45        );
46        NoNull::new(arr.into())
47    }
48}
49
50impl<T> FromIteratorReversed<Option<T::Native>> for ChunkedArray<T>
51where
52    T: PolarsNumericType,
53{
54    fn from_trusted_len_iter_rev<I: TrustedLen<Item = Option<T::Native>>>(iter: I) -> Self {
55        let arr: PrimitiveArray<T::Native> = iter.collect_reversed();
56        arr.into()
57    }
58}
59
60impl<T> FromIteratorReversed<T::Native> for NoNull<ChunkedArray<T>>
61where
62    T: PolarsNumericType,
63{
64    fn from_trusted_len_iter_rev<I: TrustedLen<Item = T::Native>>(iter: I) -> Self {
65        let arr: PrimitiveArray<T::Native> = iter.collect_reversed();
66        NoNull::new(arr.into())
67    }
68}
69
70impl FromIteratorReversed<Option<bool>> for BooleanChunked {
71    fn from_trusted_len_iter_rev<I: TrustedLen<Item = Option<bool>>>(iter: I) -> Self {
72        let arr: BooleanArray = iter.collect_reversed();
73        arr.into()
74    }
75}
76
77impl FromIteratorReversed<bool> for NoNull<BooleanChunked> {
78    fn from_trusted_len_iter_rev<I: TrustedLen<Item = bool>>(iter: I) -> Self {
79        let arr: BooleanArray = iter.collect_reversed();
80        NoNull::new(arr.into())
81    }
82}
83
84impl<Ptr> FromTrustedLenIterator<Ptr> for ListChunked
85where
86    Ptr: Borrow<Series>,
87{
88    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
89        let iter = iter.into_iter();
90        iter.collect()
91    }
92}
93
94impl FromTrustedLenIterator<Option<Series>> for ListChunked {
95    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Series>>>(iter: I) -> Self {
96        let iter = iter.into_iter();
97        iter.collect()
98    }
99}
100
101impl FromTrustedLenIterator<Option<bool>> for ChunkedArray<BooleanType> {
102    fn from_iter_trusted_length<I: IntoIterator<Item = Option<bool>>>(iter: I) -> Self
103    where
104        I::IntoIter: TrustedLen,
105    {
106        let iter = iter.into_iter();
107        let arr: BooleanArray = iter.collect_trusted();
108        arr.into()
109    }
110}
111
112impl FromTrustedLenIterator<bool> for BooleanChunked {
113    fn from_iter_trusted_length<I: IntoIterator<Item = bool>>(iter: I) -> Self
114    where
115        I::IntoIter: TrustedLen,
116    {
117        let iter = iter.into_iter();
118        let arr: BooleanArray = iter.collect_trusted();
119        arr.into()
120    }
121}
122
123impl FromTrustedLenIterator<bool> for NoNull<BooleanChunked> {
124    fn from_iter_trusted_length<I: IntoIterator<Item = bool>>(iter: I) -> Self {
125        let iter = iter.into_iter();
126        iter.collect()
127    }
128}
129impl<Ptr> FromTrustedLenIterator<Ptr> for StringChunked
130where
131    Ptr: PolarsAsRef<str>,
132{
133    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
134        let iter = iter.into_iter();
135        iter.collect()
136    }
137}
138
139impl<Ptr> FromTrustedLenIterator<Option<Ptr>> for StringChunked
140where
141    Ptr: AsRef<str>,
142{
143    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
144        let iter = iter.into_iter();
145        iter.collect()
146    }
147}
148
149impl<Ptr> FromTrustedLenIterator<Ptr> for BinaryChunked
150where
151    Ptr: PolarsAsRef<[u8]>,
152{
153    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
154        let iter = iter.into_iter();
155        iter.collect()
156    }
157}
158
159impl<Ptr> FromTrustedLenIterator<Option<Ptr>> for BinaryChunked
160where
161    Ptr: AsRef<[u8]>,
162{
163    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
164        let iter = iter.into_iter();
165        iter.collect()
166    }
167}
168
169impl<Ptr> FromTrustedLenIterator<Ptr> for BinaryOffsetChunked
170where
171    Ptr: PolarsAsRef<[u8]>,
172{
173    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
174        let arr = BinaryArray::from_iter_values(iter.into_iter());
175        ChunkedArray::with_chunk(PlSmallStr::EMPTY, arr)
176    }
177}
178
179impl<Ptr> FromTrustedLenIterator<Option<Ptr>> for BinaryOffsetChunked
180where
181    Ptr: AsRef<[u8]>,
182{
183    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
184        let iter = iter.into_iter();
185        let arr = BinaryArray::from_iter(iter);
186        ChunkedArray::with_chunk(PlSmallStr::EMPTY, arr)
187    }
188}
189
190#[cfg(feature = "object")]
191impl<T: PolarsObject> FromTrustedLenIterator<Option<T>> for ObjectChunked<T> {
192    fn from_iter_trusted_length<I: IntoIterator<Item = Option<T>>>(iter: I) -> Self {
193        let iter = iter.into_iter();
194        iter.collect()
195    }
196}
197
198#[cfg(test)]
199mod test {
200    use super::*;
201
202    #[test]
203    fn test_reverse_collect() {
204        let ca: NoNull<Int32Chunked> = (0..5).collect_reversed();
205        let arr = ca.downcast_iter().next().unwrap();
206        let s = arr.values().as_slice();
207        assert_eq!(s, &[4, 3, 2, 1, 0]);
208
209        let ca: Int32Chunked = (0..5)
210            .map(|val| match val % 2 == 0 {
211                true => Some(val),
212                false => None,
213            })
214            .collect_reversed();
215        assert_eq!(Vec::from(&ca), &[Some(4), None, Some(2), None, Some(0)]);
216    }
217}