performance[rust]: improve predicate pushdown + with_columns (#5029)

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

@@ -1077,7 +1077,7 @@ impl ProjectionPushDown {
 
                     // Make sure that columns selected with_columns are available
                     // only if not empty. If empty we already select everything.
-                    for expression in &exprs {
+                    for expression in &pruned_with_cols {
                         add_expr_to_accumulated(
                             *expression,
                             &mut acc_projections,

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

@@ -461,10 +461,35 @@ fn test_with_column_prune() -> PolarsResult<()> {
         "c1" => [0],
         "c2" => [0],
     ]?;
+    let (mut expr_arena, mut lp_arena) = get_arenas();
+
+    // only a single expression pruned and only one column selection
+    let q = df
+        .clone()
+        .lazy()
+        .with_columns([col("c0"), col("c1").alias("c4")])
+        .select([col("c1"), col("c4")]);
+    let lp = q.optimize(&mut lp_arena, &mut expr_arena).unwrap();
+    (&lp_arena).iter(lp).for_each(|(_, lp)| {
+        use ALogicalPlan::*;
+        match lp {
+            DataFrameScan { projection, .. } => {
+                let projection = projection.as_ref().unwrap();
+                let projection = projection.as_slice();
+                assert_eq!(projection.len(), 1);
+                let name = &projection[0];
+                assert_eq!(name, "c1");
+            }
+            HStack { exprs, .. } => {
+                assert_eq!(exprs.len(), 1);
+            }
+            _ => {}
+        };
+    });
 
+    // whole `with_columns` pruned
     let q = df.lazy().with_column(col("c0")).select([col("c1")]);
 
-    let (mut expr_arena, mut lp_arena) = get_arenas();
     let lp = q.clone().optimize(&mut lp_arena, &mut expr_arena).unwrap();
 
     // check if with_column is pruned