polars_ops/chunked_array/binary/
cast_binary_to_numerical.rs

1use arrow::array::{Array, BinaryViewArray, PrimitiveArray};
2use arrow::datatypes::ArrowDataType;
3use arrow::types::NativeType;
4use polars_error::PolarsResult;
5
6/// Trait for casting bytes to a primitive type
7pub trait Cast {
8    fn cast_le(val: &[u8]) -> Option<Self>
9    where
10        Self: Sized;
11    fn cast_be(val: &[u8]) -> Option<Self>
12    where
13        Self: Sized;
14}
15macro_rules! impl_cast {
16    ($primitive_type:ident) => {
17        impl Cast for $primitive_type {
18            fn cast_le(val: &[u8]) -> Option<Self> {
19                Some($primitive_type::from_le_bytes(val.try_into().ok()?))
20            }
21
22            fn cast_be(val: &[u8]) -> Option<Self> {
23                Some($primitive_type::from_be_bytes(val.try_into().ok()?))
24            }
25        }
26    };
27}
28
29impl_cast!(i8);
30impl_cast!(i16);
31impl_cast!(i32);
32impl_cast!(i64);
33impl_cast!(i128);
34impl_cast!(u8);
35impl_cast!(u16);
36impl_cast!(u32);
37impl_cast!(u64);
38impl_cast!(u128);
39impl_cast!(f32);
40impl_cast!(f64);
41
42/// Casts a [`BinaryArray`] to a [`PrimitiveArray`], making any uncastable value a Null.
43pub(super) fn cast_binview_to_primitive<T>(
44    from: &BinaryViewArray,
45    to: &ArrowDataType,
46    is_little_endian: bool,
47) -> PrimitiveArray<T>
48where
49    T: Cast + NativeType,
50{
51    let iter = from.iter().map(|x| {
52        x.and_then::<T, _>(|x| {
53            if is_little_endian {
54                T::cast_le(x)
55            } else {
56                T::cast_be(x)
57            }
58        })
59    });
60
61    PrimitiveArray::<T>::from_trusted_len_iter(iter).to(to.clone())
62}
63
64/// Casts a [`BinaryArray`] to a [`PrimitiveArray`], making any uncastable value a Null.
65pub(super) fn cast_binview_to_primitive_dyn<T>(
66    from: &dyn Array,
67    to: &ArrowDataType,
68    is_little_endian: bool,
69) -> PolarsResult<Box<dyn Array>>
70where
71    T: Cast + NativeType,
72{
73    let from = from.as_any().downcast_ref().unwrap();
74
75    Ok(Box::new(cast_binview_to_primitive::<T>(
76        from,
77        to,
78        is_little_endian,
79    )))
80}