more simplify

Signed-off-by: SHIMA Tatsuya <ts1s1andn@gmail.com>
apache · Oct 13, 2022 · 943d603 · 943d603
1 parent 4e8850a
commit 943d603
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/r/R/dplyr-collect.R b/r/R/dplyr-collect.R
@@ -62,17 +62,14 @@ restore_dplyr_features <- function(df, query) {
   # An arrow_dplyr_query holds some attributes that Arrow doesn't know about
   # After calling collect(), make sure these features are carried over
 
-  if (length(dplyr::group_vars(df))) {
-    stop("df must be ungrouped", call. = FALSE)
-  }
-
   if (length(dplyr::group_vars(query))) {
     if (is.data.frame(df)) {
       # Preserve groupings, if present
       df <- dplyr::group_by(
         df,
         !!!syms(dplyr::group_vars(query)),
-        .drop = dplyr::group_by_drop_default(query)
+        .drop = dplyr::group_by_drop_default(query),
+        .add = FALSE
       )
     } else {
       # This is a Table, via compute() or collect(as_data_frame = FALSE)

diff --git a/r/R/dplyr.R b/r/R/dplyr.R
@@ -49,8 +49,9 @@ arrow_dplyr_query <- function(.data) {
   if (inherits(.data, "data.frame")) {
     .data <- Table$create(.data)
   }
-  # If .data is a Table, it must be ungrouped (ARROW-17737)
-  if (length(group_vars.ArrowTabular(.data))) {
+  # ARROW-17737: If .data is a Table, remove groups from metadata
+  # (we've already grabbed the groups above)
+  if (inherits(.data, "ArrowTabular")) {
     .data <- ungroup.ArrowTabular(.data)
   }