fix: Fix decimal arithmetic schema (pola-rs#20398)

Author: ritchie46

Diff for: crates/polars-core/src/chunked_array/arithmetic/decimal.rs

@@ -4,7 +4,7 @@ impl Add for &DecimalChunked {
     type Output = PolarsResult<DecimalChunked>;
 
     fn add(self, rhs: Self) -> Self::Output {
-        let scale = self.scale().max(rhs.scale());
+        let scale = _get_decimal_scale_add_sub(self.scale(), rhs.scale());
         let lhs = self.to_scale(scale)?;
         let rhs = rhs.to_scale(scale)?;
         Ok((&lhs.0 + &rhs.0).into_decimal_unchecked(None, scale))
@@ -15,7 +15,7 @@ impl Sub for &DecimalChunked {
     type Output = PolarsResult<DecimalChunked>;
 
     fn sub(self, rhs: Self) -> Self::Output {
-        let scale = self.scale().max(rhs.scale());
+        let scale = _get_decimal_scale_add_sub(self.scale(), rhs.scale());
         let lhs = self.to_scale(scale)?;
         let rhs = rhs.to_scale(scale)?;
         Ok((&lhs.0 - &rhs.0).into_decimal_unchecked(None, scale))
@@ -26,7 +26,7 @@ impl Mul for &DecimalChunked {
     type Output = PolarsResult<DecimalChunked>;
 
     fn mul(self, rhs: Self) -> Self::Output {
-        let scale = self.scale() + rhs.scale();
+        let scale = _get_decimal_scale_mul(self.scale(), rhs.scale());
         Ok((&self.0 * &rhs.0).into_decimal_unchecked(None, scale))
     }
 }
@@ -35,9 +35,22 @@ impl Div for &DecimalChunked {
     type Output = PolarsResult<DecimalChunked>;
 
     fn div(self, rhs: Self) -> Self::Output {
-        // Follow postgres and MySQL adding a fixed scale increment of 4
-        let scale = self.scale() + 4;
+        let scale = _get_decimal_scale_div(self.scale());
         let lhs = self.to_scale(scale + rhs.scale())?;
         Ok((&lhs.0 / &rhs.0).into_decimal_unchecked(None, scale))
     }
 }
+
+// Used by polars-plan to determine schema.
+pub fn _get_decimal_scale_add_sub(scale_left: usize, scale_right: usize) -> usize {
+    scale_left.max(scale_right)
+}
+
+pub fn _get_decimal_scale_mul(scale_left: usize, scale_right: usize) -> usize {
+    scale_left + scale_right
+}
+
+pub fn _get_decimal_scale_div(scale_left: usize) -> usize {
+    // Follow postgres and MySQL adding a fixed scale increment of 4
+    scale_left + 4
+}

Diff for: crates/polars-core/src/chunked_array/arithmetic/mod.rs

@@ -6,6 +6,8 @@ mod numeric;
 use std::ops::{Add, Div, Mul, Rem, Sub};
 
 use arrow::compute::utils::combine_validities_and;
+#[cfg(feature = "dtype-decimal")]
+pub use decimal::{_get_decimal_scale_add_sub, _get_decimal_scale_div, _get_decimal_scale_mul};
 use num_traits::{Num, NumCast, ToPrimitive};
 pub use numeric::ArithmeticChunked;
 

Diff for: crates/polars-plan/src/plans/aexpr/schema.rs

@@ -1,3 +1,7 @@
+#[cfg(feature = "dtype-decimal")]
+use polars_core::chunked_array::arithmetic::{
+    _get_decimal_scale_add_sub, _get_decimal_scale_div, _get_decimal_scale_mul,
+};
 use recursive::recursive;
 
 use super::*;
@@ -500,6 +504,11 @@ fn get_arithmetic_field(
                         other_dtype.leaf_dtype(),
                     )?)
                 },
+                #[cfg(feature = "dtype-decimal")]
+                (Decimal(_, Some(scale_left)), Decimal(_, Some(scale_right))) => {
+                    let scale = _get_decimal_scale_add_sub(*scale_left, *scale_right);
+                    Decimal(None, Some(scale))
+                },
                 (left, right) => try_get_supertype(left, right)?,
             }
         },
@@ -549,6 +558,11 @@ fn get_arithmetic_field(
                         other_dtype.leaf_dtype(),
                     )?)
                 },
+                #[cfg(feature = "dtype-decimal")]
+                (Decimal(_, Some(scale_left)), Decimal(_, Some(scale_right))) => {
+                    let scale = _get_decimal_scale_add_sub(*scale_left, *scale_right);
+                    Decimal(None, Some(scale))
+                },
                 (left, right) => try_get_supertype(left, right)?,
             }
         },
@@ -581,6 +595,23 @@ fn get_arithmetic_field(
                         polars_bail!(InvalidOperation: "{} not allowed on {} and {}", op, left_field.dtype, right_type)
                     },
                 },
+                #[cfg(feature = "dtype-decimal")]
+                (Decimal(_, Some(scale_left)), Decimal(_, Some(scale_right))) => {
+                    let scale = match op {
+                        Operator::Multiply => _get_decimal_scale_mul(*scale_left, *scale_right),
+                        Operator::Divide | Operator::TrueDivide => {
+                            _get_decimal_scale_div(*scale_left)
+                        },
+                        _ => {
+                            debug_assert!(false);
+                            *scale_left
+                        },
+                    };
+                    let dtype = Decimal(None, Some(scale));
+                    left_field.coerce(dtype);
+                    return Ok(left_field);
+                },
+
                 (l @ List(a), r @ List(b))
                     if ![a, b]
                         .into_iter()
@@ -684,6 +715,11 @@ fn get_truediv_field(
             })
         },
         (Float32, _) => Float32,
+        #[cfg(feature = "dtype-decimal")]
+        (Decimal(_, Some(scale_left)), Decimal(_, _)) => {
+            let scale = _get_decimal_scale_div(*scale_left);
+            Decimal(None, Some(scale))
+        },
         (dt, _) if dt.is_numeric() => Float64,
         #[cfg(feature = "dtype-duration")]
         (Duration(_), Duration(_)) => Float64,

Diff for: py-polars/tests/unit/datatypes/test_decimal.py

@@ -540,3 +540,16 @@ def test_decimal_round() -> None:
         expected_s = pl.Series("a", [round(v, decimals) for v in values], dtype)
 
         assert_series_equal(got_s, expected_s)
+
+
+def test_decimal_arithmetic_schema() -> None:
+    q = pl.LazyFrame({"x": [1.0]}, schema={"x": pl.Decimal(15, 2)})
+
+    q1 = q.select(pl.col.x * pl.col.x)
+    assert q1.collect_schema() == q1.collect().schema
+    q1 = q.select(pl.col.x / pl.col.x)
+    assert q1.collect_schema() == q1.collect().schema
+    q1 = q.select(pl.col.x - pl.col.x)
+    assert q1.collect_schema() == q1.collect().schema
+    q1 = q.select(pl.col.x + pl.col.x)
+    assert q1.collect_schema() == q1.collect().schema