polars_core/chunked_array/ops/
any_value.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2#[cfg(feature = "dtype-categorical")]
3use polars_utils::sync::SyncPtr;
4
5#[cfg(feature = "object")]
6use crate::chunked_array::object::extension::polars_extension::PolarsExtension;
7use crate::prelude::*;
8use crate::series::implementations::null::NullChunked;
9use crate::utils::index_to_chunked_index;
10
11#[inline]
12#[allow(unused_variables)]
13pub(crate) unsafe fn arr_to_any_value<'a>(
14    arr: &'a dyn Array,
15    idx: usize,
16    dtype: &'a DataType,
17) -> AnyValue<'a> {
18    debug_assert!(idx < arr.len());
19    if arr.is_null(idx) {
20        return AnyValue::Null;
21    }
22
23    macro_rules! downcast_and_pack {
24        ($casttype:ident, $variant:ident) => {{
25            let arr = &*(arr as *const dyn Array as *const $casttype);
26            let v = arr.value_unchecked(idx);
27            AnyValue::$variant(v)
28        }};
29    }
30    macro_rules! downcast {
31        ($casttype:ident) => {{
32            let arr = &*(arr as *const dyn Array as *const $casttype);
33            arr.value_unchecked(idx)
34        }};
35    }
36    match dtype {
37        DataType::String => downcast_and_pack!(Utf8ViewArray, String),
38        DataType::Binary => downcast_and_pack!(BinaryViewArray, Binary),
39        DataType::Boolean => downcast_and_pack!(BooleanArray, Boolean),
40        DataType::UInt8 => downcast_and_pack!(UInt8Array, UInt8),
41        DataType::UInt16 => downcast_and_pack!(UInt16Array, UInt16),
42        DataType::UInt32 => downcast_and_pack!(UInt32Array, UInt32),
43        DataType::UInt64 => downcast_and_pack!(UInt64Array, UInt64),
44        DataType::Int8 => downcast_and_pack!(Int8Array, Int8),
45        DataType::Int16 => downcast_and_pack!(Int16Array, Int16),
46        DataType::Int32 => downcast_and_pack!(Int32Array, Int32),
47        DataType::Int64 => downcast_and_pack!(Int64Array, Int64),
48        DataType::Int128 => downcast_and_pack!(Int128Array, Int128),
49        DataType::Float32 => downcast_and_pack!(Float32Array, Float32),
50        DataType::Float64 => downcast_and_pack!(Float64Array, Float64),
51        DataType::List(dt) => {
52            let v: ArrayRef = downcast!(LargeListArray);
53            if dt.is_primitive() {
54                let s = Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, vec![v], dt);
55                AnyValue::List(s)
56            } else {
57                let s = Series::from_chunks_and_dtype_unchecked(
58                    PlSmallStr::EMPTY,
59                    vec![v],
60                    &dt.to_physical(),
61                )
62                .from_physical_unchecked(dt)
63                .unwrap();
64                AnyValue::List(s)
65            }
66        },
67        #[cfg(feature = "dtype-array")]
68        DataType::Array(dt, width) => {
69            let v: ArrayRef = downcast!(FixedSizeListArray);
70            if dt.is_primitive() {
71                let s = Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, vec![v], dt);
72                AnyValue::Array(s, *width)
73            } else {
74                let s = Series::from_chunks_and_dtype_unchecked(
75                    PlSmallStr::EMPTY,
76                    vec![v],
77                    &dt.to_physical(),
78                )
79                .from_physical_unchecked(dt)
80                .unwrap();
81                AnyValue::Array(s, *width)
82            }
83        },
84        #[cfg(feature = "dtype-categorical")]
85        DataType::Categorical(rev_map, _) => {
86            let arr = &*(arr as *const dyn Array as *const UInt32Array);
87            let v = arr.value_unchecked(idx);
88            AnyValue::Categorical(v, rev_map.as_ref().unwrap().as_ref(), SyncPtr::new_null())
89        },
90        #[cfg(feature = "dtype-categorical")]
91        DataType::Enum(rev_map, _) => {
92            let arr = &*(arr as *const dyn Array as *const UInt32Array);
93            let v = arr.value_unchecked(idx);
94            AnyValue::Enum(v, rev_map.as_ref().unwrap().as_ref(), SyncPtr::new_null())
95        },
96        #[cfg(feature = "dtype-struct")]
97        DataType::Struct(flds) => {
98            let arr = &*(arr as *const dyn Array as *const StructArray);
99            AnyValue::Struct(idx, arr, flds)
100        },
101        #[cfg(feature = "dtype-datetime")]
102        DataType::Datetime(tu, tz) => {
103            let arr = &*(arr as *const dyn Array as *const Int64Array);
104            let v = arr.value_unchecked(idx);
105            AnyValue::Datetime(v, *tu, tz.as_ref())
106        },
107        #[cfg(feature = "dtype-date")]
108        DataType::Date => {
109            let arr = &*(arr as *const dyn Array as *const Int32Array);
110            let v = arr.value_unchecked(idx);
111            AnyValue::Date(v)
112        },
113        #[cfg(feature = "dtype-duration")]
114        DataType::Duration(tu) => {
115            let arr = &*(arr as *const dyn Array as *const Int64Array);
116            let v = arr.value_unchecked(idx);
117            AnyValue::Duration(v, *tu)
118        },
119        #[cfg(feature = "dtype-time")]
120        DataType::Time => {
121            let arr = &*(arr as *const dyn Array as *const Int64Array);
122            let v = arr.value_unchecked(idx);
123            AnyValue::Time(v)
124        },
125        #[cfg(feature = "dtype-decimal")]
126        DataType::Decimal(precision, scale) => {
127            let arr = &*(arr as *const dyn Array as *const Int128Array);
128            let v = arr.value_unchecked(idx);
129            AnyValue::Decimal(v, scale.unwrap_or_else(|| unreachable!()))
130        },
131        #[cfg(feature = "object")]
132        DataType::Object(_) => {
133            // We should almost never hit this. The only known exception is when we put objects in
134            // structs. Any other hit should be considered a bug.
135            let arr = arr.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap();
136            PolarsExtension::arr_to_av(arr, idx)
137        },
138        DataType::Null => AnyValue::Null,
139        DataType::BinaryOffset => downcast_and_pack!(LargeBinaryArray, Binary),
140        dt => panic!("not implemented for {dt:?}"),
141    }
142}
143
144#[cfg(feature = "dtype-struct")]
145impl<'a> AnyValue<'a> {
146    pub fn _iter_struct_av(&self) -> impl Iterator<Item = AnyValue> {
147        match self {
148            AnyValue::Struct(idx, arr, flds) => {
149                let idx = *idx;
150                unsafe {
151                    arr.values().iter().zip(*flds).map(move |(arr, fld)| {
152                        // The dictionary arrays categories don't have to map to the rev-map in the dtype
153                        // so we set the array pointer with values of the dictionary array.
154                        #[cfg(feature = "dtype-categorical")]
155                        {
156                            use arrow::legacy::is_valid::IsValid as _;
157                            if let Some(arr) = arr.as_any().downcast_ref::<DictionaryArray<u32>>() {
158                                let keys = arr.keys();
159                                let values = arr.values();
160                                let values =
161                                    values.as_any().downcast_ref::<Utf8ViewArray>().unwrap();
162                                let arr = &*(keys as *const dyn Array as *const UInt32Array);
163
164                                if arr.is_valid_unchecked(idx) {
165                                    let v = arr.value_unchecked(idx);
166                                    match fld.dtype() {
167                                        DataType::Categorical(Some(rev_map), _) => {
168                                            AnyValue::Categorical(
169                                                v,
170                                                rev_map,
171                                                SyncPtr::from_const(values),
172                                            )
173                                        },
174                                        DataType::Enum(Some(rev_map), _) => {
175                                            AnyValue::Enum(v, rev_map, SyncPtr::from_const(values))
176                                        },
177                                        _ => unimplemented!(),
178                                    }
179                                } else {
180                                    AnyValue::Null
181                                }
182                            } else {
183                                arr_to_any_value(&**arr, idx, fld.dtype())
184                            }
185                        }
186
187                        #[cfg(not(feature = "dtype-categorical"))]
188                        {
189                            arr_to_any_value(&**arr, idx, fld.dtype())
190                        }
191                    })
192                }
193            },
194            _ => unreachable!(),
195        }
196    }
197
198    pub fn _materialize_struct_av(&'a self, buf: &mut Vec<AnyValue<'a>>) {
199        let iter = self._iter_struct_av();
200        buf.extend(iter)
201    }
202}
203
204macro_rules! get_any_value_unchecked {
205    ($self:ident, $index:expr) => {{
206        let (chunk_idx, idx) = $self.index_to_chunked_index($index);
207        debug_assert!(chunk_idx < $self.chunks.len());
208        let arr = &**$self.chunks.get_unchecked(chunk_idx);
209        debug_assert!(idx < arr.len());
210        arr_to_any_value(arr, idx, $self.dtype())
211    }};
212}
213
214macro_rules! get_any_value {
215    ($self:ident, $index:expr) => {{
216        if $index >= $self.len() {
217            polars_bail!(oob = $index, $self.len());
218        }
219        // SAFETY:
220        // bounds are checked
221        Ok(unsafe { $self.get_any_value_unchecked($index) })
222    }};
223}
224
225impl<T> ChunkAnyValue for ChunkedArray<T>
226where
227    T: PolarsNumericType,
228{
229    #[inline]
230    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
231        get_any_value_unchecked!(self, index)
232    }
233
234    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
235        get_any_value!(self, index)
236    }
237}
238
239impl ChunkAnyValue for BooleanChunked {
240    #[inline]
241    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
242        get_any_value_unchecked!(self, index)
243    }
244
245    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
246        get_any_value!(self, index)
247    }
248}
249
250impl ChunkAnyValue for StringChunked {
251    #[inline]
252    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
253        get_any_value_unchecked!(self, index)
254    }
255
256    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
257        get_any_value!(self, index)
258    }
259}
260
261impl ChunkAnyValue for BinaryChunked {
262    #[inline]
263    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
264        get_any_value_unchecked!(self, index)
265    }
266
267    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
268        get_any_value!(self, index)
269    }
270}
271
272impl ChunkAnyValue for BinaryOffsetChunked {
273    #[inline]
274    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
275        get_any_value_unchecked!(self, index)
276    }
277
278    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
279        get_any_value!(self, index)
280    }
281}
282
283impl ChunkAnyValue for ListChunked {
284    #[inline]
285    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
286        get_any_value_unchecked!(self, index)
287    }
288
289    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
290        get_any_value!(self, index)
291    }
292}
293
294#[cfg(feature = "dtype-array")]
295impl ChunkAnyValue for ArrayChunked {
296    #[inline]
297    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
298        get_any_value_unchecked!(self, index)
299    }
300
301    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
302        get_any_value!(self, index)
303    }
304}
305
306#[cfg(feature = "object")]
307impl<T: PolarsObject> ChunkAnyValue for ObjectChunked<T> {
308    #[inline]
309    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue {
310        match self.get_object_unchecked(index) {
311            None => AnyValue::Null,
312            Some(v) => AnyValue::Object(v),
313        }
314    }
315
316    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue> {
317        get_any_value!(self, index)
318    }
319}
320
321impl ChunkAnyValue for NullChunked {
322    #[inline]
323    unsafe fn get_any_value_unchecked(&self, _index: usize) -> AnyValue {
324        AnyValue::Null
325    }
326
327    fn get_any_value(&self, _index: usize) -> PolarsResult<AnyValue> {
328        Ok(AnyValue::Null)
329    }
330}
331
332#[cfg(feature = "dtype-struct")]
333impl ChunkAnyValue for StructChunked {
334    /// Gets AnyValue from LogicalType
335    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
336        polars_ensure!(i < self.len(), oob = i, self.len());
337        unsafe { Ok(self.get_any_value_unchecked(i)) }
338    }
339
340    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
341        let (chunk_idx, idx) = index_to_chunked_index(self.chunks.iter().map(|c| c.len()), i);
342        if let DataType::Struct(flds) = self.dtype() {
343            // SAFETY: we already have a single chunk and we are
344            // guarded by the type system.
345            unsafe {
346                let arr = &**self.chunks.get_unchecked(chunk_idx);
347                let arr = &*(arr as *const dyn Array as *const StructArray);
348
349                if arr.is_null_unchecked(idx) {
350                    AnyValue::Null
351                } else {
352                    AnyValue::Struct(idx, arr, flds)
353                }
354            }
355        } else {
356            unreachable!()
357        }
358    }
359}
polars_core/chunked_array/ops/any_value.rs

polars_core/chunked_array/ops/
any_value.rs
