fix fill_nan dtypes (#2933)

pola-rs · Mar 18, 2022 · 2b80e6b · 2b80e6b
1 parent 407a4e2
commit 2b80e6b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/builder.rs b/polars/polars-lazy/src/logical_plan/builder.rs
@@ -238,16 +238,17 @@ impl LogicalPlanBuilder {
 
     pub fn fill_nan(self, fill_value: Expr) -> Self {
         let schema = self.0.schema();
+
         let exprs = schema
-            .iter_names()
-            .map(|name| {
-                when(col(name).is_nan())
-                    .then(fill_value.clone())
-                    .otherwise(col(name))
-                    .alias(name)
+            .iter()
+            .filter_map(|(name, dtype)| match dtype {
+                DataType::Float32 | DataType::Float64 => {
+                    Some(col(name).fill_nan(fill_value.clone()).alias(name))
+                }
+                _ => None,
             })
             .collect();
-        self.project_local(exprs)
+        self.with_columns(exprs)
     }
 
     pub fn with_columns(self, exprs: Vec<Expr>) -> Self {

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -1612,6 +1612,13 @@ def test_fill_nan() -> None:
     df = pl.DataFrame({"a": [1, 2], "b": [3.0, float("nan")]})
     assert df.fill_nan(4).frame_equal(pl.DataFrame({"a": [1, 2], "b": [3, 4]}))
     assert df["b"].fill_nan(5.0).to_list() == [3.0, 5.0]
+    df = pl.DataFrame(
+        {
+            "a": [1.0, np.nan, 3.0],
+            "b": [datetime(1, 2, 2), datetime(2, 2, 2), datetime(3, 2, 2)],
+        }
+    )
+    assert df.fill_nan(2.0).dtypes == [pl.Float64, pl.Datetime]
 
 
 def test_shift_and_fill() -> None: