allow join on same cat source (#3363)

pola-rs · May 11, 2022 · 2d94240 · 2d94240
1 parent eaabf69
commit 2d94240
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/polars/polars-core/src/chunked_array/logical/categorical/builder.rs b/polars/polars-core/src/chunked_array/logical/categorical/builder.rs
@@ -91,6 +91,9 @@ impl RevMapping {
     pub fn same_src(&self, other: &Self) -> bool {
         match (self, other) {
             (RevMapping::Global(_, _, l), RevMapping::Global(_, _, r)) => *l == *r,
+            (RevMapping::Local(l), RevMapping::Local(r)) => {
+                std::ptr::eq(l as *const Utf8Array<_>, r as *const Utf8Array<_>)
+            }
             _ => false,
         }
     }

diff --git a/py-polars/tests/test_joins.py b/py-polars/tests/test_joins.py
@@ -45,3 +45,16 @@ def test_semi_anti_join() -> None:
         "b": ["c"],
         "payload": [30],
     }
+
+
+def test_join_same_cat_src() -> None:
+    df = pl.DataFrame(
+        data={"column": ["a", "a", "b"], "more": [1, 2, 3]},
+        columns=[("column", pl.Categorical), ("more", pl.Int32)],
+    )
+    df_agg = df.groupby("column").agg(pl.col("more").mean())
+    assert df.join(df_agg, on="column").to_dict(False) == {
+        "column": ["a", "a", "b"],
+        "more": [1, 2, 3],
+        "more_right": [1.5, 1.5, 3.0],
+    }