fix: Preserve name when casting to Enum (#14320)

pola-rs · Feb 6, 2024 · 6fc0395 · 6fc0395
1 parent 6ab6550
commit 6fc0395
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/crates/polars-core/src/chunked_array/logical/categorical/mod.rs b/crates/polars-core/src/chunked_array/logical/categorical/mod.rs
@@ -373,7 +373,8 @@ impl LogicalType for CategoricalChunked {
                 Ok(self
                     .to_enum(categories, *hash)?
                     .set_ordering(*ordering, true)
-                    .into_series())
+                    .into_series()
+                    .with_name(self.name()))
             },
             DataType::Enum(None, _) => {
                 polars_bail!(ComputeError: "can not cast to enum without categories present")

diff --git a/py-polars/tests/unit/operations/test_cast.py b/py-polars/tests/unit/operations/test_cast.py
@@ -578,6 +578,10 @@ def test_strict_cast_string_and_binary(
             _cast_lit_t(*args)  # type: ignore[arg-type]
 
 
+@pytest.mark.parametrize(
+    "dtype_in",
+    [(pl.Categorical), (pl.Enum(["1"]))],
+)
 @pytest.mark.parametrize(
     "dtype_out",
     [
@@ -593,11 +597,15 @@ def test_strict_cast_string_and_binary(
         (pl.Datetime),
         (pl.Time),
         (pl.Duration),
-        (pl.Enum(["1"])),
+        (pl.String),
+        (pl.Categorical),
+        (pl.Enum(["1", "2"])),
     ],
 )
-def test_cast_categorical_name_retention(dtype_out: PolarsDataType) -> None:
-    assert pl.Series("a", ["1"], dtype=pl.Categorical).cast(dtype_out).name == "a"
+def test_cast_categorical_name_retention(
+    dtype_in: PolarsDataType, dtype_out: PolarsDataType
+) -> None:
+    assert pl.Series("a", ["1"], dtype=dtype_in).cast(dtype_out).name == "a"
 
 
 def test_cast_date_to_time() -> None: