Skip to content

Commit

Permalink
fix(rust, python): take input_schema to create physical expr for sele…
Browse files Browse the repository at this point in the history
…ction (pola-rs#10571)
  • Loading branch information
reswqa authored Aug 17, 2023
1 parent 9b3007e commit 629c25b
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 2 deletions.
10 changes: 8 additions & 2 deletions crates/polars-lazy/src/physical_plan/planner/lp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -165,10 +165,16 @@ pub fn create_physical_plan(
Ok(Box::new(executors::SliceExec { input, offset, len }))
},
Selection { input, predicate } => {
let input_schema = lp_arena.get(input).schema(lp_arena).into_owned();
let input = create_physical_plan(input, lp_arena, expr_arena)?;
let mut state = ExpressionConversionState::default();
let predicate =
create_physical_expr(predicate, Context::Default, expr_arena, None, &mut state)?;
let predicate = create_physical_expr(
predicate,
Context::Default,
expr_arena,
Some(&input_schema),
&mut state,
)?;
Ok(Box::new(executors::FilterExec::new(
predicate,
input,
Expand Down
16 changes: 16 additions & 0 deletions py-polars/tests/unit/operations/test_filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,3 +145,19 @@ def test_clear_window_cache_after_filter_10499() -> None:
assert df.lazy().filter((pl.col("a").null_count() < pl.count()).over("b")).filter(
((pl.col("a") == 0).sum() < pl.count()).over("b")
).collect().to_dict(False) == {"a": [3, None, 5, 0, 9, 10], "b": [2, 2, 3, 3, 5, 5]}


def test_agg_function_of_filter_10565() -> None:
df_int = pl.DataFrame(data={"a": []}, schema={"a": pl.Int16})
assert df_int.filter(pl.col("a").n_unique().over("a") == 1).to_dict(False) == {
"a": []
}

df_str = pl.DataFrame(data={"a": []}, schema={"a": pl.Utf8})
assert df_str.filter(pl.col("a").n_unique().over("a") == 1).to_dict(False) == {
"a": []
}

assert df_str.lazy().filter(pl.col("a").n_unique().over("a") == 1).collect(
predicate_pushdown=False
).to_dict(False) == {"a": []}

0 comments on commit 629c25b

Please sign in to comment.