Python & Pandas: How to query if a list-type column contains something?
PythonPandasPython Problem Overview
I have a dataframe, which contains info about movies. It has a column called genre, which contains a list of genres it belongs to. For example:
df['genre']
## returns
0 ['comedy', 'sci-fi']
1 ['action', 'romance', 'comedy']
2 ['documentary']
3 ['crime','horror']
...
I want to know how can I query the dataframe, so it returns the movie belongs to a cerain genre?
For example, something may like df['genre'].contains('comedy') returns 0 or 1.
I know for a list, I can do things like:
'comedy' in ['comedy', 'sci-fi']
However, in pandas, I didn't find something similar, the only thing I know is df['genre'].str.contains(), but it didn't work for the list type.
Python Solutions
Solution 1 - Python
You can use apply for create mask and then boolean indexing:
mask = df.genre.apply(lambda x: 'comedy' in x)
df1 = df[mask]
print (df1)
genre
0 [comedy, sci-fi]
1 [action, romance, comedy]
Solution 2 - Python
using sets
df.genre.map(set(['comedy']).issubset)
0 True
1 True
2 False
3 False
dtype: bool
df.genre[df.genre.map(set(['comedy']).issubset)]
0 [comedy, sci-fi]
1 [action, romance, comedy]
dtype: object
presented in a way I like better
comedy = set(['comedy'])
iscomedy = comedy.issubset
df[df.genre.map(iscomedy)]
more efficient
comedy = set(['comedy'])
iscomedy = comedy.issubset
df[[iscomedy(l) for l in df.genre.values.tolist()]]
using str in two passes
slow! and not perfectly accurate!
df[df.genre.str.join(' ').str.contains('comedy')]
Solution 3 - Python
According to the source code, you can use .str.contains(..., regex=False).
Solution 4 - Python
You need to set regex=False and .str.contains will work for list values as you would expect:
In : df['genre'].str.contains('comedy', regex=False)
Out:
0 True
1 True
2 False
3 False
Name: genre, dtype: bool
Solution 5 - Python
A complete example:
import pandas as pd
data = pd.DataFrame([[['foo', 'bar']],
[['bar', 'baz']]], columns=['list_column'])
print(data)
list_column
0 [foo, bar]
1 [bar, baz]
filtered_data = data.loc[ lambda df: df.list_column.apply( lambda l: 'foo' in l )]
print(filtered_data)
list_column
0 [foo, bar]