fix: Don't prune alias in AnonymousFunction subtree

pola-rs · Apr 3, 2024 · d783730 · d783730
1 parent c8640cb
commit d783730
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/crates/polars-plan/src/logical_plan/conversion.rs b/crates/polars-plan/src/logical_plan/conversion.rs
@@ -210,11 +210,14 @@ fn to_aexpr_impl(expr: Expr, arena: &mut Arena<AExpr>, state: &mut ConversionSta
             function,
             output_type,
             options,
-        } => AExpr::AnonymousFunction {
-            input: to_aexprs(input, arena, state),
-            function,
-            output_type,
-            options,
+        } => {
+            state.prune_alias = false;
+            AExpr::AnonymousFunction {
+                input: to_aexprs(input, arena, state),
+                function,
+                output_type,
+                options,
+            }
         },
         Expr::Function {
             input,

diff --git a/py-polars/tests/unit/test_schema.py b/py-polars/tests/unit/test_schema.py
@@ -643,3 +643,22 @@ def test_schema_boolean_sum_horizontal() -> None:
 def test_struct_alias_prune_15401() -> None:
     df = pl.DataFrame({"a": []}, schema={"a": pl.Struct({"b": pl.Int8})})
     assert df.select(pl.col("a").alias("c").struct.field("b")).columns == ["b"]
+
+
+def test_alias_prune_in_fold_15438() -> None:
+    df = pl.DataFrame({"x": [1, 2], "expected_result": ["first", "second"]}).select(
+        actual_result=pl.fold(
+            acc=pl.lit("other", dtype=pl.Utf8),
+            function=lambda acc, x: pl.when(x).then(pl.lit(x.name)).otherwise(acc),  # type: ignore[arg-type, return-value]
+            exprs=[
+                (pl.col("x") == 1).alias("first"),
+                (pl.col("x") == 2).alias("second"),
+            ],
+        )
+    )
+    expected = pl.DataFrame(
+        {
+            "actual_result": ["first", "second"],
+        }
+    )
+    assert_frame_equal(df, expected)