Allow modification of grouping_vars (#4712)

Fixes #4709

Author: hadley

NEWS.md

@@ -1,12 +1,15 @@
 # dplyr 0.9.0 (in development)
 
+* `mutate()` and `summarise()` can now modify grouping variables (#4709).
+
 * Grouped data frames now have `[[<-`, `[<-` and `$<-` methods that will 
   re-generate the underlying grouping. Note that modifying grouping variables 
   in multiple steps (i.e. `df$grp1 <- 1; df$grp2 <- 1`) will be inefficient
   since the data frame will be regrouped after each modification.
 
 * `[.grouped_df` now regroups to respect any grouping columns that have
   been removed (#4708).
+
 * `as.tbl()` and `tbl_df()` have been formally deprecated.
   Please use `as_tibble()` instead.
 

R/mutate.R

@@ -45,9 +45,6 @@
 #' The former normalises `mass` by the global average whereas the
 #' latter normalises by the averages within gender levels.
 #'
-#' Note that you can't overwrite a grouping variable within
-#' `mutate()`.
-#'
 #' `mutate()` does not evaluate the expressions when the group is empty.
 #'
 #' @section Scoped mutation and transmutation:
@@ -259,6 +256,12 @@ mutate_finish <- function(.data, new_columns) {
     out[[new_column_names[i]]] <- if (!inherits(new_columns[[i]], "rlang_zap")) new_columns[[i]]
   }
 
+  # Re-group if needed
+  groups <- group_vars(.data)
+  if (any(groups %in% names(new_columns))) {
+    out <- grouped_df(out, intersect(groups, names(.data)))
+  }
+
   # copy back attributes
   # TODO: challenge that with some vctrs theory
   atts <- attributes(.data)

R/summarise.R

@@ -129,12 +129,14 @@ summarise.tbl_df <- function(.data, ...) {
 
   }
 
-  grouping <- group_keys(.data)
+  out <- group_keys(.data)
   if (!identical(.size, 1L)) {
-    grouping <- vec_slice(grouping, rep(seq2(1L, nrow(grouping)), .size))
+    out <- vec_slice(out, rep(seq2(1L, nrow(out)), .size))
   }
 
-  out <- add_column(grouping, !!!summaries)
+  for (i in seq_along(summaries)) {
+    out[[names(summaries)[i]]] <- summaries[[i]]
+  }
 
   if (is_grouped_df(.data) && length(group_vars(.data)) > 1) {
     out <- grouped_df(out, head(group_vars(.data), -1), group_by_drop_default(.data))

R/tbl-df.r

@@ -61,9 +61,6 @@ DataMask <- R6Class("DataMask",
     },
 
     add = function(name, chunks) {
-      if (name %in% group_vars(private$data)) {
-        abort(glue("Column `{name}` can't be modified because it's a grouping variable"))
-      }
       force(chunks)
       env_bind_active(private$bindings, !!name := function() {
         .subset2(chunks, private$current_group)

man/mutate.Rd

@@ -106,23 +106,6 @@ test_that("can use a purrr-style lambda", {
   expect_identical(summarise_at(mtcars, vars(1:2), ~ mean(.x)), summarise(mtcars, mpg = mean(mpg), cyl = mean(cyl)))
 })
 
-test_that("mutate_at and transmute_at refuses to mutate a grouping variable (#3351, #3480)", {
-  tbl <- tibble(gr1 = rep(1:2, 4), gr2 = rep(1:2, each = 4), x = 1:8) %>%
-    group_by(gr1)
-
-  expect_error(
-    mutate_at(tbl, vars(gr1), sqrt),
-    "Column `gr1` can't be modified because it's a grouping variable",
-    fixed = TRUE
-  )
-
-  expect_error(
-    transmute_at(tbl, vars(gr1), sqrt),
-    "Column `gr1` can't be modified because it's a grouping variable",
-    fixed = TRUE
-  )
-})
-
 test_that("mutate and transmute variants does not mutate grouping variable (#3351, #3480)", {
   tbl <- tibble(gr1 = rep(1:2, 4), gr2 = rep(1:2, each = 4), x = 1:8) %>%
     group_by(gr1)

tests/testthat/test-colwise-mutate.R

@@ -54,12 +54,13 @@ test_that("mutate can rename variables (#137)", {
   expect_equal(res$cyl2, res$cyl)
 })
 
-test_that("mutate refuses to modify grouping vars (#143)", {
-  expect_error(
-    mutate(group_by(mtcars, am), am = am + 2),
-    "Column `am` can't be modified because it's a grouping variable",
-    fixed = TRUE
-  )
+test_that("mutate regroups after modifying grouping vars", {
+  df <- tibble(x = 1:2, y = 2)
+  gf <- group_by(df, x)
+
+  out <- gf %>% mutate(x = 1)
+  expect_equal(out$x, c(1, 1))
+  expect_equal(nrow(group_data(out)), 1)
 })
 
 test_that("mutate handles constants (#152)", {

tests/testthat/test-mutate.r

@@ -50,13 +50,12 @@ test_that("summarise can refer to factor variables that were just created (#2217
   )
 })
 
-test_that("summarise refuses to modify grouping variable (#143)", {
-  df <- data.frame(a = c(1, 2, 1, 2), b = c(1, 1, 2, 2), x = 1:4)
-  ds <- group_by(df, a, b)
-  expect_error(
-    summarise(ds, a = mean(x), a = b + 1),
-    "Column `a` can't be modified because it's a grouping variable"
-  )
+test_that("summarise can modify grouping variables", {
+  df <- tibble(a = c(1, 2, 1, 2), b = c(1, 1, 2, 2))
+  gf <- group_by(df, a)
+
+  out <- summarise(gf, a = min(a) + 1)
+  expect_equal(out$a, c(2, 3))
 })
 
 test_that("summarise gives proper errors (#153)", {