feat: Support describe on decimal

pola-rs · Jul 10, 2024 · d461f6a · d461f6a
1 parent daf2e49
commit d461f6a
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 1 deletion.
diff --git a/crates/polars-core/src/series/implementations/decimal.rs b/crates/polars-core/src/series/implementations/decimal.rs
@@ -24,6 +24,20 @@ impl SeriesWrap<DecimalChunked> {
         f(&self.0)
     }
 
+    fn scale_factor(&self) -> u128 {
+        10u128.pow(self.0.scale() as u32)
+    }
+
+    fn apply_scale(&self, mut scalar: Scalar) -> Scalar {
+        debug_assert_eq!(scalar.dtype(), &DataType::Float64);
+        let v = scalar
+            .value()
+            .try_extract::<f64>()
+            .expect("should be f64 scalar");
+        scalar.update((v / self.scale_factor() as f64).into());
+        scalar
+    }
+
     fn agg_helper<F: Fn(&Int128Chunked) -> Series>(&self, f: F) -> Series {
         let agg_s = f(&self.0);
         match agg_s.dtype() {
@@ -370,6 +384,35 @@ impl SeriesTrait for SeriesWrap<DecimalChunked> {
             Scalar::new(self.dtype().clone(), av)
         }))
     }
+
+    fn mean(&self) -> Option<f64> {
+        self.0.mean().map(|v| v / self.scale_factor() as f64)
+    }
+
+    fn median(&self) -> Option<f64> {
+        self.0.median().map(|v| v / self.scale_factor() as f64)
+    }
+    fn median_reduce(&self) -> PolarsResult<Scalar> {
+        Ok(self.apply_scale(self.0.median_reduce()))
+    }
+
+    fn std(&self, ddof: u8) -> Option<f64> {
+        self.0.std(ddof).map(|v| v / self.scale_factor() as f64)
+    }
+    fn std_reduce(&self, ddof: u8) -> PolarsResult<Scalar> {
+        Ok(self.apply_scale(self.0.std_reduce(ddof)))
+    }
+
+    fn quantile_reduce(
+        &self,
+        quantile: f64,
+        interpol: QuantileInterpolOptions,
+    ) -> PolarsResult<Scalar> {
+        self.0
+            .quantile_reduce(quantile, interpol)
+            .map(|v| self.apply_scale(v))
+    }
+
     fn as_any(&self) -> &dyn Any {
         &self.0
     }

diff --git a/crates/polars-core/src/series/mod.rs b/crates/polars-core/src/series/mod.rs
@@ -821,7 +821,7 @@ impl Series {
                 let val = self.mean().map(|m| m as f32);
                 Scalar::new(self.dtype().clone(), val.into())
             },
-            dt if dt.is_numeric() || matches!(dt, DataType::Boolean) => {
+            dt if dt.is_numeric() || dt.is_decimal() || dt.is_bool() => {
                 let val = self.mean();
                 Scalar::new(DataType::Float64, val.into())
             },

diff --git a/py-polars/tests/unit/datatypes/test_decimal.py b/py-polars/tests/unit/datatypes/test_decimal.py
@@ -326,10 +326,40 @@ def test_decimal_aggregations() -> None:
         sum=pl.sum("a"),
         min=pl.min("a"),
         max=pl.max("a"),
+        mean=pl.mean("a"),
+        median=pl.median("a"),
     ).to_dict(as_series=False) == {
         "sum": [D("9110.33")],
         "min": [D("0.10")],
         "max": [D("9000.12")],
+        "mean": [2277.5825],
+        "median": [55.055],
+    }
+
+    assert df.describe().to_dict(as_series=False) == {
+        "statistic": [
+            "count",
+            "null_count",
+            "mean",
+            "std",
+            "min",
+            "25%",
+            "50%",
+            "75%",
+            "max",
+        ],
+        "g": [4.0, 0.0, 1.5, 0.5773502691896257, 1.0, 1.0, 2.0, 2.0, 2.0],
+        "a": [
+            4.0,
+            0.0,
+            2277.5825,
+            4481.916846516863,
+            0.1,
+            10.1,
+            100.01,
+            100.01,
+            9000.12,
+        ],
     }