polars_core/chunked_array/logical/
decimal.rs

1use std::borrow::Cow;
2
3use super::*;
4use crate::chunked_array::cast::cast_chunks;
5use crate::prelude::*;
6
7pub type DecimalChunked = Logical<DecimalType, Int128Type>;
8
9impl Int128Chunked {
10    #[inline]
11    pub fn into_decimal_unchecked(self, precision: Option<usize>, scale: usize) -> DecimalChunked {
12        let mut dt = DecimalChunked::new_logical(self);
13        dt.2 = Some(DataType::Decimal(precision, Some(scale)));
14        dt
15    }
16
17    pub fn into_decimal(
18        self,
19        precision: Option<usize>,
20        scale: usize,
21    ) -> PolarsResult<DecimalChunked> {
22        // TODO: if precision is None, do we check that the value fits within precision of 38?...
23        if let Some(precision) = precision {
24            let precision_max = 10_i128.pow(precision as u32);
25            if let Some((min, max)) = self.min_max() {
26                let max_abs = max.abs().max(min.abs());
27                polars_ensure!(
28                    max_abs < precision_max,
29                    ComputeError: "decimal precision {} can't fit values with {} digits",
30                    precision,
31                    max_abs.to_string().len()
32                );
33            }
34        }
35        Ok(self.into_decimal_unchecked(precision, scale))
36    }
37}
38
39impl LogicalType for DecimalChunked {
40    fn dtype(&self) -> &DataType {
41        self.2.as_ref().unwrap()
42    }
43
44    #[inline]
45    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
46        polars_ensure!(i < self.len(), oob = i, self.len());
47        Ok(unsafe { self.get_any_value_unchecked(i) })
48    }
49
50    #[inline]
51    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
52        match self.0.get_unchecked(i) {
53            Some(v) => AnyValue::Decimal(v, self.scale()),
54            None => AnyValue::Null,
55        }
56    }
57
58    fn cast_with_options(
59        &self,
60        dtype: &DataType,
61        cast_options: CastOptions,
62    ) -> PolarsResult<Series> {
63        let mut dtype = Cow::Borrowed(dtype);
64        if let DataType::Decimal(to_precision, to_scale) = dtype.as_ref() {
65            let from_precision = self.precision();
66            let from_scale = self.scale();
67
68            let to_precision = to_precision.or(from_precision);
69            let to_scale = to_scale.unwrap_or(from_scale);
70
71            if to_precision == from_precision && to_scale == from_scale {
72                return Ok(self.clone().into_series());
73            }
74
75            dtype = Cow::Owned(DataType::Decimal(to_precision, Some(to_scale)));
76        }
77
78        let arrow_dtype = self.dtype().to_arrow(CompatLevel::newest());
79        let chunks = self
80            .chunks
81            .iter()
82            .map(|arr| {
83                arr.as_any()
84                    .downcast_ref::<PrimitiveArray<i128>>()
85                    .unwrap()
86                    .clone()
87                    .to(arrow_dtype.clone())
88                    .to_boxed()
89            })
90            .collect::<Vec<_>>();
91        let chunks = cast_chunks(&chunks, dtype.as_ref(), cast_options)?;
92        Series::try_from((self.name().clone(), chunks))
93    }
94}
95
96impl DecimalChunked {
97    pub fn precision(&self) -> Option<usize> {
98        match self.2.as_ref().unwrap() {
99            DataType::Decimal(precision, _) => *precision,
100            _ => unreachable!(),
101        }
102    }
103
104    pub fn scale(&self) -> usize {
105        match self.2.as_ref().unwrap() {
106            DataType::Decimal(_, scale) => scale.unwrap_or_else(|| unreachable!()),
107            _ => unreachable!(),
108        }
109    }
110
111    pub fn to_scale(&self, scale: usize) -> PolarsResult<Cow<'_, Self>> {
112        if self.scale() == scale {
113            return Ok(Cow::Borrowed(self));
114        }
115
116        let mut precision = self.precision();
117        if let Some(ref mut precision) = precision {
118            if self.scale() < scale {
119                *precision += scale;
120                *precision = (*precision).min(38);
121            }
122        }
123
124        let s = self.cast_with_options(
125            &DataType::Decimal(precision, Some(scale)),
126            CastOptions::NonStrict,
127        )?;
128        Ok(Cow::Owned(s.decimal().unwrap().clone()))
129    }
130}