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::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::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(cats, mapping) => {
85            with_match_categorical_physical_type!(cats.physical(), |$C| {
86                type A = <$C as PolarsDataType>::Array;
87                let arr = &*(arr as *const dyn Array as *const A);
88                let cat_id = arr.value_unchecked(idx).as_cat();
89                AnyValue::Categorical(cat_id, mapping)
90            })
91        },
92        #[cfg(feature = "dtype-categorical")]
93        DataType::Enum(fcats, mapping) => {
94            with_match_categorical_physical_type!(fcats.physical(), |$C| {
95                type A = <$C as PolarsDataType>::Array;
96                let arr = &*(arr as *const dyn Array as *const A);
97                let cat_id = arr.value_unchecked(idx).as_cat();
98                AnyValue::Enum(cat_id, mapping)
99            })
100        },
101        #[cfg(feature = "dtype-struct")]
102        DataType::Struct(flds) => {
103            let arr = &*(arr as *const dyn Array as *const StructArray);
104            AnyValue::Struct(idx, arr, flds)
105        },
106        #[cfg(feature = "dtype-datetime")]
107        DataType::Datetime(tu, tz) => {
108            let arr = &*(arr as *const dyn Array as *const Int64Array);
109            let v = arr.value_unchecked(idx);
110            AnyValue::Datetime(v, *tu, tz.as_ref())
111        },
112        #[cfg(feature = "dtype-date")]
113        DataType::Date => {
114            let arr = &*(arr as *const dyn Array as *const Int32Array);
115            let v = arr.value_unchecked(idx);
116            AnyValue::Date(v)
117        },
118        #[cfg(feature = "dtype-duration")]
119        DataType::Duration(tu) => {
120            let arr = &*(arr as *const dyn Array as *const Int64Array);
121            let v = arr.value_unchecked(idx);
122            AnyValue::Duration(v, *tu)
123        },
124        #[cfg(feature = "dtype-time")]
125        DataType::Time => {
126            let arr = &*(arr as *const dyn Array as *const Int64Array);
127            let v = arr.value_unchecked(idx);
128            AnyValue::Time(v)
129        },
130        #[cfg(feature = "dtype-decimal")]
131        DataType::Decimal(precision, scale) => {
132            let arr = &*(arr as *const dyn Array as *const Int128Array);
133            let v = arr.value_unchecked(idx);
134            AnyValue::Decimal(v, *precision, *scale)
135        },
136        #[cfg(feature = "dtype-extension")]
137        DataType::Extension(typ, storage) => arr_to_any_value(arr, idx, storage),
138        #[cfg(feature = "object")]
139        DataType::Object(_) => {
140            // We should almost never hit this. The only known exception is when we put objects in
141            // structs. Any other hit should be considered a bug.
142            let arr = arr.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap();
143            PolarsExtension::arr_to_av(arr, idx)
144        },
145        DataType::Null => AnyValue::Null,
146        DataType::BinaryOffset => downcast_and_pack!(LargeBinaryArray, Binary),
147        dt => panic!("not implemented for {dt:?}"),
148    }
149}
150
151#[cfg(feature = "dtype-struct")]
152impl<'a> AnyValue<'a> {
153    pub fn _iter_struct_av(&self) -> impl Iterator<Item = AnyValue<'_>> {
154        let AnyValue::Struct(idx, arr, flds) = self else {
155            unreachable!()
156        };
157        unsafe {
158            arr.values()
159                .iter()
160                .zip(*flds)
161                .map(move |(arr, fld)| arr_to_any_value(&**arr, *idx, fld.dtype()))
162        }
163    }
164
165    pub fn _materialize_struct_av(&'a self, buf: &mut Vec<AnyValue<'a>>) {
166        let iter = self._iter_struct_av();
167        buf.extend(iter)
168    }
169}
170
171macro_rules! get_any_value_unchecked {
172    ($self:ident, $index:expr) => {{
173        let (chunk_idx, idx) = $self.index_to_chunked_index($index);
174        debug_assert!(chunk_idx < $self.chunks.len());
175        let arr = &**$self.chunks.get_unchecked(chunk_idx);
176        debug_assert!(idx < arr.len());
177        arr_to_any_value(arr, idx, $self.dtype())
178    }};
179}
180
181macro_rules! get_any_value {
182    ($self:ident, $index:expr) => {{
183        if $index >= $self.len() {
184            polars_bail!(oob = $index, $self.len());
185        }
186        // SAFETY:
187        // bounds are checked
188        Ok(unsafe { $self.get_any_value_unchecked($index) })
189    }};
190}
191
192impl<T> ChunkAnyValue for ChunkedArray<T>
193where
194    T: PolarsNumericType,
195{
196    #[inline]
197    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
198        get_any_value_unchecked!(self, index)
199    }
200
201    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
202        get_any_value!(self, index)
203    }
204}
205
206impl ChunkAnyValue for BooleanChunked {
207    #[inline]
208    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
209        get_any_value_unchecked!(self, index)
210    }
211
212    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
213        get_any_value!(self, index)
214    }
215}
216
217impl ChunkAnyValue for StringChunked {
218    #[inline]
219    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
220        get_any_value_unchecked!(self, index)
221    }
222
223    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
224        get_any_value!(self, index)
225    }
226}
227
228impl ChunkAnyValue for BinaryChunked {
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 BinaryOffsetChunked {
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 ListChunked {
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
261#[cfg(feature = "dtype-array")]
262impl ChunkAnyValue for ArrayChunked {
263    #[inline]
264    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
265        get_any_value_unchecked!(self, index)
266    }
267
268    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
269        get_any_value!(self, index)
270    }
271}
272
273#[cfg(feature = "object")]
274impl<T: PolarsObject> ChunkAnyValue for ObjectChunked<T> {
275    #[inline]
276    unsafe fn get_any_value_unchecked(&self, index: usize) -> AnyValue<'_> {
277        match self.get_object_unchecked(index) {
278            None => AnyValue::Null,
279            Some(v) => AnyValue::Object(v),
280        }
281    }
282
283    fn get_any_value(&self, index: usize) -> PolarsResult<AnyValue<'_>> {
284        get_any_value!(self, index)
285    }
286}
287
288impl ChunkAnyValue for NullChunked {
289    #[inline]
290    unsafe fn get_any_value_unchecked(&self, _index: usize) -> AnyValue<'_> {
291        AnyValue::Null
292    }
293
294    fn get_any_value(&self, _index: usize) -> PolarsResult<AnyValue<'_>> {
295        Ok(AnyValue::Null)
296    }
297}
298
299#[cfg(feature = "dtype-struct")]
300impl ChunkAnyValue for StructChunked {
301    /// Gets AnyValue from LogicalType
302    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
303        polars_ensure!(i < self.len(), oob = i, self.len());
304        unsafe { Ok(self.get_any_value_unchecked(i)) }
305    }
306
307    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
308        let (chunk_idx, idx) = index_to_chunked_index(self.chunks.iter().map(|c| c.len()), i);
309        if let DataType::Struct(flds) = self.dtype() {
310            // SAFETY: we already have a single chunk and we are
311            // guarded by the type system.
312            unsafe {
313                let arr = &**self.chunks.get_unchecked(chunk_idx);
314                let arr = &*(arr as *const dyn Array as *const StructArray);
315
316                if arr.is_null_unchecked(idx) {
317                    AnyValue::Null
318                } else {
319                    AnyValue::Struct(idx, arr, flds)
320                }
321            }
322        } else {
323            unreachable!()
324        }
325    }
326}
polars_core/chunked_array/ops/any_value.rs

polars_core/chunked_array/ops/
any_value.rs
