ARROW-17738: [R] dplyr::compute should convert from grouped arrow_dpl…

…yr_query to arrow Table (#14160)

Authored-by: SHIMA Tatsuya <ts1s1andn@gmail.com>
Signed-off-by: Dewey Dunnington <dewey@voltrondata.com>

r/R/dplyr-collect.R

@@ -68,9 +68,7 @@ restore_dplyr_features <- function(df, query) {
       )
     } else {
       # This is a Table, via compute() or collect(as_data_frame = FALSE)
-      df <- as_adq(df)
-      df$group_by_vars <- query$group_by_vars
-      df$drop_empty_groups <- query$drop_empty_groups
+      df$metadata$r$attributes$.group_vars <- query$group_by_vars
     }
   }
   df
@@ -80,7 +78,10 @@ collapse.arrow_dplyr_query <- function(x, ...) {
   # Figure out what schema will result from the query
   x$schema <- implicit_schema(x)
   # Nest inside a new arrow_dplyr_query (and keep groups)
-  restore_dplyr_features(arrow_dplyr_query(x), x)
+  out <- arrow_dplyr_query(x)
+  out$group_by_vars <- x$group_by_vars
+  out$drop_empty_groups <- x$drop_empty_groups
+  out
 }
 collapse.Dataset <- collapse.ArrowTabular <- collapse.RecordBatchReader <- function(x, ...) {
   arrow_dplyr_query(x)

r/tests/testthat/test-dataset-dplyr.R

@@ -284,13 +284,9 @@ test_that("compute()/collect(as_data_frame=FALSE)", {
     group_by(fct) %>%
     compute()
 
-  # the group_by() prevents compute() from returning a Table...
-  expect_s3_class(tab5, "arrow_dplyr_query")
-
-  # ... but $.data is a Table...
-  expect_r6_class(tab5$.data, "Table")
-  # ... and the mutate() was evaluated
-  expect_true("negint" %in% names(tab5$.data))
+  expect_r6_class(tab5, "Table")
+  # mutate() was evaluated
+  expect_true("negint" %in% names(tab5))
 })
 
 test_that("head/tail on query on dataset", {

r/tests/testthat/test-dplyr-query.R

@@ -119,7 +119,7 @@ test_that("collect(as_data_frame=FALSE)", {
     filter(int > 5) %>%
     group_by(int) %>%
     collect(as_data_frame = FALSE)
-  expect_s3_class(b4, "arrow_dplyr_query")
+  expect_r6_class(b4, "Table")
   expect_equal(
     as.data.frame(b4),
     expected %>%
@@ -156,7 +156,7 @@ test_that("compute()", {
     filter(int > 5) %>%
     group_by(int) %>%
     compute()
-  expect_s3_class(b4, "arrow_dplyr_query")
+  expect_r6_class(b4, "Table")
   expect_equal(
     as.data.frame(b4),
     expected %>%
@@ -579,3 +579,43 @@ test_that("needs_projection unit tests", {
     tab %>% relocate(lgl)
   ))
 })
+
+test_that("compute() on a grouped query returns a Table with groups in metadata", {
+  tab1 <- tbl %>%
+      arrow_table() %>%
+      group_by(int) %>%
+      compute()
+  expect_r6_class(tab1, "Table")
+  expect_equal(
+    as.data.frame(tab1),
+    tbl %>%
+      group_by(int)
+  )
+  expect_equal(
+    collect(tab1),
+    tbl %>%
+      group_by(int)
+  )
+})
+
+test_that("collect() is identical to compute() %>% collect()", {
+  tab1 <- tbl %>%
+    arrow_table()
+  adq1 <- tab1 %>%
+    group_by(int)
+
+  expect_equal(
+    tab1 %>%
+      compute() %>%
+      collect(),
+    tab1 %>%
+      collect()
+  )
+  expect_equal(
+    adq1 %>%
+      compute() %>%
+      collect(),
+    adq1 %>%
+      collect()
+  )
+})