From 789effe1511163e87dbb483584922b0b804437a5 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Fri, 20 Jan 2023 11:20:22 +0100
Subject: [PATCH] fix(rust, python): expand all nested wildcards in functions

---
 .../src/logical_plan/projection.rs            |  3 +-
 py-polars/tests/unit/test_arity.py            | 33 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 py-polars/tests/unit/test_arity.py

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"],
+    }