python: Series.fill_nan

pola-rs · Jan 16, 2022 · 913d102 · 913d102
1 parent 1a5a294
commit 913d102
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 2 deletions.
diff --git a/py-polars/docs/source/reference/series.rst b/py-polars/docs/source/reference/series.rst
@@ -171,6 +171,7 @@ Manipulation/ selection
     Series.ceil
     Series.set_at_idx
     Series.fill_null
+    Series.fill_nan
     Series.zip_with
     Series.interpolate
     Series.clip

diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -938,7 +938,7 @@ def fill_null(self, fill_value: Union[int, float, bool, str, "Expr"]) -> "Expr":
 
     def fill_nan(self, fill_value: Union[str, int, float, bool, "Expr"]) -> "Expr":
         """
-        Fill none value with a fill value
+        Fill floating point NaN value with a fill value
         """
         fill_value = expr_to_lit_or_expr(fill_value, str_to_lit=True)
         return wrap_expr(self._pyexpr.fill_nan(fill_value._pyexpr))

diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -3169,7 +3169,7 @@ def fill_null(self, strategy: Union[str, "pli.Expr", Any]) -> "DataFrame":
 
     def fill_nan(self, fill_value: Union["pli.Expr", int, float]) -> "DataFrame":
         """
-        Fill None/missing values by a an Expression evaluation.
+        Fill floating point NaN values by an Expression evaluation.
 
         Warnings
         --------

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -2142,6 +2142,16 @@ def __copy__(self) -> "Series":
     def __deepcopy__(self, memodict: Any = {}) -> "Series":
         return self.clone()
 
+    def fill_nan(
+        self, fill_value: Union[str, int, float, bool, "pli.Expr"]
+    ) -> "Series":
+        """
+        Fill floating point NaN value with a fill value
+        """
+        return (
+            self.to_frame().select(pli.col(self.name).fill_nan(fill_value)).to_series()
+        )
+
     def fill_null(self, strategy: Union[str, int, "pli.Expr"]) -> "Series":
         """
         Fill null values with a filling strategy.

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -1564,6 +1564,7 @@ def test_fill_null() -> None:
 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]
 
 
 def test_shift_and_fill() -> None: