polars_core/chunked_array/object/
mod.rs

1use std::any::Any;
2use std::fmt::{Debug, Display};
3use std::hash::Hash;
4
5use arrow::bitmap::utils::{BitmapIter, ZipValidity};
6use arrow::bitmap::Bitmap;
7use arrow::buffer::Buffer;
8use polars_utils::total_ord::TotalHash;
9
10use crate::prelude::*;
11
12pub mod builder;
13#[cfg(feature = "object")]
14pub(crate) mod extension;
15mod is_valid;
16mod iterator;
17pub mod registry;
18
19pub use extension::set_polars_allow_extension;
20
21#[derive(Debug, Clone)]
22pub struct ObjectArray<T>
23where
24    T: PolarsObject,
25{
26    values: Buffer<T>,
27    validity: Option<Bitmap>,
28}
29
30/// Trimmed down object safe polars object
31pub trait PolarsObjectSafe: Any + Debug + Send + Sync + Display {
32    fn type_name(&self) -> &'static str;
33
34    fn as_any(&self) -> &dyn Any;
35
36    fn to_boxed(&self) -> Box<dyn PolarsObjectSafe>;
37
38    fn equal(&self, other: &dyn PolarsObjectSafe) -> bool;
39}
40
41impl PartialEq for &dyn PolarsObjectSafe {
42    fn eq(&self, other: &Self) -> bool {
43        self.equal(*other)
44    }
45}
46
47/// Values need to implement this so that they can be stored into a Series and DataFrame
48pub trait PolarsObject:
49    Any + Debug + Clone + Send + Sync + Default + Display + Hash + TotalHash + PartialEq + Eq + TotalEq
50{
51    /// This should be used as type information. Consider this a part of the type system.
52    fn type_name() -> &'static str;
53}
54
55impl<T: PolarsObject> PolarsObjectSafe for T {
56    fn type_name(&self) -> &'static str {
57        T::type_name()
58    }
59
60    fn as_any(&self) -> &dyn Any {
61        self
62    }
63
64    fn to_boxed(&self) -> Box<dyn PolarsObjectSafe> {
65        Box::new(self.clone())
66    }
67
68    fn equal(&self, other: &dyn PolarsObjectSafe) -> bool {
69        let Some(other) = other.as_any().downcast_ref::<T>() else {
70            return false;
71        };
72        self == other
73    }
74}
75
76pub type ObjectValueIter<'a, T> = std::slice::Iter<'a, T>;
77
78impl<T> ObjectArray<T>
79where
80    T: PolarsObject,
81{
82    pub fn values_iter(&self) -> ObjectValueIter<'_, T> {
83        self.values.iter()
84    }
85
86    /// Returns an iterator of `Option<&T>` over every element of this array.
87    pub fn iter(&self) -> ZipValidity<&T, ObjectValueIter<'_, T>, BitmapIter> {
88        ZipValidity::new_with_validity(self.values_iter(), self.validity.as_ref())
89    }
90
91    /// Get a value at a certain index location
92    pub fn value(&self, index: usize) -> &T {
93        &self.values[index]
94    }
95
96    pub fn get(&self, index: usize) -> Option<&T> {
97        if self.is_valid(index) {
98            Some(unsafe { self.value_unchecked(index) })
99        } else {
100            None
101        }
102    }
103
104    /// Get a value at a certain index location
105    ///
106    /// # Safety
107    ///
108    /// This does not any bound checks. The caller needs to ensure the index is within
109    /// the size of the array.
110    pub unsafe fn value_unchecked(&self, index: usize) -> &T {
111        self.values.get_unchecked(index)
112    }
113
114    /// Check validity
115    ///
116    /// # Safety
117    /// No bounds checks
118    #[inline]
119    pub unsafe fn is_valid_unchecked(&self, i: usize) -> bool {
120        if let Some(b) = &self.validity {
121            b.get_bit_unchecked(i)
122        } else {
123            true
124        }
125    }
126
127    /// Check validity
128    ///
129    /// # Safety
130    /// No bounds checks
131    #[inline]
132    pub unsafe fn is_null_unchecked(&self, i: usize) -> bool {
133        !self.is_valid_unchecked(i)
134    }
135
136    /// Returns this array with a new validity.
137    /// # Panic
138    /// Panics iff `validity.len() != self.len()`.
139    #[must_use]
140    #[inline]
141    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
142        self.set_validity(validity);
143        self
144    }
145
146    /// Sets the validity of this array.
147    /// # Panics
148    /// This function panics iff `validity.len() != self.len()`.
149    #[inline]
150    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
151        if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
152            panic!("validity must be equal to the array's length")
153        }
154        self.validity = validity;
155    }
156}
157
158impl<T> Array for ObjectArray<T>
159where
160    T: PolarsObject,
161{
162    fn as_any(&self) -> &dyn Any {
163        self
164    }
165
166    fn dtype(&self) -> &ArrowDataType {
167        &ArrowDataType::FixedSizeBinary(size_of::<T>())
168    }
169
170    fn slice(&mut self, offset: usize, length: usize) {
171        assert!(
172            offset + length <= self.len(),
173            "the offset of the new Buffer cannot exceed the existing length"
174        );
175        unsafe { self.slice_unchecked(offset, length) }
176    }
177
178    unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
179        self.validity = self
180            .validity
181            .take()
182            .map(|bitmap| bitmap.sliced_unchecked(offset, length))
183            .filter(|bitmap| bitmap.unset_bits() > 0);
184        self.values.slice_unchecked(offset, length);
185    }
186
187    fn split_at_boxed(&self, offset: usize) -> (Box<dyn Array>, Box<dyn Array>) {
188        let (lhs, rhs) = Splitable::split_at(self, offset);
189        (Box::new(lhs), Box::new(rhs))
190    }
191
192    unsafe fn split_at_boxed_unchecked(&self, offset: usize) -> (Box<dyn Array>, Box<dyn Array>) {
193        let (lhs, rhs) = unsafe { Splitable::split_at_unchecked(self, offset) };
194        (Box::new(lhs), Box::new(rhs))
195    }
196
197    fn len(&self) -> usize {
198        self.values.len()
199    }
200
201    fn validity(&self) -> Option<&Bitmap> {
202        self.validity.as_ref()
203    }
204
205    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
206        Box::new(self.clone().with_validity(validity))
207    }
208
209    fn to_boxed(&self) -> Box<dyn Array> {
210        Box::new(self.clone())
211    }
212
213    fn as_any_mut(&mut self) -> &mut dyn Any {
214        unimplemented!()
215    }
216
217    fn null_count(&self) -> usize {
218        match &self.validity {
219            None => 0,
220            Some(validity) => validity.unset_bits(),
221        }
222    }
223}
224
225impl<T: PolarsObject> Splitable for ObjectArray<T> {
226    fn check_bound(&self, offset: usize) -> bool {
227        offset <= self.len()
228    }
229
230    unsafe fn _split_at_unchecked(&self, offset: usize) -> (Self, Self) {
231        let (left_values, right_values) = unsafe { self.values.split_at_unchecked(offset) };
232        let (left_validity, right_validity) = unsafe { self.validity.split_at_unchecked(offset) };
233        (
234            Self {
235                values: left_values,
236                validity: left_validity,
237            },
238            Self {
239                values: right_values,
240                validity: right_validity,
241            },
242        )
243    }
244}
245
246impl<T: PolarsObject> StaticArray for ObjectArray<T> {
247    type ValueT<'a> = &'a T;
248    type ZeroableValueT<'a> = Option<&'a T>;
249    type ValueIterT<'a> = ObjectValueIter<'a, T>;
250
251    #[inline]
252    unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_> {
253        self.value_unchecked(idx)
254    }
255
256    fn values_iter(&self) -> Self::ValueIterT<'_> {
257        self.values_iter()
258    }
259
260    fn iter(&self) -> ZipValidity<Self::ValueT<'_>, Self::ValueIterT<'_>, BitmapIter> {
261        self.iter()
262    }
263
264    fn with_validity_typed(self, validity: Option<Bitmap>) -> Self {
265        self.with_validity(validity)
266    }
267
268    fn full_null(length: usize, _dtype: ArrowDataType) -> Self {
269        ObjectArray {
270            values: vec![T::default(); length].into(),
271            validity: Some(Bitmap::new_with_value(false, length)),
272        }
273    }
274}
275
276impl<T: PolarsObject> ParameterFreeDtypeStaticArray for ObjectArray<T> {
277    fn get_dtype() -> ArrowDataType {
278        ArrowDataType::FixedSizeBinary(size_of::<T>())
279    }
280}
281
282impl<T> ObjectChunked<T>
283where
284    T: PolarsObject,
285{
286    /// Get a hold to an object that can be formatted or downcasted via the Any trait.
287    ///
288    /// # Safety
289    ///
290    /// No bounds checks
291    pub unsafe fn get_object_unchecked(&self, index: usize) -> Option<&dyn PolarsObjectSafe> {
292        let (chunk_idx, idx) = self.index_to_chunked_index(index);
293        self.get_object_chunked_unchecked(chunk_idx, idx)
294    }
295
296    pub(crate) unsafe fn get_object_chunked_unchecked(
297        &self,
298        chunk: usize,
299        index: usize,
300    ) -> Option<&dyn PolarsObjectSafe> {
301        let chunks = self.downcast_chunks();
302        let arr = chunks.get_unchecked(chunk);
303        if arr.is_valid_unchecked(index) {
304            Some(arr.value(index))
305        } else {
306            None
307        }
308    }
309
310    /// Get a hold to an object that can be formatted or downcasted via the Any trait.
311    pub fn get_object(&self, index: usize) -> Option<&dyn PolarsObjectSafe> {
312        if index < self.len() {
313            unsafe { self.get_object_unchecked(index) }
314        } else {
315            None
316        }
317    }
318}
319
320impl<T: PolarsObject> From<Vec<T>> for ObjectArray<T> {
321    fn from(values: Vec<T>) -> Self {
322        Self {
323            values: values.into(),
324            validity: None,
325        }
326    }
327}
polars_core/chunked_array/object/mod.rs

polars_core/chunked_array/object/
mod.rs
