fix(rust, python): fix rolling_float in case closure returns None (#5180

)
pola-rs · Oct 12, 2022 · c30d0e4 · c30d0e4
1 parent eea6dca
commit c30d0e4
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/rolling_window.rs b/polars/polars-core/src/chunked_array/ops/rolling_window.rs
@@ -227,8 +227,18 @@ mod inner_mod {
 
                 let out = f(&mut heap_container);
                 match out {
-                    Some(v) => unsafe { values.push_unchecked(v) },
-                    None => unsafe { unset_bit_raw(validity_ptr, offset + window_size - 1) },
+                    Some(v) => {
+                        // Safety: we have pre-allocated
+                        unsafe { values.push_unchecked(v) }
+                    }
+                    None => {
+                        // safety: we allocated enough for both the `values` vec
+                        // and the `validity_ptr`
+                        unsafe {
+                            values.push_unchecked(T::Native::default());
+                            unset_bit_raw(validity_ptr, offset + window_size - 1);
+                        }
+                    }
                 }
             }
             let arr = PrimitiveArray::from_data(

diff --git a/py-polars/tests/unit/test_rolling.py b/py-polars/tests/unit/test_rolling.py
@@ -342,3 +342,25 @@ def test_overlapping_groups_4628() -> None:
             [None, 40, 50],
         ],
     }
+
+
+def test_rolling_skew_lagging_null_5179() -> None:
+    s = pl.Series([None, 3, 4, 1, None, None, None, None, 3, None, 5, 4, 7, 2, 1, None])
+    assert s.rolling_skew(3).fill_nan(-1.0).to_list() == [
+        None,
+        None,
+        0.0,
+        -0.3818017741606059,
+        0.0,
+        -1.0,
+        None,
+        None,
+        -1.0,
+        -1.0,
+        0.0,
+        0.0,
+        0.38180177416060695,
+        0.23906314692954517,
+        0.6309038567106234,
+        0.0,
+    ]