CODEWITHSUNDEEP
pythonpandasdataframe
ajrlewis
ajrlewis

Reputation: 3058

Query data frame using list of columns and list of values

I have the following simple data frame:

stores = ['a', 'a', 'b', 'b', 'b']
brands = ['Nike', 'Nike', 'Adidas', 'Nike', 'Adidas']
colours = ['Black', 'Black', 'White', 'Black', 'Black']
data = dict(stores=stores, brands=brands, colours=colours)
df = pd.DataFrame(data, columns=data.keys())

I'd like to query this using a list of columns and a corresponding list of values. For e.g.

columns = ['stores', 'brands']
values = ['a', 'Nike']
df[columns == values]

Is this possible?

Upvotes: 1

Views: 86

Answers (2)

rnorris
rnorris

Reputation: 2092

You can do this in a similar way to thesilkworm's answer using only pandas:

query = " & ".join([c + " == '" + v + "'" for c,v in zip(columns, values)])
df.query(query)

Output using the above code:

>>> query = " & ".join([c + " == '" + v + "'" for c,v in zip(columns, values)])
>>> query
"stores == 'a' & brands == 'Nike'"
>>> df.query(query)
  stores brands colours
0      a   Nike   Black
1      a   Nike   Black

Note the inclusion of single quotes around v in the list comprehension. These are important, since we're comparing a string value. For more info, see the query documentation for pandas.

Upvotes: 1

sjw
sjw

Reputation: 6543

This should be possible using numpy.logical_and with reduce for an arbitrary number of conditions:

import numpy as np

df[np.logical_and.reduce([df[col] == val for col, val in zip(columns, values)])]

Results:

  stores brands colours
0      a   Nike   Black
1      a   Nike   Black

Upvotes: 2

Related Questions