fix(rust, python): projection_node always do projection locally if no…

polars/polars-lazy/polars-plan/src/logical_plan/optimizer/projection_pushdown.rs

@@ -282,6 +282,7 @@ impl ProjectionPushDown {
 
         match logical_plan {
             Projection { expr, input, .. } => {
+                let mut local_projection = Vec::with_capacity(expr.len());
                 // A projection can consist of a chain of expressions followed by an alias.
                 // We want to do the chain locally because it can have complicated side effects.
                 // The only thing we push down is the root name of the projection.
@@ -327,6 +328,10 @@ impl ProjectionPushDown {
                             }
                         }
                     }
+                    // do local as we still need the effect of the projection
+                    // e.g. a projection is more than selecting a column, it can
+                    // also be a function/ complicated expression
+                    local_projection.push(*e);
 
                     add_expr_to_accumulated(
                         *e,
@@ -344,27 +349,6 @@ impl ProjectionPushDown {
                     lp_arena,
                     expr_arena,
                 )?;
-                let lp = lp_arena.get(input);
-
-                let mut local_projection = Vec::with_capacity(expr.len());
-
-                // the projections should all be done at the latest projection node to keep the same schema order
-                if projections_seen == 0
-                    || expr.iter().any(|node| has_aexpr_alias(*node, expr_arena))
-                {
-                    let schema = lp.schema(lp_arena);
-                    for node in expr {
-                        // Due to the pushdown, a lot of projections cannot be done anymore at the final
-                        // node and should be skipped
-                        if expr_arena
-                            .get(node)
-                            .to_field(&schema, Context::Default, expr_arena)
-                            .is_ok()
-                        {
-                            local_projection.push(node);
-                        }
-                    }
-                }
 
                 let builder = ALogicalPlanBuilder::new(input, expr_arena, lp_arena);
                 let lp = self.finish_node(local_projection, builder);

polars/polars-lazy/polars-plan/src/utils.rs

@@ -104,10 +104,6 @@ where
     arena.iter(current_node).any(|(_node, e)| matches(e))
 }
 
-pub(crate) fn has_aexpr_alias(current_node: Node, arena: &Arena<AExpr>) -> bool {
-    has_aexpr(current_node, arena, |e| matches!(e, AExpr::Alias(_, _)))
-}
-
 pub fn has_aexpr_window(current_node: Node, arena: &Arena<AExpr>) -> bool {
     has_aexpr(current_node, arena, |e| matches!(e, AExpr::Window { .. }))
 }

polars/tests/it/lazy/projection_queries.rs

@@ -118,3 +118,36 @@ fn test_many_aliasing_projections_5070() -> PolarsResult<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_projection_5086() -> PolarsResult<()> {
+    let df = df![
+        "a" => ["a", "a", "a", "b"],
+        "b" => [1, 0, 1, 0],
+        "c" => [0, 1, 2, 0],
+    ]?;
+
+    let out = df
+        .lazy()
+        .select([
+            col("a"),
+            col("b").take("c").cumsum(false).over([col("a")]).gt(lit(0)),
+        ])
+        .select([
+            col("a"),
+            col("b")
+                .xor(col("b").shift(1).over([col("a")]))
+                .fill_null(lit(true))
+                .alias("keep"),
+        ])
+        .collect()?;
+
+    let expected = df![
+        "a" => ["a", "a", "a", "b"],
+        "keep" => [true, false, false, true]
+    ]?;
+
+    assert!(out.frame_equal(&expected));
+
+    Ok(())
+}