From 63e073e702570639746e1e95f26c6fa199263137 Mon Sep 17 00:00:00 2001
From: Simon Lin <simonlin.rqmmw@slmail.me>
Date: Mon, 9 Dec 2024 16:49:14 +1100
Subject: [PATCH] c

---
 .../src/plans/optimizer/predicate_pushdown/utils.rs  |  3 ++-
 .../optimizer/projection_pushdown/projection.rs      |  5 +++--
 .../src/plans/optimizer/slice_pushdown_lp.rs         |  4 +++-
 py-polars/tests/unit/operations/test_slice.py        | 12 ++++++++++++
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/crates/polars-plan/src/plans/optimizer/predicate_pushdown/utils.rs b/crates/polars-plan/src/plans/optimizer/predicate_pushdown/utils.rs
index 0a10c671446a..8db68f5a46e8 100644
--- a/crates/polars-plan/src/plans/optimizer/predicate_pushdown/utils.rs
+++ b/crates/polars-plan/src/plans/optimizer/predicate_pushdown/utils.rs
@@ -146,7 +146,8 @@ pub fn pushdown_eligibility(
     expr_arena: &mut Arena<AExpr>,
     scratch: &mut UnitVec<Node>,
 ) -> PolarsResult<(PushdownEligibility, PlHashMap<PlSmallStr, PlSmallStr>)> {
-    assert!(scratch.is_empty());
+    debug_assert!(scratch.is_empty());
+    scratch.clear();
     let ae_nodes_stack = scratch;
 
     let mut alias_to_col_map =
diff --git a/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs b/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs
index 6327238c4c21..1a595c1e62bb 100644
--- a/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs
+++ b/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs
@@ -71,13 +71,13 @@ pub(super) fn process_projection(
         } else {
             // Select the last column projection.
             let mut name = None;
-            for (_, plan) in (&*lp_arena).iter(input) {
+            'outer: for (_, plan) in (&*lp_arena).iter(input) {
                 match plan {
                     IR::Select { expr: exprs, .. } | IR::HStack { exprs, .. } => {
                         for e in exprs {
                             if !e.is_scalar(expr_arena) {
                                 name = Some(e.output_name());
-                                break;
+                                break 'outer;
                             }
                         }
                     },
@@ -122,6 +122,7 @@ pub(super) fn process_projection(
                 None
             }
         };
+
         if let Some(expr) = expr {
             add_expr_to_accumulated(expr, &mut acc_projections, &mut projected_names, expr_arena);
             local_projection.push(exprs.pop().unwrap());
diff --git a/crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs b/crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs
index 76ed0b8e3029..70e7ab84c0a7 100644
--- a/crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs
+++ b/crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs
@@ -25,6 +25,7 @@ mod inner {
         }
 
         pub fn nodes_scratch_mut(&mut self) -> &mut UnitVec<Node> {
+            self.scratch.clear();
             &mut self.scratch
         }
     }
@@ -49,7 +50,8 @@ fn can_pushdown_slice_past_projections(
     arena: &Arena<AExpr>,
     scratch: &mut UnitVec<Node>,
 ) -> (bool, bool) {
-    assert!(scratch.is_empty());
+    debug_assert!(scratch.is_empty());
+    scratch.clear();
 
     let mut can_pushdown_and_any_expr_has_column = false;
 
diff --git a/py-polars/tests/unit/operations/test_slice.py b/py-polars/tests/unit/operations/test_slice.py
index b148b2443421..9d1fb82eaa69 100644
--- a/py-polars/tests/unit/operations/test_slice.py
+++ b/py-polars/tests/unit/operations/test_slice.py
@@ -299,3 +299,15 @@ def test_slice_after_sort_with_nulls_20079() -> None:
     out = df.sort("a", nulls_last=False).slice(0, 10).collect()
     expected = pl.DataFrame({"a": [None, None, 1.2]})
     assert_frame_equal(out, expected)
+
+
+def test_slice_pushdown_panic_20216() -> None:
+    col = pl.col("A")
+
+    df = pl.LazyFrame([{"A": "1/1"}])
+    df = df.with_columns(col.str.split("/"))
+    df = df.with_columns(
+        pl.when(col.is_not_null()).then(col.list.get(0)).otherwise(None)
+    )
+
+    df.collect()