Repository: pandas-dev/pandas
Pull request: ENH: Enable indexing with nullable Boolean
closes #31503 

- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry

I think from the discussion in https://github.com/pandas-dev/pandas/issues/31503 that this is something people want to allow.

Problem statement:
API: query / boolean selection with nullable dtypes with NAs
#### Code Sample

Create a dataframe with nullable integer, string, and float data types.

```python
>>> df = df = pd.DataFrame({'a': [1, 3, np.nan], 'b': ['rain', 'shine', None], 
                                                    'a_float': [1.1, 3.3, np.nan]})
>>> df = df.convert_dtypes()
>>> df
```

|    | a    | b     |   a_float |
|---:|:-----|:------|----------:|
|  0 | 1    | rain  |       1.1 |
|  1 | 3    | shine |       3.3 |
|  2 | `<NA>` | `<NA>`  |     nan   |

Verify data types and attempt to use `query`

```python
>>> df.dtypes
a            Int64
b            string
a_float      float64

>>> df.query('a > 2') # same as df[df['a'] > 2]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/Documents/Code Practice/pandas-dev/pandas/pandas/core/frame.py in query(self, expr, inplace, **kwargs)
   3227         try:
-> 3228             new_data = self.loc[res]
   3229         except ValueError:

~/Documents/Code Practice/pandas-dev/pandas/pandas/core/indexing.py in __getitem__(self, key)
   1683             maybe_callable = com.apply_if_callable(key, self.obj)
-> 1684             return self._getitem_axis(maybe_callable, axis=axis)
   1685 

~/Documents/Code Practice/pandas-dev/pandas/pandas/core/indexing.py in _getitem_axis(self, key, axis)
   1798             return self._get_slice_axis(key, axis=axis)
-> 1799         elif com.is_bool_indexer(key):
   1800             return self._getbool_axis(key, axis=axis)

~/Documents/Code Practice/pandas-dev/pandas/pandas/core/common.py in is_bool_indexer(key)
    133                 if np.any(key.isna()):
--> 134                     raise ValueError(na_msg)
    135             return True

ValueError: cannot mask with array containing NA / NaN values

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-52-e5d239635d7b> in <module>
----> 1 df.query('a > 2')

~/Documents/Code Practice/pandas-dev/pandas/pandas/core/frame.py in query(self, expr, inplace, **kwargs)
   3230             # when res is multi-dimensional loc raises, but this is sometimes a
   3231             # valid query
-> 3232             new_data = self[res]
   3233 
   3234         if inplace:

~/Documents/Code Practice/pandas-dev/pandas/pandas/core/frame.py in __getitem__(self, key)
   2784 
   2785         # Do we have a (boolean) 1d indexer?
-> 2786         if com.is_bool_indexer(key):
   2787             return self._getitem_bool_array(key)
   2788 

~/Documents/Code Practice/pandas-dev/pandas/pandas/core/common.py in is_bool_indexer(key)
    132             if is_extension_array_dtype(key.dtype):
    133                 if np.any(key.isna()):
--> 134                     raise ValueError(na_msg)
    135             return True
    136     elif isinstance(key, list):

ValueError: cannot mask with array containing NA / NaN values

>>> df.query('a_float > 2')
```
|    |   a | b     |   a_float |
|---:|----:|:------|----------:|
|  1 |   3 | shine |       3.3 |

Using `query` with strings works...

```python
>>> df.query('b == "rain"')
```
|    |   a | b    |   a_float |
|---:|----:|:-----|----------:|
|  0 |   1 | rain |       1.1 |


...but fails for boolean selection
```python
>>> df[df['b'] == 'rain']
ValueError: cannot mask with array containing NA / NaN values
```

strings also fail for inequalities
```
>>> df.query('b >= "rain"') # also df.query('b > "rain"')
ValueError: cannot mask with array containing NA / NaN values
```

#### Problem description

The `query` method  behaves differently for nullable integers, strings, and floats. Here's my summary of how I think they work with the `query method` assuming there are missing values in the columns.

* nullable integers - fails
* strings - works with equality, fails with inequality
* float - works for all

I find it extremely difficult to use if the behavior for all of these types are different for query and boolean selection.

#### Expected Output

I think I would prefer to have both query and boolean selection working like they do with floats, where missing values evaluate as False in a condition. And even if there are missing values in the boolen mask itself, treat those as False. This would harmonize the behavior for all data types. 

This would leave it up to the user to check for missing values. I believe SQL where clauses work in such a manner (missing values in conditions evaluate as False).

#### Output of ``pd.show_versions()``

<details>

INSTALLED VERSIONS
------------------
commit           : None
python           : 3.8.1.final.0
python-bits      : 64
OS               : Darwin
OS-release       : 19.2.0
machine          : x86_64
processor        : i386
byteorder        : little
LC_ALL           : None
LANG             : en_US.UTF-8
LOCALE           : en_US.UTF-8

pandas           : 0+untagged.1.gce8af21.dirty
numpy            : 1.18.1
pytz             : 2019.3
dateutil         : 2.8.1
pip              : 20.0.2
setuptools       : 45.1.0.post20200127
Cython           : 0.29.14
pytest           : None
hypothesis       : None
sphinx           : None
blosc            : None
feather          : None
xlsxwriter       : None
lxml.etree       : 4.5.0
html5lib         : None
pymysql          : None
psycopg2         : None
jinja2           : 2.10.3
IPython          : 7.11.1
pandas_datareader: None
bs4              : None
bottleneck       : None
fastparquet      : None
gcsfs            : None
lxml.etree       : 4.5.0
matplotlib       : 3.1.1
numexpr          : None
odfpy            : None
openpyxl         : None
pandas_gbq       : None
pyarrow          : None
pytables         : None
pytest           : None
pyxlsb           : None
s3fs             : None
scipy            : 1.3.1
sqlalchemy       : 1.3.13
tables           : None
tabulate         : 0.8.3
xarray           : None
xlrd             : None
xlwt             : None
xlsxwriter       : None
numba            : None

</details>

Hints from issue discussion:
Can you include the full tracebacks? That helps show where things are failing. I'm seeing different exceptions in an environment with numexpr.


```pytb
In [17]: df.query('a > 2')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-17-e5d239635d7b> in <module>
----> 1 df.query('a > 2')

~/sandbox/pandas/pandas/core/frame.py in query(self, expr, inplace, **kwargs)
   3223         kwargs["level"] = kwargs.pop("level", 0) + 1
   3224         kwargs["target"] = None
-> 3225         res = self.eval(expr, **kwargs)
   3226
   3227         try:

~/sandbox/pandas/pandas/core/frame.py in eval(self, expr, inplace, **kwargs)
   3338         kwargs["resolvers"] = kwargs.get("resolvers", ()) + tuple(resolvers)
   3339
-> 3340         return _eval(expr, inplace=inplace, **kwargs)
   3341
   3342     def select_dtypes(self, include=None, exclude=None) -> "DataFrame":

~/sandbox/pandas/pandas/core/computation/eval.py in eval(expr, parser, engine, truediv, local_dict, global_dict, resolvers, level, target, inplace)
    328         eng = _engines[engine]
    329         eng_inst = eng(parsed_expr)
--> 330         ret = eng_inst.evaluate()
    331
    332         if parsed_expr.assigner is None:

~/sandbox/pandas/pandas/core/computation/engines.py in evaluate(self)
     71
     72         # make sure no names in resolvers and locals/globals clash
---> 73         res = self._evaluate()
     74         return reconstruct_object(
     75             self.result_type, res, self.aligned_axes, self.expr.terms.return_type

~/sandbox/pandas/pandas/core/computation/engines.py in _evaluate(self)
    112         scope = env.full_scope
    113         _check_ne_builtin_clash(self.expr)
--> 114         return ne.evaluate(s, local_dict=scope)
    115
    116

~/Envs/pandas-dev/lib/python3.7/site-packages/numexpr/necompiler.py in evaluate(ex, local_dict, global_dict, out, order, casting, **kwargs)
    820     # Create a signature
    821     signature = [(name, getType(arg)) for (name, arg) in
--> 822                  zip(names, arguments)]
    823
    824     # Look up numexpr if possible.

~/Envs/pandas-dev/lib/python3.7/site-packages/numexpr/necompiler.py in <listcomp>(.0)
    819
    820     # Create a signature
--> 821     signature = [(name, getType(arg)) for (name, arg) in
    822                  zip(names, arguments)]
    823

~/Envs/pandas-dev/lib/python3.7/site-packages/numexpr/necompiler.py in getType(a)
    701     if kind == 'S':
    702         return bytes
--> 703     raise ValueError("unknown type %s" % a.dtype.name)
    704
    705

ValueError: unknown type object

```


> I think I would prefer to have both query and boolean selection working like they do with floats, where missing values evaluate as False in a condition.

We've defined NA such that it propagates in comparison operations. There was a sizable discussion on it if you go back to the original issue, and we reached consensus that propagating was the best.

```python
In [10]: 1 == pd.NA
Out[10]: <NA>
```

Not saying that can't change, since NA is experimental, but it'll be a high hurdle.
> We've defined NA such that it propagates in comparison operations. 

And in addition also the indexing behaviour with missing values.

So the behaviour you see for boolean masking with missing values is the consequence of two things:

* The comparison operation step propagates missing values
* The indexing step raises on missing values

And both are open for discussion (as Tom noted, there is a reason this is experimental). But I personally think it is rather the second that we should maybe discuss. The alternative under consideration for that one was to not error on missing values in a boolean mask, but to interpret the boolean value as False (skip it in masking). That would *mostly* give the same behaviour as we had before.

We decided to take the "raise on missing values" behaviour for now, as it is the most conservative: the user can always do `fillna(True)` or `fillna(False)` on the boolean mask to be explicit in which behaviour they want. 
(well, when using `query` this is actually not that straightforward ..)

But I think in this area we explicitly said: let's get some more experience with it and re-evaluate later. So feedback on this is certainly welcome. See https://github.com/pandas-dev/pandas/issues/28778 for some of the earlier discussion.
@TomAugspurger - Edited post to include traceback

This is going to be extraordinarily confusing for users and an absolute pain to have all these data types with all these different behaviors when filtering with missing values. 

How is anybody going to reasonably use nullable integers to filter their data? Convert back to float first? 

Users are going to be so confused that nullable integers filter differently than floats and strings only filter for equality, but object columns work (even with pd.NA in them).

```python
>>> df1 = pd.DataFrame({'a': ['rain', 'shine', pd.NA]}) # dtype is object
>>> df1.query('a == "rain"') # works
>>> df1.query('a <= "rain"') # works
```
> Using query with strings works...  df.query('b == "rain"')

This seems a bug, it should be consistent with boolean selection `df[df['b'] == 'rain']`

> How is anybody going to reasonably use nullable integers to filter their data? Convert back to float first?

There is no need to convert to a float. You can still filter with missing values, but for now you need to be explicit how to treat missing values (with fillna), as I explained above. But to be clear, this is not set in stone. Feedback is very valuable, but so let's discuss the specifics.
And it's important to distinguish the two operations here. We have

```python
mask = integer_array == 1  # comparison propagates NA
filtered = df[mask]        # indexing raises on NA
```

Changing the behavior of NA in comparisons just (or primarily) for indexing isn't warranted IMO. There are situations where I'm doing a comparison operation not followed by indexing, and I'd like the NA to propagate.

So IMO the focus should be on the behavior of NA in indexing. @tdpetrou would you be able to go through the NA indexing discussion and summarize the arguments for the various options (IIRC, the 3 were raise -- which we implemented, propagate NA, and treat NA as False).?

edit: Whoops, I didn't see https://github.com/pandas-dev/pandas/issues/31503#issuecomment-580831243 until now, and realized I basically repeated Joris 😄 
My two cents: I think a sensible default would be to treat NA as False when filtering (so we only select values that are _definitely_ true). This would be the same as a typical where clause in SQL, and in my opinion is generally what people want when they filter data.
@TomAugspurger Agreed - NA comparisons should remain NA. For indexing, NA's should be treated as False. Problems solved. It would dramatically simplify things for the user. There are many reasons for this. 

• Matches behavior of old pandas types
• Matches SQL
• Prevents you from having to fill in missing values (which you often don't know)  - people would fill these in with garbage just to make the filter work or just convert to float.

Right now, filtering depends on the data type, which to me, is just absurd. The users are going to get so confused and frustrated. 

The workflow right now for filtering nullable ints is one of two things:
```python
>>> df.fillna({'a': garbage_number}).query('a > 2') # or
>>> df.astype({'a': 'float'}).query('a > 2')
```

There are so many instances where you just want to filter on the known data without having to go through these huge unnecessary headaches.

```
Try to keep things respectful... We put a good amount of discussion into this. My recollection is that raising seemed the safest thing to do for now, in part because it matches the behavior of NaN, and in part because it left open the door to specifying a behavior later on.

---

> • Matches behavior of old pandas types

Can you clarify that? When indexing with NaNs, NumPy (or pandas) will raise, right?
> My two cents: I think a sensible default would be to treat NA as False when filtering (so we only select values that are definitely true).

This does make a ton of sense for boolean masking. Keep in mind, though, that our `__getitem__` also allows fancy indexing (list of integers). What is the expected behavior of the following?

```python
In [3]: s = pd.Series([1, 2, 3])

In [4]: m = pd.array([1, 1, None])

In [5]: s[m]
```

It's not at all clear to me how the `NA` should be treated there. When you index with a list of integers, you'd expect the output shape to match the shape of the indexer, right? Should that value be `NA`? Should it be excluded?

Would it be strange to treat NA in an IntegerArray differently than NA in a boolean array? We don't have a typed NA (`NA[bool], NA[int]`), though it was discussed.
@dsaxton Thanks for the input!

>> Matches behavior of old pandas types
>
> Can you clarify that? When indexing with NaNs, NumPy (or pandas) will raise, right?

@TomAugspurger but you typically didn't get booleans with missing values, as the comparison didn't result in them. So for the combined workflow (comparison + boolean indexing), it mostly matches old behaviour


> Right now, filtering depends on the data type, which to me, is just absurd.

@tdpetrou this is not absurd. We are adding those dtypes because we want to improve things, and yes, that means they don't behave exactly the same as the old ones. There is a reason they are labeled as experimental, and are not yet the default. 
This means we are going to have a period where we have multiple dtypes yes. And I agree this will give a period where this can cause more confusion. But IMO this is unavoidable if we want to improve things. Better documentation might help? Feedback on this is also certainly welcome as someone looking at this with a fresh eye.

Commit message: Delete another test

Patch to review:
diff --git a/doc/source/whatsnew/v1.1.0.rst b/doc/source/whatsnew/v1.1.0.rst
index dfb5ae56cbe1f..52bd13eff3bc1 100644
--- a/doc/source/whatsnew/v1.1.0.rst
+++ b/doc/source/whatsnew/v1.1.0.rst
@@ -212,6 +212,7 @@ Sparse
 ExtensionArray
 ^^^^^^^^^^^^^^
 
+- Enabled filtering with nullable Boolean arrays. (:issue:`31503`)
 -
 -
 
diff --git a/pandas/core/common.py b/pandas/core/common.py
index 745a56ce2be7f..e7b1d1097a085 100644
--- a/pandas/core/common.py
+++ b/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:
diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py
index 77003719360d9..77ce84509a25b 100755
--- a/pandas/core/indexing.py
+++ b/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)
diff --git a/pandas/tests/arrays/categorical/test_indexing.py b/pandas/tests/arrays/categorical/test_indexing.py
index 85d5a6a3dc3ac..652ae2a9dde87 100644
--- a/pandas/tests/arrays/categorical/test_indexing.py
+++ b/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):
     """
diff --git a/pandas/tests/arrays/test_boolean.py b/pandas/tests/arrays/test_boolean.py
index cb9b07db4a0df..b3dc4dbc11b3d 100644
--- a/pandas/tests/arrays/test_boolean.py
+++ b/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)
diff --git a/pandas/tests/extension/base/getitem.py b/pandas/tests/extension/base/getitem.py
index 8615a8df22dcc..50be6494a8e63 100644
--- a/pandas/tests/extension/base/getitem.py
+++ b/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])],
diff --git a/pandas/tests/indexing/test_na_indexing.py b/pandas/tests/indexing/test_na_indexing.py
index befe4fee8ecf8..f1ccaacfec049 100644
--- a/pandas/tests/indexing/test_na_indexing.py
+++ b/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]
diff --git a/pandas/tests/series/indexing/test_boolean.py b/pandas/tests/series/indexing/test_boolean.py
index 28f3c0f7429f8..b78dd2fc84f12 100644
--- a/pandas/tests/series/indexing/test_boolean.py
+++ b/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


Code review comments:
diff_hunk: @@ -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):
path: pandas/tests/indexing/test_na_indexing.py
text: @user1:
Same: adapt for treating NA as false.

diff_hunk: @@ -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):
path: pandas/tests/extension/base/getitem.py
text: @user1:
Rather than removing, can you adapt this test to treat the NAs as False?

diff_hunk: @@ -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)
path: pandas/core/indexing.py
text: @user1:
This should likely be `key.to_numpy(fill_value=False)`.

I think that will avoid the full scans of the values from the `isna` and `.any()`.

@author:
Yes, that's much nicer