nealrichardson commented on a change in pull request #10992:
URL: https://github.com/apache/arrow/pull/10992#discussion_r698766675
##########
File path: r/R/dplyr-collect.R
##########
@@ -47,16 +36,69 @@ collect.ArrowTabular <- function(x, as_data_frame = TRUE,
...) {
x
}
}
-collect.Dataset <- function(x, ...) dplyr::collect(arrow_dplyr_query(x), ...)
+collect.Dataset <- function(x, ...) dplyr::collect(as_adq(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
pull.arrow_dplyr_query <- function(.data, var = -1) {
- .data <- arrow_dplyr_query(.data)
+ .data <- as_adq(.data)
var <- vars_pull(names(.data), !!enquo(var))
.data$selected_columns <- set_names(.data$selected_columns[var], var)
dplyr::collect(.data)[[1]]
}
pull.Dataset <- pull.ArrowTabular <- pull.arrow_dplyr_query
+
+# TODO: Correctly handle group_vars after summarize; also in collapse()
+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(query$group_by_vars) > 0) {
+ # Preserve groupings, if present
+ if (is.data.frame(df)) {
+ df <- dplyr::grouped_df(
+ df,
+ dplyr::group_vars(query),
+ drop = dplyr::group_by_drop_default(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
+}
+
+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
+ arrow_dplyr_query(x)
+}
Review comment:
Is there something more it should do? There may be more subtlety in how
the ExecPlan gets built at evaluation time, but for the purposes of getting a
"fresh" arrow_dplyr_query to chain more things on, I think this gets that job
done. It's really about making sure the order of ExecNodes stays faithful to
the user's request.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
