diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs b/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs index cff472cd08cc..d08666c61786 100644 --- a/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs +++ b/polars/polars-lazy/polars-plan/src/logical_plan/projection.rs @@ -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, }); diff --git a/py-polars/tests/unit/test_arity.py b/py-polars/tests/unit/test_arity.py new file mode 100644 index 000000000000..fa97ca057715 --- /dev/null +++ b/py-polars/tests/unit/test_arity.py @@ -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"], + }