polars_core/chunked_array/builder/list/
primitive.rs

1use super::*;
2
3pub struct ListPrimitiveChunkedBuilder<T>
4where
5    T: PolarsNumericType,
6{
7    pub builder: LargePrimitiveBuilder<T::Native>,
8    field: Field,
9    fast_explode: bool,
10}
11
12impl<T> ListPrimitiveChunkedBuilder<T>
13where
14    T: PolarsNumericType,
15{
16    pub fn new(
17        name: PlSmallStr,
18        capacity: usize,
19        values_capacity: usize,
20        inner_type: DataType,
21    ) -> Self {
22        debug_assert!(
23            inner_type.to_physical().is_primitive_numeric(),
24            "inner type must be primitive, got {}",
25            inner_type
26        );
27        let values = MutablePrimitiveArray::<T::Native>::with_capacity(values_capacity);
28        let builder = LargePrimitiveBuilder::<T::Native>::new_with_capacity(values, capacity);
29        let field = Field::new(name, DataType::List(Box::new(inner_type)));
30
31        Self {
32            builder,
33            field,
34            fast_explode: true,
35        }
36    }
37
38    pub fn new_with_values_type(
39        name: PlSmallStr,
40        capacity: usize,
41        values_capacity: usize,
42        values_type: DataType,
43        logical_type: DataType,
44    ) -> Self {
45        let values = MutablePrimitiveArray::<T::Native>::with_capacity_from(
46            values_capacity,
47            values_type.to_arrow(CompatLevel::newest()),
48        );
49        let builder = LargePrimitiveBuilder::<T::Native>::new_with_capacity(values, capacity);
50        let field = Field::new(name, DataType::List(Box::new(logical_type)));
51        Self {
52            builder,
53            field,
54            fast_explode: true,
55        }
56    }
57
58    #[inline]
59    pub fn append_slice(&mut self, items: &[T::Native]) {
60        let values = self.builder.mut_values();
61        values.extend_from_slice(items);
62        self.builder.try_push_valid().unwrap();
63
64        if items.is_empty() {
65            self.fast_explode = false;
66        }
67    }
68
69    #[inline]
70    pub fn append_opt_slice(&mut self, opt_v: Option<&[T::Native]>) {
71        match opt_v {
72            Some(items) => self.append_slice(items),
73            None => {
74                self.builder.push_null();
75            },
76        }
77    }
78    /// Appends from an iterator over values
79    #[inline]
80    pub fn append_values_iter_trusted_len<I: Iterator<Item = T::Native> + TrustedLen>(
81        &mut self,
82        iter: I,
83    ) {
84        let values = self.builder.mut_values();
85
86        if iter.size_hint().0 == 0 {
87            self.fast_explode = false;
88        }
89        // SAFETY:
90        // trusted len, trust the type system
91        values.extend_values(iter);
92        self.builder.try_push_valid().unwrap();
93    }
94
95    #[inline]
96    pub fn append_values_iter<I: Iterator<Item = T::Native>>(&mut self, iter: I) {
97        let values = self.builder.mut_values();
98
99        if iter.size_hint().0 == 0 {
100            self.fast_explode = false;
101        }
102        values.extend_values(iter);
103        self.builder.try_push_valid().unwrap();
104    }
105
106    /// Appends from an iterator over values
107    #[inline]
108    pub fn append_iter<I: Iterator<Item = Option<T::Native>> + TrustedLen>(&mut self, iter: I) {
109        let values = self.builder.mut_values();
110
111        if iter.size_hint().0 == 0 {
112            self.fast_explode = false;
113        }
114        // SAFETY:
115        // trusted len, trust the type system
116        unsafe { values.extend_trusted_len_unchecked(iter) };
117        self.builder.try_push_valid().unwrap();
118    }
119}
120
121impl<T> ListBuilderTrait for ListPrimitiveChunkedBuilder<T>
122where
123    T: PolarsNumericType,
124{
125    #[inline]
126    fn append_null(&mut self) {
127        self.fast_explode = false;
128        self.builder.push_null();
129    }
130
131    #[inline]
132    fn append_series(&mut self, s: &Series) -> PolarsResult<()> {
133        if s.is_empty() {
134            self.fast_explode = false;
135        }
136        let physical = s.to_physical_repr();
137        let ca = physical.unpack::<T>().map_err(|_| {
138            polars_err!(SchemaMismatch: "cannot build list with different dtypes 
139
140Expected {}, got {}.", self.field.dtype(), s.dtype())
141        })?;
142        let values = self.builder.mut_values();
143
144        ca.downcast_iter().for_each(|arr| {
145            if arr.null_count() == 0 {
146                values.extend_from_slice(arr.values().as_slice())
147            } else {
148                // SAFETY:
149                // Arrow arrays are trusted length iterators.
150                unsafe { values.extend_trusted_len_unchecked(arr.into_iter()) }
151            }
152        });
153        // overflow of i64 is far beyond polars capable lengths.
154        unsafe { self.builder.try_push_valid().unwrap_unchecked() };
155        Ok(())
156    }
157
158    fn field(&self) -> &Field {
159        &self.field
160    }
161
162    fn inner_array(&mut self) -> ArrayRef {
163        self.builder.as_box()
164    }
165
166    fn fast_explode(&self) -> bool {
167        self.fast_explode
168    }
169}
polars_core/chunked_array/builder/list/primitive.rs

polars_core/chunked_array/builder/list/
primitive.rs
