polars_core/chunked_array/ops/
any_value.rs

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

polars_core/chunked_array/ops/
any_value.rs
