From 01fadd9172044fb466f0043d1fbe48bb1a4a3432 Mon Sep 17 00:00:00 2001 From: nameexhaustion Date: Mon, 9 Dec 2024 21:00:53 +1100 Subject: [PATCH] fix: Fix incorrect lazy `select(len())` with some select orderings (#20222) --- .../src/plans/optimizer/projection_pushdown/projection.rs | 4 ++-- py-polars/tests/unit/test_scalar.py | 7 +++++++ 2 files changed, 9 insertions(+), 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 index 6327238c4c21..79b461cb90aa 100644 --- 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 index 868cb72e6991..094109b0c4d3 100644 --- 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