From 2d6ceb8da1413efd55c98f7f155f8b630780a016 Mon Sep 17 00:00:00 2001
From: Lawrence Mitchell <lmitchell@nvidia.com>
Date: Wed, 3 Jul 2024 10:35:40 +0000
Subject: [PATCH] fix(rust): avoid panic when projecting solitary count into
 empty frame

When the schema of a frame is empty, a count expression is still valid
since it does not refer to any columns. In this scenario, selecting
the "last" column in the schema would previously panic, since there
are none. Fix this by handling this edge case explicitly.

- Closes #16904
---
 .../projection_pushdown/projection.rs         | 31 ++++++++++++-------
 py-polars/tests/unit/test_projections.py      | 11 +++++++
 2 files changed, 30 insertions(+), 12 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 93aa71fb68bc..f8b338f9d20a 100644
--- a/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs
+++ b/crates/polars-plan/src/plans/optimizer/projection_pushdown/projection.rs
@@ -62,18 +62,25 @@ pub(super) fn process_projection(
     // the whole file while we only want the count
     if exprs.len() == 1 && is_count(exprs[0].node(), expr_arena) {
         let input_schema = lp_arena.get(input).schema(lp_arena);
-        // simply select the last column
-        // NOTE: the first can be the inserted index column, so that might not work
-        let (first_name, _) = input_schema.try_get_at_index(input_schema.len() - 1)?;
-        let expr = expr_arena.add(AExpr::Column(ColumnName::from(first_name.as_str())));
-        if !acc_projections.is_empty() {
-            check_double_projection(
-                &exprs[0],
-                expr_arena,
-                &mut acc_projections,
-                &mut projected_names,
-            );
-        }
+        let expr = if input_schema.is_empty() {
+            // If the input schema is empty, we should just project
+            // ourselves
+            exprs[0].node()
+        } else {
+            // simply select the last column
+            // NOTE: the first can be the inserted index column, so that might not work
+            let (first_name, _) = input_schema.try_get_at_index(input_schema.len() - 1)?;
+            let expr = expr_arena.add(AExpr::Column(ColumnName::from(first_name.as_str())));
+            if !acc_projections.is_empty() {
+                check_double_projection(
+                    &exprs[0],
+                    expr_arena,
+                    &mut acc_projections,
+                    &mut projected_names,
+                );
+            }
+            expr
+        };
         add_expr_to_accumulated(expr, &mut acc_projections, &mut projected_names, expr_arena);
         local_projection.push(exprs.pop().unwrap());
         proj_pd.is_count_star = true;
diff --git a/py-polars/tests/unit/test_projections.py b/py-polars/tests/unit/test_projections.py
index 802061265baa..307ea6a0349b 100644
--- a/py-polars/tests/unit/test_projections.py
+++ b/py-polars/tests/unit/test_projections.py
@@ -511,3 +511,14 @@ def test_projection_pushdown_semi_anti_no_selection(
     assert "PROJECT 1/2" in (
         q_a.join(q_b, left_on="a", right_on="b", how=how).explain()
     )
+
+
+def test_projection_empty_frame_len_16904() -> None:
+    df = pl.LazyFrame({})
+
+    q = df.select(pl.len())
+
+    assert "PROJECT */0" in q.explain()
+
+    expect = pl.DataFrame({"len": [0]}, schema_overrides={"len": pl.UInt32()})
+    assert_frame_equal(q.collect(), expect)