fix outer join categorical (#2760)

pola-rs · Feb 24, 2022 · f2ad684 · f2ad684
1 parent 4044d75
commit f2ad684
Show file tree

Hide file tree

Showing 3 changed files with 118 additions and 27 deletions.
diff --git a/polars/polars-core/src/frame/hash_join/mod.rs b/polars/polars-core/src/frame/hash_join/mod.rs
@@ -1431,17 +1431,18 @@ impl DataFrame {
                 )
             },
         );
+
         let mut s = s_left
             .to_physical_repr()
             .zip_outer_join_column(&s_right.to_physical_repr(), &opt_join_tuples);
         s.rename(s_left.name());
         let s = match s_left.dtype() {
             #[cfg(feature = "dtype-categorical")]
             DataType::Categorical(_) => {
-                let ca_or = s_left.categorical().unwrap();
+                let ca_left = s_left.categorical().unwrap();
+                let new_rev_map = ca_left.merge_categorical_map(s_right.categorical().unwrap());
                 let logical = s.u32().unwrap().clone();
-                CategoricalChunked::from_cats_and_rev_map(logical, ca_or.get_rev_map().clone())
-                    .into_series()
+                CategoricalChunked::from_cats_and_rev_map(logical, new_rev_map).into_series()
             }
             dt @ DataType::Datetime(_, _)
             | dt @ DataType::Time

diff --git a/py-polars/Cargo.lock b/py-polars/Cargo.lock
diff --git a/py-polars/tests/test_categorical.py b/py-polars/tests/test_categorical.py
@@ -25,6 +25,23 @@ def test_categorical_outer_join() -> None:
     expected = pl.DataFrame({"val1": [1], "key1": [42], "key2": ["bar"], "val2": [2]})
 
     assert out.frame_equal(expected)
+    with pl.StringCache():
+        dfa = pl.DataFrame(
+            [
+                pl.Series("key", ["foo", "bar"], dtype=pl.Categorical),
+                pl.Series("val1", [3, 1]),
+            ]
+        )
+        dfb = pl.DataFrame(
+            [
+                pl.Series("key", ["bar", "baz"], dtype=pl.Categorical),
+                pl.Series("val2", [6, 8]),
+            ]
+        )
+
+    df = dfa.join(dfb, on="key", how="outer")
+    # the cast is important to test the rev map
+    assert df["key"].cast(pl.Utf8).to_list() == ["bar", "baz", "foo"]
 
 
 def test_read_csv_categorical() -> None: