polars_core/chunked_array/logical/
decimal.rs

1use std::borrow::Cow;
2
3use arrow::bitmap::Bitmap;
4use polars_compute::decimal::{dec128_fits, dec128_rescale, dec128_verify_prec_scale};
5
6use super::*;
7use crate::chunked_array::cast::cast_chunks;
8use crate::prelude::arity::{unary_elementwise, unary_kernel};
9use crate::prelude::*;
10
11pub type DecimalChunked = Logical<DecimalType, Int128Type>;
12
13impl Int128Chunked {
14    #[inline]
15    pub fn into_decimal_unchecked(self, precision: usize, scale: usize) -> DecimalChunked {
16        // SAFETY: no invalid states (from a safety perspective).
17        unsafe { DecimalChunked::new_logical(self, DataType::Decimal(precision, scale)) }
18    }
19
20    pub fn into_decimal(self, precision: usize, scale: usize) -> PolarsResult<DecimalChunked> {
21        dec128_verify_prec_scale(precision, scale)?;
22        if let Some((min, max)) = self.min_max() {
23            let max_abs = max.abs().max(min.abs());
24            polars_ensure!(
25                dec128_fits(max_abs, precision),
26                ComputeError: "decimal precision {} can't fit values with {} digits",
27                precision,
28                max_abs.to_string().len()
29            );
30        }
31        Ok(self.into_decimal_unchecked(precision, scale))
32    }
33}
34
35impl LogicalType for DecimalChunked {
36    fn dtype(&self) -> &DataType {
37        &self.dtype
38    }
39
40    #[inline]
41    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
42        polars_ensure!(i < self.len(), oob = i, self.len());
43        Ok(unsafe { self.get_any_value_unchecked(i) })
44    }
45
46    #[inline]
47    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
48        match self.phys.get_unchecked(i) {
49            Some(v) => AnyValue::Decimal(v, self.precision(), self.scale()),
50            None => AnyValue::Null,
51        }
52    }
53
54    fn cast_with_options(
55        &self,
56        dtype: &DataType,
57        cast_options: CastOptions,
58    ) -> PolarsResult<Series> {
59        if let DataType::Decimal(to_prec, to_scale) = dtype {
60            return Ok(self
61                .with_prec_scale(*to_prec, *to_scale, cast_options.is_strict())?
62                .into_owned()
63                .into_series());
64        }
65
66        match dtype {
67            DataType::Decimal(to_prec, to_scale) => {
68                return Ok(self
69                    .with_prec_scale(*to_prec, *to_scale, cast_options.is_strict())?
70                    .into_owned()
71                    .into_series());
72            },
73
74            dt if dt.is_primitive_numeric()
75                | matches!(dt, DataType::String | DataType::Boolean) =>
76            {
77                // Normally we don't set the Arrow logical type, but now we temporarily set it so
78                // we can re-use the compute cast kernels.
79                let arrow_dtype = self.dtype().to_arrow(CompatLevel::newest());
80                let chunks = self
81                    .physical()
82                    .chunks
83                    .iter()
84                    .map(|arr| {
85                        arr.as_any()
86                            .downcast_ref::<PrimitiveArray<i128>>()
87                            .unwrap()
88                            .clone()
89                            .to(arrow_dtype.clone())
90                            .to_boxed()
91                    })
92                    .collect::<Vec<_>>();
93                let chunks = cast_chunks(&chunks, dtype, cast_options)?;
94                Series::try_from((self.name().clone(), chunks))
95            },
96
97            dt => {
98                polars_bail!(
99                    InvalidOperation:
100                    "casting from {:?} to {:?} not supported",
101                    self.dtype(), dt
102                )
103            },
104        }
105    }
106}
107
108impl DecimalChunked {
109    pub fn precision(&self) -> usize {
110        match &self.dtype {
111            DataType::Decimal(precision, _) => *precision,
112            _ => unreachable!(),
113        }
114    }
115
116    pub fn scale(&self) -> usize {
117        match &self.dtype {
118            DataType::Decimal(_, scale) => *scale,
119            _ => unreachable!(),
120        }
121    }
122
123    pub fn with_prec_scale(
124        &self,
125        prec: usize,
126        scale: usize,
127        strict: bool,
128    ) -> PolarsResult<Cow<'_, Self>> {
129        if self.precision() == prec && self.scale() == scale {
130            return Ok(Cow::Borrowed(self));
131        }
132
133        dec128_verify_prec_scale(prec, scale)?;
134        let phys = if self.scale() == scale {
135            if prec >= self.precision() {
136                // Increasing precision is always allowed.
137                self.phys.clone()
138            } else if strict {
139                if let Some((min, max)) = self.phys.min_max() {
140                    let max_abs = max.abs().max(min.abs());
141                    polars_ensure!(
142                        dec128_fits(max_abs, prec),
143                        ComputeError: "decimal precision {} can't fit values with {} digits",
144                        prec,
145                        max_abs.to_string().len()
146                    );
147                }
148                self.phys.clone()
149            } else {
150                unary_kernel(&self.phys, |arr| {
151                    let new_valid: Bitmap = arr
152                        .iter()
153                        .map(|opt_x| {
154                            if let Some(x) = opt_x {
155                                dec128_fits(*x, prec)
156                            } else {
157                                false
158                            }
159                        })
160                        .collect();
161                    arr.clone().with_validity_typed(Some(new_valid))
162                })
163            }
164        } else {
165            let old_s = self.scale();
166            unary_elementwise(&self.phys, |x| dec128_rescale(x?, old_s, prec, scale))
167        };
168
169        let ca = unsafe { DecimalChunked::new_logical(phys, DataType::Decimal(prec, scale)) };
170        Ok(Cow::Owned(ca))
171    }
172
173    /// Converts self to a physical representation with the given precision and
174    /// scale, returning the given sentinel value instead for values which don't
175    /// fit in the given precision and scale. This can be useful for comparisons.
176    pub fn into_phys_with_prec_scale_or_sentinel(
177        &self,
178        prec: usize,
179        scale: usize,
180        sentinel: i128,
181    ) -> Int128Chunked {
182        if self.precision() <= prec && self.scale() == scale {
183            return self.phys.clone();
184        }
185
186        let old_s = self.scale();
187        unary_elementwise(&self.phys, |x| {
188            Some(dec128_rescale(x?, old_s, prec, scale).unwrap_or(sentinel))
189        })
190    }
191}
polars_core/chunked_array/logical/decimal.rs

polars_core/chunked_array/logical/
decimal.rs
