diff --git a/doc/source/whatsnew/v3.0.0.rst b/doc/source/whatsnew/v3.0.0.rst index 7fddb36185789..5d1668514c0a6 100644 --- a/doc/source/whatsnew/v3.0.0.rst +++ b/doc/source/whatsnew/v3.0.0.rst @@ -482,6 +482,7 @@ Other - Bug in :class:`DataFrame` when passing a ``dict`` with a NA scalar and ``columns`` that would always return ``np.nan`` (:issue:`57205`) - Bug in :func:`eval` where the names of the :class:`Series` were not preserved when using ``engine="numexpr"``. (:issue:`10239`) - Bug in :func:`unique` on :class:`Index` not always returning :class:`Index` (:issue:`57043`) +- Bug in :func:`mask` to handle NaN values in condition of function. (:issue:`56844`) - Bug in :meth:`DataFrame.eval` and :meth:`DataFrame.query` which caused an exception when using NumPy attributes via ``@`` notation, e.g., ``df.eval("@np.floor(a)")``. (:issue:`58041`) - Bug in :meth:`DataFrame.eval` and :meth:`DataFrame.query` which did not allow to use ``tan`` function. (:issue:`55091`) - Bug in :meth:`DataFrame.sort_index` when passing ``axis="columns"`` and ``ignore_index=True`` and ``ascending=False`` not returning a :class:`RangeIndex` columns (:issue:`57293`) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 3da047eaa2d4e..a61b42712e999 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -9999,6 +9999,9 @@ def mask( cond = common.apply_if_callable(cond, self) other = common.apply_if_callable(other, self) + if isinstance(cond, (ABCDataFrame, ABCSeries)): + cond = cond.fillna(False) + # see gh-21891 if not hasattr(cond, "__invert__"): cond = np.array(cond) diff --git a/pandas/tests/frame/indexing/test_mask.py b/pandas/tests/frame/indexing/test_mask.py index 264e27c9c122e..014f038158033 100644 --- a/pandas/tests/frame/indexing/test_mask.py +++ b/pandas/tests/frame/indexing/test_mask.py @@ -150,3 +150,12 @@ def test_mask_inplace_no_other(): df.mask(cond, inplace=True) expected = DataFrame({"a": [np.nan, 2], "b": ["x", np.nan]}) tm.assert_frame_equal(df, expected) + + +def test_mask_with_NA(): + df = DataFrame({"A": [0, 1, 2]}) + cond = Series([True, False, pd.NA], dtype=pd.BooleanDtype()) + + result = df.mask(cond, other=100) + expected = DataFrame({"A": [100, 1, 2]}) + tm.assert_frame_equal(result, expected)