polars_core/chunked_array/builder/
fixed_size_list.rs

1use arrow::types::NativeType;
2use polars_utils::pl_str::PlSmallStr;
3
4use crate::prelude::*;
5
6pub(crate) struct FixedSizeListNumericBuilder<T: NativeType> {
7    inner: Option<MutableFixedSizeListArray<MutablePrimitiveArray<T>>>,
8    width: usize,
9    name: PlSmallStr,
10    logical_dtype: DataType,
11}
12
13impl<T: NativeType> FixedSizeListNumericBuilder<T> {
14    /// # Safety
15    ///
16    /// The caller must ensure that the physical numerical type match logical type.
17    pub(crate) unsafe fn new(
18        name: PlSmallStr,
19        width: usize,
20        capacity: usize,
21        logical_dtype: DataType,
22    ) -> Self {
23        let mp = MutablePrimitiveArray::<T>::with_capacity(capacity * width);
24        let inner = Some(MutableFixedSizeListArray::new(mp, width));
25        Self {
26            inner,
27            width,
28            name,
29            logical_dtype,
30        }
31    }
32}
33
34pub trait FixedSizeListBuilder {
35    /// # Safety
36    ///
37    /// `arr` must have at least `(offset + 1) * width` valid elements of the
38    /// builder's expected inner type
39    unsafe fn push_unchecked(&mut self, arr: &dyn Array, offset: usize);
40    /// # Safety
41    ///
42    /// The builder must have been properly initialized
43    unsafe fn push_null(&mut self);
44    fn finish(&mut self) -> ArrayChunked;
45}
46
47impl<T: NativeType> FixedSizeListBuilder for FixedSizeListNumericBuilder<T> {
48    #[inline]
49    unsafe fn push_unchecked(&mut self, arr: &dyn Array, offset: usize) {
50        let start = offset * self.width;
51        let end = start + self.width;
52        let arr = arr
53            .as_any()
54            .downcast_ref::<PrimitiveArray<T>>()
55            .unwrap_unchecked();
56        let inner = self.inner.as_mut().unwrap_unchecked();
57
58        let values = arr.values().as_slice();
59        let validity = arr.validity();
60        if let Some(validity) = validity {
61            let iter = (start..end).map(|i| {
62                if validity.get_bit_unchecked(i) {
63                    Some(*values.get_unchecked(i))
64                } else {
65                    None
66                }
67            });
68            inner.push_unchecked(Some(iter))
69        } else {
70            let iter = (start..end).map(|i| Some(*values.get_unchecked(i)));
71            inner.push_unchecked(Some(iter))
72        }
73    }
74
75    #[inline]
76    unsafe fn push_null(&mut self) {
77        let inner = self.inner.as_mut().unwrap_unchecked();
78        inner.push_null()
79    }
80
81    fn finish(&mut self) -> ArrayChunked {
82        let arr: FixedSizeListArray = self.inner.take().unwrap().into();
83        // SAFETY: physical type matches the logical
84        unsafe {
85            ChunkedArray::from_chunks_and_dtype(
86                self.name.clone(),
87                vec![Box::new(arr)],
88                DataType::Array(Box::new(self.logical_dtype.clone()), self.width),
89            )
90        }
91    }
92}
93
94pub(crate) struct AnonymousOwnedFixedSizeListBuilder {
95    inner: fixed_size_list::AnonymousBuilder,
96    name: PlSmallStr,
97    inner_dtype: Option<DataType>,
98}
99
100impl AnonymousOwnedFixedSizeListBuilder {
101    pub(crate) fn new(
102        name: PlSmallStr,
103        width: usize,
104        capacity: usize,
105        inner_dtype: Option<DataType>,
106    ) -> Self {
107        let inner = fixed_size_list::AnonymousBuilder::new(capacity, width);
108        Self {
109            inner,
110            name,
111            inner_dtype,
112        }
113    }
114}
115
116impl FixedSizeListBuilder for AnonymousOwnedFixedSizeListBuilder {
117    #[inline]
118    unsafe fn push_unchecked(&mut self, arr: &dyn Array, offset: usize) {
119        let arr = arr.sliced_unchecked(offset * self.inner.width, self.inner.width);
120        self.inner.push(arr)
121    }
122
123    #[inline]
124    unsafe fn push_null(&mut self) {
125        self.inner.push_null()
126    }
127
128    fn finish(&mut self) -> ArrayChunked {
129        let arr = std::mem::take(&mut self.inner)
130            .finish(
131                self.inner_dtype
132                    .as_ref()
133                    .map(|dt| dt.to_arrow(CompatLevel::newest()))
134                    .as_ref(),
135            )
136            .unwrap();
137        ChunkedArray::with_chunk(self.name.clone(), arr)
138    }
139}
140
141pub fn get_fixed_size_list_builder(
142    inner_type_logical: &DataType,
143    capacity: usize,
144    width: usize,
145    name: PlSmallStr,
146) -> PolarsResult<Box<dyn FixedSizeListBuilder>> {
147    let phys_dtype = inner_type_logical.to_physical();
148
149    let builder = if phys_dtype.is_primitive_numeric() {
150        with_match_physical_numeric_type!(phys_dtype, |$T| {
151        // SAFETY: physical type match logical type
152        unsafe {
153            Box::new(FixedSizeListNumericBuilder::<$T>::new(name, width, capacity,inner_type_logical.clone())) as Box<dyn FixedSizeListBuilder>
154        }
155        })
156    } else {
157        Box::new(AnonymousOwnedFixedSizeListBuilder::new(
158            name,
159            width,
160            capacity,
161            Some(inner_type_logical.clone()),
162        ))
163    };
164    Ok(builder)
165}
polars_core/chunked_array/builder/fixed_size_list.rs

polars_core/chunked_array/builder/
fixed_size_list.rs
