ENH: Enable indexing with nullable Boolean

doc/source/whatsnew/v1.1.0.rst

@@ -36,6 +36,11 @@ For example:
    ser["2014"]
    ser.loc["May 2015"]
 
+Indexing with Nullable Boolean Arrays
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Previously indexing with a nullable Boolean array would raise a ``ValueError``, however this is now permitted with ``NA`` being treated as ``False``. (:issue:`31503`)
+
 .. _whatsnew_110.enhancements.other:
 
 Other enhancements

pandas/core/common.py

@@ -114,24 +114,16 @@ def is_bool_indexer(key: Any) -> bool:
     check_array_indexer : Check that `key` is a valid array to index,
         and convert to an ndarray.
     """
-    na_msg = "cannot mask with array containing NA / NaN values"
     if isinstance(key, (ABCSeries, np.ndarray, ABCIndex)) or (
         is_array_like(key) and is_extension_array_dtype(key.dtype)
     ):
         if key.dtype == np.object_:
             key = np.asarray(values_from_object(key))
 
             if not lib.is_bool_array(key):
-                if isna(key).any():
-                    raise ValueError(na_msg)
                 return False
             return True
         elif is_bool_dtype(key.dtype):
-            # an ndarray with bool-dtype by definition has no missing values.
-            # So we only need to check for NAs in ExtensionArrays
-            if is_extension_array_dtype(key.dtype):
-                if np.any(key.isna()):
-                    raise ValueError(na_msg)
             return True
     elif isinstance(key, list):
         try:

pandas/core/indexers.py

@@ -10,6 +10,7 @@
 from pandas.core.dtypes.common import (
     is_array_like,
     is_bool_dtype,
+    is_extension_array_dtype,
     is_integer_dtype,
     is_list_like,
 )
@@ -333,14 +334,11 @@ def check_array_indexer(array: AnyArrayLike, indexer: Any) -> Any:
     ...
     IndexError: Boolean index has wrong length: 3 instead of 2.
 
-    A ValueError is raised when the mask cannot be converted to
-    a bool-dtype ndarray.
+    NA values are treated as False.
 
     >>> mask = pd.array([True, pd.NA])
     >>> pd.api.indexers.check_array_indexer(arr, mask)
-    Traceback (most recent call last):
-    ...
-    ValueError: Cannot mask with a boolean indexer containing NA values
+    array([ True, False])
 
     A numpy boolean mask will get passed through (if the length is correct):
 
@@ -392,10 +390,17 @@ def check_array_indexer(array: AnyArrayLike, indexer: Any) -> Any:
 
     dtype = indexer.dtype
     if is_bool_dtype(dtype):
-        try:
-            indexer = np.asarray(indexer, dtype=bool)
-        except ValueError:
-            raise ValueError("Cannot mask with a boolean indexer containing NA values")
+        if is_extension_array_dtype(dtype):
+            indexer = indexer.to_numpy(dtype=bool, na_value=False)
+        else:
+            try:
+                indexer = np.asarray(indexer, dtype=bool)
+            except ValueError:
+                msg = (
+                    "Cannot mask with a non-ExtensionArray boolean indexer "
+                    " containing missing values"
+                )
+                raise ValueError(msg)
 
         # GH26658
         if len(indexer) != len(array):

pandas/core/indexing.py

@@ -8,6 +8,8 @@
 from pandas.util._decorators import Appender
 
 from pandas.core.dtypes.common import (
+    is_bool_dtype,
+    is_extension_array_dtype,
     is_float,
     is_integer,
     is_iterator,
@@ -2222,6 +2224,8 @@ def check_bool_indexer(index: Index, key) -> np.ndarray:
                 "the indexed object do not match)."
             )
         result = result.astype(bool)._values
+    elif is_extension_array_dtype(key) and is_bool_dtype(key):
+        result = key.to_numpy(dtype=bool, na_value=False)
     else:
         # key might be sparse / object-dtype bool, check_array_indexer needs bool array
         result = np.asarray(result, dtype=bool)

pandas/tests/arrays/categorical/test_indexing.py

@@ -239,15 +239,14 @@ def test_mask_with_boolean(index):
     tm.assert_series_equal(result, expected)
 
 
-@pytest.mark.parametrize("index", [True, False])
-def test_mask_with_boolean_raises(index):
+def test_mask_with_boolean_na_treated_as_false():
     s = Series(range(3))
     idx = Categorical([True, False, None])
-    if index:
-        idx = CategoricalIndex(idx)
 
-    with pytest.raises(ValueError, match="NA / NaN"):
-        s[idx]
+    result = s[idx]
+    expected = s[idx.fillna(False)]
+
+    tm.assert_series_equal(result, expected)
 
 
 @pytest.fixture

pandas/tests/arrays/test_boolean.py

@@ -929,3 +929,23 @@ def test_diff():
     result = s.diff()
     expected = pd.Series(expected)
     tm.assert_series_equal(result, expected)
+
+
+def test_nullable_boolean_mask_series():
+    s = pd.Series([1, 2, 3])
+    mask = pd.array([True, True, None], dtype="boolean")
+
+    result = s[mask]
+    expected = s.iloc[:2]
+
+    tm.assert_series_equal(result, expected)
+
+
+def test_nullable_boolean_mask_frame():
+    df = pd.DataFrame({"a": [1, 2, 3]})
+    mask = pd.array([True, True, None], dtype="boolean")
+
+    result = df[mask]
+    expected = df.iloc[:2, :]
+
+    tm.assert_frame_equal(result, expected)

pandas/tests/extension/base/getitem.py

@@ -158,21 +158,21 @@ def test_getitem_boolean_array_mask(self, data):
         result = pd.Series(data)[mask]
         self.assert_series_equal(result, expected)
 
-    def test_getitem_boolean_array_mask_raises(self, data):
+    def test_getitem_boolean_na_treated_as_false(self, data):
         mask = pd.array(np.zeros(data.shape, dtype="bool"), dtype="boolean")
         mask[:2] = pd.NA
 
-        msg = (
-            "Cannot mask with a boolean indexer containing NA values|"
-            "cannot mask with array containing NA / NaN values"
-        )
-        with pytest.raises(ValueError, match=msg):
-            data[mask]
+        result = data[mask]
+        expected = data[mask.fillna(False)]
+
+        tm.assert_frame_equal(result, expected)
 
         s = pd.Series(data)
 
-        with pytest.raises(ValueError):
-            s[mask]
+        result = s[mask]
+        expected = s[mask.fillna(False)]
+
+        tm.assert_series_equal(result, expected)
 
     @pytest.mark.parametrize(
         "idx",

pandas/tests/indexing/test_na_indexing.py

@@ -62,18 +62,28 @@ def test_series_mask_boolean(values, dtype, mask, box_mask, frame):
 
 
 @pytest.mark.parametrize("frame", [True, False])
-def test_indexing_with_na_raises(frame):
+def test_na_treated_as_false(frame):
     s = pd.Series([1, 2, 3], name="name")
 
     if frame:
         s = s.to_frame()
+
     mask = pd.array([True, False, None], dtype="boolean")
-    match = "cannot mask with array containing NA / NaN values"
-    with pytest.raises(ValueError, match=match):
-        s[mask]
 
-    with pytest.raises(ValueError, match=match):
-        s.loc[mask]
+    result = s[mask]
+    expected = s[mask.fillna(False)]
+
+    result_loc = s.loc[mask]
+    expected_loc = s.loc[mask.fillna(False)]
 
-    with pytest.raises(ValueError, match=match):
-        s.iloc[mask]
+    result_iloc = s.iloc[mask]
+    expected_iloc = s.iloc[mask.fillna(False)]
+
+    if frame:
+        tm.assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result_loc, expected_loc)
+        tm.assert_frame_equal(result_iloc, expected_iloc)
+    else:
+        tm.assert_series_equal(result, expected)
+        tm.assert_series_equal(result_loc, expected_loc)
+        tm.assert_series_equal(result_iloc, expected_iloc)

pandas/tests/series/indexing/test_boolean.py

@@ -70,14 +70,6 @@ def test_getitem_boolean_object(string_series):
     s2[mask] = 5
     tm.assert_series_equal(cop, s2)
 
-    # nans raise exception
-    omask[5:10] = np.nan
-    msg = "cannot mask with array containing NA / NaN values"
-    with pytest.raises(ValueError, match=msg):
-        s[omask]
-    with pytest.raises(ValueError, match=msg):
-        s[omask] = 5
-
 
 def test_getitem_setitem_boolean_corner(datetime_series):
     ts = datetime_series