feat: Implement <group_by/lazy_group_by>$len() (#1638)

Author: etiennebacher

NEWS.md

@@ -8,6 +8,7 @@
 - `<lazyframe>$remove()` and `<dataframe>$remove()` as a complement to
   `$filter()` (#1632).
 - New method `<expr>$is_close()` (#1637).
+- New methods `<group_by>$len()` and `<lazy_group_by>$len()` (#1638).
 
 ## polars 1.5.0
 

R/dataframe-group_by-general.R

@@ -207,3 +207,23 @@ groupby__n_unique <- function() {
   self$agg(pl$all()$n_unique()) |>
     wrap()
 }
+
+#' @inherit lazygroupby__len title params
+#' @inherit as_polars_df return
+#' @examples
+#' df <- pl$DataFrame(
+#'   a = c("Apple", "Apple", "Orange"),
+#'   b = c(1, NA, 2)
+#' )
+#' df$group_by("a")$len()
+#'
+#' df$group_by("a")$len("n")
+groupby__len <- function(name = NULL) {
+  wrap({
+    len_expr <- pl$len()
+    if (!is.null(name)) {
+      len_expr <- len_expr$alias(name)
+    }
+    self$agg(len_expr)
+  })
+}

R/lazyframe-group_by.R

@@ -207,3 +207,26 @@ lazygroupby__n_unique <- function() {
   self$agg(pl$all()$n_unique()) |>
     wrap()
 }
+
+#' Return the number of rows in each group
+#'
+#' @param name Assign a name to the resulting column. If `NULL`, defaults to
+#' `"len"`.
+#' @inherit as_polars_lf return
+#' @examples
+#' lf <- pl$LazyFrame(
+#'   a = c("Apple", "Apple", "Orange"),
+#'   b = c(1, NA, 2)
+#' )
+#' lf$group_by("a")$len()$collect()
+#'
+#' lf$group_by("a")$len("n")$collect()
+lazygroupby__len <- function(name = NULL) {
+  wrap({
+    len_expr <- pl$len()
+    if (!is.null(name)) {
+      len_expr <- len_expr$alias(name)
+    }
+    self$agg(len_expr)
+  })
+}

altdoc/mkdocs.yml

@@ -74,6 +74,7 @@ nav:
         - GroupBy:
             - agg: man/groupby__agg.md
             - head: man/groupby__head.md
+            - len: man/groupby__len.md
             - mean: man/groupby__mean.md
             - median: man/groupby__median.md
             - min: man/groupby__min.md
@@ -155,6 +156,7 @@ nav:
         - GroupBy:
             - agg: man/lazygroupby__agg.md
             - head: man/lazygroupby__head.md
+            - len: man/lazygroupby__len.md
             - mean: man/lazygroupby__mean.md
             - median: man/lazygroupby__median.md
             - min: man/lazygroupby__min.md

man/groupby__len.Rd

@@ -628,3 +628,27 @@
        [1] "mpg"  "cyl"  "disp" "hp"   "drat" "wt"   "qsec" "vs"   "am"   "gear"
       [11] "carb"
 
+# group_by() + len()
+
+    Code
+      df$group_by("a", .maintain_order = TRUE)$len(1)
+    Condition
+      Error:
+      ! Evaluation failed in `$len()`.
+      Caused by error in `len_expr$alias()`:
+      ! Evaluation failed in `$alias()`.
+      Caused by error:
+      ! Argument `name` must be character, not double
+
+---
+
+    Code
+      df$group_by("a", .maintain_order = TRUE)$len(TRUE)
+    Condition
+      Error:
+      ! Evaluation failed in `$len()`.
+      Caused by error in `len_expr$alias()`:
+      ! Evaluation failed in `$alias()`.
+      Caused by error:
+      ! Argument `name` must be character, not logical
+

man/lazygroupby__len.Rd

@@ -2616,3 +2616,29 @@ test_that("active bindings", {
   expect_snapshot(as_polars_lf(mtcars)$width)
   expect_snapshot(as_polars_lf(mtcars)$columns)
 })
+
+test_that("group_by() + len()", {
+  df <- pl$DataFrame(
+    a = c("Apple", "Apple", "Orange"),
+    b = c(1, NA, 2)
+  )
+
+  expect_query_equal(
+    .input$group_by("a", .maintain_order = TRUE)$len(),
+    df,
+    pl$DataFrame(a = c("Apple", "Orange"), len = c(2, 1))$cast(len = pl$UInt32)
+  )
+  expect_query_equal(
+    .input$group_by("a", .maintain_order = TRUE)$len("n"),
+    df,
+    pl$DataFrame(a = c("Apple", "Orange"), n = c(2, 1))$cast(n = pl$UInt32)
+  )
+  expect_snapshot(
+    df$group_by("a", .maintain_order = TRUE)$len(1),
+    error = TRUE
+  )
+  expect_snapshot(
+    df$group_by("a", .maintain_order = TRUE)$len(TRUE),
+    error = TRUE
+  )
+})