Add n function to all grouped sources.

Closes tidyverse#10

R/manip-grouped-df.r

@@ -57,6 +57,8 @@ summarise.grouped_df <- function(.data, ...) {
   calls <- named_dots(...)
   if (is.lazy(.data)) .data <- build_index(.data)
   v <- make_view(.data, parent.frame())
+  v$add_function("n", function() length(rows))
+
   ngrps <- length(attr(.data, "index"))
 
   output_summary <- function(j) {

R/manip-grouped-dt.r

@@ -57,7 +57,14 @@ summarise.grouped_dt <- function(.data, ...) {
   }
 
   cols <- named_dots(...)
-  list_call <- as.call(c(quote(list), named_dots(...)))
+  # Replace n() with .N
+  for (i in seq_along(cols)) {
+    if (identical(cols[[i]], quote(n()))) {
+      cols[[i]] <- quote(.N)
+    }
+  }
+
+  list_call <- as.call(c(quote(list), cols))
   call <- substitute(data[, list_call, by = vars])
 
   env <- new.env(parent = parent.frame(), size = 1L)

R/translate-sql-base.r

@@ -38,5 +38,7 @@ base_sql$is.null <- function(x) {
 base_sql$c     <- function(...) escape(c(...))
 base_sql$`:`   <- function(from, to) escape(from:to)
 
+base_sql$n     <- sql_prefix("count")
+
 senv <- new.env(parent = emptyenv())
 senv$pi <- structure("PI()", class = "sql")

R/view.r

@@ -15,7 +15,14 @@ view <- function(data, index, parent = parent.frame()) {
     rows
   }
 
+  # Tools to manage extra functions
+  fun_env <- new.env(parent = parent)
+  add_function <- function(name, fun) {
+    fun_env[[name]] <- fun
+  }
+
   # Tools to manage active bindings
+  grp_env <- new.env(parent = fun_env, size = nrow(data))
   add_binding <- function(name, fun) {
     makeActiveBinding(name, fun, grp_env)
   }
@@ -27,7 +34,6 @@ view <- function(data, index, parent = parent.frame()) {
     }
   }
 
-  grp_env <- new.env(parent, size = nrow(data))
   for (name in names(data)) {
     add_binding(name, from_data(name))
   }
@@ -36,5 +42,6 @@ view <- function(data, index, parent = parent.frame()) {
     eval(expr, grp_env)
   }
 
-  list(set_group = set_group, eval = local_eval, add_binding = add_binding)
+  list(set_group = set_group, eval = local_eval, add_function = add_function, 
+    add_binding = add_binding)
 }

inst/tests/test-group-size.r → inst/tests/test-equiv-grouped.r

@@ -1,5 +1,6 @@
-context("Group size")
+context("Equivalence (grouped)")
 
+library(data.table)
 data("baseball", package = "plyr")
 player_df <- group_by(baseball, id)
 player_dt <- group_by(data.table(baseball), id)
@@ -16,4 +17,13 @@ test_that("group size the same regardless of data source", {
 
   expect_equal(gs_dt, gs_df)
   expect_equal(gs_db, gs_df)
-})
+})
+
+test_that("n that same regardless of source", {
+  count_df <- summarise(player_df, count = n())
+  count_dt <- ungroup(summarise(player_dt, count = n()))
+  count_db <- summarise(player_db, count = n())
+
+  expect_equal(count_dt$n, count_df$n)
+  expect_equal(count_db$n, count_df$n)
+})