polars_core/chunked_array/ops/
any_value.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2
3use crate::prelude::*;
4use crate::series::implementations::null::NullChunked;
5use crate::utils::index_to_chunked_index;
6
7/// # Safety
8/// `idx` MUST be in-bounds for `arr` and `dtype` has to match the data stored in `arr`.
9#[inline]
10#[allow(unused_variables)]
11pub(crate) unsafe fn arr_to_any_value<'a>(
12    arr: &'a dyn Array,
13    idx: usize,
14    dtype: &'a DataType,
15) -> AnyValue<'a> {
16    debug_assert!(idx < arr.len());
17    if arr.is_null(idx) {
18        return AnyValue::Null;
19    }
20
21    macro_rules! downcast_and_pack {
22        ($casttype:ident, $variant:ident) => {{
23            let arr = &*(arr as *const dyn Array as *const $casttype);
24            let v = arr.value_unchecked(idx);
25            AnyValue::$variant(v)
26        }};
27    }
28    macro_rules! downcast {
29        ($casttype:ident) => {{
30            let arr = &*(arr as *const dyn Array as *const $casttype);
31            arr.value_unchecked(idx)
32        }};
33    }
34    match dtype {
35        DataType::String => downcast_and_pack!(Utf8ViewArray, String),
36        DataType::Binary => downcast_and_pack!(BinaryViewArray, Binary),
37        DataType::Boolean => downcast_and_pack!(BooleanArray, Boolean),
38        DataType::UInt8 => downcast_and_pack!(UInt8Array, UInt8),
39        DataType::UInt16 => downcast_and_pack!(UInt16Array, UInt16),
40        DataType::UInt32 => downcast_and_pack!(UInt32Array, UInt32),
41        DataType::UInt64 => downcast_and_pack!(UInt64Array, UInt64),
42        DataType::UInt128 => downcast_and_pack!(UInt128Array, UInt128),
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::Float16 => downcast_and_pack!(Float16Array, Float16),
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(cats, mapping) => {
86            with_match_categorical_physical_type!(cats.physical(), |$C| {
87                type A = <$C as PolarsDataType>::Array;
88                let arr = &*(arr as *const dyn Array as *const A);
89                let cat_id = arr.value_unchecked(idx).as_cat();
90                AnyValue::Categorical(cat_id, mapping)
91            })
92        },
93        #[cfg(feature = "dtype-categorical")]
94        DataType::Enum(fcats, mapping) => {
95            with_match_categorical_physical_type!(fcats.physical(), |$C| {
96                type A = <$C as PolarsDataType>::Array;
97                let arr = &*(arr as *const dyn Array as *const A);
98                let cat_id = arr.value_unchecked(idx).as_cat();
99                AnyValue::Enum(cat_id, mapping)
100            })
101        },
102        #[cfg(feature = "dtype-struct")]
103        DataType::Struct(flds) => {
104            let arr = &*(arr as *const dyn Array as *const StructArray);
105            AnyValue::Struct(idx, arr, flds)
106        },
107        #[cfg(feature = "dtype-datetime")]
108        DataType::Datetime(tu, tz) => {
109            let arr = &*(arr as *const dyn Array as *const Int64Array);
110            let v = arr.value_unchecked(idx);
111            AnyValue::Datetime(v, *tu, tz.as_ref())
112        },
113        #[cfg(feature = "dtype-date")]
114        DataType::Date => {
115            let arr = &*(arr as *const dyn Array as *const Int32Array);
116            let v = arr.value_unchecked(idx);
117            AnyValue::Date(v)
118        },
119        #[cfg(feature = "dtype-duration")]
120        DataType::Duration(tu) => {
121            let arr = &*(arr as *const dyn Array as *const Int64Array);
122            let v = arr.value_unchecked(idx);
123            AnyValue::Duration(v, *tu)
124        },
125        #[cfg(feature = "dtype-time")]
126        DataType::Time => {
127            let arr = &*(arr as *const dyn Array as *const Int64Array);
128            let v = arr.value_unchecked(idx);
129            AnyValue::Time(v)
130        },
131        #[cfg(feature = "dtype-decimal")]
132        DataType::Decimal(precision, scale) => {
133            let arr = &*(arr as *const dyn Array as *const Int128Array);
134            let v = arr.value_unchecked(idx);
135            AnyValue::Decimal(v, *precision, *scale)
136        },
137        #[cfg(feature = "dtype-map")]
138        DataType::Map(_, _) => {
139            let entries_dtype = dtype.map_entries_dtype().unwrap();
140            let v: ArrayRef = downcast!(LargeListArray);
141            let s = Series::from_chunks_and_dtype_unchecked(
142                PlSmallStr::EMPTY,
143                vec![v],
144                &entries_dtype.to_physical(),
145            )
146            .from_physical_unchecked(&entries_dtype)
147            .unwrap();
148            AnyValue::Map(s)
149        },
150        #[cfg(feature = "dtype-extension")]
151        DataType::Extension(typ, storage) => arr_to_any_value(arr, idx, storage),
152        #[cfg(feature = "object")]
153        DataType::Object(_) => {
154            use crate::chunked_array::object::registry::get_object_array_getter;
155            get_object_array_getter()(arr, idx).unwrap()
156        },
157        DataType::Null => AnyValue::Null,
158        DataType::BinaryOffset => downcast_and_pack!(LargeBinaryArray, Binary),
159        dt => panic!("not implemented for {dt:?}"),
160    }
161}
162
163#[cfg(feature = "dtype-struct")]
164impl<'a> AnyValue<'a> {
165    pub fn _iter_struct_av(&self) -> impl Iterator<Item = AnyValue<'_>> {
166        let AnyValue::Struct(idx, arr, flds) = self else {
167            unreachable!()
168        };
169        unsafe {
170            arr.values()
171                .iter()
172                .zip(*flds)
173                .map(move |(arr, fld)| arr_to_any_value(&**arr, *idx, fld.dtype()))
174        }
175    }
176
177    pub fn _materialize_struct_av(&'a self, buf: &mut Vec<AnyValue<'a>>) {
178        let iter = self._iter_struct_av();
179        buf.extend(iter)
180    }
181}
182
183macro_rules! get_any_value_unchecked {
184    ($self:ident, $index:expr) => {{
185        let (chunk_idx, idx) = $self.index_to_chunked_index($index);
186        debug_assert!(chunk_idx < $self.chunks.len());
187        let arr = &**$self.chunks.get_unchecked(chunk_idx);
188        debug_assert!(idx < arr.len());
189        arr_to_any_value(arr, idx, $self.dtype())
190    }};
191}
192
193macro_rules! get_any_value {
194    ($self:ident, $index:expr) => {{
195        if $index >= $self.len() {
196            polars_bail!(oob = $index, $self.len());
197        }
198        // SAFETY:
199        // bounds are checked
200        Ok(unsafe { $self.get_any_value_unchecked($index) })
201    }};
202}
203
204impl<T> ChunkAnyValue for ChunkedArray<T>
205where
206    T: PolarsNumericType,
207{
208    #[inline]
209    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
210        get_any_value_unchecked!(self, index)
211    }
212
213    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
214        get_any_value!(self, index)
215    }
216}
217
218impl ChunkAnyValue for BooleanChunked {
219    #[inline]
220    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
221        get_any_value_unchecked!(self, index)
222    }
223
224    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
225        get_any_value!(self, index)
226    }
227}
228
229impl ChunkAnyValue for StringChunked {
230    #[inline]
231    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
232        get_any_value_unchecked!(self, index)
233    }
234
235    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
236        get_any_value!(self, index)
237    }
238}
239
240impl ChunkAnyValue for BinaryChunked {
241    #[inline]
242    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
243        get_any_value_unchecked!(self, index)
244    }
245
246    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
247        get_any_value!(self, index)
248    }
249}
250
251impl ChunkAnyValue for BinaryOffsetChunked {
252    #[inline]
253    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
254        get_any_value_unchecked!(self, index)
255    }
256
257    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
258        get_any_value!(self, index)
259    }
260}
261
262impl ChunkAnyValueBypassValidity for BinaryOffsetChunked {
263    #[inline]
264    unsafe fn get_any_value_bypass_validity(&self, index: usize) -> AnyValue<'_> {
265        debug_assert!(index < self.len());
266        let (chunk_idx, idx) = self.index_to_chunked_index(index);
267        debug_assert!(chunk_idx < self.chunks.len());
268        let arr = &**self.chunks.get_unchecked(chunk_idx);
269        let arr = &*(arr as *const dyn Array as *const LargeBinaryArray);
270        let v = arr.value_unchecked(idx);
271        AnyValue::Binary(v)
272    }
273}
274
275impl ChunkAnyValue for ListChunked {
276    #[inline]
277    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
278        get_any_value_unchecked!(self, index)
279    }
280
281    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
282        get_any_value!(self, index)
283    }
284}
285
286#[cfg(feature = "dtype-array")]
287impl ChunkAnyValue for ArrayChunked {
288    #[inline]
289    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
290        get_any_value_unchecked!(self, index)
291    }
292
293    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
294        get_any_value!(self, index)
295    }
296}
297
298#[cfg(feature = "object")]
299impl<T: PolarsObject> ChunkAnyValue for ObjectChunked<T> {
300    #[inline]
301    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
302        match self.get_object_unchecked(index) {
303            None => AnyValue::Null,
304            Some(v) => AnyValue::Object(v),
305        }
306    }
307
308    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
309        get_any_value!(self, index)
310    }
311}
312
313impl ChunkAnyValue for NullChunked {
314    #[inline]
315    unsafe fn get_any_value_unchecked(&self, _index: usize) -> AnyValue<'_> {
316        AnyValue::Null
317    }
318
319    fn get_any_value(&self, _index: usize) -> PolarsResult<AnyValue<'_>> {
320        Ok(AnyValue::Null)
321    }
322}
323
324#[cfg(feature = "dtype-struct")]
325impl ChunkAnyValue for StructChunked {
326    /// Gets AnyValue from LogicalType
327    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
328        polars_ensure!(i < self.len(), oob = i, self.len());
329        unsafe { Ok(self.get_any_value_unchecked(i)) }
330    }
331
332    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
333        let (chunk_idx, idx) = index_to_chunked_index(self.chunks.iter().map(|c| c.len()), i);
334        if let DataType::Struct(flds) = self.dtype() {
335            // SAFETY: we already have a single chunk and we are
336            // guarded by the type system.
337            unsafe {
338                let arr = &**self.chunks.get_unchecked(chunk_idx);
339                let arr = &*(arr as *const dyn Array as *const StructArray);
340
341                if arr.is_null_unchecked(idx) {
342                    AnyValue::Null
343                } else {
344                    AnyValue::Struct(idx, arr, flds)
345                }
346            }
347        } else {
348            unreachable!()
349        }
350    }
351}
polars_core/chunked_array/ops/any_value.rs

polars_core/chunked_array/ops/
any_value.rs
