Boolean Indexing in Pandas

This is the third post in the series on indexing and selecting data in pandas. If you haven’t read the others yet, see the first post that covers the basics of selecting based on index or relative numerical indexing, and the second post, that talks about slicing. In this post, I’m going to talk about boolean indexing which is the way that I usually select subsets of data when I work with pandas.

What is boolean indexing?

For those familiar with NumPy, this may already be second nature, but for beginners it is not so obvious. Boolean indexing works for a given array by passing a boolean vector into the indexing operator ([]), returning all values that are True.

One thing to note, this array needs to be the same length as the array dimension being indexed.

Let’s look at an example.

>>> import pandas as pd
>>> import numpy as np
>>>
>>> a = np.arange(5)
>>> a
array([0, 1, 2, 3, 4])

Now we can select the first, second and last elements of our array using a list of array indices.

>>> a[[0, 1, 4]]
array([0, 1, 4])

Boolean indexing can do the same, by creating a boolean array of the same size as the entire array, with elements 0, 1 and 4 set to True, all others False.

>>> mask = np.array([True, True, False, False, True])
>>> a[mask]
array([0, 1, 4])

Boolean operators

So now we know how to index our array with a single boolean array. But building that array by hand is a pain, so what you will usually end up doing is applying operations to the original array that return a boolean array themselves.

For example, to select all elements less than 3:

>>> a[a < 3]
array([0, 1, 2])

or all even elements:

>>> a[a % 2 == 0]
array([0, 2, 4])

And we can combine these using expressions, to AND them (&) or OR them (|). With these operators, we can select the same elements from our first example.

>>> a[(a < 2) | (a >= 4)]
array([0, 1, 4])

Another very helpful operators is the inverse or not operator, (~). Remember to watch your parentheses.

>>> a[~((a < 2) | (a >= 4))]
array([2, 3])

On to pandas

In pandas, boolean indexing works pretty much like in NumPy, especially in a Series. You pass in a vector the same length as the Series. Note that this vector doesn’t have to have an index, but if you use a Series as the argument, it does have an index so you need to be aware of how your index aligns. A common method of using boolean indexing is to apply functions to the original Series so your index will always match.

Series

>>> s = pd.Series(np.arange(5), index=list("abcde"))
>>> s
a 0
b 1
c 2
d 3
e 4
dtype: int64
>>> s[[True, True, False, False, True]] # this vector is just a list of boolean values
a 0
b 1
e 4
dtype: int64
>>> s[np.array([True, True, False, False, True])] # this vector is a NumPy array of boolean values
a 0
b 1
e 4
dtype: int64

But, since our index is not the default (i.e. not a RangeIndex), if we use another Series of the same length, it will not work. It needs a matching index, and the default index created below doesn’t match our character index. So we have to specify the index to match.

>>> try:
...     s[pd.Series([True, True, False, False, True])]
... except Exception as ie:
...     print(ie)

Unalignable boolean Series provided as indexer (index of the boolean Series and of the indexed object do not match).
>>> s[pd.Series([True, True, False, False, True], index=list("abcde"))]
a 0
b 1
e 4
dtype: int64

But instead of making a new Series, we’ll just base all of our expressions on our source data Series or DataFrame, then they’ll share an index.

>>> # just like before with NumPy
>>> s[(s < 2) | (s > 3)]
a 0
b 1
e 4
dtype: int64

Make note that you need to surround each expression with parentheses because the Python parser will apply the boolean operators incorrectly. For the example above, it would apply it as s < (2 | s ) < 3. You’ll realize you’re forgetting parentheses when you get complaints about the boolean operators being applied to a series. See?

>>> s[s < 2 | s > 3]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mcw/.pyenv/versions/pandas/lib/python3.8/site-packages/pandas/core/generic.py", line 1329, in __nonzero__
    raise ValueError(
ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

DataFrame

We can also do boolean indexing on DataFrames. A popular way to create the boolean vector is to use one or more of the columns of the DataFrame.

>>> df = pd.DataFrame({'x': np.arange(5), 'y': np.arange(5, 10)})
>>> df[df['x'] < 3]
   x y
0 0 5
1 1 6
2 2 7

You can also supply multiple conditions, just like before with Series. (Remember those parentheses!)

>>> df[(df['x'] < 3) & (df['y'] > 5)]
   x y
1 1 6
2 2 7

.loc, .iloc, and []

If you remember from previous posts, pandas has three primary ways to index the containers. The indexing operator ([]) is sort of a hybrid of using the index labels or location based offsets. .loc is meant for using the index labels, .iloc is for integer based indexing. The good news is that all of them accept boolean arrays, and return subsets of the underlying container.

>>> mask = (df['x'] < 3) & (df['y'] > 5)
>>> mask
0 False
1 True
2 True
3 False
4 False
dtype: bool
>>> df[mask]
   x y
1 1 6
2 2 7
>>> df.loc[mask]
   x y
1 1 6
2 2 7

Note that .iloc is a little different than the others, if you pass in this mask, you’ll get an exception.

>>> try:
...     df.iloc[mask]
... except NotImplementedError as nie:
...     print(nie)
...
iLocation based boolean indexing on an integer type is not available

This is by design, .iloc is only intended to take positional arguments. However, our mask is a Series with an index, so it is rejected. You can still pass in a boolean vector, but just pass in the vector itself without the index.

>>> df.iloc[mask.to_numpy()]
   x y
1 1 6
2 2 7
>>> # or
>>> df.iloc[mask.values]
   x y
1 1 6
2 2 7

Examples!

I think one of the most helpful things when thinking about boolean indexing is to see some examples. You are only limited by what you can express by grouping together any combination of expressions on your data. You do this by carefully grouping your boolean expressions and using parentheses wisely. It can also help to break the problem into pieces as you work on it.

In this series I’ve been grabbing data from the Chicago Data Portal. This time, I thought the list of lobbyists might be interesting. Due to the need for lobbyists to re-register every year, there’s some repeating data. Let’s take a look.

>>> # you should be able to grab this dataset as an unauthenticated user, but you can be rate limited
>>> lbys = pd.read_json("https://data.cityofchicago.org/resource/tq3e-t5yq.json")
>>> lbys.dtypes
year int64
lobbyist_id int64
salutation object
first_name object
last_name object
address_1 object
city object
state object
zip object
country object
email object
phone object
fax object
employer_id int64
employer_name object
created_date object
middle_initial object
address_2 object
suffix object
dtype: object
>>> lbys['created_date'] = pd.to_datetime(lbys['created_date'])
>>> # I'll drop the personally identifiable data, just to be nice
>>> lbys = lbys.drop(['email', 'phone', 'fax', 'last_name'], axis=1)
>>>
>>> lbys.head(3)
   year lobbyist_id salutation first_name address_1 ... employer_name created_date middle_initial address_2 suffix
0 2020 18883 MR. PERICLES 6969 W WABANSIA AVE ... THE PERICLES ORGANIZATION 2020-08-07 NaN NaN NaN
1 2019 18883 NaN PERICLES 6969 W WABANSIA AVE ... THE PERICLES ORGANIZATION 2020-01-21 NaN NaN NaN
2 2018 18883 NaN PERICLES 6969 W WABANSIA AVE ... THE PERICLES ORGANIZATION 2018-12-12 NaN NaN NaN

[3 rows x 15 columns]

In terms of examples, there’s not really too much complexity to deal with, but here’s a few to give you an idea what boolean indexing looks like.

>>> lbys[lbys['year'] == 2020] # all lobbyists registered in 2020
     year lobbyist_id salutation first_name ... created_date middle_initial address_2 suffix
0 2020 18883 MR. PERICLES ... 2020-08-07 NaN NaN NaN
4 2020 17521 MR. STEVE ... 2020-01-15 NaN NaN NaN
.. ... ... ... ... ... ... ... ... ...
998 2020 24740 MRS. CHRISONIA ... 2020-01-16 D. SUITE 1700 NaN
999 2020 15081 MS. LIZ ... 2020-05-29 NaN STE. 900 NaN

[125 rows x 15 columns]
>>> lbys[(lbys['year'] == 2020) & (lbys['city'] == 'CHICAGO')] # lobbyists registered in 2020 from Chicago
     year lobbyist_id salutation first_name ... created_date middle_initial address_2 suffix
0 2020 18883 MR. PERICLES ... 2020-08-07 NaN NaN NaN
4 2020 17521 MR. STEVE ... 2020-01-15 NaN NaN NaN
.. ... ... ... ... ... ... ... ... ...
998 2020 24740 MRS. CHRISONIA ... 2020-01-16 D. SUITE 1700 NaN
999 2020 15081 MS. LIZ ... 2020-05-29 NaN STE. 900 NaN

[76 rows x 15 columns]
>>> # let's get the most popular employer for 2020
>>> pop_emp_id = lbys[lbys['year'] == 2020].groupby('employer_id').count().sort_values(by='lobbyist_id', ascending=False).index[0]
>>> # who works for them?
>>> lbys[(lbys['employer_id'] == pop_emp_id) & (lbys['year'] == 2020)]
     year lobbyist_id salutation first_name address_1 ... employer_name created_date middle_initial address_2 suffix
619 2020 24484 NaN BRIAN 1330 W FULTON ST ... STERLING BAY, LLC AND ITS AFFILIATES 2020-01-15 NaN STE 800 NaN
654 2020 24106 NaN HOWARD 1330 W. FULTON ST ... STERLING BAY, LLC AND ITS AFFILIATES 2020-01-17 NaN SUITE 800 NaN
982 2020 23828 MS. SHELLY 1330 W. FULTON ST ... STERLING BAY, LLC AND ITS AFFILIATES 2020-01-15 NaN SUITE 800 NaN

[3 rows x 15 columns]

If we only want to deal with 2020 data, we can just make a new smaller DataFrame with that data.

lbys = lbys[lbys['year'] == 2020]

Boolean indexing with isin

A helpful method that is often paired with boolean indexing is Series.isin. It returns a boolean vector with all rows that match one of the elements in the arguments.

>>> lbys['state'].tail()
995 MO
996 MO
997 IL
998 IL
999 IL
Name: state, dtype: object
>>> lbys['state'].isin(['IL']).tail()
995 False
996 False
997 True
998 True
999 True
Name: state, dtype: bool
>>> lbys[lbys['state'].isin(['WI', 'IA', 'MO', 'KY', 'IN'])] # lobbyists from bordering states
     year lobbyist_id salutation first_name ... created_date middle_initial address_2 suffix
164 2020 24879 NaN DAN ... 2020-06-15 NaN NaN NaN
786 2012 4644 MR. BRIAN ... 2012-05-25 NaN NaN NaN
788 2019 15543 MR. LORENZO ... 2019-01-10 NaN NaN NaN
789 2018 15543 MR. LORENZO ... 2018-01-17 NaN NaN NaN
790 2017 15543 MR. LORENZO ... 2017-01-19 NaN NaN NaN
791 2020 15543 MR. LORENZO ... 2020-01-10 NaN NaN NaN
792 2016 15543 MR. LORENZO ... 2016-10-13 NaN NaN NaN
994 2016 10222 MS. MARILYN ... 2016-01-13 NaN MO1-800-14-40 NaN
995 2015 10222 MS. MARILYN ... 2015-01-14 NaN MO1-800-14-40 NaN
996 2014 10222 MS. MARILYN ... 2014-05-06 NaN MO1-800-14-40 NaN

[10 rows x 15 columns]

I’ll wrap it up with a slightly more complicated expression.

>>> lbys[
... ~(lbys['state'].isin(['WI', 'IA', 'MO', 'KY', 'IN'])) & # lobbyists NOT from bordering states
... (lbys['state'] != 'IL') & # and NOT from IL
... (lbys['created_date'] >= '2020-07-01') # created in the last half of the year
... ]
     year lobbyist_id salutation first_name ... created_date middle_initial address_2 suffix
26 2021 24967 NaN JON ... 2020-12-27 NaN SUITE 104 NaN
27 2020 24967 NaN JON ... 2020-10-21 NaN SUITE 104 NaN
34 2021 24901 NaN JOSEPH ... 2020-12-20 NaN NaN NaN
160 2021 23782 MR. ALAN ... 2021-01-01 P SUITE 404-352 NaN
227 2020 24909 MS. LAKEITHA ... 2020-07-14 NaN NaN NaN
332 2021 22341 NaN CAROLINE ... 2020-12-23 NaN NaN NaN
511 2021 5361 NaN JAY ... 2020-12-27 NaN SUITE 450 NaN
609 2021 4561 NaN BRET ... 2020-12-18 NaN NaN NaN
660 2021 13801 MR. JOHN ... 2020-12-27 NaN MC 482-C30-C76 NaN
700 2020 24925 NaN TAMI ... 2020-08-06 NaN NaN NaN
862 2020 24969 MR. ALEX ... 2020-08-31 NaN NaN NaN
951 2021 6164 MS. GABRIELLE ... 2020-12-21 NaN 73RD FLOOR NaN

[12 rows x 15 columns]

I also find it helpful to sometimes create a variable for storing the mask. So for the above example, instead of having to parse the entire expression when reading the code, it can be helpful to have expressive variable names for the parts of the indexing expression.

>>> non_bordering = ~(lbys['state'].isin(['WI', 'IA', 'MO', 'KY', 'IN']))
>>> non_illinois = (lbys['state'] != 'IL')
>>>
>>> # more readable maybe?
>>> lbys[non_bordering & non_illinois & (lbys['created_date'] >= '2020-07-01')]

Often when building a complex expression, it can be helpful to build it in pieces, so assigning parts of the mask to variables can make a much more complicated expression easier to read, at the cost of extra variables to deal with. In general, I use variables in the mask if I have to reuse them multiple times, but if only used once, I do the entire expression in place.

In summary, boolean indexing is really quite simple, but powerful. I’ll be looking at a few other ways to select data in pandas in upcoming posts.