polars_compute/cast/
decimal_to.rs

1use arrow::array::*;
2use arrow::datatypes::ArrowDataType;
3use arrow::types::NativeType;
4use num_traits::{AsPrimitive, Float, NumCast};
5use polars_error::PolarsResult;
6
7#[inline]
8fn decimal_to_decimal_impl<F: Fn(i128) -> Option<i128>>(
9    from: &PrimitiveArray<i128>,
10    op: F,
11    to_precision: usize,
12    to_scale: usize,
13) -> PrimitiveArray<i128> {
14    let upper_bound_for_precision = 10_i128.saturating_pow(to_precision as u32);
15    let lower_bound_for_precision = upper_bound_for_precision.saturating_neg();
16
17    let values = from.iter().map(|x| {
18        x.and_then(|x| {
19            op(*x).and_then(|x| {
20                if x >= upper_bound_for_precision || x <= lower_bound_for_precision {
21                    None
22                } else {
23                    Some(x)
24                }
25            })
26        })
27    });
28    PrimitiveArray::<i128>::from_trusted_len_iter(values)
29        .to(ArrowDataType::Decimal(to_precision, to_scale))
30}
31
32/// Returns a [`PrimitiveArray<i128>`] with the cast values. Values are `None` on overflow
33pub fn decimal_to_decimal(
34    from: &PrimitiveArray<i128>,
35    to_precision: usize,
36    to_scale: usize,
37) -> PrimitiveArray<i128> {
38    let (from_precision, from_scale) =
39        if let ArrowDataType::Decimal(p, s) = from.dtype().to_logical_type() {
40            (*p, *s)
41        } else {
42            panic!("internal error: i128 is always a decimal")
43        };
44
45    if to_scale == from_scale && to_precision >= from_precision {
46        // fast path
47        return from
48            .clone()
49            .to(ArrowDataType::Decimal(to_precision, to_scale));
50    }
51    // todo: other fast paths include increasing scale and precision by so that
52    // a number will never overflow (validity is preserved)
53
54    if from_scale > to_scale {
55        let factor = 10_i128.pow((from_scale - to_scale) as u32);
56        decimal_to_decimal_impl(
57            from,
58            |x: i128| x.checked_div(factor),
59            to_precision,
60            to_scale,
61        )
62    } else {
63        let factor = 10_i128.pow((to_scale - from_scale) as u32);
64        decimal_to_decimal_impl(
65            from,
66            |x: i128| x.checked_mul(factor),
67            to_precision,
68            to_scale,
69        )
70    }
71}
72
73pub(super) fn decimal_to_decimal_dyn(
74    from: &dyn Array,
75    to_precision: usize,
76    to_scale: usize,
77) -> PolarsResult<Box<dyn Array>> {
78    let from = from.as_any().downcast_ref().unwrap();
79    Ok(Box::new(decimal_to_decimal(from, to_precision, to_scale)))
80}
81
82/// Returns a [`PrimitiveArray<i128>`] with the cast values. Values are `None` on overflow
83pub fn decimal_to_float<T>(from: &PrimitiveArray<i128>) -> PrimitiveArray<T>
84where
85    T: NativeType + Float,
86    f64: AsPrimitive<T>,
87{
88    let (_, from_scale) = if let ArrowDataType::Decimal(p, s) = from.dtype().to_logical_type() {
89        (*p, *s)
90    } else {
91        panic!("internal error: i128 is always a decimal")
92    };
93
94    let div = 10_f64.powi(from_scale as i32);
95    let values = from
96        .values()
97        .iter()
98        .map(|x| (*x as f64 / div).as_())
99        .collect();
100
101    PrimitiveArray::<T>::new(T::PRIMITIVE.into(), values, from.validity().cloned())
102}
103
104pub(super) fn decimal_to_float_dyn<T>(from: &dyn Array) -> PolarsResult<Box<dyn Array>>
105where
106    T: NativeType + Float,
107    f64: AsPrimitive<T>,
108{
109    let from = from.as_any().downcast_ref().unwrap();
110    Ok(Box::new(decimal_to_float::<T>(from)))
111}
112
113/// Returns a [`PrimitiveArray<i128>`] with the cast values. Values are `None` on overflow
114pub fn decimal_to_integer<T>(from: &PrimitiveArray<i128>) -> PrimitiveArray<T>
115where
116    T: NativeType + NumCast,
117{
118    let (_, from_scale) = if let ArrowDataType::Decimal(p, s) = from.dtype().to_logical_type() {
119        (*p, *s)
120    } else {
121        panic!("internal error: i128 is always a decimal")
122    };
123
124    let factor = 10_i128.pow(from_scale as u32);
125    let values = from.iter().map(|x| x.and_then(|x| T::from(*x / factor)));
126
127    PrimitiveArray::from_trusted_len_iter(values)
128}
129
130pub(super) fn decimal_to_integer_dyn<T>(from: &dyn Array) -> PolarsResult<Box<dyn Array>>
131where
132    T: NativeType + NumCast,
133{
134    let from = from.as_any().downcast_ref().unwrap();
135    Ok(Box::new(decimal_to_integer::<T>(from)))
136}
137
138/// Returns a [`Utf8Array`] where every element is the utf8 representation of the decimal.
139#[cfg(feature = "dtype-decimal")]
140pub(super) fn decimal_to_utf8view(from: &PrimitiveArray<i128>) -> Utf8ViewArray {
141    use arrow::compute::decimal::DecimalFmtBuffer;
142
143    let (_, from_scale) = if let ArrowDataType::Decimal(p, s) = from.dtype().to_logical_type() {
144        (*p, *s)
145    } else {
146        panic!("internal error: i128 is always a decimal")
147    };
148
149    let mut mutable = MutableBinaryViewArray::with_capacity(from.len());
150    let mut fmt_buf = DecimalFmtBuffer::new();
151    for &x in from.values().iter() {
152        mutable.push_value_ignore_validity(fmt_buf.format(x, from_scale, false))
153    }
154
155    mutable.freeze().with_validity(from.validity().cloned())
156}
157
158#[cfg(feature = "dtype-decimal")]
159pub(super) fn decimal_to_utf8view_dyn(from: &dyn Array) -> Utf8ViewArray {
160    let from = from.as_any().downcast_ref().unwrap();
161    decimal_to_utf8view(from)
162}