fix(rust): keep track of removed indices in CWC (#16443)

crates/polars-plan/src/logical_plan/optimizer/cluster_with_columns.rs

@@ -157,6 +157,7 @@ pub fn optimize(root: Node, lp_arena: &mut Arena<IR>, expr_arena: &Arena<AExpr>)
         let mut has_seen_unpushable = false;
         let mut needs_simple_projection = false;
 
+        let mut already_removed = 0;
         *current_exprs.exprs_mut() = std::mem::take(current_exprs.exprs_mut())
             .into_iter()
             .zip(pushable.iter())
@@ -166,8 +167,11 @@ pub fn optimize(root: Node, lp_arena: &mut Arena<IR>, expr_arena: &Arena<AExpr>)
                     needs_simple_projection = has_seen_unpushable;
 
                     input_exprs.exprs_mut().push(expr);
-                    let (column, datatype) = new_current_schema.shift_remove_index(i).unwrap();
+                    let (column, datatype) = new_current_schema
+                        .shift_remove_index(i - already_removed)
+                        .unwrap();
                     new_input_schema.with_column(column, datatype);
+                    already_removed += 1;
 
                     None
                 } else {

py-polars/tests/unit/test_cwc.py

@@ -1,5 +1,6 @@
 # Tests for the optimization pass cluster WITH_COLUMNS
 
+
 import polars as pl
 
 
@@ -134,3 +135,22 @@ def test_cwc_with_internal_aliases() -> None:
         """[[(col("a")) == (2)].cast(Boolean).alias("c"), [(col("b")) * (3)].alias("d")]"""
         in explain
     )
+
+
+def test_issue_16436() -> None:
+    df = pl.DataFrame(
+        {
+            "x": [1.12, 2.21, 4.2, 3.21],
+            "y": [2.11, 3.32, 2.1, 6.12],
+        }
+    )
+
+    df = (
+        df.lazy()
+        .with_columns((pl.col("y") / pl.col("x")).alias("z"))
+        .with_columns(
+            pl.when(pl.col("z").is_infinite()).then(0).otherwise(pl.col("z")).alias("z")
+        )
+        .fill_nan(0)
+        .collect()
+    )