c

pola-rs · Dec 9, 2024 · 63e073e · 63e073e
1 parent 58a38af
commit 63e073e
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 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
@@ -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
@@ -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
@@ -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
@@ -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()