diff --git a/crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs b/crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs index 666de0d3ee3d..01ea7004101c 100644 --- a/crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs +++ b/crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs @@ -45,7 +45,8 @@ mod inner { } } - pub(super) fn nodes_scratch_mut(&mut self) -> &mut UnitVec { + /// Returns shared scratch space after clearing. + pub(super) fn empty_nodes_scratch_mut(&mut self) -> &mut UnitVec { self.nodes_scratch.clear(); &mut self.nodes_scratch } @@ -115,7 +116,7 @@ impl PredicatePushDown<'_> { &[], &acc_predicates, expr_arena, - self.nodes_scratch_mut(), + self.empty_nodes_scratch_mut(), )?; let local_predicates = match eligibility { @@ -296,7 +297,7 @@ impl PredicatePushDown<'_> { &[predicate.clone()], &acc_predicates, expr_arena, - self.nodes_scratch_mut(), + self.empty_nodes_scratch_mut(), )? .0 { @@ -663,7 +664,7 @@ impl PredicatePushDown<'_> { for v in acc_predicates.values() { let ae = expr_arena.get(v.node()); assert!(permits_filter_pushdown( - self.nodes_scratch_mut(), + self.empty_nodes_scratch_mut(), ae, expr_arena )); @@ -677,7 +678,7 @@ impl PredicatePushDown<'_> { for v in acc_predicates.values() { let ae = expr_arena.get(v.node()); assert!(permits_filter_pushdown( - self.nodes_scratch_mut(), + self.empty_nodes_scratch_mut(), ae, expr_arena )); 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 8db68f5a46e8..02859e14a6a6 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,6 @@ pub fn pushdown_eligibility( expr_arena: &mut Arena, scratch: &mut UnitVec, ) -> PolarsResult<(PushdownEligibility, PlHashMap)> { - debug_assert!(scratch.is_empty()); scratch.clear(); let ae_nodes_stack = scratch; diff --git a/crates/polars-plan/src/plans/optimizer/slice_pushdown_expr.rs b/crates/polars-plan/src/plans/optimizer/slice_pushdown_expr.rs index cb879388c27c..8d9bff3ea868 100644 --- a/crates/polars-plan/src/plans/optimizer/slice_pushdown_expr.rs +++ b/crates/polars-plan/src/plans/optimizer/slice_pushdown_expr.rs @@ -29,7 +29,7 @@ impl OptimizationRule for SlicePushDown { let out = match expr_arena.get(*input) { ae @ Alias(..) | ae @ Cast { .. } => { let ae = ae.clone(); - let scratch = self.nodes_scratch_mut(); + let scratch = self.empty_nodes_scratch_mut(); ae.nodes(scratch); let input = scratch[0]; let new_input = pushdown(input, offset, length, expr_arena); 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 70e7ab84c0a7..5bd92ffbf97c 100644 --- a/crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs +++ b/crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs @@ -24,7 +24,8 @@ mod inner { } } - pub fn nodes_scratch_mut(&mut self) -> &mut UnitVec { + /// Returns shared scratch space after clearing. + pub fn empty_nodes_scratch_mut(&mut self) -> &mut UnitVec { self.scratch.clear(); &mut self.scratch } @@ -50,7 +51,6 @@ fn can_pushdown_slice_past_projections( arena: &Arena, scratch: &mut UnitVec, ) -> (bool, bool) { - debug_assert!(scratch.is_empty()); scratch.clear(); let mut can_pushdown_and_any_expr_has_column = false; @@ -496,15 +496,16 @@ impl SlicePushDown { // [Pushdown] // these nodes will be pushed down. // State is None, we can continue - m @(Select {..}, None) | - m @ (SimpleProjection {..}, _) + m @ (Select {..}, None) + | m @ (HStack {..}, None) + | m @ (SimpleProjection {..}, _) => { let (lp, state) = m; self.pushdown_and_continue(lp, state, lp_arena, expr_arena) } // there is state, inspect the projection to determine how to deal with it (Select {input, expr, schema, options}, Some(_)) => { - if can_pushdown_slice_past_projections(&expr, expr_arena, self.nodes_scratch_mut()).1 { + if can_pushdown_slice_past_projections(&expr, expr_arena, self.empty_nodes_scratch_mut()).1 { let lp = Select {input, expr, schema, options}; self.pushdown_and_continue(lp, state, lp_arena, expr_arena) } @@ -515,7 +516,7 @@ impl SlicePushDown { } } (HStack {input, exprs, schema, options}, _) => { - let (can_pushdown, can_pushdown_and_any_expr_has_column) = can_pushdown_slice_past_projections(&exprs, expr_arena, self.nodes_scratch_mut()); + let (can_pushdown, can_pushdown_and_any_expr_has_column) = can_pushdown_slice_past_projections(&exprs, expr_arena, self.empty_nodes_scratch_mut()); if can_pushdown_and_any_expr_has_column || ( // If the schema length is greater then an input column is being projected, so diff --git a/py-polars/tests/unit/operations/test_slice.py b/py-polars/tests/unit/operations/test_slice.py index 79ae4317976c..a3d4d95c0b43 100644 --- a/py-polars/tests/unit/operations/test_slice.py +++ b/py-polars/tests/unit/operations/test_slice.py @@ -304,10 +304,9 @@ def test_slice_after_sort_with_nulls_20079() -> None: 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) - ) + q = pl.LazyFrame({"A": "1/1"}) + q = q.with_columns(col.str.split("/")) + q = q.with_columns(pl.when(col.is_not_null()).then(col.list.get(0)).otherwise(None)) - assert_frame_equal(df.collect(), pl.DataFrame({"A": ["1"]})) + assert_frame_equal(q.slice(0, 1).collect(), pl.DataFrame({"A": ["1"]})) + assert_frame_equal(q.collect(), pl.DataFrame({"A": ["1"]}))