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."
+  )
+})