Silently add grouping vars whenever selecting.

#170

R/manip-df.r

@@ -49,15 +49,9 @@ select.tbl_df <- function(.data, ...) {
 
 #' @export
 select.grouped_df <- function(.data, ...) {
-  vars <- select_vars(names(.data), ..., env = parent.frame())
+  vars <- select_vars(names(.data), ..., env = parent.frame(),
+    include = as.character(groups(.data)))
 
-  # Don't remove grouping vars!
-  missing <- setdiff(as.character(groups(.data)), vars)
-  if (length(missing) > 0) {
-    vars <- append(vars, structure(missing, names = missing)) 
-    warning("grouping variables added implicitely: ",
-      paste0(missing, collapse = ","), call. = FALSE)
-  }
   select_impl(.data, vars)
 }
 

R/manip-grouped-dt.r

@@ -117,7 +117,8 @@ arrange.grouped_dt <- function(.data, ...) {
 #' @rdname manip_grouped_dt
 #' @export
 select.grouped_dt <- function(.data, ...) {
-  vars <- select_vars(names(.data), ..., env = parent.frame())
+  vars <- select_vars(names(.data), ..., env = parent.frame(),
+    include = as.character(groups(.data)))
   out <- .data[, vars, drop = FALSE, with = FALSE]
   setnames(out, names(vars))
 

R/manip-sql.r

@@ -12,7 +12,8 @@ arrange.tbl_sql <- function(.data, ...) {
 
 #' @export
 select.tbl_sql <- function(.data, ...) {
-  vars <- select_vars(tbl_vars(.data), ..., env = parent.frame())
+  vars <- select_vars(tbl_vars(.data), ..., env = parent.frame(),
+    include = as.character(groups(.data)))
   # Index into variables so that select can be applied multiple times
   # and after a mutate.
   idx <- match(vars, tbl_vars(.data))

R/manip.r

@@ -130,6 +130,8 @@ select <- function(.data, ...) UseMethod("select")
 #'
 #' @param vars A character vector of existing column names.
 #' @param ... Expressions to compute
+#' @param include Character vector of column names that must always be
+#'   included in output.
 #' @export
 #' @keywords internal
 #' @return A named character vector. Values are existing column names,
@@ -156,7 +158,7 @@ select <- function(.data, ...) UseMethod("select")
 #' # Rename variables
 #' select_vars(names(iris), petal_length = Petal.Length)
 #' select_vars(names(iris), petal = starts_with("Petal"))
-select_vars <- function(vars, ..., env = parent.frame()) {
+select_vars <- function(vars, ..., env = parent.frame(), include = character()) {
   args <- dots(...)
   if (length(args) == 0) return(setNames(vars, vars))
 
@@ -222,10 +224,15 @@ select_vars <- function(vars, ..., env = parent.frame()) {
   excl <- abs(ind[ind < 0])
   incl <- incl[match(incl, excl, 0L) == 0L]
 
+  # Add variables that must be included
+  sel <- setNames(vars[incl], names(incl))
+  sel <- c(setdiff(include, sel), sel)
+
   # Ensure all output vars named
-  unnamed <- names2(incl) == ""
-  names(incl)[unnamed] <- vars[incl][unnamed]
-  setNames(vars[incl], names(incl))
+  unnamed <- names2(sel) == ""
+  names(sel)[unnamed] <- sel[unnamed]
+
+  sel
 }
 
 

inst/tests/test-group-by.r

@@ -139,9 +139,10 @@ test_that("original data table not modified by grouping", {
 })
 
 test_that("select(group_by(.)) implicitely adds grouping variables (#170)", {
-  expect_warning( 
-    res <- group_by(mtcars, vs) %.% select(mpg) %.% summarise( mpg = mean( mpg ) )
-  )
-  expect_equal(nrow(res), 2)
+  res <- mtcars %.% group_by(vs) %.% select(mpg)
+  expect_equal(names(res), c("vs", "mpg"))
+
+  res <- mtcars %.% tbl_dt() %.% group_by(vs) %.% select(mpg)
+  expect_equal(names(res), c("vs", "mpg"))
 })
 

inst/tests/test-select.r

@@ -64,7 +64,8 @@ test_that("num_range selects numeric ranges", {
 
 # Data table -------------------------------------------------------------------
 
-test_that("select changes columns in copy of data table", {dt <- data.table(x = 1:4, y = letters[1:4])
+test_that("select changes columns in copy of data table", {
+  dt <- data.table(x = 1:4, y = letters[1:4])
 
   expect_equal(names(select(dt, x, z = y)), c("x", "z"))
   expect_equal(names(dt), c("x", "y"))
@@ -94,3 +95,8 @@ test_that("select renames variables (#317)", {
   expect_equal(tbl_vars(first %.% select(A)), "A")
   expect_equal(tbl_vars(first %.% select(B = A)), "B")
 })
+
+test_that("select preserves grouping vars", {
+  first <- tbls$sqlite %.% group_by(b) %.% select(a)
+  expect_equal(tbl_vars(first), c("b", "a"))
+})