perf: Native streaming int_range with len or count (#24280)

Co-authored-by: Orson Peters <orsonpeters@gmail.com>

Author: coastalwhite

crates/polars-plan/src/plans/aexpr/builder.rs

@@ -381,6 +381,22 @@ impl AExprBuilder {
         nc.gt(idx_zero, arena)
     }
 
+    pub fn drop_nulls(self, arena: &mut Arena<AExpr>) -> Self {
+        Self::function(
+            vec![self.expr_ir_retain_name(arena)],
+            IRFunctionExpr::DropNulls,
+            arena,
+        )
+    }
+
+    pub fn drop_nans(self, arena: &mut Arena<AExpr>) -> Self {
+        Self::function(
+            vec![self.expr_ir_retain_name(arena)],
+            IRFunctionExpr::DropNans,
+            arena,
+        )
+    }
+
     pub fn eq(self, other: impl IntoAExprBuilder, arena: &mut Arena<AExpr>) -> Self {
         self.binary_op(other, Operator::Eq, arena)
     }
@@ -465,6 +481,10 @@ impl AExprBuilder {
         ExprIR::new(self.node(), OutputName::Alias(name.into()))
     }
 
+    pub fn expr_ir_retain_name(self, arena: &Arena<AExpr>) -> ExprIR {
+        ExprIR::from_node(self.node(), arena)
+    }
+
     pub fn expr_ir_unnamed(self) -> ExprIR {
         self.expr_ir(PlSmallStr::EMPTY)
     }

crates/polars-stream/src/physical_plan/lower_expr.rs

@@ -1196,6 +1196,92 @@ fn lower_exprs_with_ctx(
                 transformed_exprs.push(ctx.expr_arena.add(AExpr::Column(value_key.clone())));
             },
 
+            // pl.row_index() maps to this.
+            #[cfg(feature = "range")]
+            AExpr::Function {
+                input: ref inner_exprs,
+                function: IRFunctionExpr::Range(IRRangeFunction::IntRange { step: 1, dtype }),
+                options: _,
+            } if {
+                let start_is_zero = match ctx.expr_arena.get(inner_exprs[0].node()) {
+                    AExpr::Literal(lit) => lit.extract_usize().ok() == Some(0),
+                    _ => false,
+                };
+                let stop_is_len = matches!(ctx.expr_arena.get(inner_exprs[1].node()), AExpr::Len);
+
+                dtype == DataType::IDX_DTYPE && start_is_zero && stop_is_len
+            } =>
+            {
+                let out_name = unique_column_name();
+                let row_idx_col_aexpr = ctx.expr_arena.add(AExpr::Column(out_name.clone()));
+                let row_idx_col_expr_ir =
+                    ExprIR::new(row_idx_col_aexpr, OutputName::ColumnLhs(out_name.clone()));
+                let row_idx_stream = build_select_stream_with_ctx(
+                    build_row_idx_stream(input, out_name, None, ctx.phys_sm),
+                    &[row_idx_col_expr_ir],
+                    ctx,
+                )?;
+                input_streams.insert(row_idx_stream);
+                transformed_exprs.push(row_idx_col_aexpr);
+            },
+
+            #[cfg(feature = "range")]
+            AExpr::Function {
+                input: ref inner_exprs,
+                function: IRFunctionExpr::Range(IRRangeFunction::IntRange { step: 1, dtype }),
+                options: _,
+            } if {
+                let start_is_zero = match ctx.expr_arena.get(inner_exprs[0].node()) {
+                    AExpr::Literal(lit) => lit.extract_usize().ok() == Some(0),
+                    _ => false,
+                };
+                let stop_is_count = matches!(
+                    ctx.expr_arena.get(inner_exprs[1].node()),
+                    AExpr::Agg(IRAggExpr::Count { .. })
+                );
+
+                start_is_zero && stop_is_count
+            } =>
+            {
+                let AExpr::Agg(IRAggExpr::Count {
+                    input: input_expr,
+                    include_nulls,
+                }) = ctx.expr_arena.get(inner_exprs[1].node())
+                else {
+                    unreachable!();
+                };
+                let (input_expr, include_nulls) = (*input_expr, *include_nulls);
+
+                let out_name = unique_column_name();
+                let mut row_idx_col_aexpr = ctx.expr_arena.add(AExpr::Column(out_name.clone()));
+                if dtype != IDX_DTYPE {
+                    row_idx_col_aexpr = AExprBuilder::new_from_node(row_idx_col_aexpr)
+                        .cast(dtype, ctx.expr_arena)
+                        .node();
+                }
+                let row_idx_col_expr_ir =
+                    ExprIR::new(row_idx_col_aexpr, OutputName::ColumnLhs(out_name.clone()));
+
+                let mut input_expr = AExprBuilder::new_from_node(input_expr);
+                if !include_nulls {
+                    input_expr = input_expr.drop_nulls(ctx.expr_arena);
+                }
+                let input_expr = input_expr.expr_ir_retain_name(ctx.expr_arena);
+
+                let row_idx_stream = build_select_stream_with_ctx(
+                    build_row_idx_stream(
+                        build_select_stream_with_ctx(input, &[input_expr], ctx)?,
+                        out_name,
+                        None,
+                        ctx.phys_sm,
+                    ),
+                    &[row_idx_col_expr_ir],
+                    ctx,
+                )?;
+                input_streams.insert(row_idx_stream);
+                transformed_exprs.push(row_idx_col_aexpr);
+            },
+
             // Lower arbitrary elementwise functions.
             ref node @ AExpr::Function {
                 input: ref inner_exprs,
@@ -1605,35 +1691,6 @@ fn lower_exprs_with_ctx(
                 transformed_exprs.push(AExprBuilder::col(out_name.clone(), ctx.expr_arena).node());
             },
 
-            // pl.row_index() maps to this.
-            #[cfg(feature = "range")]
-            AExpr::Function {
-                input: ref inner_exprs,
-                function: IRFunctionExpr::Range(IRRangeFunction::IntRange { step: 1, dtype }),
-                options: _,
-            } if {
-                let start_is_zero = match ctx.expr_arena.get(inner_exprs[0].node()) {
-                    AExpr::Literal(lit) => lit.extract_usize().ok() == Some(0),
-                    _ => false,
-                };
-                let stop_is_len = matches!(ctx.expr_arena.get(inner_exprs[1].node()), AExpr::Len);
-
-                dtype == DataType::IDX_DTYPE && start_is_zero && stop_is_len
-            } =>
-            {
-                let out_name = unique_column_name();
-                let row_idx_col_aexpr = ctx.expr_arena.add(AExpr::Column(out_name.clone()));
-                let row_idx_col_expr_ir =
-                    ExprIR::new(row_idx_col_aexpr, OutputName::ColumnLhs(out_name.clone()));
-                let row_idx_stream = build_select_stream_with_ctx(
-                    build_row_idx_stream(input, out_name, None, ctx.phys_sm),
-                    &[row_idx_col_expr_ir],
-                    ctx,
-                )?;
-                input_streams.insert(row_idx_stream);
-                transformed_exprs.push(row_idx_col_aexpr);
-            },
-
             AExpr::Slice {
                 input: inner,
                 offset,

py-polars/tests/unit/functions/range/test_int_range.py

@@ -286,3 +286,42 @@ def test_int_ranges_non_numeric_input_should_error() -> None:
         _ = df.select(pl.int_ranges("start", "end"))
 
     assert "conversion from `str` to `i64` failed" in str(excinfo.value)
+
+
+def test_int_range_len_count() -> None:
+    values = [1, 2, None, 4, 5, 6]
+
+    lf = pl.Series("a", values).to_frame().lazy()
+
+    def irange(e: pl.Expr) -> pl.LazyFrame:
+        return lf.select(r=pl.int_range(0, e, dtype=pl.get_index_type()))
+
+    q = irange(pl.len())
+    assert_series_equal(
+        q.collect().to_series(),
+        pl.Series("r", [0, 1, 2, 3, 4, 5], pl.get_index_type()),
+    )
+
+    q = irange(pl.col.a.len())
+    assert_series_equal(
+        q.collect().to_series(),
+        pl.Series("r", [0, 1, 2, 3, 4, 5], pl.get_index_type()),
+    )
+
+    q = irange(pl.col.a.filter(pl.col.a.ne_missing(4)).len())
+    assert_series_equal(
+        q.collect().to_series(),
+        pl.Series("r", [0, 1, 2, 3, 4], pl.get_index_type()),
+    )
+
+    q = irange(pl.col.a.count())
+    assert_series_equal(
+        q.collect().to_series(),
+        pl.Series("r", [0, 1, 2, 3, 4], pl.get_index_type()),
+    )
+
+    q = irange(pl.col.a.filter(pl.col.a.ne_missing(4)).count())
+    assert_series_equal(
+        q.collect().to_series(),
+        pl.Series("r", [0, 1, 2, 3], pl.get_index_type()),
+    )