fix[rust]: fix predicate pushdown of unions #4780

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/utils.rs

@@ -243,7 +243,11 @@ where
                     if projection_roots.len() == 1 {
                         // we were able to rename the alias column with the root column name
                         // before pushing down the predicate
-                        rename_aexpr_root_names(predicate, expr_arena, projection_roots[0].clone());
+                        let predicate = rename_aexpr_root_names(
+                            predicate,
+                            expr_arena,
+                            projection_roots[0].clone(),
+                        );
 
                         insert_and_combine_predicate(
                             acc_predicates,

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

@@ -2089,20 +2089,3 @@ fn test_partitioned_gb_ternary() -> Result<()> {
 
     Ok(())
 }
-
-#[test]
-fn test_foo() -> Result<()> {
-    let df = df![
-        "a" => ["a", "b"]
-    ]?;
-
-    let df = df
-        .lazy()
-        .select([all().cast(DataType::Categorical(None)).list()])
-        .collect()?;
-
-    let out = concat_df(&[df.clone(), df.clone()])?;
-    dbg!(out.agg_chunks());
-
-    Ok(())
-}

polars/polars-lazy/src/utils.rs

@@ -225,15 +225,20 @@ pub(crate) fn aexpr_to_root_nodes(root: Node, arena: &Arena<AExpr>) -> Vec<Node>
 /// In some cases we can have multiple roots.
 /// For instance in predicate pushdown the predicates are combined by their root column
 /// When combined they may be a binary expression with the same root columns
-pub(crate) fn rename_aexpr_root_names(node: Node, arena: &mut Arena<AExpr>, new_name: Arc<str>) {
-    let roots = aexpr_to_root_nodes(node, arena);
-
-    for node in roots {
-        arena.replace_with(node, |ae| match ae {
-            AExpr::Column(_) => AExpr::Column(new_name.clone()),
-            _ => panic!("should be only a column"),
-        });
-    }
+pub(crate) fn rename_aexpr_root_names(
+    node: Node,
+    arena: &mut Arena<AExpr>,
+    new_name: Arc<str>,
+) -> Node {
+    // we convert to expression as we cannot easily copy the aexpr.
+    let mut new_expr = node_to_expr(node, arena);
+    new_expr.mutate().apply(|e| {
+        if let Expr::Column(name) = e {
+            *name = new_name.clone()
+        }
+        true
+    });
+    to_aexpr(new_expr, arena)
 }
 
 /// Rename the root of the expression from `current` to `new` and assign to new node in arena.

py-polars/tests/unit/test_df.py

@@ -2153,3 +2153,21 @@ def test_set() -> None:
     # we cannot index with any type, such as bool
     with pytest.raises(ValueError):
         df[True] = 1  # type: ignore[index]
+
+
+def test_union_with_aliases_4770() -> None:
+    lf = pl.DataFrame(
+        {
+            "a": [1, None],
+            "b": [3, 4],
+        }
+    ).lazy()
+
+    lf = pl.concat(
+        [
+            lf.select([pl.col("a").alias("x")]),
+            lf.select([pl.col("b").alias("x")]),
+        ]
+    ).filter(pl.col("x").is_not_null())
+
+    assert lf.collect()["x"].to_list() == [1, 3, 4]