polars_core/chunked_array/builder/list/
mod.rs

1mod anonymous;
2mod binary;
3mod boolean;
4#[cfg(feature = "dtype-categorical")]
5mod categorical;
6mod dtypes;
7mod null;
8mod primitive;
9
10pub use anonymous::*;
11use arrow::legacy::array::list::AnonymousBuilder;
12use arrow::legacy::array::null::MutableNullArray;
13pub use binary::*;
14pub use boolean::*;
15#[cfg(feature = "dtype-categorical")]
16use categorical::*;
17use dtypes::*;
18pub use null::*;
19pub use primitive::*;
20
21use super::*;
22#[cfg(feature = "object")]
23use crate::chunked_array::object::registry::get_object_builder;
24
25pub trait ListBuilderTrait {
26    fn append_opt_series(&mut self, opt_s: Option<&Series>) -> PolarsResult<()> {
27        match opt_s {
28            Some(s) => return self.append_series(s),
29            None => self.append_null(),
30        }
31        Ok(())
32    }
33    fn append_series(&mut self, s: &Series) -> PolarsResult<()>;
34    fn append_null(&mut self);
35
36    fn field(&self) -> &Field {
37        unimplemented!()
38    }
39
40    fn inner_array(&mut self) -> ArrayRef {
41        unimplemented!()
42    }
43
44    fn fast_explode(&self) -> bool {
45        unimplemented!()
46    }
47
48    fn finish(&mut self) -> ListChunked {
49        let arr = self.inner_array();
50
51        let mut ca = ListChunked::new_with_compute_len(Arc::new(self.field().clone()), vec![arr]);
52        if self.fast_explode() {
53            ca.set_fast_explode()
54        }
55        ca
56    }
57}
58
59impl<S: ?Sized> ListBuilderTrait for Box<S>
60where
61    S: ListBuilderTrait,
62{
63    fn append_opt_series(&mut self, opt_s: Option<&Series>) -> PolarsResult<()> {
64        (**self).append_opt_series(opt_s)
65    }
66
67    fn append_series(&mut self, s: &Series) -> PolarsResult<()> {
68        (**self).append_series(s)
69    }
70
71    fn append_null(&mut self) {
72        (**self).append_null()
73    }
74
75    fn finish(&mut self) -> ListChunked {
76        (**self).finish()
77    }
78}
79
80type LargePrimitiveBuilder<T> = MutableListArray<i64, MutablePrimitiveArray<T>>;
81type LargeListBinViewBuilder<T> = MutableListArray<i64, MutableBinaryViewArray<T>>;
82type LargeListBooleanBuilder = MutableListArray<i64, MutableBooleanArray>;
83type LargeListNullBuilder = MutableListArray<i64, MutableNullArray>;
84
85pub fn get_list_builder(
86    inner_type_logical: &DataType,
87    value_capacity: usize,
88    list_capacity: usize,
89    name: PlSmallStr,
90) -> Box<dyn ListBuilderTrait> {
91    match inner_type_logical {
92        #[cfg(feature = "dtype-categorical")]
93        DataType::Categorical(Some(rev_map), ordering) => {
94            return create_categorical_chunked_listbuilder(
95                name,
96                *ordering,
97                list_capacity,
98                value_capacity,
99                rev_map.clone(),
100            );
101        },
102        #[cfg(feature = "dtype-categorical")]
103        DataType::Enum(Some(rev_map), ordering) => {
104            let list_builder = ListEnumCategoricalChunkedBuilder::new(
105                name,
106                *ordering,
107                list_capacity,
108                value_capacity,
109                (**rev_map).clone(),
110            );
111            return Box::new(list_builder);
112        },
113        _ => {},
114    }
115
116    let physical_type = inner_type_logical.to_physical();
117
118    match &physical_type {
119        #[cfg(feature = "object")]
120        DataType::Object(_) => {
121            let builder = get_object_builder(PlSmallStr::EMPTY, 0).get_list_builder(
122                name,
123                value_capacity,
124                list_capacity,
125            );
126            Box::new(builder)
127        },
128        #[cfg(feature = "dtype-struct")]
129        DataType::Struct(_) => Box::new(AnonymousOwnedListBuilder::new(
130            name,
131            list_capacity,
132            Some(inner_type_logical.clone()),
133        )),
134        DataType::Null => Box::new(ListNullChunkedBuilder::new(name, list_capacity)),
135        DataType::List(_) => Box::new(AnonymousOwnedListBuilder::new(
136            name,
137            list_capacity,
138            Some(inner_type_logical.clone()),
139        )),
140        #[cfg(feature = "dtype-array")]
141        DataType::Array(..) => Box::new(AnonymousOwnedListBuilder::new(
142            name,
143            list_capacity,
144            Some(inner_type_logical.clone()),
145        )),
146        #[cfg(feature = "dtype-decimal")]
147        DataType::Decimal(_, _) => Box::new(
148            ListPrimitiveChunkedBuilder::<Int128Type>::new_with_values_type(
149                name,
150                list_capacity,
151                value_capacity,
152                physical_type,
153                inner_type_logical.clone(),
154            ),
155        ),
156        _ => {
157            macro_rules! get_primitive_builder {
158                ($type:ty) => {{
159                    let builder = ListPrimitiveChunkedBuilder::<$type>::new(
160                        name,
161                        list_capacity,
162                        value_capacity,
163                        inner_type_logical.clone(),
164                    );
165                    Box::new(builder)
166                }};
167            }
168            macro_rules! get_bool_builder {
169                () => {{
170                    let builder =
171                        ListBooleanChunkedBuilder::new(name, list_capacity, value_capacity);
172                    Box::new(builder)
173                }};
174            }
175            macro_rules! get_string_builder {
176                () => {{
177                    let builder =
178                        ListStringChunkedBuilder::new(name, list_capacity, 5 * value_capacity);
179                    Box::new(builder)
180                }};
181            }
182            macro_rules! get_binary_builder {
183                () => {{
184                    let builder =
185                        ListBinaryChunkedBuilder::new(name, list_capacity, 5 * value_capacity);
186                    Box::new(builder)
187                }};
188            }
189            match_dtype_to_logical_apply_macro!(
190                physical_type,
191                get_primitive_builder,
192                get_string_builder,
193                get_binary_builder,
194                get_bool_builder
195            )
196        },
197    }
198}