Add some missing grouped_df methods (#4710)

* Drop groups when subsetting. Fixes #4708

* Implement subset assignment methods for grouped_df. Fixes #4004

* Add news

Author: hadley

NAMESPACE

@@ -1,7 +1,10 @@
 # Generated by roxygen2: do not edit by hand
 
+S3method("$<-",grouped_df)
 S3method("[",fun_list)
 S3method("[",grouped_df)
+S3method("[<-",grouped_df)
+S3method("[[<-",grouped_df)
 S3method(anti_join,data.frame)
 S3method(anti_join,tbl_df)
 S3method(arrange,data.frame)

NEWS.md

@@ -1,5 +1,12 @@
 # dplyr 0.9.0 (in development)
 
+* 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/grouped-df.r

@@ -204,20 +204,46 @@ as_tibble.grouped_df <- function(x, ...) {
 #' @importFrom tibble is_tibble
 #' @export
 `[.grouped_df` <- function(x, i, j, drop = FALSE) {
-  y <- NextMethod()
+  out <- NextMethod()
 
-  if (isTRUE(drop) && !is_tibble(y)) {
-    return(y)
+  if (!is.data.frame(out)) {
+    return(out)
   }
 
-  group_names <- group_vars(x)
-  if (!all(group_names %in% names(y))) {
-    as_tibble(y)
+  if (drop) {
+    as_tibble(out)
   } else {
-    grouped_df(y, group_names, group_by_drop_default(x))
+    groups <- intersect(names(out), group_vars(x))
+    if ((missing(i) || nargs() == 2) && identical(groups, group_vars(x))) {
+      new_grouped_df(out, group_data(x))
+    } else {
+      grouped_df(out, groups, group_by_drop_default(x))
+    }
+  }
+}
+
+#' @export
+`$<-.grouped_df` <- function(x, name, ..., value) {
+  out <- NextMethod()
+  if (name %in% group_vars(x)) {
+    grouped_df(out, intersect(names(out), group_vars(x)), group_by_drop_default(x))
+  } else {
+    out
   }
 }
 
+#' @export
+`[<-.grouped_df` <- function(x, i, j, ..., value) {
+  out <- NextMethod()
+  grouped_df(out, intersect(names(out), group_vars(x)), group_by_drop_default(x))
+}
+
+#' @export
+`[[<-.grouped_df` <- function(x, ..., value) {
+  out <- NextMethod()
+  grouped_df(out, intersect(names(out), group_vars(x)), group_by_drop_default(x))
+}
+
 #' @method rbind grouped_df
 #' @export
 rbind.grouped_df <- function(...) {

R/select.R

@@ -148,16 +148,6 @@ select_impl <- function(.data, vars) {
   if (is_grouped_df(.data)) {
     # we might have to alter the names of the groups metadata
     groups <- attr(.data, "groups")
-
-    # check grouped metadata
-    group_names <- names(groups)[seq_len(ncol(groups) - 1L)]
-    if (any(test <- ! group_names %in% vars)) {
-      abort(
-        glue("{col} not found in groups metadata. Probably a corrupt grouped_df object.", col = group_names[test[1L]]),
-        "dplyr_select_corrupt_grouped_df"
-      )
-    }
-
     group_vars <- c(vars[vars %in% names(groups)], .rows = ".rows")
     groups <- select_impl(groups, group_vars)
 

tests/testthat/test-group-by.r

@@ -319,12 +319,6 @@ test_that("tbl_sum gets the right number of groups", {
   expect_equal(res, c("A tibble" = "4 x 1", "Groups" = "x [2]"))
 })
 
-test_that("grouped data frames support drop=TRUE (#3714)", {
-  expect_is(group_by(iris, Species)[ , "Sepal.Width", drop=TRUE], "numeric")
-
-  expect_is(group_by(iris, Species)[ , c("Species", "Sepal.Width"), drop=TRUE], "grouped_df")
-})
-
 test_that("group_by ignores empty quosures (3780)", {
   empty <- quo()
   expect_equal(group_by(mtcars, cyl), group_by(mtcars, cyl, !!empty))

tests/testthat/test-grouped-df.r

@@ -6,3 +6,53 @@ test_that("can selectively ungroup", {
   expect_equal(gf %>% ungroup(x) %>% group_vars(), "y")
   expect_error(gf %>% ungroup(z) %>% group_vars(), "z")
 })
+
+
+test_that("[ method can remove grouping vars", {
+  df <- tibble(x = 1, y = 2, z = 3)
+  gf <- group_by(df, x, y)
+
+  expect_equal(gf, gf)
+  expect_equal(gf[1], group_by(df[1], x))
+  expect_equal(gf[3], df[3])
+})
+
+test_that("[ method reuses group_data() if possible", {
+  df <- tibble(x = 1, y = 2, z = 3)
+  gf <- group_by(df, x, y)
+
+  expect_reference(group_data(gf), group_data(gf[1:2]))
+  expect_reference(group_data(gf), group_data(gf[, 1:2]))
+})
+
+test_that("[ supports drop=TRUE (#3714)", {
+  df <- tibble(x = 1, y = 2)
+  gf <- group_by(df, x)
+
+  expect_type(gf[, "y", drop = TRUE], "double")
+  expect_s3_class(gf[, c("x", "y"), drop = TRUE], "tbl_df")
+})
+
+test_that("$<-, [[<-, and [<- update grouping data if needed", {
+  df <- tibble(x = 1, y = 2)
+  gf <- group_by(df, x)
+
+  expect_equal(group_data(`$<-`(gf, "x", 2))$x, 2)
+  expect_equal(group_data(`$<-`(gf, "y", 2))$x, 1)
+
+  expect_equal(group_data({gf2 <- gf; gf2[[1]] <- 3; gf2})$x, 3)
+  expect_equal(group_data(`[<-`(gf, 1, "x", 4))$x, 4)
+})
+
+test_that("can remove grouping cols with subset assignment", {
+  df <- tibble(x = 1, y = 2)
+  gf1 <- gf2 <- gf3 <- group_by(df, x, y)
+
+  gf1$x <- NULL
+  gf2[["x"]] <- NULL
+  gf3[, "x"] <- NULL
+
+  expect_named(group_data(gf1), c("y", ".rows"))
+  expect_named(group_data(gf2), c("y", ".rows"))
+  expect_named(group_data(gf3), c("y", ".rows"))
+})

tests/testthat/test-select.r

@@ -64,13 +64,6 @@ test_that("select can be before group_by (#309)", {
   expect_equal(names(dfagg), c("id", "year", "var1"))
 })
 
-test_that("rename errors with invalid grouped data frame (#640)", {
-  df <- tibble(a = 1:3, b = 2:4, d = 3:5) %>% group_by(a, b)
-  df$a <- NULL
-  expect_error(df %>% rename(e = d))
-  expect_error(df %>% rename(e = b))
-})
-
 test_that("rename() handles data pronoun", {
   expect_identical(rename(tibble(x = 1), y = .data$x), tibble(y = 1))
 })