Enum filtering has broken. about polars

[cmdlineluser]

Do you happen to know what version this worked on?

Going backwards, I get a different error on 0.20.21

TypeError: invalid literal value: 'State.VIC'

But it fails on every version I've tried.

[orlp]

I'm not even sure if this is actually a bug.

[mcrumiller]

We do not automatically convert from python Enum types to polars Enum Series; you can put in a feature request for this. Note that the dtype of your dataframe is simply a string:

>>> data
shape: (6, 1)
┌─────────────────┐
│ state           │
│ ---             │
│ str             │
╞═════════════════╡
│ victoria        │
│ victoria        │
│ victoria        │
│ new south wales │
│ new south wales │
│ new south wales │
└─────────────────┘

type(State.VIC) is <enum 'State'>, so polars is trying to filter a string column based on an object, and doesn't like it. State.VIC.value is a string, and so the filter works.

So tl;dr you are not actually performing enum filtering, but trying to filter a string column based on a python Enum object, which polars does not recognize.

[Andre-Medina]

Sorry guys, I just re-tested on previous versions and found I missed something key. The enum class was a string enum, not just an enum.

from enum import Enum
import polars as pl

class State(str, Enum):  # NOTE: missed the `str` before the enum
    
    VIC = "victoria"
    NSW = "new south wales"

data = pl.DataFrame({
    'state': [State.VIC] * 3 + [State.NSW] * 3  
})

print(data)

prints out: (note: correctly converts the 'enum' to a string)

shape: (6, 1)
┌─────────────────┐
│ state           │
│ ---             │
│ str             │
╞═════════════════╡
│ victoria        │
│ victoria        │
│ victoria        │
│ new south wales │
│ new south wales │
│ new south wales │
└─────────────────┘

Filtering if on 0.20.31:

>>> print(data.filter(pl.col('state') == State.VIC)) # Does not filter
shape: (0, 1)
┌───────┐
│ state │
│ ---   │
│ str   │
╞═══════╡
└───────┘
>>> print(data.filter(pl.col('state') == State.VIC.value)) # Adding the .value, filters correctly
shape: (3, 1)
┌──────────┐
│ state    │
│ ---      │
│ str      │
╞══════════╡
│ victoria │
│ victoria │
│ victoria │
└──────────┘

However, back in 0.20.25

>>> print(data.filter(pl.col('state') == State.VIC)) # With or without the .value, filters correctly 
shape: (3, 1)
┌──────────┐
│ state    │
│ ---      │
│ str      │
╞══════════╡
│ victoria │
│ victoria │
│ victoria │
└──────────┘

This string enum also works in other libraries, so there's precedent for it to also work in polars:

assert State.VIC == "victoria"
my_dict = {}
my_dict[State.VIC] = 'foo'
print(my_dict['victoria']) # prints out 'foo'

# Or in pandas
import pandas as pd
data_pd = pd.DataFrame({'state': [State.VIC] * 3 + [State.NSW] * 3})
print(data_pd.loc[lambda df: df['state'] == State.VIC])
# Prints:
#        state
# 0  State.VIC
# 1  State.VIC
# 2  State.VIC