DOC/TST: Indexing with NA raises

asv_bench/benchmarks/indexing.py

@@ -131,6 +131,7 @@ def setup(self):
         self.col_scalar = columns[10]
         self.bool_indexer = self.df[self.col_scalar] > 0
         self.bool_obj_indexer = self.bool_indexer.astype(object)
+        self.boolean_indexer = (self.df[self.col_scalar] > 0).astype("boolean")
 
     def time_loc(self):
         self.df.loc[self.idx_scalar, self.col_scalar]
@@ -144,6 +145,9 @@ def time_boolean_rows(self):
     def time_boolean_rows_object(self):
         self.df[self.bool_obj_indexer]
 
+    def time_boolean_rows_boolean(self):
+        self.df[self.bool_obj_indexer]
+
 
 class DataFrameNumericIndexing:
     def setup(self):

doc/source/user_guide/boolean.rst

@@ -14,6 +14,29 @@ Nullable Boolean Data Type
 
 .. versionadded:: 1.0.0
 
+
+.. _boolean.indexing:
+
+Indexing with NA values
+-----------------------
+
+pandas does not allow indexing with NA values. Attempting to do so
+will raise a ``ValueError``.
+
+.. ipython:: python
+   :okexcept:
+
+   s = pd.Series([1, 2, 3])
+   mask = pd.array([True, False, None])
+   s[mask]
+
+The missing values will need to be explicitly filled with True or False prior
+to using the array as a mask.
+
+.. ipython:: python
+
+   s[mask.fillna(False)]
+
 .. _boolean.kleene:
 
 Kleene Logical Operations

pandas/core/arrays/boolean.py

@@ -289,6 +289,13 @@ def _from_factorized(cls, values, original: "BooleanArray"):
     def _formatter(self, boxed=False):
         return str
 
+    @property
+    def _hasnans(self):
+        # Note: this is expensive right now! The hope is that we can
+        # make this faster by having an optional mask, but not have to change
+        # source code using it..
+        return self._mask.any()
+
     def __getitem__(self, item):
         if is_integer(item):
             if self._mask[item]:
@@ -311,7 +318,7 @@ def _coerce_to_ndarray(self, dtype=None, na_value: "Scalar" = libmissing.NA):
         if dtype is None:
             dtype = object
         if is_bool_dtype(dtype):
-            if not self.isna().any():
+            if not self._hasnans:
                 return self._data
             else:
                 raise ValueError(
@@ -485,7 +492,7 @@ def astype(self, dtype, copy=True):
 
         if is_bool_dtype(dtype):
             # astype_nansafe converts np.nan to True
-            if self.isna().any():
+            if self._hasnans:
                 raise ValueError("cannot convert float NaN to bool")
             else:
                 return self._data.astype(dtype, copy=copy)
@@ -497,7 +504,7 @@ def astype(self, dtype, copy=True):
             )
         # for integer, error if there are missing values
         if is_integer_dtype(dtype):
-            if self.isna().any():
+            if self._hasnans:
                 raise ValueError("cannot convert NA to integer")
         # for float dtype, ensure we use np.nan before casting (numpy cannot
         # deal with pd.NA)

pandas/tests/indexing/test_loc.py

@@ -373,6 +373,9 @@ def test_loc_index(self):
         result = df.loc[mask.values]
         tm.assert_frame_equal(result, expected)
 
+        result = df.loc[pd.array(mask, dtype="boolean")]
+        tm.assert_frame_equal(result, expected)
+
     def test_loc_general(self):
 
         df = DataFrame(

pandas/tests/indexing/test_na_indexing.py

@@ -0,0 +1,79 @@
+import pytest
+
+import pandas as pd
+import pandas.util.testing as tm
+
+
+@pytest.mark.parametrize(
+    "values, dtype",
+    [
+        ([1, 2, 3], "int64"),
+        ([1.0, 2.0, 3.0], "float64"),
+        (["a", "b", "c"], "object"),
+        (["a", "b", "c"], "string"),
+        ([1, 2, 3], "datetime64[ns]"),
+        ([1, 2, 3], "datetime64[ns, CET]"),
+        ([1, 2, 3], "timedelta64[ns]"),
+        (["2000", "2001", "2002"], "Period[D]"),
+        ([1, 0, 3], "Sparse"),
+        ([pd.Interval(0, 1), pd.Interval(1, 2), pd.Interval(3, 4)], "interval"),
+    ],
+)
+@pytest.mark.parametrize(
+    "mask", [[True, False, False], [True, True, True], [False, False, False]]
+)
+@pytest.mark.parametrize("box_mask", [True, False])
+@pytest.mark.parametrize("frame", [True, False])
+def test_series_mask_boolean(values, dtype, mask, box_mask, frame):
+    ser = pd.Series(values, dtype=dtype, index=["a", "b", "c"])
+    if frame:
+        ser = ser.to_frame()
+    mask = pd.array(mask, dtype="boolean")
+    if box_mask:
+        mask = pd.Series(mask, index=ser.index)
+
+    expected = ser[mask.astype("bool")]
+
+    result = ser[mask]
+    tm.assert_equal(result, expected)
+
+    if not box_mask:
+        # Series.iloc[Series[bool]] isn't allowed
+        result = ser.iloc[mask]
+        tm.assert_equal(result, expected)
+
+    result = ser.loc[mask]
+    tm.assert_equal(result, expected)
+
+    # empty
+    mask = mask[:0]
+    ser = ser.iloc[:0]
+    expected = ser[mask.astype("bool")]
+    result = ser[mask]
+    tm.assert_equal(result, expected)
+
+    if not box_mask:
+        # Series.iloc[Series[bool]] isn't allowed
+        result = ser.iloc[mask]
+        tm.assert_equal(result, expected)
+
+    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 index with vector 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]