polars_core/chunked_array/ops/
any_value.rs

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

polars_core/chunked_array/ops/
any_value.rs
