polars_core/chunked_array/
collect.rs

1//! Methods for collecting into a ChunkedArray.
2//!
3//! For types that don't have dtype parameters:
4//! iter.(try_)collect_ca(_trusted) (name)
5//!
6//! For all types:
7//! iter.(try_)collect_ca(_trusted)_like (other_df)  Copies name/dtype from other_df
8//! iter.(try_)collect_ca(_trusted)_with_dtype (name, df)
9//!
10//! The try variants work on iterators of Results, the trusted variants do not
11//! check the length of the iterator.
12
13use std::sync::Arc;
14
15use arrow::trusted_len::TrustedLen;
16use polars_utils::pl_str::PlSmallStr;
17
18use crate::chunked_array::ChunkedArray;
19use crate::datatypes::{
20    ArrayCollectIterExt, ArrayFromIter, ArrayFromIterDtype, DataType, Field, PolarsDataType,
21};
22use crate::prelude::CompatLevel;
23
24pub trait ChunkedCollectIterExt<T: PolarsDataType>: Iterator + Sized {
25    #[inline]
26    fn collect_ca_with_dtype(self, name: PlSmallStr, dtype: DataType) -> ChunkedArray<T>
27    where
28        T::Array: ArrayFromIterDtype<Self::Item>,
29    {
30        let field = Arc::new(Field::new(name, dtype.clone()));
31        let arr = self.collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
32        ChunkedArray::from_chunk_iter_and_field(field, [arr])
33    }
34
35    #[inline]
36    fn collect_ca_like(self, name_dtype_src: &ChunkedArray<T>) -> ChunkedArray<T>
37    where
38        T::Array: ArrayFromIterDtype<Self::Item>,
39    {
40        let field = Arc::clone(&name_dtype_src.field);
41        let arr = self.collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
42        ChunkedArray::from_chunk_iter_and_field(field, [arr])
43    }
44
45    #[inline]
46    fn collect_ca_trusted_with_dtype(self, name: PlSmallStr, dtype: DataType) -> ChunkedArray<T>
47    where
48        T::Array: ArrayFromIterDtype<Self::Item>,
49        Self: TrustedLen,
50    {
51        let field = Arc::new(Field::new(name, dtype.clone()));
52        let arr = self.collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
53        ChunkedArray::from_chunk_iter_and_field(field, [arr])
54    }
55
56    #[inline]
57    fn collect_ca_trusted_like(self, name_dtype_src: &ChunkedArray<T>) -> ChunkedArray<T>
58    where
59        T::Array: ArrayFromIterDtype<Self::Item>,
60        Self: TrustedLen,
61    {
62        let field = Arc::clone(&name_dtype_src.field);
63        let arr = self.collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
64        ChunkedArray::from_chunk_iter_and_field(field, [arr])
65    }
66
67    #[inline]
68    fn try_collect_ca_with_dtype<U, E>(
69        self,
70        name: PlSmallStr,
71        dtype: DataType,
72    ) -> Result<ChunkedArray<T>, E>
73    where
74        T::Array: ArrayFromIterDtype<U>,
75        Self: Iterator<Item = Result<U, E>>,
76    {
77        let field = Arc::new(Field::new(name, dtype.clone()));
78        let arr = self.try_collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
79        Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
80    }
81
82    #[inline]
83    fn try_collect_ca_like<U, E>(
84        self,
85        name_dtype_src: &ChunkedArray<T>,
86    ) -> Result<ChunkedArray<T>, E>
87    where
88        T::Array: ArrayFromIterDtype<U>,
89        Self: Iterator<Item = Result<U, E>>,
90    {
91        let field = Arc::clone(&name_dtype_src.field);
92        let arr = self.try_collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
93        Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
94    }
95
96    #[inline]
97    fn try_collect_ca_trusted_with_dtype<U, E>(
98        self,
99        name: PlSmallStr,
100        dtype: DataType,
101    ) -> Result<ChunkedArray<T>, E>
102    where
103        T::Array: ArrayFromIterDtype<U>,
104        Self: Iterator<Item = Result<U, E>> + TrustedLen,
105    {
106        let field = Arc::new(Field::new(name, dtype.clone()));
107        let arr =
108            self.try_collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
109        Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
110    }
111
112    #[inline]
113    fn try_collect_ca_trusted_like<U, E>(
114        self,
115        name_dtype_src: &ChunkedArray<T>,
116    ) -> Result<ChunkedArray<T>, E>
117    where
118        T::Array: ArrayFromIterDtype<U>,
119        Self: Iterator<Item = Result<U, E>> + TrustedLen,
120    {
121        let field = Arc::clone(&name_dtype_src.field);
122        let arr =
123            self.try_collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
124        Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
125    }
126}
127
128impl<T: PolarsDataType, I: Iterator> ChunkedCollectIterExt<T> for I {}
129
130pub trait ChunkedCollectInferIterExt<T: PolarsDataType>: Iterator + Sized {
131    #[inline]
132    fn collect_ca(self, name: PlSmallStr) -> ChunkedArray<T>
133    where
134        T::Array: ArrayFromIter<Self::Item>,
135    {
136        let field = Arc::new(Field::new(name, T::get_dtype()));
137        let arr = self.collect_arr();
138        ChunkedArray::from_chunk_iter_and_field(field, [arr])
139    }
140
141    #[inline]
142    fn collect_ca_trusted(self, name: PlSmallStr) -> ChunkedArray<T>
143    where
144        T::Array: ArrayFromIter<Self::Item>,
145        Self: TrustedLen,
146    {
147        let field = Arc::new(Field::new(name, T::get_dtype()));
148        let arr = self.collect_arr_trusted();
149        ChunkedArray::from_chunk_iter_and_field(field, [arr])
150    }
151
152    #[inline]
153    fn try_collect_ca<U, E>(self, name: PlSmallStr) -> Result<ChunkedArray<T>, E>
154    where
155        T::Array: ArrayFromIter<U>,
156        Self: Iterator<Item = Result<U, E>>,
157    {
158        let field = Arc::new(Field::new(name, T::get_dtype()));
159        let arr = self.try_collect_arr()?;
160        Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
161    }
162
163    #[inline]
164    fn try_collect_ca_trusted<U, E>(self, name: PlSmallStr) -> Result<ChunkedArray<T>, E>
165    where
166        T::Array: ArrayFromIter<U>,
167        Self: Iterator<Item = Result<U, E>> + TrustedLen,
168    {
169        let field = Arc::new(Field::new(name, T::get_dtype()));
170        let arr = self.try_collect_arr_trusted()?;
171        Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
172    }
173}
174
175impl<T: PolarsDataType, I: Iterator> ChunkedCollectInferIterExt<T> for I {}
polars_core/chunked_array/collect.rs

polars_core/chunked_array/
collect.rs
