polars_core/chunked_array/logical/
decimal.rs

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

polars_core/chunked_array/logical/
decimal.rs
