ENH: Enable indexing with nullable Boolean

doc/source/whatsnew/v1.1.0.rst

@@ -36,6 +36,31 @@ 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`)
+
+*pandas 1.0.x*
+
+.. code-block:: python
+
+    >>> s = pd.Series([1, 2, 3, 4])
+    >>> mask = pd.array([True, True, False, None], dtype="boolean")
+    >>> s[mask]
+    Traceback (most recent call last):
+    ...
+    ValueError: cannot mask with array containing NA / NaN values
+
+*pandas 1.1.0*
+
+.. code-block:: python
+
+    >>> s[mask]
+    0    1
+    1    2
+    dtype: int64
+
 .. _whatsnew_110.enhancements.other:
 
 Other enhancements

pandas/core/arrays/datetimelike.py

@@ -520,7 +520,9 @@ def __getitem__(self, key):
         if com.is_bool_indexer(key):
             # first convert to boolean, because check_array_indexer doesn't
             # allow object dtype
-            key = np.asarray(key, dtype=bool)
+            if is_object_dtype(key):
+                key = np.asarray(key, dtype=bool)
+
             key = check_array_indexer(self, key)
             if key.all():
                 key = slice(0, None, None)

pandas/core/common.py

@@ -118,24 +118,19 @@ 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):
+                na_msg = "Cannot mask with non-boolean array containing NA / NaN values"
                 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,10 @@ def check_array_indexer(array: AnyArrayLike, indexer: Any) -> Any:
 
     dtype = indexer.dtype
     if is_bool_dtype(dtype):
-        try:
+        if is_extension_array_dtype(dtype):
+            indexer = indexer.to_numpy(dtype=bool, na_value=False)
+        else:
             indexer = np.asarray(indexer, dtype=bool)
-        except ValueError:
-            raise ValueError("Cannot mask with a boolean indexer containing NA values")
 
         # GH26658
         if len(indexer) != len(array):

pandas/core/indexing.py

@@ -13,6 +13,7 @@
     is_iterator,
     is_list_like,
     is_numeric_dtype,
+    is_object_dtype,
     is_scalar,
     is_sequence,
 )
@@ -2207,10 +2208,12 @@ def check_bool_indexer(index: Index, key) -> np.ndarray:
                 "the indexed object do not match)."
             )
         result = result.astype(bool)._values
-    else:
+    elif is_object_dtype(key):
         # key might be sparse / object-dtype bool, check_array_indexer needs bool array
         result = np.asarray(result, dtype=bool)
         result = check_array_indexer(index, result)
+    else:
+        result = check_array_indexer(index, result)
 
     return result
 

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/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)]
+
+        self.assert_extension_array_equal(result, expected)
 
         s = pd.Series(data)
 
-        with pytest.raises(ValueError):
-            s[mask]
+        result = s[mask]
+        expected = s[mask.fillna(False)]
+
+        self.assert_series_equal(result, expected)
 
     @pytest.mark.parametrize(
         "idx",

pandas/tests/indexing/test_check_indexer.py

@@ -31,15 +31,16 @@ def test_valid_input(indexer, expected):
     tm.assert_numpy_array_equal(result, expected)
 
 
-@pytest.mark.parametrize(
-    "indexer", [[True, False, None], pd.array([True, False, None], dtype="boolean")],
-)
-def test_bool_raise_missing_values(indexer):
-    array = np.array([1, 2, 3])
+def test_boolean_na_returns_indexer():
+    # TODO: What to do now about list input that contains
+    # bools and None?
+    arr = np.array([1, 2, 3])
+    indexer = pd.array([True, False, None], dtype="boolean")
 
-    msg = "Cannot mask with a boolean indexer containing NA values"
-    with pytest.raises(ValueError, match=msg):
-        check_array_indexer(array, indexer)
+    result = check_array_indexer(arr, indexer)
+    expected = np.array([True, False, False], dtype=bool)
+
+    tm.assert_numpy_array_equal(result, expected)
 
 
 @pytest.mark.parametrize(

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

@@ -72,7 +72,7 @@ def test_getitem_boolean_object(string_series):
 
     # nans raise exception
     omask[5:10] = np.nan
-    msg = "cannot mask with array containing NA / NaN values"
+    msg = "Cannot mask with non-boolean array containing NA / NaN values"
     with pytest.raises(ValueError, match=msg):
         s[omask]
     with pytest.raises(ValueError, match=msg):