ARROW-11754: [R] Support dplyr::compute()

r/R/arrow-package.R

@@ -32,7 +32,7 @@
     c(
       "select", "filter", "collect", "summarise", "group_by", "groups",
       "group_vars", "group_by_drop_default", "ungroup", "mutate", "transmute",
-      "arrange", "rename", "pull", "relocate"
+      "arrange", "rename", "pull", "relocate", "compute"
     )
   )
   for (cl in c("Dataset", "ArrowTabular", "arrow_dplyr_query")) {

r/R/dplyr.R

@@ -619,9 +619,19 @@ collect.arrow_dplyr_query <- function(x, as_data_frame = TRUE, ...) {
     restore_dplyr_features(tab, x)
   }
 }
-collect.ArrowTabular <- as.data.frame.ArrowTabular
+collect.ArrowTabular <- function(x, as_data_frame = TRUE, ...) {
+  if (as_data_frame) {
+    as.data.frame(x, ...)
+  } else {
+    x
+  }
+}
 collect.Dataset <- function(x, ...) dplyr::collect(arrow_dplyr_query(x), ...)
 
+compute.arrow_dplyr_query <- function(x, ...) dplyr::collect(x, as_data_frame = FALSE)
+compute.ArrowTabular <- function(x, ...) x
+compute.Dataset <- compute.arrow_dplyr_query
+
 ensure_group_vars <- function(x) {
   if (inherits(x, "arrow_dplyr_query")) {
     # Before pulling data from Arrow, make sure all group vars are in the projection
@@ -670,7 +680,7 @@ restore_dplyr_features <- function(df, query) {
         drop = dplyr::group_by_drop_default(query)
       )
     } else {
-      # This is a Table, via collect(as_data_frame = FALSE)
+      # This is a Table, via compute() or collect(as_data_frame = FALSE)
       df <- arrow_dplyr_query(df)
       df$group_by_vars <- query$group_by_vars
       df$drop_empty_groups <- query$drop_empty_groups

r/tests/testthat/test-dataset.R

@@ -1068,6 +1068,59 @@ See $.data for the source Arrow object",
   )
 })
 
+test_that("compute()/collect(as_data_frame=FALSE)", {
+  skip_if_not_available("parquet")
+  ds <- open_dataset(dataset_dir)
+
+  tab1 <- ds %>% compute()
+  expect_is(tab1, "Table")
+
+  tab2 <- ds %>% collect(as_data_frame = FALSE)
+  expect_is(tab2, "Table")
+
+  tab3 <-  ds %>%
+    mutate(negint = -int) %>%
+    filter(negint > - 100) %>%
+    arrange(chr) %>%
+    select(negint) %>%
+    compute()
+
+  expect_is(tab3, "Table")
+
+  expect_equal(
+    tab3 %>% collect(),
+    tibble(negint = -1:-10)
+  )
+
+  tab4 <-  ds %>%
+    mutate(negint = -int) %>%
+    filter(negint > - 100) %>%
+    arrange(chr) %>%
+    select(negint) %>%
+    collect(as_data_frame = FALSE)
+
+  expect_is(tab3, "Table")
+
+  expect_equal(
+    tab4 %>% collect(),
+    tibble(negint = -1:-10)
+  )
+
+  tab5 <- ds %>%
+    mutate(negint = -int) %>%
+    group_by(fct) %>%
+    compute()
+
+  # the group_by() prevents compute() from returning a Table...
+  expect_is(tab5, "arrow_dplyr_query")
+
+  # ... but $.data is a Table...
+  expect_is(tab5$.data, "Table")
+  # ... and the mutate() was evaluated
+  expect_true("negint" %in% names(tab5$.data))
+
+})
+
 test_that("head/tail", {
   skip_if_not_available("parquet")
   ds <- open_dataset(dataset_dir)

r/tests/testthat/test-dplyr.R

@@ -178,6 +178,10 @@ test_that("pull", {
 test_that("collect(as_data_frame=FALSE)", {
   batch <- record_batch(tbl)
 
+  b1 <- batch %>% collect(as_data_frame = FALSE)
+
+  expect_is(b1, "RecordBatch")
+
   b2 <- batch %>%
     select(int, chr) %>%
     filter(int > 5) %>%
@@ -208,6 +212,43 @@ test_that("collect(as_data_frame=FALSE)", {
     )
 })
 
+test_that("compute()", {
+  batch <- record_batch(tbl)
+
+  b1 <- batch %>% compute()
+
+  expect_is(b1, "RecordBatch")
+
+  b2 <- batch %>%
+    select(int, chr) %>%
+    filter(int > 5) %>%
+    compute()
+
+  expect_is(b2, "RecordBatch")
+  expected <- tbl[tbl$int > 5 & !is.na(tbl$int), c("int", "chr")]
+  expect_equal(as.data.frame(b2), expected)
+
+  b3 <- batch %>%
+    select(int, strng = chr) %>%
+    filter(int > 5) %>%
+    compute()
+  expect_is(b3, "RecordBatch")
+  expect_equal(as.data.frame(b3), set_names(expected, c("int", "strng")))
+
+  b4 <- batch %>%
+    select(int, strng = chr) %>%
+    filter(int > 5) %>%
+    group_by(int) %>%
+    compute()
+  expect_is(b4, "arrow_dplyr_query")
+  expect_equal(
+    as.data.frame(b4),
+    expected %>%
+      rename(strng = chr) %>%
+      group_by(int)
+  )
+})
+
 test_that("head", {
   batch <- record_batch(tbl)