feat: Extend allowed types of the categories arg of dy.Enum
#138
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #138 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 49 49
Lines 2812 2818 +6
=========================================
+ Hits 2812 2818 +6 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
jjurm
changed the title
categories arg of dy.Enumcategories arg of dy.Enum
jjurm
requested review from
borchero,
AndreasAlbertQC and
delsner
as code owners
There was a problem hiding this comment.
Neat, looks great, thanks @jjurm! I think it's fine to represent the categories as series, I can't come up with a downside and we mirror the polars API this way. Would love to hear a second opinion on this by @AndreasAlbertQC or @delsner though 😄
There was a problem hiding this comment.
Thanks for the PR @jjurm! Personally, I'd have gone for list, but I also don't think it's super important 🤷
| values=[getattr(v, "value", v) for v in categories.__members__.values()] | ||
| ) | ||
| elif not isinstance(categories, pl.Series): | ||
| categories = pl.Series(values=categories) |
There was a problem hiding this comment.
My take is that I would not use a polars object when I can just use vanilla python without drawbacks, which seems to be the case here. 🤷 Why not just list?
There was a problem hiding this comment.
Most of the reason is to mirror Polars, so we can pass categories around without intermediate conversions, e.g.:
- A call to the
.dtypeproperty creates thepl.Enumtype which would convert lists to Series on every call. This conversion is skipped when the categories already are apl.Series. - Storing categories as a Series allows us to defer the category comparison logic to Polars, therefore ensuring a more consistent behavior with Polars in the future (sure, this is more of a cosmetic reason)
| @@ -72,7 +81,7 @@ def dtype(self) -> pl.DataType: | |||
| def validate_dtype(self, dtype: PolarsDataType) -> bool: | |||
| if not isinstance(dtype, pl.Enum): | |||
| return False | |||
| return self.categories == dtype.categories.to_list() | |||
| return self.categories.equals(dtype.categories) | |||
There was a problem hiding this comment.
Does not need to change in this PR, but while we are talking: I always feel like this should be a set comparison so that enums with the same categories but different order are considered the same. wdyt?
There was a problem hiding this comment.
Enum uses category order for comparisons and sorts, so two Enums with the same labels but different order are not equivalent. Therefore I’d keep it order-sensitive on purpose.
There was a problem hiding this comment.
Fair point. I tend to think of enums just in terms of "what values are allowed" and never in terms of sorting, but that might be my personal bias :)
There was a problem hiding this comment.
For reference, Polars Enum is also order-sensitive.
| @@ -22,7 +24,7 @@ class Enum(Column): | |||
|
|
|||
| def __init__( | |||
| self, | |||
| categories: Sequence[str], | |||
| categories: pl.Series | Iterable[str] | type[enum.Enum], | |||
There was a problem hiding this comment.
IMO fine to accept Series here, independently of whether we use it internally as well 👍
Motivation
This makes the
categoriesargument accept the same types as inpl.Enum, making it easier to convert from polars dtypes to dataframely dtypes.Fixes #136.
Changes
dy.Enum.categoriesis now stored as apl.Series- happy to think of another solution if you don't think we should convert all sequences to series. However, this is what polars does internally anyway oncedy.Enum.dtypeis called.