adjust for kleene logic in drop_na (#3529)

fixes #3525
pola-rs · May 29, 2022 · efa6dac · efa6dac
1 parent d5a843e
commit efa6dac
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/polars/polars-lazy/src/dsl/mod.rs b/polars/polars-lazy/src/dsl/mod.rs
@@ -30,7 +30,6 @@ use polars_core::export::arrow::{array::BooleanArray, bitmap::MutableBitmap};
 use polars_core::prelude::*;
 
 use std::fmt::Debug;
-use std::ops::Not;
 use std::{
     ops::{Add, Div, Mul, Rem, Sub},
     sync::Arc,
@@ -301,12 +300,12 @@ impl Expr {
             |s| match s.dtype() {
                 DataType::Float32 => {
                     let ca = s.f32()?;
-                    let mask = ca.is_nan().not();
+                    let mask = ca.is_not_nan().fill_null(FillNullStrategy::One)?;
                     ca.filter(&mask).map(|ca| ca.into_series())
                 }
                 DataType::Float64 => {
                     let ca = s.f64()?;
-                    let mask = ca.is_nan().not();
+                    let mask = ca.is_not_nan().fill_null(FillNullStrategy::One)?;
                     ca.filter(&mask).map(|ca| ca.into_series())
                 }
                 _ => Ok(s),

diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py
@@ -1540,3 +1540,14 @@ def test_cumulative_eval() -> None:
     expr = pl.element().first() - pl.element().last() ** 2
     expected = pl.Series("values", [None, -3.0, -8.0, -15.0, -24.0])
     verify_series_and_expr_api(s, expected, "cumulative_eval", expr)
+
+
+def test_drop_nan_ignore_null_3525() -> None:
+    df = pl.DataFrame({"a": [1.0, float("NaN"), 2.0, None, 3.0, 4.0]})
+    assert df.select(pl.col("a").drop_nans()).to_series().to_list() == [
+        2.0,
+        2.0,
+        None,
+        3.0,
+        4.0,
+    ]