diff --git a/r/R/dplyr-collect.R b/r/R/dplyr-collect.R index ca018a87860..1714af8a509 100644 --- a/r/R/dplyr-collect.R +++ b/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) diff --git a/r/tests/testthat/test-dataset-dplyr.R b/r/tests/testthat/test-dataset-dplyr.R index f3ec858c6ce..04b1a460dc4 100644 --- a/r/tests/testthat/test-dataset-dplyr.R +++ b/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", { diff --git a/r/tests/testthat/test-dplyr-query.R b/r/tests/testthat/test-dplyr-query.R index b08016f2170..c40815df69d 100644 --- a/r/tests/testthat/test-dplyr-query.R +++ b/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() + ) +})