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_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(T::get_dtype().to_arrow(CompatLevel::newest()), values, None);
42        NoNull::new(arr.into())
43    }
44}
45
46impl<T> FromIteratorReversed<Option<T::Native>> for ChunkedArray<T>
47where
48    T: PolarsNumericType,
49{
50    fn from_trusted_len_iter_rev<I: TrustedLen<Item = Option<T::Native>>>(iter: I) -> Self {
51        let arr: PrimitiveArray<T::Native> = iter.collect_reversed();
52        arr.into()
53    }
54}
55
56impl<T> FromIteratorReversed<T::Native> for NoNull<ChunkedArray<T>>
57where
58    T: PolarsNumericType,
59{
60    fn from_trusted_len_iter_rev<I: TrustedLen<Item = T::Native>>(iter: I) -> Self {
61        let arr: PrimitiveArray<T::Native> = iter.collect_reversed();
62        NoNull::new(arr.into())
63    }
64}
65
66impl FromIteratorReversed<Option<bool>> for BooleanChunked {
67    fn from_trusted_len_iter_rev<I: TrustedLen<Item = Option<bool>>>(iter: I) -> Self {
68        let arr: BooleanArray = iter.collect_reversed();
69        arr.into()
70    }
71}
72
73impl FromIteratorReversed<bool> for NoNull<BooleanChunked> {
74    fn from_trusted_len_iter_rev<I: TrustedLen<Item = bool>>(iter: I) -> Self {
75        let arr: BooleanArray = iter.collect_reversed();
76        NoNull::new(arr.into())
77    }
78}
79
80impl<Ptr> FromTrustedLenIterator<Ptr> for ListChunked
81where
82    Ptr: Borrow<Series>,
83{
84    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
85        let iter = iter.into_iter();
86        iter.collect()
87    }
88}
89
90impl FromTrustedLenIterator<Option<Series>> for ListChunked {
91    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Series>>>(iter: I) -> Self {
92        let iter = iter.into_iter();
93        iter.collect()
94    }
95}
96
97impl FromTrustedLenIterator<Option<bool>> for ChunkedArray<BooleanType> {
98    fn from_iter_trusted_length<I: IntoIterator<Item = Option<bool>>>(iter: I) -> Self
99    where
100        I::IntoIter: TrustedLen,
101    {
102        let iter = iter.into_iter();
103        let arr: BooleanArray = iter.collect_trusted();
104        arr.into()
105    }
106}
107
108impl FromTrustedLenIterator<bool> for BooleanChunked {
109    fn from_iter_trusted_length<I: IntoIterator<Item = bool>>(iter: I) -> Self
110    where
111        I::IntoIter: TrustedLen,
112    {
113        let iter = iter.into_iter();
114        let arr: BooleanArray = iter.collect_trusted();
115        arr.into()
116    }
117}
118
119impl FromTrustedLenIterator<bool> for NoNull<BooleanChunked> {
120    fn from_iter_trusted_length<I: IntoIterator<Item = bool>>(iter: I) -> Self {
121        let iter = iter.into_iter();
122        iter.collect()
123    }
124}
125impl<Ptr> FromTrustedLenIterator<Ptr> for StringChunked
126where
127    Ptr: PolarsAsRef<str>,
128{
129    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
130        let iter = iter.into_iter();
131        iter.collect()
132    }
133}
134
135impl<Ptr> FromTrustedLenIterator<Option<Ptr>> for StringChunked
136where
137    Ptr: AsRef<str>,
138{
139    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
140        let iter = iter.into_iter();
141        iter.collect()
142    }
143}
144
145impl<Ptr> FromTrustedLenIterator<Ptr> for BinaryChunked
146where
147    Ptr: PolarsAsRef<[u8]>,
148{
149    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
150        let iter = iter.into_iter();
151        iter.collect()
152    }
153}
154
155impl<Ptr> FromTrustedLenIterator<Option<Ptr>> for BinaryChunked
156where
157    Ptr: AsRef<[u8]>,
158{
159    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
160        let iter = iter.into_iter();
161        iter.collect()
162    }
163}
164
165impl<Ptr> FromTrustedLenIterator<Ptr> for BinaryOffsetChunked
166where
167    Ptr: PolarsAsRef<[u8]>,
168{
169    fn from_iter_trusted_length<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
170        let arr = BinaryArray::from_iter_values(iter.into_iter());
171        ChunkedArray::with_chunk(PlSmallStr::EMPTY, arr)
172    }
173}
174
175impl<Ptr> FromTrustedLenIterator<Option<Ptr>> for BinaryOffsetChunked
176where
177    Ptr: AsRef<[u8]>,
178{
179    fn from_iter_trusted_length<I: IntoIterator<Item = Option<Ptr>>>(iter: I) -> Self {
180        let iter = iter.into_iter();
181        let arr = BinaryArray::from_iter(iter);
182        ChunkedArray::with_chunk(PlSmallStr::EMPTY, arr)
183    }
184}
185
186#[cfg(feature = "object")]
187impl<T: PolarsObject> FromTrustedLenIterator<Option<T>> for ObjectChunked<T> {
188    fn from_iter_trusted_length<I: IntoIterator<Item = Option<T>>>(iter: I) -> Self {
189        let iter = iter.into_iter();
190        iter.collect()
191    }
192}
193
194#[cfg(test)]
195mod test {
196    use super::*;
197
198    #[test]
199    fn test_reverse_collect() {
200        let ca: NoNull<Int32Chunked> = (0..5).collect_reversed();
201        let arr = ca.downcast_iter().next().unwrap();
202        let s = arr.values().as_slice();
203        assert_eq!(s, &[4, 3, 2, 1, 0]);
204
205        let ca: Int32Chunked = (0..5)
206            .map(|val| match val % 2 == 0 {
207                true => Some(val),
208                false => None,
209            })
210            .collect_reversed();
211        assert_eq!(Vec::from(&ca), &[Some(4), None, Some(2), None, Some(0)]);
212    }
213}