Skip to main content

polars_core/chunked_array/builder/list/
mod.rs

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