polars_core/chunked_array/builder/
mod.rs

1mod boolean;
2#[cfg(feature = "dtype-categorical")]
3pub mod categorical;
4#[cfg(feature = "dtype-array")]
5pub mod fixed_size_list;
6pub mod list;
7mod null;
8mod primitive;
9mod string;
10
11use std::sync::Arc;
12
13use arrow::array::*;
14use arrow::bitmap::Bitmap;
15pub use boolean::*;
16#[cfg(feature = "dtype-categorical")]
17pub use categorical::*;
18#[cfg(feature = "dtype-array")]
19pub(crate) use fixed_size_list::*;
20pub use list::*;
21pub use null::*;
22pub use primitive::*;
23pub use string::*;
24
25use crate::chunked_array::to_primitive;
26use crate::prelude::*;
27use crate::utils::{NoNull, get_iter_capacity};
28
29// N: the value type; T: the sentinel type
30pub trait ChunkedBuilder<N, T: PolarsDataType> {
31    fn append_value(&mut self, val: N);
32    fn append_null(&mut self);
33
34    #[inline]
35    fn append_option(&mut self, opt_val: Option<N>) {
36        match opt_val {
37            Some(v) => self.append_value(v),
38            None => self.append_null(),
39        }
40    }
41
42    fn finish(self) -> ChunkedArray<T>;
43
44    fn shrink_to_fit(&mut self);
45}
46
47// Used in polars/src/chunked_array/apply.rs:24 to collect from aligned vecs and null bitmaps
48impl<T> FromIterator<(Vec<T::Native>, Option<Bitmap>)> for ChunkedArray<T>
49where
50    T: PolarsNumericType,
51{
52    fn from_iter<I: IntoIterator<Item = (Vec<T::Native>, Option<Bitmap>)>>(iter: I) -> Self {
53        let chunks = iter
54            .into_iter()
55            .map(|(values, opt_buffer)| to_primitive::<T>(values, opt_buffer));
56        ChunkedArray::from_chunk_iter(PlSmallStr::EMPTY, chunks)
57    }
58}
59
60pub trait NewChunkedArray<T, N> {
61    fn from_slice(name: PlSmallStr, v: &[N]) -> Self;
62    fn from_slice_options(name: PlSmallStr, opt_v: &[Option<N>]) -> Self;
63
64    /// Create a new ChunkedArray from an iterator.
65    fn from_iter_options(name: PlSmallStr, it: impl Iterator<Item = Option<N>>) -> Self;
66
67    /// Create a new ChunkedArray from an iterator.
68    fn from_iter_values(name: PlSmallStr, it: impl Iterator<Item = N>) -> Self;
69}
70
71impl<T> NewChunkedArray<T, T::Native> for ChunkedArray<T>
72where
73    T: PolarsNumericType,
74{
75    fn from_slice(name: PlSmallStr, v: &[T::Native]) -> Self {
76        let arr =
77            PrimitiveArray::from_slice(v).to(T::get_static_dtype().to_arrow(CompatLevel::newest()));
78        ChunkedArray::with_chunk(name, arr)
79    }
80
81    fn from_slice_options(name: PlSmallStr, opt_v: &[Option<T::Native>]) -> Self {
82        Self::from_iter_options(name, opt_v.iter().copied())
83    }
84
85    fn from_iter_options(
86        name: PlSmallStr,
87        it: impl Iterator<Item = Option<T::Native>>,
88    ) -> ChunkedArray<T> {
89        let mut builder = PrimitiveChunkedBuilder::new(name, get_iter_capacity(&it));
90        it.for_each(|opt| builder.append_option(opt));
91        builder.finish()
92    }
93
94    /// Create a new ChunkedArray from an iterator.
95    fn from_iter_values(name: PlSmallStr, it: impl Iterator<Item = T::Native>) -> ChunkedArray<T> {
96        let ca: NoNull<ChunkedArray<_>> = it.collect();
97        let mut ca = ca.into_inner();
98        ca.rename(name);
99        ca
100    }
101}
102
103impl NewChunkedArray<BooleanType, bool> for BooleanChunked {
104    fn from_slice(name: PlSmallStr, v: &[bool]) -> Self {
105        Self::from_iter_values(name, v.iter().copied())
106    }
107
108    fn from_slice_options(name: PlSmallStr, opt_v: &[Option<bool>]) -> Self {
109        Self::from_iter_options(name, opt_v.iter().copied())
110    }
111
112    fn from_iter_options(
113        name: PlSmallStr,
114        it: impl Iterator<Item = Option<bool>>,
115    ) -> ChunkedArray<BooleanType> {
116        let mut builder = BooleanChunkedBuilder::new(name, get_iter_capacity(&it));
117        it.for_each(|opt| builder.append_option(opt));
118        builder.finish()
119    }
120
121    /// Create a new ChunkedArray from an iterator.
122    fn from_iter_values(
123        name: PlSmallStr,
124        it: impl Iterator<Item = bool>,
125    ) -> ChunkedArray<BooleanType> {
126        let mut ca: ChunkedArray<_> = it.collect();
127        ca.rename(name);
128        ca
129    }
130}
131
132impl<S> NewChunkedArray<StringType, S> for StringChunked
133where
134    S: AsRef<str>,
135{
136    fn from_slice(name: PlSmallStr, v: &[S]) -> Self {
137        let arr = Utf8ViewArray::from_slice_values(v);
138        ChunkedArray::with_chunk(name, arr)
139    }
140
141    fn from_slice_options(name: PlSmallStr, opt_v: &[Option<S>]) -> Self {
142        let arr = Utf8ViewArray::from_slice(opt_v);
143        ChunkedArray::with_chunk(name, arr)
144    }
145
146    fn from_iter_options(name: PlSmallStr, it: impl Iterator<Item = Option<S>>) -> Self {
147        let arr = MutableBinaryViewArray::from_iterator(it).freeze();
148        ChunkedArray::with_chunk(name, arr)
149    }
150
151    /// Create a new ChunkedArray from an iterator.
152    fn from_iter_values(name: PlSmallStr, it: impl Iterator<Item = S>) -> Self {
153        let arr = MutableBinaryViewArray::from_values_iter(it).freeze();
154        ChunkedArray::with_chunk(name, arr)
155    }
156}
157
158impl<B> NewChunkedArray<BinaryType, B> for BinaryChunked
159where
160    B: AsRef<[u8]>,
161{
162    fn from_slice(name: PlSmallStr, v: &[B]) -> Self {
163        let arr = BinaryViewArray::from_slice_values(v);
164        ChunkedArray::with_chunk(name, arr)
165    }
166
167    fn from_slice_options(name: PlSmallStr, opt_v: &[Option<B>]) -> Self {
168        let arr = BinaryViewArray::from_slice(opt_v);
169        ChunkedArray::with_chunk(name, arr)
170    }
171
172    fn from_iter_options(name: PlSmallStr, it: impl Iterator<Item = Option<B>>) -> Self {
173        let arr = MutableBinaryViewArray::from_iterator(it).freeze();
174        ChunkedArray::with_chunk(name, arr)
175    }
176
177    /// Create a new ChunkedArray from an iterator.
178    fn from_iter_values(name: PlSmallStr, it: impl Iterator<Item = B>) -> Self {
179        let arr = MutableBinaryViewArray::from_values_iter(it).freeze();
180        ChunkedArray::with_chunk(name, arr)
181    }
182}
183
184#[cfg(test)]
185mod test {
186    use super::*;
187
188    #[test]
189    fn test_primitive_builder() {
190        let mut builder =
191            PrimitiveChunkedBuilder::<UInt32Type>::new(PlSmallStr::from_static("foo"), 6);
192        let values = &[Some(1), None, Some(2), Some(3), None, Some(4)];
193        for val in values {
194            builder.append_option(*val);
195        }
196        let ca = builder.finish();
197        assert_eq!(Vec::from(&ca), values);
198    }
199
200    #[test]
201    fn test_list_builder() {
202        let mut builder = ListPrimitiveChunkedBuilder::<Int32Type>::new(
203            PlSmallStr::from_static("a"),
204            10,
205            5,
206            DataType::Int32,
207        );
208
209        // Create a series containing two chunks.
210        let mut s1 =
211            Int32Chunked::from_slice(PlSmallStr::from_static("a"), &[1, 2, 3]).into_series();
212        let s2 = Int32Chunked::from_slice(PlSmallStr::from_static("b"), &[4, 5, 6]).into_series();
213        s1.append(&s2).unwrap();
214
215        builder.append_series(&s1).unwrap();
216        builder.append_series(&s2).unwrap();
217        let ls = builder.finish();
218        if let AnyValue::List(s) = ls.get_any_value(0).unwrap() {
219            // many chunks are aggregated to one in the ListArray
220            assert_eq!(s.len(), 6)
221        } else {
222            panic!()
223        }
224        if let AnyValue::List(s) = ls.get_any_value(1).unwrap() {
225            assert_eq!(s.len(), 3)
226        } else {
227            panic!()
228        }
229
230        // Test list collect.
231        let out = [&s1, &s2].iter().copied().collect::<ListChunked>();
232        assert_eq!(out.get_as_series(0).unwrap().len(), 6);
233        assert_eq!(out.get_as_series(1).unwrap().len(), 3);
234
235        let mut builder = ListPrimitiveChunkedBuilder::<Int32Type>::new(
236            PlSmallStr::from_static("a"),
237            10,
238            5,
239            DataType::Int32,
240        );
241        builder.append_series(&s1).unwrap();
242        builder.append_null();
243
244        let out = builder.finish();
245        let out = out
246            .explode(ExplodeOptions {
247                empty_as_null: true,
248                keep_nulls: true,
249            })
250            .unwrap();
251        assert_eq!(out.len(), 7);
252        assert_eq!(out.get(6).unwrap(), AnyValue::Null);
253    }
254}
polars_core/chunked_array/builder/mod.rs

polars_core/chunked_array/builder/
mod.rs
