diff --git a/R/dataframe__frame.R b/R/dataframe__frame.R index 9b4e67dbd..ba2097900 100644 --- a/R/dataframe__frame.R +++ b/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 diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R index cde638634..379cba413 100644 --- a/R/extendr-wrappers.R +++ b/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 diff --git a/R/lazyframe__lazy.R b/R/lazyframe__lazy.R index e82f8466b..45f253fb4 100644 --- a/R/lazyframe__lazy.R +++ b/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 diff --git a/man/DataFrame_with_row_count.Rd b/man/DataFrame_with_row_count.Rd new file mode 100644 index 000000000..1f9860a51 --- /dev/null +++ b/man/DataFrame_with_row_count.Rd @@ -0,0 +1,20 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/dataframe__frame.R +\name{DataFrame_with_row_count} +\alias{DataFrame_with_row_count} +\title{Eager with_row_count} +\usage{ +DataFrame_with_row_count(name, offset = NULL) +} +\arguments{ +\item{name}{string name of the created column} + +\item{offset}{positive integer offset for the start of the counter} +} +\value{ +A new \code{DataFrame} object with a counter column in front +} +\description{ +Add a new column at index 0 that counts the rows +} +\keyword{DataFrame} diff --git a/man/LazyFrame_with_row_count.Rd b/man/LazyFrame_with_row_count.Rd new file mode 100644 index 000000000..1029aa8d7 --- /dev/null +++ b/man/LazyFrame_with_row_count.Rd @@ -0,0 +1,20 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/lazyframe__lazy.R +\name{LazyFrame_with_row_count} +\alias{LazyFrame_with_row_count} +\title{Lazy with_row_count} +\usage{ +LazyFrame_with_row_count(name, offset = NULL) +} +\arguments{ +\item{name}{string name of the created column} + +\item{offset}{positive integer offset for the start of the counter} +} +\value{ +A new \code{LazyFrame} object with a counter column in front +} +\description{ +Add a new column at index 0 that counts the rows +} +\keyword{LazyFrame} diff --git a/man/nanoarrow.Rd b/man/nanoarrow.Rd index 3ecaf02a4..7af2018a2 100644 --- a/man/nanoarrow.Rd +++ b/man/nanoarrow.Rd @@ -16,13 +16,13 @@ \alias{as_record_batch_reader.DataFrame} \title{polars to nanoarrow and arrow} \usage{ -\method{as_nanoarrow_array_stream}{DataFrame}(x, ..., schema = NULL) +as_nanoarrow_array_stream.DataFrame(x, ..., schema = NULL) -\method{infer_nanoarrow_schema}{DataFrame}(x, ...) +infer_nanoarrow_schema.DataFrame(x, ...) -\method{as_arrow_table}{DataFrame}(x, ...) +as_arrow_table.DataFrame(x, ...) -\method{as_record_batch_reader}{DataFrame}(x, ..., schema = NULL) +as_record_batch_reader.DataFrame(x, ..., schema = NULL) } \arguments{ \item{x}{a polars DataFrame} diff --git a/src/rust/src/lazy/dataframe.rs b/src/rust/src/lazy/dataframe.rs index c8a622788..0a09e2bcc 100644 --- a/src/rust/src/lazy/dataframe.rs +++ b/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 { + 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, diff --git a/src/rust/src/rdataframe/mod.rs b/src/rust/src/rdataframe/mod.rs index 3d3b09bc9..b9090fa52 100644 --- a/src/rust/src/rdataframe/mod.rs +++ b/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 { + 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() diff --git a/tests/testthat/test-dataframe.R b/tests/testthat/test-dataframe.R index 04ff8fbcd..5bfe3b74b 100644 --- a/tests/testthat/test-dataframe.R +++ b/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)))) +}) diff --git a/tests/testthat/test-lazy.R b/tests/testthat/test-lazy.R index 1d545cc60..a698621c2 100644 --- a/tests/testthat/test-lazy.R +++ b/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)))) +})