polars_core/chunked_array/arithmetic/
decimal.rs

1use std::cell::Cell;
2
3use polars_arrow::bitmap::Bitmap;
4use polars_arrow::compute::utils::combine_validities_and;
5use polars_compute::arithmetic::DivRounding;
6use polars_compute::decimal::{
7    DEC128_MAX_PREC, dec128_add_scaled, dec128_div_scaled, dec128_int_div_scaled,
8    dec128_mul_scaled, dec128_rem_scaled, dec128_rescale, dec128_sub_scaled, i64_binary_values,
9    i64_unary_values,
10};
11
12use super::*;
13use crate::prelude::arity::{
14    apply_binary_kernel_broadcast, broadcast_binary_elementwise_values,
15    broadcast_try_binary_elementwise,
16};
17
18impl DecimalChunked {
19    /// Applies `kernel(l, left_scale, r, right_scale, scale)` elementwise, producing a
20    /// `Decimal(38, scale)`. The kernel returns `None` if the result doesn't fit, or on
21    /// division by zero.
22    fn apply_scaled_kernel(
23        &self,
24        rhs: &Self,
25        scale: usize,
26        op: &str,
27        kernel: impl Fn(i128, usize, i128, usize, usize) -> Option<i128>,
28    ) -> PolarsResult<Self> {
29        let left_s = self.scale();
30        let right_s = rhs.scale();
31        let phys = broadcast_try_binary_elementwise(
32            self.physical(),
33            rhs.physical(),
34            |opt_l, opt_r| {
35                let (Some(l), Some(r)) = (opt_l, opt_r) else {
36                    return PolarsResult::Ok(None);
37                };
38                let ret = kernel(l, left_s, r, right_s, scale).ok_or_else(|| {
39                    // Only division and remainder can fail with a zero operand.
40                    if r == 0 {
41                        polars_err!(ComputeError: "division by zero Decimal")
42                    } else {
43                        polars_err!(
44                            ComputeError: "overflow in decimal {op}: result doesn't fit Decimal({DEC128_MAX_PREC}, {scale})"
45                        )
46                    }
47                })?;
48                Ok(Some(ret))
49            },
50        )?;
51        Ok(phys.into_decimal_unchecked(DEC128_MAX_PREC, scale))
52    }
53
54    /// [`Self::apply_scaled_kernel`] for kernels that only fail on overflow. It first
55    /// computes all slots without per-row validity, and only goes row by row if a slot
56    /// failed, as null slots can hold any value.
57    fn apply_scaled_kernel_values(
58        &self,
59        rhs: &Self,
60        scale: usize,
61        op: &str,
62        kernel: impl Fn(i128, usize, i128, usize, usize) -> Option<i128>,
63    ) -> PolarsResult<Self> {
64        let left_s = self.scale();
65        let right_s = rhs.scale();
66        let mut failed = false;
67        let phys = broadcast_binary_elementwise_values(self.physical(), rhs.physical(), |l, r| {
68            let ret = kernel(l, left_s, r, right_s, scale);
69            failed |= ret.is_none();
70            ret.unwrap_or(0)
71        });
72        if failed {
73            return self.apply_scaled_kernel(rhs, scale, op, kernel);
74        }
75        Ok(phys.into_decimal_unchecked(DEC128_MAX_PREC, scale))
76    }
77
78    /// Applies `op` to all slots as i64s. Returns `None` if a slot, null ones included,
79    /// doesn't fit an i64.
80    fn apply_i64_values(
81        &self,
82        rhs: &Self,
83        scale: usize,
84        op: impl Fn(i64, i64) -> i128,
85    ) -> Option<Self> {
86        let failed = Cell::new(false);
87        let to_arr = |values: Option<Vec<i128>>, validity: Option<Bitmap>| {
88            let Some(values) = values else {
89                // The output is not used if a chunk failed.
90                failed.set(true);
91                return Int128Array::new_empty(ArrowDataType::Int128);
92            };
93            Int128Array::from_vec(values).with_validity(validity)
94        };
95        let phys = apply_binary_kernel_broadcast(
96            self.physical(),
97            rhs.physical(),
98            |l, r| {
99                let values = i64_binary_values(l.values(), r.values(), &op);
100                to_arr(values, combine_validities_and(l.validity(), r.validity()))
101            },
102            |l, r| {
103                let values = i64::try_from(l)
104                    .ok()
105                    .and_then(|l| i64_unary_values(r.values(), |r| op(l, r)));
106                to_arr(values, r.validity().cloned())
107            },
108            |l, r| {
109                let values = i64::try_from(r)
110                    .ok()
111                    .and_then(|r| i64_unary_values(l.values(), |l| op(l, r)));
112                to_arr(values, l.validity().cloned())
113            },
114        );
115        (!failed.get()).then(|| phys.into_decimal_unchecked(DEC128_MAX_PREC, scale))
116    }
117
118    /// A single non-null value at `scale`, if it has another scale and fits.
119    fn scalar_with_scale(&self, scale: usize) -> Option<Self> {
120        if self.len() != 1 || self.scale() == scale {
121            return None;
122        }
123        let value = dec128_rescale(
124            self.physical().get(0)?,
125            self.scale(),
126            DEC128_MAX_PREC,
127            scale,
128        )?;
129        Some(
130            Int128Chunked::from_slice(self.name().clone(), &[value])
131                .into_decimal_unchecked(DEC128_MAX_PREC, scale),
132        )
133    }
134
135    /// Applies an addition or subtraction kernel at the larger scale. A single
136    /// value is brought to that scale once instead of in every row.
137    fn add_sub(
138        &self,
139        rhs: &Self,
140        op: &str,
141        i64_op: impl Fn(i64, i64) -> i128,
142        kernel: impl Fn(i128, usize, i128, usize, usize) -> Option<i128>,
143    ) -> PolarsResult<Self> {
144        let scale = self.scale().max(rhs.scale());
145        let lhs_scalar = self.scalar_with_scale(scale);
146        let rhs_scalar = rhs.scalar_with_scale(scale);
147        let lhs = lhs_scalar.as_ref().unwrap_or(self);
148        let rhs = rhs_scalar.as_ref().unwrap_or(rhs);
149        if lhs.scale() == rhs.scale()
150            && let Some(out) = lhs.apply_i64_values(rhs, scale, i64_op)
151        {
152            return Ok(out);
153        }
154        lhs.apply_scaled_kernel_values(rhs, scale, op, kernel)
155    }
156
157    /// Multiplies with the result rounded to `scale`.
158    pub fn mul_with_scale(&self, rhs: &Self, scale: usize) -> PolarsResult<Self> {
159        if self.scale() + rhs.scale() == scale
160            && let Some(out) = self.apply_i64_values(rhs, scale, |l, r| l as i128 * r as i128)
161        {
162            return Ok(out);
163        }
164        self.apply_scaled_kernel_values(rhs, scale, "multiplication", dec128_mul_scaled)
165    }
166
167    /// Divides with the result rounded to `scale`.
168    pub fn div_with_scale(&self, rhs: &Self, scale: usize) -> PolarsResult<Self> {
169        self.apply_scaled_kernel(rhs, scale, "division", dec128_div_scaled)
170    }
171
172    /// The exact remainder, with the sign of `rhs` when flooring (as `%`), else of `self`.
173    pub fn rem_with(&self, rhs: &Self, rounding: DivRounding) -> PolarsResult<Self> {
174        let scale = self.scale().max(rhs.scale());
175        self.apply_scaled_kernel(rhs, scale, "remainder", |l, sl, r, sr, s| {
176            dec128_rem_scaled(l, sl, r, sr, s, rounding)
177        })
178    }
179
180    /// The integer quotient, rounded per `rounding`.
181    pub fn int_div(&self, rhs: &Self, rounding: DivRounding) -> PolarsResult<Self> {
182        self.int_div_with_scale(rhs, self.scale().max(rhs.scale()), rounding)
183    }
184
185    /// [`Self::int_div`] with the quotient at `scale`.
186    pub fn int_div_with_scale(
187        &self,
188        rhs: &Self,
189        scale: usize,
190        rounding: DivRounding,
191    ) -> PolarsResult<Self> {
192        self.apply_scaled_kernel(rhs, scale, "integer division", |l, sl, r, sr, s| {
193            dec128_int_div_scaled(l, sl, r, sr, s, rounding)
194        })
195    }
196}
197
198impl Add for &DecimalChunked {
199    type Output = PolarsResult<DecimalChunked>;
200
201    fn add(self, rhs: Self) -> Self::Output {
202        self.add_sub(
203            rhs,
204            "addition",
205            |l, r| l as i128 + r as i128,
206            dec128_add_scaled,
207        )
208    }
209}
210
211impl Sub for &DecimalChunked {
212    type Output = PolarsResult<DecimalChunked>;
213
214    fn sub(self, rhs: Self) -> Self::Output {
215        self.add_sub(
216            rhs,
217            "subtraction",
218            |l, r| l as i128 - r as i128,
219            dec128_sub_scaled,
220        )
221    }
222}
223
224impl Mul for &DecimalChunked {
225    type Output = PolarsResult<DecimalChunked>;
226
227    fn mul(self, rhs: Self) -> Self::Output {
228        self.mul_with_scale(rhs, self.scale().max(rhs.scale()))
229    }
230}
231
232impl Div for &DecimalChunked {
233    type Output = PolarsResult<DecimalChunked>;
234
235    fn div(self, rhs: Self) -> Self::Output {
236        self.div_with_scale(rhs, self.scale().max(rhs.scale()))
237    }
238}
polars_core/chunked_array/arithmetic/decimal.rs

polars_core/chunked_array/arithmetic/
decimal.rs
