join expression argument fix

pola-rs · Aug 6, 2022 · 18027ff · 18027ff
1 parent 093d55c
commit 18027ff
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/alp.rs b/polars/polars-lazy/src/logical_plan/alp.rs
@@ -799,10 +799,9 @@ impl<'a> ALogicalPlanBuilder<'a> {
 
         let right_names: PlHashSet<_> = right_on
             .iter()
-            .map(|e| match self.expr_arena.get(*e) {
-                AExpr::Alias(_, name) => name.clone(),
-                AExpr::Column(name) => name.clone(),
-                _ => panic!("could not determine join column names"),
+            .map(|e| {
+                aexpr_to_root_column_name(*e, self.expr_arena)
+                    .expect("could not determine join column names")
             })
             .collect();
 

diff --git a/py-polars/tests/test_joins.py b/py-polars/tests/test_joins.py
@@ -301,3 +301,6 @@ def test_join_on_cast() -> None:
         "row_nr": [1, 2, 3, 5],
         "a": [-2, 3, 3, 10],
     }
+    assert df_a.lazy().join(
+        df_b.lazy(), on=pl.col("a").cast(pl.Int64)
+    ).collect().to_dict(False) == {"row_nr": [1, 2, 3, 5], "a": [-2, 3, 3, 10]}