polars_core/chunked_array/object/
builder.rs

1use polars_arrow::array::builder::{ArrayBuilder, ShareStrategy};
2use polars_arrow::bitmap::BitmapBuilder;
3use polars_arrow::datatypes::ExtensionType;
4use polars_utils::vec::PushUnchecked;
5
6use super::*;
7use crate::chunked_array::object::registry::run_with_gil;
8use crate::utils::get_iter_capacity;
9
10pub struct ObjectChunkedBuilder<T> {
11    field: Field,
12    bitmask_builder: BitmapBuilder,
13    values: Vec<T>,
14}
15
16impl<T> ObjectChunkedBuilder<T>
17where
18    T: PolarsObject,
19{
20    pub fn field(&self) -> &Field {
21        &self.field
22    }
23    pub fn new(name: PlSmallStr, capacity: usize) -> Self {
24        ObjectChunkedBuilder {
25            field: Field::new(name, DataType::Object(T::type_name())),
26            values: Vec::with_capacity(capacity),
27            bitmask_builder: BitmapBuilder::with_capacity(capacity),
28        }
29    }
30
31    /// Appends a value of type `T` into the builder
32    #[inline]
33    pub fn append_value(&mut self, v: T) {
34        self.values.push(v);
35        self.bitmask_builder.push(true);
36    }
37
38    /// Appends a null slot into the builder
39    #[inline]
40    pub fn append_null(&mut self) {
41        self.values.push(T::default());
42        self.bitmask_builder.push(false);
43    }
44
45    #[inline]
46    pub fn append_option(&mut self, opt: Option<T>) {
47        match opt {
48            Some(s) => self.append_value(s),
49            None => self.append_null(),
50        }
51    }
52
53    pub fn finish(mut self) -> ObjectChunked<T> {
54        let null_bitmap: Option<Bitmap> = self.bitmask_builder.into_opt_validity();
55
56        let len = self.values.len();
57        let null_count = null_bitmap
58            .as_ref()
59            .map(|validity| validity.unset_bits())
60            .unwrap_or(0);
61
62        let arr = Box::new(ObjectArray {
63            dtype: ArrowDataType::FixedSizeBinary(size_of::<T>()),
64            values: self.values.into(),
65            validity: null_bitmap,
66        });
67
68        self.field.dtype = get_object_type::<T>();
69
70        unsafe { ChunkedArray::new_with_dims(Arc::new(self.field), vec![arr], len, null_count) }
71    }
72}
73
74/// Initialize a polars Object data type. The type has got information needed to
75/// construct new objects.
76pub(crate) fn get_object_type<T: PolarsObject>() -> DataType {
77    DataType::Object(T::type_name())
78}
79
80impl<T> Default for ObjectChunkedBuilder<T>
81where
82    T: PolarsObject,
83{
84    fn default() -> Self {
85        ObjectChunkedBuilder::new(PlSmallStr::EMPTY, 0)
86    }
87}
88
89impl<T> NewChunkedArray<ObjectType<T>, T> for ObjectChunked<T>
90where
91    T: PolarsObject,
92{
93    fn from_slice(name: PlSmallStr, v: &[T]) -> Self {
94        Self::from_iter_values(name, v.iter().cloned())
95    }
96
97    fn from_slice_options(name: PlSmallStr, opt_v: &[Option<T>]) -> Self {
98        let mut builder = ObjectChunkedBuilder::<T>::new(name, opt_v.len());
99        opt_v
100            .iter()
101            .cloned()
102            .for_each(|opt| builder.append_option(opt));
103        builder.finish()
104    }
105
106    fn from_iter_options(
107        name: PlSmallStr,
108        it: impl Iterator<Item = Option<T>>,
109    ) -> ObjectChunked<T> {
110        let mut builder = ObjectChunkedBuilder::new(name, get_iter_capacity(&it));
111        it.for_each(|opt| builder.append_option(opt));
112        builder.finish()
113    }
114
115    /// Create a new ChunkedArray from an iterator.
116    fn from_iter_values(name: PlSmallStr, it: impl Iterator<Item = T>) -> ObjectChunked<T> {
117        let mut builder = ObjectChunkedBuilder::new(name, get_iter_capacity(&it));
118        it.for_each(|v| builder.append_value(v));
119        builder.finish()
120    }
121}
122
123impl<T> ObjectChunked<T>
124where
125    T: PolarsObject,
126{
127    pub fn new_from_vec(name: PlSmallStr, v: Vec<T>) -> Self {
128        let field = Arc::new(Field::new(name, DataType::Object(T::type_name())));
129        let len = v.len();
130        let arr = Box::new(ObjectArray {
131            dtype: ArrowDataType::FixedSizeBinary(size_of::<T>()),
132            values: v.into(),
133            validity: None,
134        });
135
136        unsafe { ObjectChunked::new_with_dims(field, vec![arr], len, 0) }
137    }
138
139    pub fn new_from_vec_and_validity(
140        name: PlSmallStr,
141        v: Vec<T>,
142        validity: Option<Bitmap>,
143    ) -> Self {
144        let field = Arc::new(Field::new(name, DataType::Object(T::type_name())));
145        let len = v.len();
146        let null_count = validity.as_ref().map(|v| v.unset_bits()).unwrap_or(0);
147        let arr = Box::new(ObjectArray {
148            dtype: ArrowDataType::FixedSizeBinary(size_of::<T>()),
149            values: v.into(),
150            validity,
151        });
152
153        unsafe { ObjectChunked::new_with_dims(field, vec![arr], len, null_count) }
154    }
155
156    pub fn new_empty(name: PlSmallStr) -> Self {
157        Self::new_from_vec(name, vec![])
158    }
159}
160
161/// Convert a Series of dtype object to an Arrow Array of FixedSizeBinary
162pub(crate) fn object_series_to_arrow_array(s: &Series) -> ArrayRef {
163    // The list builder knows how to create an arrow array
164    // we simply piggy back on that code.
165
166    // SAFETY: 0..len is in bounds
167    let list_s = unsafe {
168        let groups = vec![[0, s.len() as IdxSize]];
169        s.agg_list(&GroupsType::new_slice(groups, false, true))
170    };
171    let arr = &list_s.chunks()[0];
172    let arr = arr.as_any().downcast_ref::<ListArray<i64>>().unwrap();
173
174    let mut arr: Box<dyn Array> = arr.values().clone();
175
176    if let ArrowDataType::Extension(ext_type) = arr.dtype()
177        && let ExtensionType {
178            name,
179            inner: ArrowDataType::FixedSizeBinary(8),
180            metadata: Some(_),
181        } = ext_type.as_ref()
182        && name == POLARS_OBJECT_EXTENSION_NAME
183    {
184        *arr.dtype_mut() = ArrowDataType::FixedSizeBinary(8);
185    } else {
186        panic!()
187    }
188
189    arr
190}
191
192impl<T: PolarsObject> ArrayBuilder for ObjectChunkedBuilder<T> {
193    fn dtype(&self) -> &ArrowDataType {
194        &ArrowDataType::FixedSizeBinary(size_of::<T>())
195    }
196
197    fn reserve(&mut self, additional: usize) {
198        self.bitmask_builder.reserve(additional);
199        self.values.reserve(additional);
200    }
201
202    fn freeze(self) -> Box<dyn Array> {
203        Box::new(ObjectArray {
204            dtype: ArrowDataType::FixedSizeBinary(size_of::<T>()),
205            values: self.values.into(),
206            validity: self.bitmask_builder.into_opt_validity(),
207        })
208    }
209
210    fn freeze_reset(&mut self) -> Box<dyn Array> {
211        Box::new(ObjectArray {
212            dtype: ArrowDataType::FixedSizeBinary(size_of::<T>()),
213            values: core::mem::take(&mut self.values).into(),
214            validity: core::mem::take(&mut self.bitmask_builder).into_opt_validity(),
215        })
216    }
217
218    fn len(&self) -> usize {
219        self.values.len()
220    }
221
222    fn extend_nulls(&mut self, length: usize) {
223        run_with_gil(|| {
224            self.values.resize(self.values.len() + length, T::default());
225        });
226        self.bitmask_builder.extend_constant(length, false);
227    }
228
229    fn subslice_extend(
230        &mut self,
231        other: &dyn Array,
232        start: usize,
233        length: usize,
234        _share: ShareStrategy,
235    ) {
236        run_with_gil(|| {
237            let other: &ObjectArray<T> = other.as_any().downcast_ref().unwrap();
238            self.values
239                .extend_from_slice(&other.values[start..start + length]);
240            self.bitmask_builder
241                .subslice_extend_from_opt_validity(other.validity(), start, length);
242        })
243    }
244
245    fn subslice_extend_repeated(
246        &mut self,
247        other: &dyn Array,
248        start: usize,
249        length: usize,
250        repeats: usize,
251        _share: ShareStrategy,
252    ) {
253        run_with_gil(|| {
254            for _ in 0..repeats {
255                let other: &ObjectArray<T> = other.as_any().downcast_ref().unwrap();
256                self.values
257                    .extend_from_slice(&other.values[start..start + length]);
258                self.bitmask_builder.subslice_extend_from_opt_validity(
259                    other.validity(),
260                    start,
261                    length,
262                );
263            }
264        })
265    }
266
267    fn subslice_extend_each_repeated(
268        &mut self,
269        other: &dyn Array,
270        start: usize,
271        length: usize,
272        repeats: usize,
273        _share: ShareStrategy,
274    ) {
275        run_with_gil(|| {
276            let other: &ObjectArray<T> = other.as_any().downcast_ref().unwrap();
277
278            self.values.reserve(length * repeats);
279            for value in other.values[start..start + length].iter() {
280                unsafe {
281                    for _ in 0..repeats {
282                        self.values.push_unchecked(value.clone());
283                    }
284                }
285            }
286        });
287
288        self.bitmask_builder
289            .subslice_extend_each_repeated_from_opt_validity(
290                other.validity(),
291                start,
292                length,
293                repeats,
294            );
295    }
296
297    unsafe fn gather_extend(&mut self, other: &dyn Array, idxs: &[IdxSize], _share: ShareStrategy) {
298        run_with_gil(|| {
299            let other: &ObjectArray<T> = other.as_any().downcast_ref().unwrap();
300            let other_values_slice = other.values.as_slice();
301            self.values.extend(
302                idxs.iter()
303                    .map(|idx| other_values_slice.get_unchecked(*idx as usize).clone()),
304            );
305        });
306        self.bitmask_builder
307            .gather_extend_from_opt_validity(other.validity(), idxs, other.len());
308    }
309
310    fn opt_gather_extend(&mut self, other: &dyn Array, idxs: &[IdxSize], _share: ShareStrategy) {
311        run_with_gil(|| {
312            let other: &ObjectArray<T> = other.as_any().downcast_ref().unwrap();
313            let other_values_slice = other.values.as_slice();
314            self.values.reserve(idxs.len());
315            unsafe {
316                for idx in idxs {
317                    let val = if (*idx as usize) < other.len() {
318                        other_values_slice.get_unchecked(*idx as usize).clone()
319                    } else {
320                        T::default()
321                    };
322                    self.values.push_unchecked(val);
323                }
324            }
325        });
326        self.bitmask_builder.opt_gather_extend_from_opt_validity(
327            other.validity(),
328            idxs,
329            other.len(),
330        );
331    }
332}
polars_core/chunked_array/object/builder.rs

polars_core/chunked_array/object/
builder.rs
