fix[rust]: rolling_max no_null dispatch (#4442)

pola-rs · Aug 16, 2022 · c171634 · c171634
1 parent fdfb001
commit c171634
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/polars/polars-time/src/chunkedarray/rolling_window/mod.rs b/polars/polars-time/src/chunkedarray/rolling_window/mod.rs
@@ -238,8 +238,8 @@ where
         check_input(options.window_size, options.min_periods)?;
         let ca = ca.rechunk();
 
-        match ca.has_validity() {
-            false => rolling_agg_fn(
+        match ca.null_count() {
+            0 => rolling_agg_fn(
                 arr.values().as_slice(),
                 options.window_size,
                 options.min_periods,

diff --git a/py-polars/tests/test_queries.py b/py-polars/tests/test_queries.py
@@ -336,3 +336,33 @@ def test_edge_cast_string_duplicates_4259() -> None:
 
     mask = df.select(["a", "b"]).is_duplicated()
     assert df.filter(pl.lit(mask)).shape == (0, 2)
+
+
+def test_query_4438() -> None:
+    df = pl.DataFrame({"x": [1, 2, 3, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 1, 1]})
+
+    q = (
+        df.lazy()
+        .with_column(pl.col("x").rolling_max(window_size=3).alias("rolling_max"))
+        .fill_null(strategy="backward")
+        .with_column(
+            pl.col("rolling_max").rolling_max(window_size=3).alias("rolling_max_2")
+        )
+    )
+    assert q.collect()["rolling_max_2"].to_list() == [
+        None,
+        None,
+        3,
+        10,
+        10,
+        10,
+        10,
+        10,
+        9,
+        8,
+        7,
+        6,
+        5,
+        4,
+        3,
+    ]