From a8f488935adb07d2d5995e8331aaee523cbc25d3 Mon Sep 17 00:00:00 2001 From: Luciano Scarpulla Date: Tue, 30 Jul 2024 18:12:20 +0200 Subject: [PATCH] feat: add `fill_null` to `DaskExpr` --- narwhals/_dask/expr.py | 3 +++ tests/dask_test.py | 23 +++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py index 1c3ecdf4b..3b64e6073 100644 --- a/narwhals/_dask/expr.py +++ b/narwhals/_dask/expr.py @@ -232,6 +232,9 @@ def sum(self) -> Self: "sum", ) + def fill_null(self, value: Any) -> DaskExpr: + return self._from_call(lambda _input, _val: _input.fillna(_val), "fillna", value) + @property def str(self: Self) -> DaskExprStringNamespace: return DaskExprStringNamespace(self) diff --git a/tests/dask_test.py b/tests/dask_test.py index f91754968..275df0116 100644 --- a/tests/dask_test.py +++ b/tests/dask_test.py @@ -479,6 +479,29 @@ def test_drop_nulls() -> None: compare_dicts(result_d, expected_d) +def test_fill_null_series() -> None: + import dask.dataframe as dd + + data = { + "a": [0.0, None, 2, 3, 4], + "b": [1.0, None, None, 5, 3], + "c": [5.0, None, 3, 2, 1], + } + df = nw.from_native(dd.from_pandas(pd.DataFrame(data))) + + expected = { + "a": [0.0, 99, 2, 3, 4], + "b": [1.0, 99, 99, 5, 3], + "c": [5.0, 99, 3, 2, 1], + } + result = df.with_columns( + a=nw.col("a").fill_null(99), + b=nw.col("b").fill_null(99), + c=nw.col("c").fill_null(99), + ) + compare_dicts(result, expected) + + def test_comparison_operations() -> None: import dask.dataframe as dd