fix: Fix ColumnNotFound due to projection between filter/cache in CSPE (#26946)

Author: nameexhaustion

crates/polars-plan/src/plans/optimizer/cse/cache_states.rs

@@ -373,10 +373,42 @@ pub(super) fn set_cache_states(
                     .block_at_cache(1);
                 let lp = pred_pd.optimize(start_lp, lp_arena, expr_arena)?;
                 lp_arena.replace(node, lp.clone());
+
+                // TODO: Drop filter column if it isn't used after the filter.
+
+                let mut updated_cache_node = node;
+
+                loop {
+                    match lp_arena.get(updated_cache_node) {
+                        IR::Cache { .. } => break,
+                        IR::SimpleProjection { input, .. } => updated_cache_node = *input,
+                        _ => unreachable!(),
+                    }
+                }
+
                 for &parents in &v.parents[1..] {
-                    let node = get_filter_node(parents, lp_arena)
+                    let filter_node = get_filter_node(parents, lp_arena)
                         .expect("expected filter; this is an optimizer bug");
-                    lp_arena.replace(node, lp.clone());
+
+                    let IR::Filter { input, .. } = lp_arena.get(filter_node) else {
+                        unreachable!()
+                    };
+
+                    let new_lp = match lp_arena.get(*input) {
+                        IR::SimpleProjection { input, columns } => {
+                            debug_assert!(matches!(lp_arena.get(*input), IR::Cache { .. }));
+                            IR::SimpleProjection {
+                                input: updated_cache_node,
+                                columns: columns.clone(),
+                            }
+                        },
+                        ir => {
+                            debug_assert!(matches!(ir, IR::Cache { .. }));
+                            lp_arena.get(updated_cache_node).clone()
+                        },
+                    };
+
+                    lp_arena.replace(filter_node, new_lp);
                 }
             } else {
                 let child = *v.children.first().unwrap();

py-polars/tests/unit/lazyframe/test_cse.py

@@ -1330,3 +1330,29 @@ def f_b(df: pl.DataFrame) -> pl.DataFrame:
         schema={"A": pl.Int32, "PART": pl.Int32, "B": pl.Int32},
     )
     assert_frame_equal(out, expected)
+
+
+def test_cspe_projection_between_filter_and_cache_26916() -> None:
+    lf = pl.LazyFrame(
+        {
+            "VendorID": [1, 1, 2, 2, 2],
+            "total_amount": [10.0, 20.0, 30.0, 40.0, 50.0],
+            "passenger_count": [1, 2, 1, 3, 2],
+        }
+    )
+
+    g1 = lf.group_by("VendorID").agg(pl.mean("total_amount"))
+    g2 = lf.group_by("VendorID").agg(pl.mean("passenger_count"))
+
+    q = g1.join(g2, "VendorID").filter(VendorID=1)
+
+    assert_frame_equal(
+        q.collect(),
+        pl.DataFrame(
+            {
+                "VendorID": 1,
+                "total_amount": 15.0,
+                "passenger_count": 1.5,
+            }
+        ),
+    )