fix(rust, python): improve numeric stability of rolling_variance (#5207)

pola-rs · Oct 14, 2022 · 5fde255 · 5fde255
1 parent 0355b7d
commit 5fde255
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/polars/polars-arrow/src/kernels/rolling/no_nulls/variance.rs b/polars/polars-arrow/src/kernels/rolling/no_nulls/variance.rs
@@ -93,6 +93,7 @@ impl<
             + NumCast
             + One
             + Zero
+            + PartialOrd
             + Sub<Output = T>,
     > RollingAggWindowNoNulls<'a, T> for VarWindow<'a, T>
 {
@@ -114,7 +115,14 @@ impl<
             T::zero()
         } else {
             // apply Bessel's correction
-            var / (count - T::one()) * count
+            let out = var / (count - T::one()) * count;
+            // variance cannot be negative.
+            // if it is negative it is due to numeric instability
+            if out < T::zero() {
+                T::zero()
+            } else {
+                out
+            }
         }
     }
 }
@@ -196,6 +204,7 @@ impl<
             + One
             + Zero
             + Sub<Output = T>
+            + PartialOrd
             + Pow<T, Output = T>,
     > RollingAggWindowNoNulls<'a, T> for StdWindow<'a, T>
 {

diff --git a/py-polars/tests/unit/test_rolling.py b/py-polars/tests/unit/test_rolling.py
@@ -364,3 +364,20 @@ def test_rolling_skew_lagging_null_5179() -> None:
         0.6309038567106234,
         0.0,
     ]
+
+
+def test_rolling_var_numerical_stability_5197() -> None:
+    s = pl.Series([*[1.2] * 4, *[3.3] * 7])
+    assert s.to_frame("a").with_columns(pl.col("a").rolling_var(5))[:, 0].to_list() == [
+        None,
+        None,
+        None,
+        None,
+        0.882,
+        1.3229999999999997,
+        1.3229999999999997,
+        0.8819999999999983,
+        0.0,
+        0.0,
+        0.0,
+    ]