Skip to content

Commit

Permalink
fix(rust, python): expand all nested wildcards in functions
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Jan 20, 2023
1 parent a7fffe3 commit 789effe
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,8 @@ fn expand_function_inputs(mut expr: Expr, schema: &Schema) -> Expr {
if options.input_wildcard_expansion =>
{
*input = rewrite_projections(input.clone(), schema, &[]).unwrap();
false
// continue iteration, there might be more functions.
true
}
_ => true,
});
Expand Down
33 changes: 33 additions & 0 deletions py-polars/tests/unit/test_arity.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import polars as pl


def test_nested_when_then_and_wildcard_expansion_6284() -> None:
df = pl.DataFrame(
{
"1": ["a", "b"],
"2": ["c", "d"],
}
)

out0 = df.with_column(
pl.when(pl.any(pl.all() == "a"))
.then("a")
.otherwise(pl.when(pl.any(pl.all() == "d")).then("d").otherwise(None))
.alias("result")
)

out1 = df.with_column(
pl.when(pl.any(pl.all() == "a"))
.then("a")
.when(pl.any(pl.all() == "d"))
.then("d")
.otherwise(None)
.alias("result")
)

assert out0.frame_equal(out1)
assert out0.to_dict(False) == {
"1": ["a", "b"],
"2": ["c", "d"],
"result": ["a", "d"],
}

0 comments on commit 789effe

Please sign in to comment.