fill_null with expression (#2106)

Author: janpipek

narwhals/_arrow/series.py

@@ -667,10 +667,10 @@ def fill_aux(
             )
 
         ser = self._native_series
-        dtype = ser.type
 
         if value is not None:
-            res_ser = self._from_native_series(pc.fill_null(ser, lit(value, dtype)))  # type: ignore[attr-defined]
+            _, value = extract_native(self, value)
+            res_ser = self._from_native_series(pc.fill_null(ser, value))  # type: ignore[attr-defined]
         elif limit is None:
             fill_func = (
                 pc.fill_null_forward if strategy == "forward" else pc.fill_null_backward

narwhals/_duckdb/expr.py

@@ -450,9 +450,10 @@ def fill_null(self: Self, value: Any, strategy: Any, limit: int | None) -> Self:
             msg = "todo"
             raise NotImplementedError(msg)
 
-        return self._from_call(
-            lambda _input: CoalesceOperator(_input, lit(value)), "fill_null"
-        )
+        def func(_input: duckdb.Expression, value: Any) -> duckdb.Expression:
+            return CoalesceOperator(_input, value)
+
+        return self._from_call(func, "fill_null", value=value)
 
     def cast(self: Self, dtype: DType | type[DType]) -> Self:
         def func(_input: duckdb.Expression) -> duckdb.Expression:

narwhals/_pandas_like/series.py

@@ -531,6 +531,7 @@ def fill_null(
     ) -> Self:
         ser = self._native_series
         if value is not None:
+            _, value = align_and_extract_native(self, value)
             res_ser = self._from_native_series(ser.fillna(value=value))
         else:
             res_ser = self._from_native_series(

narwhals/expr.py

@@ -12,6 +12,7 @@
 from narwhals._expression_parsing import ExprMetadata
 from narwhals._expression_parsing import apply_n_ary_operation
 from narwhals._expression_parsing import combine_metadata
+from narwhals._expression_parsing import extract_compliant
 from narwhals.dtypes import _validate_dtype
 from narwhals.exceptions import LengthChangingExprError
 from narwhals.expr_cat import ExprCatNamespace
@@ -1333,7 +1334,7 @@ def fill_null(
         """Fill null values with given value.
 
         Arguments:
-            value: Value used to fill null values.
+            value: Value or expression used to fill null values.
             strategy: Strategy used to fill null values.
             limit: Number of consecutive null values to fill when using the 'forward' or 'backward' strategy.
 
@@ -1352,34 +1353,37 @@ def fill_null(
             ...     {
             ...         "a": [2, None, None, 3],
             ...         "b": [2.0, float("nan"), float("nan"), 3.0],
+            ...         "c": [1, 2, 3, 4],
             ...     }
             ... )
             >>> df = nw.from_native(df_native)
             >>> df.with_columns(
-            ...     nw.col("a", "b").fill_null(0).name.suffix("_nulls_filled")
+            ...     nw.col("a", "b").fill_null(0).name.suffix("_filled"),
+            ...     nw.col("a").fill_null(nw.col("c")).name.suffix("_filled_with_c"),
             ... )
-            ┌────────────────────────────────────────────────┐
-            |               Narwhals DataFrame               |
-            |------------------------------------------------|
-            |shape: (4, 4)                                   |
-            |┌──────┬─────┬────────────────┬────────────────┐|
-            |│ a    ┆ b   ┆ a_nulls_filled ┆ b_nulls_filled │|
-            |│ ---  ┆ --- ┆ ---            ┆ ---            │|
-            |│ i64  ┆ f64 ┆ i64            ┆ f64            │|
-            |╞══════╪═════╪════════════════╪════════════════╡|
-            |│ 2    ┆ 2.0 ┆ 2              ┆ 2.0            │|
-            |│ null ┆ NaN ┆ 0              ┆ NaN            │|
-            |│ null ┆ NaN ┆ 0              ┆ NaN            │|
-            |│ 3    ┆ 3.0 ┆ 3              ┆ 3.0            │|
-            |└──────┴─────┴────────────────┴────────────────┘|
-            └────────────────────────────────────────────────┘
+            ┌────────────────────────────────────────────────────────────┐
+            |                     Narwhals DataFrame                     |
+            |------------------------------------------------------------|
+            |shape: (4, 6)                                               |
+            |┌──────┬─────┬─────┬──────────┬──────────┬─────────────────┐|
+            |│ a    ┆ b   ┆ c   ┆ a_filled ┆ b_filled ┆ a_filled_with_c │|
+            |│ ---  ┆ --- ┆ --- ┆ ---      ┆ ---      ┆ ---             │|
+            |│ i64  ┆ f64 ┆ i64 ┆ i64      ┆ f64      ┆ i64             │|
+            |╞══════╪═════╪═════╪══════════╪══════════╪═════════════════╡|
+            |│ 2    ┆ 2.0 ┆ 1   ┆ 2        ┆ 2.0      ┆ 2               │|
+            |│ null ┆ NaN ┆ 2   ┆ 0        ┆ NaN      ┆ 2               │|
+            |│ null ┆ NaN ┆ 3   ┆ 0        ┆ NaN      ┆ 3               │|
+            |│ 3    ┆ 3.0 ┆ 4   ┆ 3        ┆ 3.0      ┆ 3               │|
+            |└──────┴─────┴─────┴──────────┴──────────┴─────────────────┘|
+            └────────────────────────────────────────────────────────────┘
 
             Using a strategy:
 
-            >>> df.with_columns(
+            >>> df.select(
+            ...     nw.col("a", "b"),
             ...     nw.col("a", "b")
             ...     .fill_null(strategy="forward", limit=1)
-            ...     .name.suffix("_nulls_forward_filled")
+            ...     .name.suffix("_nulls_forward_filled"),
             ... )
             ┌────────────────────────────────────────────────────────────────┐
             |                       Narwhals DataFrame                       |
@@ -1408,7 +1412,9 @@ def fill_null(
             raise ValueError(msg)
         return self._from_callable(
             lambda plx: self._to_compliant_expr(plx).fill_null(
-                value=value, strategy=strategy, limit=limit
+                value=extract_compliant(plx, value, str_as_lit=True),
+                strategy=strategy,
+                limit=limit,
             )
         )
 

narwhals/series.py

@@ -1328,7 +1328,9 @@ def fill_null(
             msg = f"strategy not supported: {strategy}"
             raise ValueError(msg)
         return self._from_compliant_series(
-            self._compliant_series.fill_null(value=value, strategy=strategy, limit=limit)
+            self._compliant_series.fill_null(
+                value=self._extract_native(value), strategy=strategy, limit=limit
+            )
         )
 
     def is_between(

tests/expr_and_series/fill_null_test.py

@@ -31,6 +31,27 @@ def test_fill_null(request: pytest.FixtureRequest, constructor: Constructor) ->
     assert_equal_data(result, expected)
 
 
+def test_fill_null_series_expression(
+    request: pytest.FixtureRequest, constructor: Constructor
+) -> None:
+    if "pyspark" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+    data = {
+        "a": [0.0, None, 2.0, 3.0, 4.0],
+        "b": [1.0, None, None, 5.0, 3.0],
+        "c": [5.0, 2.0, None, 2.0, 1.0],
+    }
+    df = nw.from_native(constructor(data))
+
+    result = df.with_columns(nw.col("a", "b").fill_null(nw.col("c")))
+    expected = {
+        "a": [0.0, 2, 2, 3, 4],
+        "b": [1.0, 2, None, 5, 3],
+        "c": [5.0, 2, None, 2, 1],
+    }
+    assert_equal_data(result, expected)
+
+
 def test_fill_null_exceptions(constructor: Constructor) -> None:
     data = {
         "a": [0.0, None, 2.0, 3.0, 4.0],