Merge pull request #3490 from tidyverse/feature-1185-group

- `group_indices()` can be used without argument in expressions in verbs (#1185).

NAMESPACE

@@ -76,8 +76,10 @@ S3method(group_by_,tbl_cube)
 S3method(group_indices,data.frame)
 S3method(group_indices,default)
 S3method(group_indices,grouped_df)
+S3method(group_indices,rowwise_df)
 S3method(group_indices_,data.frame)
 S3method(group_indices_,grouped_df)
+S3method(group_indices_,rowwise_df)
 S3method(group_size,data.frame)
 S3method(group_size,grouped_df)
 S3method(group_size,rowwise_df)
@@ -427,6 +429,7 @@ importFrom(glue,glue)
 importFrom(magrittr,"%>%")
 importFrom(methods,is)
 importFrom(pkgconfig,get_config)
+importFrom(rlang,dots_n)
 importFrom(stats,lag)
 importFrom(stats,setNames)
 importFrom(stats,update)

NEWS.md

@@ -110,6 +110,8 @@
 
 * `transmute()` no longer prints a message when including a group variable.
 
+* `group_indices()` can be used, without argument inside a dplyr expression (#1185).
+
 ## Documentation
 
 * Improved documentation for set operations (#3238, @edublancas).

R/group-indices.R

@@ -37,9 +37,23 @@ group_indices_.data.frame <- function(.data, ..., .dots = list()) {
   group_indices(.data, !!!dots)
 }
 
+#' @export
+group_indices.rowwise_df <- function(.data, ...) {
+  if (dots_n(...)) {
+    warn("group_indices_.rowwise_df ignores extra arguments")
+  }
+  seq_len(nrow(.data))
+}
+#' @export
+group_indices_.rowwise_df <- function(.data, ..., .dots = list()) {
+  dots <- compat_lazy_dots(.dots, caller_env(), ...)
+  group_indices(.data, !!!dots)
+}
+
+#' @importFrom rlang dots_n
 #' @export
 group_indices.grouped_df <- function(.data, ...) {
-  if (length(list(...))) {
+  if (dots_n(...)) {
     warn("group_indices_.grouped_df ignores extra arguments")
   }
   grouped_indices_grouped_df_impl(.data)

inst/include/dplyr/HybridHandlerMap.h

@@ -14,6 +14,7 @@ void install_window_handlers(HybridHandlerMap& handlers);
 void install_offset_handlers(HybridHandlerMap& handlers);
 void install_in_handlers(HybridHandlerMap& handlers);
 void install_debug_handlers(HybridHandlerMap& handlers);
+void install_group_handlers(HybridHandlerMap& handlers);
 
 bool hybridable(RObject arg);
 

inst/include/tools/SlicingIndex.h

@@ -82,7 +82,7 @@ class NaturalSlicingIndex : public SlicingIndex {
   }
 
   virtual int group() const {
-    return -1;
+    return 0 ;
   }
 
   virtual bool is_identity(SEXP x) const {
@@ -111,7 +111,7 @@ class OffsetSlicingIndex : public SlicingIndex {
   }
 
   inline int group() const {
-    return -1;
+    return 0;
   }
 
 private:

src/hybrid.cpp

@@ -72,6 +72,7 @@ HybridHandlerMap& get_handlers() {
     install_offset_handlers(handlers);
     install_in_handlers(handlers);
     install_debug_handlers(handlers);
+    install_group_handlers(handlers);
   }
   return handlers;
 }

src/hybrid_group.cpp

@@ -0,0 +1,44 @@
+#include "pch.h"
+#include <dplyr/main.h>
+
+#include <dplyr/HybridHandlerMap.h>
+
+#include <dplyr/Result/ILazySubsets.h>
+
+#include <dplyr/Result/Result.h>
+
+using namespace Rcpp;
+using namespace dplyr;
+
+
+class GroupHybrid : public Result {
+public:
+  SEXP process(const RowwiseDataFrame&) {
+    Rprintf("GroupHybrid::RowwiseDataFrame()\n") ;
+    return R_NilValue ;
+  }
+
+  SEXP process(const GroupedDataFrame&) {
+    Rprintf("GroupHybrid::GroupedDataFrame()\n") ;
+    return R_NilValue ;
+  }
+
+  SEXP process(const SlicingIndex& i) {
+    return IntegerVector(i.size(), i.group() + 1);
+  }
+
+};
+
+Result* group_indices_prototype(SEXP call, const ILazySubsets&, int nargs) {
+  // if there are arguments, give up on hybrid evaluation
+  // and let R handle things, perhaps it will be a group_indices(data, ...)
+  if (nargs != 0)
+    return 0;
+
+  return new GroupHybrid();
+}
+
+void install_group_handlers(HybridHandlerMap& handlers) {
+  Environment ns_dplyr = Environment::namespace_env("dplyr");
+  handlers[ Rf_install("group_indices") ] = HybridHandler(group_indices_prototype, HybridHandler::DPLYR, ns_dplyr["group_indices"]);
+}

tests/testthat/test-group-indices.R

@@ -31,3 +31,37 @@ test_that("group indices are updated correctly for joined grouped data frames (#
   res <- inner_join(d1, d2, by = "x")
   expect_equal(group_indices(res), res$x)
 })
+
+test_that("group_indices() works for rowwise data (#3491)", {
+  df <- rowwise(data.frame(x = 1:10))
+  expect_equal(group_indices(df), 1:10)
+})
+
+test_that("group_indices() warns when passed extra arguments on grouped or rowwise data", {
+  df <- rowwise(data.frame(x = 1:10))
+  expect_warning(idx <- group_indices(df, x))
+  expect_equal(idx, 1:10)
+
+  expect_warning(idx <- group_indices(group_by(df,x), x))
+  expect_equal(idx, 1:10)
+})
+
+test_that("group_indices() can be used inside mutate (#1185)", {
+  df <- data_frame(v1 = c(3, 3, 2, 2, 3, 1), v2 = 1:6) %>% group_by(v1)
+  expect_identical(
+    pull(mutate(df, g = group_indices())),
+    group_indices(df)
+  )
+
+  df <- data_frame(v1 = c(3, 3, 2, 2, 3, 1), v2 = 1:6)
+  expect_identical(
+    pull(mutate(df, g = group_indices())),
+    group_indices(df)
+  )
+
+  df <- rowwise(data_frame(v1 = c(3, 3, 2, 2, 3, 1), v2 = 1:6))
+  expect_identical(
+    pull(mutate(df, g = group_indices())),
+    group_indices(df)
+  )
+})