ARROW-15016: [R] show_exec_plan for an arrow_dplyr_query

r/NAMESPACE

@@ -352,6 +352,7 @@ export(s3_bucket)
 export(schema)
 export(set_cpu_count)
 export(set_io_thread_count)
+export(show_exec_plan)
 export(starts_with)
 export(string)
 export(struct)

r/R/arrow-package.R

@@ -41,7 +41,8 @@
  "group_vars", "group_by_drop_default", "ungroup", "mutate", "transmute",
  "arrange", "rename", "pull", "relocate", "compute", "collapse",
  "distinct", "left_join", "right_join", "inner_join", "full_join",
- "semi_join", "anti_join", "count", "tally", "rename_with", "union", "union_all", "glimpse"
+ "semi_join", "anti_join", "count", "tally", "rename_with", "union",
+ "union_all", "glimpse", "show_query", "explain"
  )
  )
  for (cl in c("Dataset", "ArrowTabular", "RecordBatchReader", "arrow_dplyr_query")) {

r/R/dplyr.R

@@ -219,6 +219,53 @@ tail.arrow_dplyr_query <- function(x, n = 6L, ...) {
  x
 }
 
+#' Show the details of an Arrow Execution Plan
+#'
+#' This is a function which gives more details about the logical query plan
+#' that will be executed when evaluating an `arrow_dplyr_query` object.
+#' It calls the C++ `ExecPlan` object's print method.
+#' Functionally, it is similar to `dplyr::explain()`. This function is used as
+#' the `dplyr::explain()` and `dplyr::show_query()` methods.
+#'
+#' @param x an `arrow_dplyr_query` to print the `ExecPlan` for.
+#'
+#' @return `x`, invisibly.
+#' @export
+#'
+#' @examplesIf arrow_with_dataset() && requireNamespace("dplyr", quietly = TRUE)
+#' library(dplyr)
+#' mtcars %>%
+#' arrow_table() %>%
+#' filter(mpg > 20) %>%
+#' mutate(x = gear/carb) %>%
+#' 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))
+ invisible(x)
+}
+
+show_query.arrow_dplyr_query <- function(x, ...) {
+ show_exec_plan(x)
+}
+
+show_query.Dataset <- show_query.ArrowTabular <- show_query.RecordBatchReader <- show_query.arrow_dplyr_query
+
+explain.arrow_dplyr_query <- function(x, ...) {
+ show_exec_plan(x)
+}
+
+explain.Dataset <- explain.ArrowTabular <- explain.RecordBatchReader <- explain.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

r/R/query-engine.R

@@ -193,6 +193,8 @@ 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()`
  assert_is(node, "ExecNode")
 
  # Sorting and head/tail (if sorted) are handled in the SinkNode,
@@ -210,6 +212,8 @@ 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
@@ -273,6 +277,39 @@ 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)
  )
 )

r/_pkgdown.yml

@@ -220,6 +220,7 @@ reference:
  - value_counts
  - list_compute_functions
  - register_scalar_function
+ - show_exec_plan
  - title: Connections to other systems
  contents:
  - to_arrow

r/src/compute-exec.cpp

@@ -60,6 +60,10 @@ std::pair<std::shared_ptr<compute::ExecPlan>, std::shared_ptr<arrow::RecordBatch
 ExecPlan_prepare(const std::shared_ptr<compute::ExecPlan>& plan,
  const std::shared_ptr<compute::ExecNode>& final_node,
  cpp11::list sort_options, cpp11::strings metadata, int64_t head = -1) {
+ // a section of this code is copied and used in ExecPlan_BuildAndShow - the 2 need
+ // to be in sync
+ // Start of chunk used in ExecPlan_BuildAndShow
+
  // For now, don't require R to construct SinkNodes.
  // Instead, just pass the node we should collect as an argument.
  arrow::AsyncGenerator<arrow::util::optional<compute::ExecBatch>> sink_gen;
@@ -88,6 +92,8 @@ ExecPlan_prepare(const std::shared_ptr<compute::ExecPlan>& plan,
  compute::SinkNodeOptions{&sink_gen});
  }
 
+ // End of chunk used in ExecPlan_BuildAndShow
+
  StopIfNotOk(plan->Validate());
 
  // If the generator is destroyed before being completely drained, inform plan
@@ -155,6 +161,46 @@ std::shared_ptr<arrow::Schema> ExecNode_output_schema(
  return node->output_schema();
 }
 
+// [[arrow::export]]
+std::string ExecPlan_BuildAndShow(const std::shared_ptr<compute::ExecPlan>& plan,
+ const std::shared_ptr<compute::ExecNode>& final_node,
+ cpp11::list sort_options, int64_t head = -1) {
+ // a section of this code is copied from ExecPlan_prepare - the 2 need to be in sync
+ // Start of chunk copied from ExecPlan_prepare
+
+ // For now, don't require R to construct SinkNodes.
+ // Instead, just pass the node we should collect as an argument.
+ arrow::AsyncGenerator<arrow::util::optional<compute::ExecBatch>> sink_gen;
+
+ // Sorting uses a different sink node; there is no general sort yet
+ if (sort_options.size() > 0) {
+ if (head >= 0) {
+ // Use the SelectK node to take only what we need
+ MakeExecNodeOrStop(
+ "select_k_sink", plan.get(), {final_node.get()},
+ compute::SelectKSinkNodeOptions{
+ arrow::compute::SelectKOptions(
+ head, std::dynamic_pointer_cast<compute::SortOptions>(
+ make_compute_options("sort_indices", sort_options))
+ ->sort_keys),
+ &sink_gen});
+ } else {
+ MakeExecNodeOrStop("order_by_sink", plan.get(), {final_node.get()},
+ compute::OrderBySinkNodeOptions{
+ *std::dynamic_pointer_cast<compute::SortOptions>(
+ make_compute_options("sort_indices", sort_options)),
+ &sink_gen});
+ }
+ } else {
+ MakeExecNodeOrStop("sink", plan.get(), {final_node.get()},
+ compute::SinkNodeOptions{&sink_gen});
+ }
+
+ // End of chunk copied from ExecPlan_prepare
+
+ return plan->ToString();
+}
+
 #if defined(ARROW_R_WITH_DATASET)
 
 #include <arrow/dataset/file_base.h>

r/tests/testthat/test-dataset-dplyr.R

@@ -340,3 +340,80 @@ test_that("dplyr method not implemented messages", {
  fixed = TRUE
  )
 })
+
+test_that("show_exec_plan(), show_query() and explain() with datasets", {
+ # show_query() and explain() are wrappers around show_exec_plan() and are not
+ # tested separately
+
+ ds <- open_dataset(dataset_dir, partitioning = schema(part = uint8()))
+
+ # minimal test
+ expect_output(
+ ds %>%
+ show_exec_plan(),
+ regexp = paste0(
+ "ExecPlan with .* nodes:.*", # boiler plate for ExecPlan
+ "ProjectNode.*", # output columns
+ "SourceNode" # entry point
+ )
+ )
+
+ # filter and select
+ expect_output(
+ ds %>%
+ select(string = chr, integer = int, part) %>%
+ filter(integer > 6L & part == 1) %>%
+ show_exec_plan(),
+ regexp = paste0(
+ "ExecPlan with .* nodes:.*", # boiler plate for ExecPlan
+ "ProjectNode.*", # output columns
+ "FilterNode.*", # filter node
+ "int > 6.*cast.*", # filtering expressions + auto-casting of part
+ "SourceNode" # entry point
+ )
+ )
+
+ # group_by and summarise
+ expect_output(
+ ds %>%
+ group_by(part) %>%
+ summarise(avg = mean(int)) %>%
+ show_exec_plan(),
+ regexp = paste0(
+ "ExecPlan with .* nodes:.*", # boiler plate for ExecPlan
+ "ProjectNode.*", # output columns
+ "GroupByNode.*", # group by node
+ "keys=.*part.*", # key for aggregations
+ "aggregates=.*hash_mean.*", # aggregations
+ "ProjectNode.*", # input columns
+ "SourceNode" # entry point
+ )
+ )
+
+ # arrange and head
+ expect_output(
+ ds %>%
+ filter(lgl) %>%
+ arrange(chr) %>%
+ show_exec_plan(),
+ regexp = paste0(
+ "ExecPlan with .* nodes:.*", # boiler plate for ExecPlan
+ "OrderBySinkNode.*chr.*ASC.*", # arrange goes via the OrderBy sink node
+ "ProjectNode.*", # output columns
+ "FilterNode.*", # filter node
+ "filter=lgl.*", # filtering expression
+ "SourceNode" # entry point
+ )
+ )
+
+ # printing the ExecPlan for a nested query would currently force the
+ # evaluation of the inner one(s), which we want to avoid => no output
+ expect_warning(
+ ds %>%
+ filter(lgl) %>%
+ arrange(chr) %>%
+ head() %>%
+ show_exec_plan(),
+ "The `ExecPlan` cannot be printed for a nested query."
+ )
+})