fix(rust, python): don't clear rev_map when categorical series is cle… (

pola-rs#9720)
c-peters · Jul 14, 2023 · f8d9f5e · f8d9f5e
1 parent f3abd78
commit f8d9f5e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/polars/polars-core/src/series/ops/null.rs b/polars/polars-core/src/series/ops/null.rs
@@ -8,7 +8,14 @@ impl Series {
                 ListChunked::full_null_with_dtype(name, size, inner_dtype).into_series()
             }
             #[cfg(feature = "dtype-categorical")]
-            DataType::Categorical(_) => CategoricalChunked::full_null(name, size).into_series(),
+            DataType::Categorical(rev_map) => {
+                let mut ca = CategoricalChunked::full_null(name, size);
+                // ensure we keep the rev-map of a cleared series
+                if let Some(rev_map) = rev_map {
+                    unsafe { ca.set_rev_map(rev_map.clone(), false) }
+                }
+                ca.into_series()
+            }
             #[cfg(feature = "dtype-date")]
             DataType::Date => Int32Chunked::full_null(name, size)
                 .into_date()

diff --git a/py-polars/tests/unit/datatypes/test_categorical.py b/py-polars/tests/unit/datatypes/test_categorical.py
@@ -401,3 +401,13 @@ def test_construct_with_null() -> None:
 
     s = pl.Series([{"struct_A": None}], dtype=pl.Struct({"struct_A": pl.Categorical}))
     assert s.to_list() == [{"struct_A": None}]
+
+
+def test_categorical_concat_string_cached() -> None:
+    with pl.StringCache():
+        df1 = pl.DataFrame({"x": ["A"]}).with_columns(pl.col("x").cast(pl.Categorical))
+        df2 = pl.DataFrame({"x": ["B"]}).with_columns(pl.col("x").cast(pl.Categorical))
+
+    out = pl.concat([df1, df2])
+    assert out.dtypes == [pl.Categorical]
+    assert out["x"].to_list() == ["A", "B"]