Implement with_row_count for DataFrame and LazyFrame (#329)

R/dataframe__frame.R

@@ -369,7 +369,16 @@ DataFrame.property_setters = new.env(parent = emptyenv())
   self
 }
 
-
+#' @title Eager with_row_count
+#' @description Add a new column at index 0 that counts the rows
+#' @keywords DataFrame
+#' @param name string name of the created column
+#' @param offset positive integer offset for the start of the counter
+#' @return A new `DataFrame` object with a counter column in front
+#' @docType NULL
+DataFrame_with_row_count = function(name, offset = NULL) {
+  .pr$DataFrame$with_row_count(self, name, offset) |> unwrap()
+}
 
 #' Get and set column names of a DataFrame
 #' @name DataFrame_columns

R/extendr-wrappers.R

@@ -1,7 +1,4 @@
 # Generated by extendr: Do not edit by hand
-
-# nolint start
-
 #
 # This file was created with the following call:
 #   .Call("wrap__make_polars_wrappers", use_symbols = TRUE, package_name = "polars")
@@ -115,6 +112,8 @@ DataFrame$set_column_from_series <- function(x) .Call(wrap__DataFrame__set_colum
 
 DataFrame$new_par_from_list <- function(robj_list) .Call(wrap__DataFrame__new_par_from_list, robj_list)
 
+DataFrame$with_row_count <- function(name, offset) .Call(wrap__DataFrame__with_row_count, self, name, offset)
+
 DataFrame$print <- function() .Call(wrap__DataFrame__print, self)
 
 DataFrame$columns <- function() .Call(wrap__DataFrame__columns, self)
@@ -933,6 +932,8 @@ LazyFrame$with_columns <- function(exprs) .Call(wrap__LazyFrame__with_columns, s
 
 LazyFrame$with_column <- function(expr) .Call(wrap__LazyFrame__with_column, self, expr)
 
+LazyFrame$with_row_count <- function(name, offset) .Call(wrap__LazyFrame__with_row_count, self, name, offset)
+
 LazyFrame$join_asof <- function(other, left_on, right_on, left_by, right_by, allow_parallel, force_parallel, suffix, strategy, tolerance, tolerance_str) .Call(wrap__LazyFrame__join_asof, self, other, left_on, right_on, left_by, right_by, allow_parallel, force_parallel, suffix, strategy, tolerance, tolerance_str)
 
 LazyFrame$join <- function(other, left_on, right_on, how, suffix, allow_parallel, force_parallel) .Call(wrap__LazyFrame__join, self, other, left_on, right_on, how, suffix, allow_parallel, force_parallel)
@@ -1099,5 +1100,3 @@ FeatureInfo$to_r <- function() .Call(wrap__FeatureInfo__to_r, self)
 #' @export
 `[[.FeatureInfo` <- `$.FeatureInfo`
 
-
-# nolint end

R/lazyframe__lazy.R

@@ -239,6 +239,17 @@ LazyFrame_with_columns = function(...) {
 #' @docType NULL
 LazyFrame_with_column = "use_extendr_wrapper"
 
+#' @title Lazy with_row_count
+#' @description Add a new column at index 0 that counts the rows
+#' @keywords LazyFrame
+#' @param name string name of the created column
+#' @param offset positive integer offset for the start of the counter
+#' @return A new `LazyFrame` object with a counter column in front
+#' @docType NULL
+LazyFrame_with_row_count = function(name, offset = NULL) {
+  .pr$LazyFrame$with_row_count(self, name, offset) |> unwrap()
+}
+
 #' @title Apply filter to LazyFrame
 #' @description Filter rows with an Expression defining a boolean column
 #' @keywords LazyFrame

src/rust/src/lazy/dataframe.rs

@@ -233,6 +233,17 @@ impl LazyFrame {
         LazyFrame(self.0.clone().with_column(expr.0.clone()))
     }
 
+    fn with_row_count(&self, name: Robj, offset: Robj) -> RResult<Self> {
+        Ok(self
+            .0
+            .clone()
+            .with_row_count(
+                robj_to!(String, name)?.as_str(),
+                robj_to!(Option, u32, offset)?,
+            )
+            .into())
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn join_asof(
         &self,

src/rust/src/rdataframe/mod.rs

@@ -10,7 +10,7 @@ use crate::rdatatype;
 use crate::rdatatype::RPolarsDataType;
 use crate::rlib;
 use crate::robj_to;
-use crate::rpolarserr::RResult;
+use crate::rpolarserr::{polars_to_rpolars_err, RResult};
 use crate::utils::extendr_concurrent::ParRObj;
 pub use lazy::dataframe::*;
 
@@ -134,6 +134,18 @@ impl DataFrame {
             .map(DataFrame)
     }
 
+    pub fn with_row_count(&self, name: Robj, offset: Robj) -> RResult<Self> {
+        Ok(self
+            .0
+            .clone()
+            .with_row_count(
+                robj_to!(String, name)?.as_str(),
+                robj_to!(Option, u32, offset)?,
+            )
+            .map_err(polars_to_rpolars_err)?
+            .into())
+    }
+
     pub fn print(&self) -> Self {
         rprintln!("{:#?}", self.0);
         self.clone()

tests/testthat/test-dataframe.R

@@ -975,3 +975,8 @@ test_that("glimpse", {
   )
   expect_true(is_string(pl$DataFrame(iris)$glimpse(return_as_string = TRUE)))
 })
+
+test_that("with_row_count", {
+  df = pl$DataFrame(mtcars)
+  expect_identical(df$with_row_count("idx", 42)$select(pl$col("idx"))$to_data_frame()$idx, as.double(42:(41+nrow(mtcars))))
+})

tests/testthat/test-lazy.R

@@ -596,3 +596,8 @@ test_that("width", {
   expect_equal(dat$width, 11)
   expect_equal(ncol(dat), 11)
 })
+
+test_that("with_row_count", {
+  lf = pl$LazyFrame(mtcars)
+  expect_identical(lf$with_row_count("idx", 42)$select(pl$col("idx"))$collect()$to_data_frame()$idx, as.double(42:(41+nrow(mtcars))))
+})