fix[rust]: ewm_mean leading nulls function (#4465)

pola-rs · Aug 17, 2022 · 828a69a · 828a69a
1 parent 90f2f51
commit 828a69a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/polars/polars-core/src/series/ops/ewm.rs b/polars/polars-core/src/series/ops/ewm.rs
@@ -19,8 +19,8 @@ fn prepare_primitive_array<T: NativeType>(
     let leading = std::cmp::max(min_periods, leading_nulls);
     if leading > 1 {
         let mut validity = MutableBitmap::with_capacity(vals.len());
-        validity.extend_constant(min_periods, false);
-        validity.extend_constant(vals.len() - min_periods, true);
+        validity.extend_constant(min_periods - 1, false);
+        validity.extend_constant(vals.len() - (min_periods - 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
@@ -1619,6 +1619,14 @@ def test_ewm_mean() -> None:
     )
 
 
+def test_ewm_mean_leading_nulls() -> None:
+    for min_periods in [1, 2, 3]:
+        assert (
+            pl.Series([1, 2, 3, 4]).ewm_mean(3, min_periods=min_periods).null_count()
+            == min_periods - 1
+        )
+
+
 def test_ewm_std_var() -> None:
     a = pl.Series("a", [2, 5, 3])