ENH: Enable indexing with nullable Boolean

doc/source/whatsnew/v1.1.0.rst

@@ -212,6 +212,7 @@ Sparse
 ExtensionArray
 ^^^^^^^^^^^^^^
 
+- Enabled filtering with nullable Boolean arrays. (:issue:`31503`)
 -
 -
 

pandas/core/common.py

@@ -114,7 +114,6 @@ 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)
     ):
@@ -127,11 +126,6 @@ def is_bool_indexer(key: Any) -> bool:
                 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/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,12 @@ 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):
+        mask = isna(key)
+        if mask.any():
+            result = np.asarray(key.fillna(False), dtype=bool)
+        else:
+            result = np.asarray(key, dtype=bool)
     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,17 +239,6 @@ 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):
-    s = Series(range(3))
-    idx = Categorical([True, False, None])
-    if index:
-        idx = CategoricalIndex(idx)
-
-    with pytest.raises(ValueError, match="NA / NaN"):
-        s[idx]
-
-
 @pytest.fixture
 def non_coercible_categorical(monkeypatch):
     """

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,22 +158,6 @@ 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):
-        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]
-
-        s = pd.Series(data)
-
-        with pytest.raises(ValueError):
-            s[mask]
-
     @pytest.mark.parametrize(
         "idx",
         [[0, 1, 2], pd.array([0, 1, 2], dtype="Int64"), np.array([0, 1, 2])],

pandas/tests/indexing/test_na_indexing.py

@@ -59,21 +59,3 @@ def test_series_mask_boolean(values, dtype, mask, box_mask, frame):
 
     result = ser.loc[mask]
     tm.assert_equal(result, expected)
-
-
-@pytest.mark.parametrize("frame", [True, False])
-def test_indexing_with_na_raises(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]
-
-    with pytest.raises(ValueError, match=match):
-        s.iloc[mask]

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