From 943d603da2a67907be33dea06890bd98ee76ac35 Mon Sep 17 00:00:00 2001
From: SHIMA Tatsuya <ts1s1andn@gmail.com>
Date: Thu, 13 Oct 2022 10:45:32 +0000
Subject: [PATCH] more simplify

Signed-off-by: SHIMA Tatsuya <ts1s1andn@gmail.com>
---
 r/R/dplyr-collect.R | 7 ++-----
 r/R/dplyr.R         | 5 +++--
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/r/R/dplyr-collect.R b/r/R/dplyr-collect.R
index 8a08f7f04482e..5640565b15385 100644
--- 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
index 2e60444b1f0ed..08fef273ac6bc 100644
--- 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)
   }