Skip to content
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

refactor: Remove debug asserts on scratch space #20224

Merged
merged 7 commits into from
Dec 9, 2024
Merged

refactor: Remove debug asserts on scratch space #20224

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3db9b5f
c
nameexhaustion Dec 9, 2024
6812e52
c
nameexhaustion Dec 9, 2024
40cf7c4
c
nameexhaustion Dec 9, 2024
a743a24
c
nameexhaustion Dec 9, 2024
989c255
c
nameexhaustion Dec 9, 2024
fe879b4
c
nameexhaustion Dec 9, 2024
2ebabe4
c
nameexhaustion Dec 9, 2024
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
11 changes: 6 additions & 5 deletions crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ mod inner {
}
}

pub(super) fn nodes_scratch_mut(&mut self) -> &mut UnitVec<Node> {
/// Returns shared scratch space after clearing.
pub(super) fn empty_nodes_scratch_mut(&mut self) -> &mut UnitVec<Node> {
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

rename and add doc

self.nodes_scratch.clear();
&mut self.nodes_scratch
}
Expand Down Expand Up @@ -115,7 +116,7 @@ impl PredicatePushDown<'_> {
&[],
&acc_predicates,
expr_arena,
self.nodes_scratch_mut(),
self.empty_nodes_scratch_mut(),
)?;

let local_predicates = match eligibility {
Expand Down Expand Up @@ -296,7 +297,7 @@ impl PredicatePushDown<'_> {
&[predicate.clone()],
&acc_predicates,
expr_arena,
self.nodes_scratch_mut(),
self.empty_nodes_scratch_mut(),
)?
.0
{
Expand Down Expand Up @@ -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
));
Expand All @@ -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
));
Expand Down
1 change: 0 additions & 1 deletion crates/polars-plan/src/plans/optimizer/predicate_pushdown/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,6 @@ pub fn pushdown_eligibility(
expr_arena: &mut Arena<AExpr>,
scratch: &mut UnitVec<Node>,
) -> PolarsResult<(PushdownEligibility, PlHashMap<PlSmallStr, PlSmallStr>)> {
debug_assert!(scratch.is_empty());
Copy link
Collaborator Author

@nameexhaustion nameexhaustion Dec 9, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

removed assert here

scratch.clear();
let ae_nodes_stack = scratch;

Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/plans/optimizer/slice_pushdown_expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
15 changes: 8 additions & 7 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 @@ -24,7 +24,8 @@ mod inner {
}
}

pub fn nodes_scratch_mut(&mut self) -> &mut UnitVec<Node> {
/// Returns shared scratch space after clearing.
pub fn empty_nodes_scratch_mut(&mut self) -> &mut UnitVec<Node> {
self.scratch.clear();
&mut self.scratch
}
Expand All @@ -50,7 +51,6 @@ fn can_pushdown_slice_past_projections(
arena: &Arena<AExpr>,
scratch: &mut UnitVec<Node>,
) -> (bool, bool) {
debug_assert!(scratch.is_empty());
Copy link
Collaborator Author

@nameexhaustion nameexhaustion Dec 9, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

removed assert here

scratch.clear();

let mut can_pushdown_and_any_expr_has_column = false;
Expand Down Expand Up @@ -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)
Copy link
Collaborator Author

@nameexhaustion nameexhaustion Dec 9, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

The original MRE of the panic didn't actually contain a slice() - I realized this was because we were always checking the projections even when we didn't need to for with_columns()

| 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)
}
Expand All @@ -514,8 +515,8 @@ impl SlicePushDown {
self.no_pushdown_restart_opt(lp, state, lp_arena, expr_arena)
}
}
(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());
(HStack {input, exprs, schema, options}, Some(_)) => {
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
Expand Down
11 changes: 5 additions & 6 deletions py-polars/tests/unit/operations/test_slice.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]}))
Loading