From c4fdebc2be0a1164425adbe4ff16a1dd61d22106 Mon Sep 17 00:00:00 2001 From: Marshall Crumiller Date: Tue, 23 Jan 2024 18:34:31 -0500 Subject: [PATCH] Allow dtype merge of list[enum] --- crates/polars-core/src/datatypes/dtype.rs | 5 ++++- py-polars/tests/unit/datatypes/test_enum.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/crates/polars-core/src/datatypes/dtype.rs b/crates/polars-core/src/datatypes/dtype.rs index 76cb154da5fa..faf9f3088fbe 100644 --- a/crates/polars-core/src/datatypes/dtype.rs +++ b/crates/polars-core/src/datatypes/dtype.rs @@ -450,7 +450,10 @@ pub fn merge_dtypes(left: &DataType, right: &DataType) -> PolarsResult merger.merge_map(rev_map_r)?; Categorical(Some(merger.finish()), *ordering) }, - (RevMapping::Local(_, idl), RevMapping::Local(_, idr)) if idl == idr => { + (RevMapping::Local(_, idl), RevMapping::Local(_, idr)) + | (RevMapping::Enum(_, idl), RevMapping::Enum(_, idr)) + if idl == idr => + { left.clone() }, _ => polars_bail!(string_cache_mismatch), diff --git a/py-polars/tests/unit/datatypes/test_enum.py b/py-polars/tests/unit/datatypes/test_enum.py index 617cbb098e50..3cfdb75be0a8 100644 --- a/py-polars/tests/unit/datatypes/test_enum.py +++ b/py-polars/tests/unit/datatypes/test_enum.py @@ -70,6 +70,24 @@ def test_nested_enum_creation() -> None: assert s.dtype == dtype +def test_nested_enum_concat() -> None: + dtype = pl.List(pl.Enum(["a", "b", "c", "d"])) + s1 = pl.Series([[None, "a"], ["b", "c"]], dtype=dtype) + s2 = pl.Series([["c", "d"], ["a", None]], dtype=dtype) + expected = pl.Series( + [ + [None, "a"], + ["b", "c"], + ["c", "d"], + ["a", None], + ], + dtype=dtype, + ) + + assert_series_equal(pl.concat((s1, s2)), expected) + assert_series_equal(s1.extend(s2), expected) + + def test_casting_to_an_enum_from_utf() -> None: dtype = pl.Enum(["a", "b", "c"]) s = pl.Series([None, "a", "b", "c"])