fix: Fix incorrect lazy select(len()) with some select orderings (p…

…ola-rs#20222)
photino · Dec 9, 2024 · 01fadd9 · 01fadd9
1 parent 4163fd3
commit 01fadd9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs b/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs
@@ -71,13 +71,13 @@ pub(super) fn process_projection(
         } else {
             // Select the last column projection.
             let mut name = None;
-            for (_, plan) in (&*lp_arena).iter(input) {
+            'outer: for (_, plan) in (&*lp_arena).iter(input) {
                 match plan {
                     IR::Select { expr: exprs, .. } | IR::HStack { exprs, .. } => {
                         for e in exprs {
                             if !e.is_scalar(expr_arena) {
                                 name = Some(e.output_name());
-                                break;
+                                break 'outer;
                             }
                         }
                     },

diff --git a/py-polars/tests/unit/test_scalar.py b/py-polars/tests/unit/test_scalar.py
@@ -63,3 +63,10 @@ def test_scalar_len_20046() -> None:
         .item()
         == 3
     )
+
+    q = pl.LazyFrame({"a": range(3)}).select(
+        pl.first("a"),
+        pl.col("a").alias("b"),
+    )
+
+    assert q.select(pl.len()).collect().item() == 3