fix[rust]: ewm correct for leading nulls (#4477)

pola-rs · Aug 18, 2022 · a50971a · a50971a
1 parent f379d0c
commit a50971a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/polars/polars-core/src/series/ops/ewm.rs b/polars/polars-core/src/series/ops/ewm.rs
@@ -16,11 +16,11 @@ fn prepare_primitive_array<T: NativeType>(
     min_periods: usize,
     leading_nulls: usize,
 ) -> PrimitiveArray<T> {
-    let leading = std::cmp::max(min_periods, leading_nulls);
-    if leading > 1 {
+    let n = min_periods + leading_nulls;
+    if n > 1 {
         let mut validity = MutableBitmap::with_capacity(vals.len());
-        validity.extend_constant(min_periods - 1, false);
-        validity.extend_constant(vals.len() - (min_periods - 1), true);
+        validity.extend_constant(n - 1, false);
+        validity.extend_constant(vals.len() - (n - 1), true);
 
         PrimitiveArray::from_data_default(vals.into(), Some(validity.into()))
     } else {

diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py
@@ -1632,6 +1632,12 @@ def test_ewm_mean_leading_nulls() -> None:
             pl.Series([1, 2, 3, 4]).ewm_mean(3, min_periods=min_periods).null_count()
             == min_periods - 1
         )
+    assert pl.Series([None, 1.0, 1.0, 1.0]).ewm_mean(
+        alpha=0.5, min_periods=1
+    ).to_list() == [None, 1.0, 1.0, 1.0]
+    assert pl.Series([None, 1.0, 1.0, 1.0]).ewm_mean(
+        alpha=0.5, min_periods=2
+    ).to_list() == [None, None, 1.0, 1.0]
 
 
 def test_ewm_std_var() -> None: