How to filter categorical column based on a list of string values

user_15115301 0
Viewed 125
df = pl.DataFrame({"a": ["foo", "bar", "ham"]}).select(pl.col('a').cast(pl.Categorical))
df
shape: (3, 1)
│ a   │
│ cat │
|-----|
│ foo │
│ bar │
│ ham │

Now apply the filter:
df.filter(pl.col('a').isin(['foo']))
df
┌─────┐
│ a   │
│ --- │
│ cat │
╞═════╡
└─────┘

This won't work. I know that expr.is_in([1]) will work, because that's an encoding for one of the values, but how do i say is_in with a pure string representation that will be translated to the appropriate categorical value?

python-polars
1 Answers

You need to set a global string cache to compare categoricals created in different columns/lists.

with pl.StringCache():
    df = pl.DataFrame({"a": ["foo", "bar", "ham"]}).select(pl.col('a').cast(pl.Categorical))
    print(df.filter(pl.col('a').is_in(['foo'])))

shape: (1, 1)
┌─────┐
│ a   │
│ --- │
│ cat │
╞═════╡
│ foo │
└─────┘
user_67170540
Related