polars_core/chunked_array/builder/list/
anonymous.rs

1use super::*;
2
3pub struct AnonymousListBuilder<'a> {
4    name: PlSmallStr,
5    builder: AnonymousBuilder<'a>,
6    fast_explode: bool,
7    inner_dtype: Option<DataType>,
8}
9
10impl Default for AnonymousListBuilder<'_> {
11    fn default() -> Self {
12        Self::new(PlSmallStr::EMPTY, 0, None)
13    }
14}
15
16impl<'a> AnonymousListBuilder<'a> {
17    pub fn new(name: PlSmallStr, capacity: usize, inner_dtype: Option<DataType>) -> Self {
18        Self {
19            name,
20            builder: AnonymousBuilder::new(capacity),
21            fast_explode: true,
22            inner_dtype,
23        }
24    }
25
26    pub fn append_opt_series(&mut self, opt_s: Option<&'a Series>) -> PolarsResult<()> {
27        match opt_s {
28            Some(s) => return self.append_series(s),
29            None => {
30                self.append_null();
31            },
32        }
33        Ok(())
34    }
35
36    pub fn append_opt_array(&mut self, opt_s: Option<&'a dyn Array>) {
37        match opt_s {
38            Some(s) => self.append_array(s),
39            None => {
40                self.append_null();
41            },
42        }
43    }
44
45    pub fn append_array(&mut self, arr: &'a dyn Array) {
46        self.builder.push(arr)
47    }
48
49    #[inline]
50    pub fn append_null(&mut self) {
51        self.fast_explode = false;
52        self.builder.push_null();
53    }
54
55    #[inline]
56    pub fn append_empty(&mut self) {
57        self.fast_explode = false;
58        self.builder.push_empty()
59    }
60
61    pub fn append_series(&mut self, s: &'a Series) -> PolarsResult<()> {
62        match (s.dtype(), &self.inner_dtype) {
63            (DataType::Null, _) => {},
64            (dt, None) => self.inner_dtype = Some(dt.clone()),
65            (dt, Some(set_dt)) => {
66                polars_bail!(ComputeError: "dtypes don't match, got {dt}, expected: {set_dt}")
67            },
68        }
69        if s.is_empty() {
70            self.append_empty();
71        } else {
72            self.builder.push_multiple(s.chunks());
73        }
74        Ok(())
75    }
76
77    pub fn finish(&mut self) -> ListChunked {
78        // Don't use self from here on out.
79        let slf = std::mem::take(self);
80        if slf.builder.is_empty() {
81            ListChunked::full_null_with_dtype(
82                slf.name.clone(),
83                0,
84                &slf.inner_dtype.unwrap_or(DataType::Null),
85            )
86        } else {
87            let inner_dtype_physical = self
88                .inner_dtype
89                .as_ref()
90                .map(|dt| dt.to_physical().to_arrow(CompatLevel::newest()));
91            let arr = slf.builder.finish(inner_dtype_physical.as_ref()).unwrap();
92
93            let list_dtype_logical = match &self.inner_dtype {
94                None => DataType::from_arrow_dtype(arr.dtype()),
95                Some(dt) => DataType::List(Box::new(dt.clone())),
96            };
97
98            let mut ca = ListChunked::with_chunk(PlSmallStr::EMPTY, arr);
99            if slf.fast_explode {
100                ca.set_fast_explode();
101            }
102            ca.field = Arc::new(Field::new(slf.name, list_dtype_logical));
103            ca
104        }
105    }
106}
107
108pub struct AnonymousOwnedListBuilder {
109    name: PlSmallStr,
110    builder: AnonymousBuilder<'static>,
111    owned: Vec<Series>,
112    inner_dtype: Option<DataType>,
113    fast_explode: bool,
114}
115
116impl Default for AnonymousOwnedListBuilder {
117    fn default() -> Self {
118        Self::new(PlSmallStr::EMPTY, 0, None)
119    }
120}
121
122impl ListBuilderTrait for AnonymousOwnedListBuilder {
123    fn append_series(&mut self, s: &Series) -> PolarsResult<()> {
124        match (s.dtype(), &self.inner_dtype) {
125            (DataType::Null, _) => {},
126            (dt, None) => self.inner_dtype = Some(dt.clone()),
127            (dt, Some(set_dt)) => {
128                polars_ensure!(dt == set_dt, ComputeError: "dtypes don't match, got {dt}, expected: {set_dt}")
129            },
130        }
131        if s.is_empty() {
132            self.append_empty();
133        } else {
134            unsafe {
135                self.builder
136                    .push_multiple(&*(s.chunks().as_ref() as *const [ArrayRef]));
137            }
138            // This make sure that the underlying ArrayRef's are not dropped.
139            self.owned.push(s.clone());
140        }
141        Ok(())
142    }
143
144    #[inline]
145    fn append_null(&mut self) {
146        self.fast_explode = false;
147        self.builder.push_null()
148    }
149
150    fn finish(&mut self) -> ListChunked {
151        let inner_dtype = std::mem::take(&mut self.inner_dtype);
152        // Don't use self from here on out.
153        let slf = std::mem::take(self);
154        let inner_dtype_physical = inner_dtype
155            .as_ref()
156            .map(|dt| dt.to_physical().to_arrow(CompatLevel::newest()));
157        let arr = slf.builder.finish(inner_dtype_physical.as_ref()).unwrap();
158
159        let list_dtype_logical = match inner_dtype {
160            None => DataType::from_arrow_dtype(arr.dtype()),
161            Some(dt) => DataType::List(Box::new(dt)),
162        };
163
164        let mut ca = ListChunked::with_chunk(PlSmallStr::EMPTY, arr);
165        if slf.fast_explode {
166            ca.set_fast_explode();
167        }
168        ca.field = Arc::new(Field::new(slf.name, list_dtype_logical));
169        ca
170    }
171}
172
173impl AnonymousOwnedListBuilder {
174    pub fn new(name: PlSmallStr, capacity: usize, inner_dtype: Option<DataType>) -> Self {
175        Self {
176            name,
177            builder: AnonymousBuilder::new(capacity),
178            owned: Vec::with_capacity(capacity),
179            inner_dtype,
180            fast_explode: true,
181        }
182    }
183
184    #[inline]
185    pub fn append_empty(&mut self) {
186        self.fast_explode = false;
187        self.builder.push_empty()
188    }
189}
polars_core/chunked_array/builder/list/anonymous.rs

polars_core/chunked_array/builder/list/
anonymous.rs
