fix: Don't cse as_struct (#19280)

pola-rs · Oct 17, 2024 · 74d726b · 74d726b
1 parent 93f8902
commit 74d726b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/crates/polars-plan/src/plans/optimizer/cse/cse_expr.rs b/crates/polars-plan/src/plans/optimizer/cse/cse_expr.rs
@@ -181,6 +181,11 @@ enum VisitRecord {
 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,
     }

diff --git a/py-polars/tests/unit/test_cse.py b/py-polars/tests/unit/test_cse.py
@@ -788,3 +788,19 @@ def test_eager_cse_during_struct_expansion_18411() -> None:
         df.select(pl.col("foo").replace(classes, counts))
         == df.select(pl.col("foo").replace(classes, counts))
     )["foo"].all()
+
+
+def test_cse_skip_as_struct_19253() -> None:
+    df = pl.LazyFrame({"x": [1, 2], "y": [4, 5]})
+
+    assert (
+        df.with_columns(
+            q1=pl.struct(pl.col.x - pl.col.y.mean()),
+            q2=pl.struct(pl.col.x - pl.col.y.mean().over("y")),
+        ).collect()
+    ).to_dict(as_series=False) == {
+        "x": [1, 2],
+        "y": [4, 5],
+        "q1": [{"x": -3.5}, {"x": -2.5}],
+        "q2": [{"x": -3.0}, {"x": -3.0}],
+    }