handle equality of categorical vs enum

pola-rs · Jan 24, 2024 · d65ddaf · d65ddaf
1 parent 01a0ac4
commit d65ddaf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/crates/polars-core/src/datatypes/dtype.rs b/crates/polars-core/src/datatypes/dtype.rs
@@ -71,9 +71,12 @@ impl PartialEq for DataType {
         use DataType::*;
         {
             match (self, other) {
-                // Don't include rev maps in comparisons
                 #[cfg(feature = "dtype-categorical")]
-                (Categorical(_, _), Categorical(_, _)) => true,
+                (Categorical(rev_l, _), Categorical(rev_r, _)) => {
+                    let is_l_enum = rev_l.as_ref().map_or(false, |x| x.is_enum());
+                    let is_r_enum = rev_r.as_ref().map_or(false, |x| x.is_enum());
+                    is_l_enum == is_r_enum
+                },
                 (Datetime(tu_l, tz_l), Datetime(tu_r, tz_r)) => tu_l == tu_r && tz_l == tz_r,
                 (List(left_inner), List(right_inner)) => left_inner == right_inner,
                 #[cfg(feature = "dtype-duration")]

diff --git a/py-polars/tests/unit/datatypes/test_enum.py b/py-polars/tests/unit/datatypes/test_enum.py
@@ -380,7 +380,7 @@ def test_enum_creating_col_expr() -> None:
         },
         schema={
             "col1": pl.Enum(["a", "b", "c"]),
-            "col2": pl.String,
+            "col2": pl.Categorical(),
             "col3": pl.Enum(["g", "h", "i"]),
         },
     )