William Wade
William Wade

Reputation: 405

Using lists in a pandas query

I am performing a query on a DataFrame:

Index Category
1     Foo
2     Bar
3     Cho
4     Foo

I would like to return the rows where the category is "Foo" or "Bar". When I use the code:

df.query("Catergory==['Foo','Bar']")

This works fine and returns:

Index Category
1     Foo
2     Bar
4     Foo

However in future I will want the filter to be changed dynamically so I wrote:

filter_list=['Foo','Bar']
df.query("Catergory==filter_list")

Which threw out the error:

UndefinedVariableError: name 'filter_list' is not defined

Other variations I tried with no success were:

df.query("Catergory"==filter_list)
df.query("Catergory=="filter_list)

Respectively producing:

ValueError: expr must be a string to be evaluated, <class 'bool'> given
SyntaxError: invalid syntax

Upvotes: 24

Views: 31434

Answers (3)

samredai
samredai

Reputation: 712

Try this:

df = pd.DataFrame({'Index':[1,2,3,4],'Category':['Foo','Bar','Cho','Foo']})
filter_list = ['Foo','Bar']

df.query(f'Category=={filter_list}')

Upvotes: 10

Erfan
Erfan

Reputation: 42916

Use isin method.

df.query('Category.isin(@filter_list)')

Upvotes: 22

Scott Boston
Scott Boston

Reputation: 153510

Use @ to reference variables in query:

filter_list=['Foo','Bar']

df.query("Category == @filter_list")

Output:

   Index Category
0      1      Foo
1      2      Bar
3      4      Foo

Upvotes: 38

Related Questions