jonkeane commented on a change in pull request #10992:
URL: https://github.com/apache/arrow/pull/10992#discussion_r698757563
##########
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:
Am I reading this right that what this does is figures out the resultant
schema, and then wrap another `arrow_dplyr_query` layer around the resultant
object? Am I missing something else that it's doing? Are we thinking about
expanding this more in the future?
--
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
