ARROW-17178: [R] Support head() in arrow_dplyr_query with user-defined function

r/R/compute.R

@@ -385,13 +385,6 @@ register_scalar_function <- function(name, fun, in_type, out_type,
  update_cache = TRUE
  )
 
- # User-defined functions require some special handling
- # in the query engine which currently require an opt-in using
- # the R_ARROW_COLLECT_WITH_UDF environment variable while this
- # behaviour is stabilized.
- # TODO(ARROW-17178) remove the need for this!
- Sys.setenv(R_ARROW_COLLECT_WITH_UDF = "true")
-
  invisible(NULL)
 }
 

r/R/dplyr.R

@@ -266,16 +266,19 @@ tail.arrow_dplyr_query <- function(x, n = 6L, ...) {
 #' show_exec_plan()
 show_exec_plan <- function(x) {
  adq <- as_adq(x)
- plan <- ExecPlan$create()
+
  # do not show the plan if we have a nested query (as this will force the
  # evaluation of the inner query/queries)
  # TODO see if we can remove after ARROW-16628
  if (is_collapsed(x) && has_head_tail(x$.data)) {
  warn("The `ExecPlan` cannot be printed for a nested query.")
  return(invisible(x))
  }
- final_node <- plan$Build(adq)
- cat(plan$BuildAndShow(final_node))
+
+ result <- as_record_batch_reader(adq)
+ cat(result$Plan()$ToString())
+ result$Close()
+
  invisible(x)
 }
 

r/R/query-engine.R

@@ -194,13 +194,11 @@ ExecPlan <- R6Class("ExecPlan",
  }
  node
  },
- Run = function(node, as_table = FALSE) {
- # a section of this code is used by `BuildAndShow()` too - the 2 need to be in sync
- # Start of chunk used in `BuildAndShow()`
+ Run = function(node) {
  assert_is(node, "ExecNode")
 
  # Sorting and head/tail (if sorted) are handled in the SinkNode,
- # created in ExecPlan_run
+ # created in ExecPlan_build
  sorting <- node$extras$sort %||% list()
  select_k <- node$extras$head %||% -1L
  has_sorting <- length(sorting) > 0
@@ -214,16 +212,7 @@ ExecPlan <- R6Class("ExecPlan",
  sorting$orders <- as.integer(sorting$orders)
  }
 
- # End of chunk used in `BuildAndShow()`
-
- # If we are going to return a Table anyway, we do this in one step and
- # entirely in one C++ call to ensure that we can execute user-defined
- # functions from the worker threads spawned by the ExecPlan. If not, we
- # use ExecPlan_run which returns a RecordBatchReader that can be
- # manipulated in R code (but that right now won't work with
- # user-defined functions).
- exec_fun <- if (as_table) ExecPlan_read_table else ExecPlan_run
- out <- exec_fun(
+ out <- ExecPlan_run(
  self,
  node,
  sorting,
@@ -240,18 +229,13 @@ ExecPlan <- R6Class("ExecPlan",
  slice_size <- node$extras$head %||% node$extras$tail
  if (!is.null(slice_size)) {
  out <- head(out, slice_size)
- # We already have everything we need for the head, so StopProducing
- self$Stop()
  }
  } else if (!is.null(node$extras$tail)) {
  # TODO(ARROW-16630): proper BottomK support
  # Reverse the row order to get back what we expect
  out <- as_arrow_table(out)
  out <- out[rev(seq_len(nrow(out))), , drop = FALSE]
- # Put back into RBR
- if (!as_table) {
- out <- as_record_batch_reader(out)
- }
+ out <- as_record_batch_reader(out)
  }
 
  # If arrange() created $temp_columns, make sure to omit them from the result
@@ -261,11 +245,7 @@ ExecPlan <- R6Class("ExecPlan",
  if (length(node$extras$sort$temp_columns) > 0) {
  tab <- as_arrow_table(out)
  tab <- tab[, setdiff(names(tab), node$extras$sort$temp_columns), drop = FALSE]
- if (!as_table) {
- out <- as_record_batch_reader(tab)
- } else {
- out <- tab
- }
+ out <- as_record_batch_reader(tab)
  }
 
  out
@@ -279,40 +259,9 @@ ExecPlan <- R6Class("ExecPlan",
  ...
  )
  },
- # SinkNodes (involved in arrange and/or head/tail operations) are created in
- # ExecPlan_run and are not captured by the regulat print method. We take a
- # similar approach to expose them before calling the print method.
- BuildAndShow = function(node) {
- # a section of this code is copied from `Run()` - the 2 need to be in sync
- # Start of chunk copied from `Run()`
-
- assert_is(node, "ExecNode")
-
- # Sorting and head/tail (if sorted) are handled in the SinkNode,
- # created in ExecPlan_run
- sorting <- node$extras$sort %||% list()
- select_k <- node$extras$head %||% -1L
- has_sorting <- length(sorting) > 0
- if (has_sorting) {
- if (!is.null(node$extras$tail)) {
- # Reverse the sort order and take the top K, then after we'll reverse
- # the resulting rows so that it is ordered as expected
- sorting$orders <- !sorting$orders
- select_k <- node$extras$tail
- }
- sorting$orders <- as.integer(sorting$orders)
- }
-
- # End of chunk copied from `Run()`
-
- ExecPlan_BuildAndShow(
- self,
- node,
- sorting,
- select_k
- )
- },
- Stop = function() ExecPlan_StopProducing(self)
+ ToString = function() {
+ ExecPlan_ToString(self)
+ }
  )
 )
 # nolint end.
@@ -396,6 +345,23 @@ ExecNode <- R6Class("ExecNode",
  )
 )
 
+ExecPlanReader <- R6Class("ExecPlanReader",
+ inherit = RecordBatchReader,
+ public = list(
+ batches = function() ExecPlanReader__batches(self),
+ read_table = function() Table__from_ExecPlanReader(self),
+ Plan = function() ExecPlanReader__Plan(self),
+ PlanStatus = function() ExecPlanReader__PlanStatus(self),
+ ToString = function() {
+ sprintf(
+ "<Status: %s>\n\n%s\n\nSee $Plan() for details.",
+ self$PlanStatus(),
+ super$ToString()
+ )
+ }
+ )
+)
+
 do_exec_plan_substrait <- function(substrait_plan) {
  if (is.string(substrait_plan)) {
  substrait_plan <- substrait__internal__SubstraitFromJSON(substrait_plan)

r/R/record-batch-reader.R

@@ -98,6 +98,7 @@ RecordBatchReader <- R6Class("RecordBatchReader",
  read_next_batch = function() RecordBatchReader__ReadNext(self),
  batches = function() RecordBatchReader__batches(self),
  read_table = function() Table__from_RecordBatchReader(self),
+ Close = function() RecordBatchReader__Close(self),
  export_to_c = function(stream_ptr) ExportRecordBatchReader(self, stream_ptr),
  ToString = function() self$schema$ToString()
  ),

r/R/table.R

@@ -328,15 +328,5 @@ as_arrow_table.RecordBatchReader <- function(x, ...) {
 #' @rdname as_arrow_table
 #' @export
 as_arrow_table.arrow_dplyr_query <- function(x, ...) {
- # See query-engine.R for ExecPlan/Nodes
- plan <- ExecPlan$create()
- final_node <- plan$Build(x)
-
- run_with_event_loop <- identical(
- Sys.getenv("R_ARROW_COLLECT_WITH_UDF", ""),
- "true"
- )
-
- result <- plan$Run(final_node, as_table = run_with_event_loop)
- as_arrow_table(result)
+ as_arrow_table(as_record_batch_reader(x))
 }

r/src/arrow_types.h

@@ -58,6 +58,8 @@ class ExecNode;
 } // namespace compute
 } // namespace arrow
 
+class ExecPlanReader;
+
 #if defined(ARROW_R_WITH_PARQUET)
 #include <parquet/type_fwd.h>
 #endif