fix(python): fix window function with nullable values; regression due… (

#5874)
pola-rs · Dec 21, 2022 · bf6bdd7 · bf6bdd7
1 parent 4af64a6
commit bf6bdd7
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 11 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/window.rs b/polars/polars-lazy/src/physical_plan/expressions/window.rs
@@ -721,6 +721,7 @@ where
     } else {
         let mut values = Vec::with_capacity(len);
         let mut validity = MutableBitmap::with_capacity(len);
+        validity.extend_constant(len, true);
         let ptr = values.as_mut_ptr() as *mut T::Native;
 
         match groups {
@@ -729,19 +730,18 @@ where
                     .zip(groups.all().iter())
                     .for_each(|(opt_v, g)| {
                         for idx in g {
-                            debug_assert!((*idx as usize) < len);
+                            let idx = *idx as usize;
+                            debug_assert!(idx < len);
                             unsafe {
-                                let valid = match opt_v {
+                                match opt_v {
                                     Some(v) => {
-                                        *ptr.add(*idx as usize) = v;
-                                        true
+                                        *ptr.add(idx) = v;
                                     }
                                     None => {
-                                        *ptr.add(*idx as usize) = T::Native::default();
-                                        false
+                                        *ptr.add(idx) = T::Native::default();
+                                        validity.set_unchecked(idx, false);
                                     }
                                 };
-                                validity.push_unchecked(valid)
                             }
                         }
                     })
@@ -753,22 +753,22 @@ where
                     for idx in start..end {
                         debug_assert!(idx < len);
                         unsafe {
-                            let valid = match opt_v {
+                            match opt_v {
                                 Some(v) => {
                                     *ptr.add(idx) = v;
-                                    true
                                 }
                                 None => {
                                     *ptr.add(idx) = T::Native::default();
-                                    false
+                                    validity.set_unchecked(idx, false);
                                 }
                             };
-                            validity.push_unchecked(valid)
                         }
                     }
                 }
             }
         }
+        // safety: we have written all slots
+        unsafe { values.set_len(len) }
         let arr = PrimitiveArray::new(
             T::get_dtype().to_physical().to_arrow(),
             values.into(),

diff --git a/py-polars/tests/unit/test_window.py b/py-polars/tests/unit/test_window.py
@@ -257,3 +257,12 @@ def test_nested_aggregation_window_expression() -> None:
         "y": [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
         "foo": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, None, None, 1],
     }
+
+
+def test_window_5868() -> None:
+    df = pl.DataFrame({"value": [None, 2], "id": [None, 1]})
+
+    assert df.with_column(pl.col("value").max().over("id")).to_dict(False) == {
+        "value": [None, 2],
+        "id": [None, 1],
+    }