fix: Ensure Function name correctness in cse (#20929)

crates/polars-plan/src/plans/optimizer/cse/cse_expr.rs

@@ -182,10 +182,6 @@ fn skip_pre_visit(ae: &AExpr, is_groupby: bool) -> bool {
     match ae {
         AExpr::Window { .. } => true,
         #[cfg(feature = "dtype-struct")]
-        AExpr::Function {
-            function: FunctionExpr::AsStruct,
-            ..
-        } => true,
         AExpr::Ternary { .. } => is_groupby,
         _ => false,
     }
@@ -724,6 +720,45 @@ impl CommonSubExprOptimizer {
                 } else {
                     out_e.set_node(out_node);
 
+                    // Ensure the function ExprIR's have the proper names.
+                    let mut scratch = vec![];
+                    let mut stack = vec![(e.node(), out_node)];
+                    while let Some((original, new)) = stack.pop() {
+                        let aes = expr_arena.get_many_mut([original, new]);
+
+                        aes[0].inputs_rev(&mut scratch);
+                        aes[1].inputs_rev(&mut scratch);
+
+                        for i in 0..scratch.len() / 2 {
+                            stack.push((scratch[i], scratch[i + 1]));
+                        }
+                        scratch.clear();
+
+                        match expr_arena.get_many_mut([original, new]) {
+                            [AExpr::Function {
+                                input: input_original,
+                                ..
+                            }, AExpr::Function {
+                                input: input_new, ..
+                            }] => {
+                                for (new, original) in input_new.iter_mut().zip(input_original) {
+                                    new.set_alias(original.output_name().clone());
+                                }
+                            },
+                            [AExpr::AnonymousFunction {
+                                input: input_original,
+                                ..
+                            }, AExpr::AnonymousFunction {
+                                input: input_new, ..
+                            }] => {
+                                for (new, original) in input_new.iter_mut().zip(input_original) {
+                                    new.set_alias(original.output_name().clone());
+                                }
+                            },
+                            _ => {},
+                        }
+                    }
+
                     // If we don't end with an alias we add an alias. Because the normal left-hand
                     // rule we apply for determining the name will not work we now refer to
                     // intermediate temporary names starting with the `CSE_REPLACED` constant.

crates/polars-plan/src/utils.rs

@@ -119,6 +119,7 @@ pub fn aexpr_output_name(node: Node, arena: &Arena<AExpr>) -> PolarsResult<PlSma
             AExpr::Alias(_, name) => return Ok(name.clone()),
             AExpr::Len => return Ok(get_len_name()),
             AExpr::Literal(val) => return Ok(val.output_column_name().clone()),
+            AExpr::Ternary { truthy, .. } => return aexpr_output_name(*truthy, arena),
             _ => {},
         }
     }

py-polars/tests/unit/test_cse.py

@@ -791,7 +791,7 @@ def test_eager_cse_during_struct_expansion_18411() -> None:
     )["foo"].all()
 
 
-def test_cse_skip_as_struct_19253() -> None:
+def test_cse_as_struct_19253() -> None:
     df = pl.LazyFrame({"x": [1, 2], "y": [4, 5]})
 
     assert (
@@ -807,6 +807,16 @@ def test_cse_skip_as_struct_19253() -> None:
     }
 
 
+@pytest.mark.may_fail_auto_streaming
+def test_cse_as_struct_value_counts_20927() -> None:
+    assert pl.DataFrame({"x": [i for i in range(1, 6) for _ in range(i)]}).select(
+        pl.struct("x").value_counts().struct.unnest()
+    ).sort("count").to_dict(as_series=False) == {
+        "x": [{"x": 1}, {"x": 2}, {"x": 3}, {"x": 4}, {"x": 5}],
+        "count": [1, 2, 3, 4, 5],
+    }
+
+
 def test_cse_union_19227() -> None:
     lf = pl.LazyFrame({"A": [1], "B": [2]})
     lf_1 = lf.select(C="A", B="B")