strategy arg for pl.Expr.fill_null

py-polars/polars/lazy/expr.py

@@ -773,8 +773,21 @@ def shift_and_fill(self, periods: int, fill_value: "Expr") -> "Expr":
 
     def fill_null(self, fill_value: Union[str, int, float, "Expr"]) -> "Expr":
         """
-        Fill none value with a fill value
-        """
+        Fill none value with a fill value or strategy
+
+        fill_value
+            Fill null strategy or a value
+                   * "backward"
+                   * "forward"
+                   * "min"
+                   * "max"
+                   * "mean"
+                   * "one"
+                   * "zero"
+        """
+        if fill_value in ["backward", "forward", "min", "max", "mean", "zero", "one"]:
+            return wrap_expr(self._pyexpr.fill_null_with_strategy(fill_value))
+
         fill_value = expr_to_lit_or_expr(fill_value, str_to_lit=True)
         return wrap_expr(self._pyexpr.fill_null(fill_value._pyexpr))
 

py-polars/src/conversion.rs

@@ -423,3 +423,16 @@ pub(crate) fn dicts_to_rows(records: &PyAny) -> PyResult<(Vec<Row>, Vec<String>)
     }
     Ok((rows, keys_first))
 }
+
+pub(crate) fn parse_strategy(strat: &str) -> FillNullStrategy {
+    match strat {
+        "backward" => FillNullStrategy::Backward,
+        "forward" => FillNullStrategy::Forward,
+        "min" => FillNullStrategy::Min,
+        "max" => FillNullStrategy::Max,
+        "mean" => FillNullStrategy::Mean,
+        "zero" => FillNullStrategy::Zero,
+        "one" => FillNullStrategy::One,
+        s => panic!("Strategy {} not supported", s),
+    }
+}

py-polars/src/dataframe.rs

@@ -308,7 +308,7 @@ impl PyDataFrame {
             "brotli" => Compression::Brotli,
             "lz4" => Compression::Lz4,
             "zstd" => Compression::Zstd,
-            s => panic!("compression {} not supported", s)
+            s => panic!("compression {} not supported", s),
         };
 
         ParquetWriter::new(buf)

py-polars/src/lazy/dsl.rs

@@ -1,6 +1,6 @@
 use crate::conversion::str_to_null_behavior;
 use crate::lazy::utils::py_exprs_to_exprs;
-use crate::prelude::str_to_rankmethod;
+use crate::prelude::{parse_strategy, str_to_rankmethod};
 use crate::series::PySeries;
 use crate::utils::{reinterpret, str_to_polarstype};
 use polars::lazy::dsl;
@@ -232,6 +232,14 @@ impl PyExpr {
         self.clone().inner.fill_null(expr.inner).into()
     }
 
+    pub fn fill_null_with_strategy(&self, strategy: &str) -> PyExpr {
+        let strat = parse_strategy(strategy);
+        self.clone()
+            .inner
+            .apply(move |s| s.fill_null(strat), GetOutput::same_type())
+            .into()
+    }
+
     pub fn fill_nan(&self, expr: PyExpr) -> PyExpr {
         self.clone().inner.fill_nan(expr.inner).into()
     }

py-polars/src/series.rs

@@ -692,16 +692,7 @@ impl PySeries {
     }
 
     pub fn fill_null(&self, strategy: &str) -> PyResult<Self> {
-        let strat = match strategy {
-            "backward" => FillNullStrategy::Backward,
-            "forward" => FillNullStrategy::Forward,
-            "min" => FillNullStrategy::Min,
-            "max" => FillNullStrategy::Max,
-            "mean" => FillNullStrategy::Mean,
-            "zero" => FillNullStrategy::Zero,
-            "one" => FillNullStrategy::One,
-            s => return Err(PyPolarsEr::Other(format!("Strategy {} not supported", s)).into()),
-        };
+        let strat = parse_strategy(strategy);
         let series = self.series.fill_null(strat).map_err(PyPolarsEr::from)?;
         Ok(PySeries::new(series))
     }

py-polars/tests/test_lazy.py

@@ -392,6 +392,11 @@ def test_fill_nan():
     assert df.lazy().fill_nan(2.0).collect()["a"] == [1.0, 2.0, 3.0]
 
 
+def test_fill_null():
+    df = pl.DataFrame({"a": [1.0, None, 3.0]})
+    assert df.select([pl.col("a").fill_null("min")])["a"][1] == 1.0
+
+
 def test_take(fruits_cars):
     df = fruits_cars