Skip to content

fix: Fix panic on lazy concat->filter->slice with CSPE #26907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ritchie46 merged 2 commits into from
Mar 13, 2026
Merged

fix: Fix panic on lazy concat->filter->slice with CSPE #26907

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eecfe8e
c
nameexhaustion Mar 13, 2026
82ac9bd
c
nameexhaustion Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions crates/polars-plan/src/plans/optimizer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,8 +201,7 @@ pub fn optimize(

if opt_flags.slice_pushdown() {
let mut slice_pushdown_opt = SlicePushDown::new();
let ir = ir_arena.take(root);
let ir = slice_pushdown_opt.optimize(ir, ir_arena, expr_arena)?;
let ir = slice_pushdown_opt.optimize(root, ir_arena, expr_arena)?;

ir_arena.replace(root, ir);

Expand Down Expand Up @@ -246,8 +245,7 @@ pub fn optimize(

if repeat_slice_pd_after_filter_pd {
let mut slice_pushdown_opt = SlicePushDown::new();
let ir = ir_arena.take(root);
let ir = slice_pushdown_opt.optimize(ir, ir_arena, expr_arena)?;
let ir = slice_pushdown_opt.optimize(root, ir_arena, expr_arena)?;

ir_arena.replace(root, ir);
}
Expand Down
55 changes: 28 additions & 27 deletions crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,10 +190,9 @@ impl SlicePushDown {
let new_inputs = inputs
.into_iter()
.map(|node| {
let alp = lp_arena.take(node);
// No state, so we do not push down the slice here.
let state = None;
let alp = self.pushdown(alp, state, lp_arena, expr_arena)?;
let alp = self.pushdown(node, state, lp_arena, expr_arena)?;
lp_arena.replace(node, alp);
Ok(node)
})
Expand All @@ -216,26 +215,37 @@ impl SlicePushDown {
let new_inputs = inputs
.into_iter()
.map(|node| {
let alp = lp_arena.take(node);
let alp = self.pushdown(alp, state, lp_arena, expr_arena)?;
let alp = self.pushdown(node, state, lp_arena, expr_arena)?;
lp_arena.replace(node, alp);
Ok(node)
})
.collect::<PolarsResult<UnitVec<_>>>()?;
Ok(lp.with_inputs(new_inputs))
}

/// This will take the `ir_node` from the `lp_arena`, replacing it with `IR::Invalid` (except if
/// `ir_node` is a `IR::Cache`).
#[recursive]
fn pushdown(
&mut self,
lp: IR,
ir_node: Node,
state: Option<State>,
lp_arena: &mut Arena<IR>,
expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<IR> {
use IR::*;

match (lp, state) {
// Don't take this, the node can be referenced multiple times in the tree.
Copy link
Copy Markdown
Collaborator Author

@nameexhaustion nameexhaustion Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fix here

if let IR::Cache { .. } = lp_arena.get(ir_node) {
return self.no_pushdown_restart_opt(
lp_arena.get(ir_node).clone(),
state,
lp_arena,
expr_arena,
);
}

match (lp_arena.take(ir_node), state) {
#[cfg(feature = "python")]
(
PythonScan { mut options },
Expand Down Expand Up @@ -305,7 +315,8 @@ impl SlicePushDown {
predicate_file_skip_applied,
};

self.pushdown(lp, None, lp_arena, expr_arena)
lp_arena.replace(ir_node, lp);
self.pushdown(ir_node, None, lp_arena, expr_arena)
} else {
let lp = Scan {
sources,
Expand Down Expand Up @@ -385,8 +396,7 @@ impl SlicePushDown {
.map(|len| State { offset: 0, len });

for input in &mut inputs {
let input_lp = lp_arena.take(*input);
let input_lp = self.pushdown(input_lp, subplan_slice, lp_arena, expr_arena)?;
let input_lp = self.pushdown(*input, subplan_slice, lp_arena, expr_arena)?;
lp_arena.replace(*input, input_lp);
}
options.slice = opt_state.map(|x| (x.offset, x.len.try_into().unwrap()));
Expand Down Expand Up @@ -440,12 +450,10 @@ impl SlicePushDown {
}

// first restart optimization in both inputs and get the updated LP
let lp_left = lp_arena.take(input_left);
let lp_left = self.pushdown(lp_left, None, lp_arena, expr_arena)?;
let lp_left = self.pushdown(input_left, None, lp_arena, expr_arena)?;
let input_left = lp_arena.add(lp_left);

let lp_right = lp_arena.take(input_right);
let lp_right = self.pushdown(lp_right, None, lp_arena, expr_arena)?;
let lp_right = self.pushdown(input_right, None, lp_arena, expr_arena)?;
let input_right = lp_arena.add(lp_right);

// then assign the slice state to the join operation
Expand Down Expand Up @@ -476,8 +484,7 @@ impl SlicePushDown {
Some(state),
) => {
// first restart optimization in inputs and get the updated LP
let input_lp = lp_arena.take(input);
let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
let input_lp = self.pushdown(input, None, lp_arena, expr_arena)?;
let input = lp_arena.add(input_lp);

if let Some(existing_slice) = &mut Arc::make_mut(&mut options).slice {
Expand Down Expand Up @@ -528,8 +535,7 @@ impl SlicePushDown {
},
(Distinct { input, mut options }, Some(state)) => {
// first restart optimization in inputs and get the updated LP
let input_lp = lp_arena.take(input);
let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
let input_lp = self.pushdown(input, None, lp_arena, expr_arena)?;
let input = lp_arena.add(input_lp);

if let Some(existing_slice) = &mut options.slice {
Expand Down Expand Up @@ -594,8 +600,7 @@ impl SlicePushDown {
assert!(slice.is_none() || slice == new_slice);

// first restart optimization in inputs and get the updated LP
let input_lp = lp_arena.take(input);
let input_lp = self.pushdown(input_lp, None, lp_arena, expr_arena)?;
let input_lp = self.pushdown(input, None, lp_arena, expr_arena)?;
let input = lp_arena.add(input_lp);

Ok(Sort {
Expand All @@ -613,8 +618,6 @@ impl SlicePushDown {
},
Some(outer_slice),
) => {
let alp = lp_arena.take(input);

// If offset is negative the length can never be greater than it.
if offset < 0 {
#[allow(clippy::unnecessary_cast)] // Necessary when IdxSize = u64.
Expand All @@ -626,10 +629,10 @@ impl SlicePushDown {
if let Some(combined) =
combine_outer_inner_slice(outer_slice, State { offset, len })
{
self.pushdown(alp, Some(combined), lp_arena, expr_arena)
self.pushdown(input, Some(combined), lp_arena, expr_arena)
} else {
let lp =
self.pushdown(alp, Some(State { offset, len }), lp_arena, expr_arena)?;
self.pushdown(input, Some(State { offset, len }), lp_arena, expr_arena)?;
let input = lp_arena.add(lp);
self.slice_node_in_optimized_plan = true;
Ok(Slice {
Expand All @@ -647,8 +650,6 @@ impl SlicePushDown {
},
None,
) => {
let alp = lp_arena.take(input);

// If offset is negative the length can never be greater than it.
if offset < 0 {
#[allow(clippy::unnecessary_cast)] // Necessary when IdxSize = u64.
Expand All @@ -658,7 +659,7 @@ impl SlicePushDown {
}

let state = Some(State { offset, len });
self.pushdown(alp, state, lp_arena, expr_arena)
self.pushdown(input, state, lp_arena, expr_arena)
},
m @ (Filter { .. }, _)
| m @ (DataFrameScan { .. }, _)
Expand Down Expand Up @@ -809,7 +810,7 @@ impl SlicePushDown {

pub fn optimize(
&mut self,
logical_plan: IR,
logical_plan: Node,
lp_arena: &mut Arena<IR>,
expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<IR> {
Expand Down
10 changes: 10 additions & 0 deletions py-polars/tests/unit/lazyframe/test_optimizations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -618,3 +618,13 @@ def test_scan_select_all_columns_no_projection_pyarrow() -> None:
ds = pad.dataset(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
plan = pl.scan_pyarrow_dataset(ds).select(pl.col("a"), pl.col("b")).explain()
assert "PROJECT */2 COLUMNS" in plan


def test_slice_pushdown_with_cache_arena_take_panic_26905() -> None:
lf = pl.LazyFrame({"x": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
q = pl.concat([lf, lf]).select(pl.all()).filter(pl.col("x") > 3).head(2)

assert_frame_equal(
q.collect(),
pl.DataFrame({"x": [4, 5]}),
)
Loading