fix zip_with categorical (#2457)

polars/polars-core/src/chunked_array/ops/zip.rs

@@ -164,7 +164,13 @@ impl ChunkZip<CategoricalType> for CategoricalChunked {
         other: &ChunkedArray<CategoricalType>,
     ) -> Result<ChunkedArray<CategoricalType>> {
         let ca: CategoricalChunked = self.deref().zip_with(mask, other.deref())?.into();
-        Ok(ca.set_state(self))
+        // first set old state
+        let mut ca = ca.set_state(self);
+        // then merge state
+        let state = ca.merge_categorical_map(other);
+        // and set this
+        ca.categorical_map = Some(state);
+        Ok(ca)
     }
 }
 

polars/polars-lazy/src/tests/arity.rs

@@ -34,3 +34,24 @@ fn test_pearson_corr() -> Result<()> {
     assert!((s.get(1).unwrap() - 0.9552238805970149).abs() < 0.000001);
     Ok(())
 }
+
+#[test]
+fn test_when_then_otherwise_categorical() -> Result<()> {
+    let df = df!["col1"=> ["a", "b", "a", "b"],
+        "col2"=> ["a", "a", "b", "b"],
+        "col3"=> ["same", "same", "same", "same"]
+    ]?;
+
+    let out = df
+        .lazy()
+        .with_column(col("*").cast(DataType::Categorical))
+        .select([when(col("col1").eq(col("col2")))
+            .then(col("col3"))
+            .otherwise(col("col1"))])
+        .collect()?;
+    let col = out.column("col3")?;
+    assert_eq!(col.dtype(), &DataType::Categorical);
+    let s = format!("{}", col);
+    assert!(s.contains("same"));
+    Ok(())
+}