polars_core/chunked_array/object/extension/
list.rs

1use arrow::offset::Offsets;
2
3use crate::chunked_array::object::builder::ObjectChunkedBuilder;
4use crate::chunked_array::object::extension::create_extension;
5use crate::prelude::*;
6
7impl<T: PolarsObject> ObjectChunked<T> {
8    pub(crate) fn get_list_builder(
9        name: PlSmallStr,
10        values_capacity: usize,
11        list_capacity: usize,
12    ) -> Box<dyn ListBuilderTrait> {
13        Box::new(ExtensionListBuilder::<T>::new(
14            name,
15            values_capacity,
16            list_capacity,
17        ))
18    }
19}
20
21pub(crate) struct ExtensionListBuilder<T: PolarsObject> {
22    values_builder: ObjectChunkedBuilder<T>,
23    offsets: Vec<i64>,
24    fast_explode: bool,
25}
26
27impl<T: PolarsObject> ExtensionListBuilder<T> {
28    pub(crate) fn new(name: PlSmallStr, values_capacity: usize, list_capacity: usize) -> Self {
29        let mut offsets = Vec::with_capacity(list_capacity + 1);
30        offsets.push(0);
31        Self {
32            values_builder: ObjectChunkedBuilder::new(name, values_capacity),
33            offsets,
34            fast_explode: true,
35        }
36    }
37}
38
39impl<T: PolarsObject> ListBuilderTrait for ExtensionListBuilder<T> {
40    fn append_series(&mut self, s: &Series) -> PolarsResult<()> {
41        let arr = s.as_any().downcast_ref::<ObjectChunked<T>>().unwrap();
42
43        for v in arr.iter() {
44            self.values_builder.append_option(v.cloned())
45        }
46        if arr.is_empty() {
47            self.fast_explode = false;
48        }
49        let len_so_far = self.offsets[self.offsets.len() - 1];
50        self.offsets.push(len_so_far + arr.len() as i64);
51        Ok(())
52    }
53
54    fn append_null(&mut self) {
55        self.values_builder.append_null();
56        let len_so_far = self.offsets[self.offsets.len() - 1];
57        self.offsets.push(len_so_far + 1);
58    }
59
60    fn finish(&mut self) -> ListChunked {
61        let values_builder = std::mem::take(&mut self.values_builder);
62        let offsets = std::mem::take(&mut self.offsets);
63        let ca = values_builder.finish();
64        let obj_arr = ca.downcast_chunks().get(0).unwrap().clone();
65
66        // SAFETY: this is safe because we just created the PolarsExtension
67        // meaning that the sentinel is heap allocated and the dereference of
68        // the pointer does not fail.
69        let mut pe = create_extension(obj_arr.into_iter_cloned());
70        unsafe { pe.set_to_series_fn::<T>() };
71        let extension_array = Box::new(pe.take_and_forget()) as ArrayRef;
72        let extension_dtype = extension_array.dtype();
73
74        let dtype = ListArray::<i64>::default_datatype(extension_dtype.clone());
75        let arr = ListArray::<i64>::new(
76            dtype,
77            // SAFETY: offsets are monotonically increasing.
78            unsafe { Offsets::new_unchecked(offsets).into() },
79            extension_array,
80            None,
81        );
82
83        let mut listarr = ListChunked::with_chunk(ca.name().clone(), arr);
84        if self.fast_explode {
85            listarr.set_fast_explode()
86        }
87        listarr
88    }
89}
polars_core/chunked_array/object/extension/list.rs

polars_core/chunked_array/object/extension/
list.rs
