diff --git a/r/NAMESPACE b/r/NAMESPACE index 0b18ace9ad3..cde81d977b7 100644 --- a/r/NAMESPACE +++ b/r/NAMESPACE @@ -101,6 +101,7 @@ S3method(dimnames,ArrowTabular) S3method(head,ArrowDatum) S3method(head,ArrowTabular) S3method(head,Dataset) +S3method(head,ExecPlanReader) S3method(head,RecordBatchReader) S3method(head,Scanner) S3method(head,arrow_dplyr_query) diff --git a/r/R/query-engine.R b/r/R/query-engine.R index b90ef382146..7e737af3ffa 100644 --- a/r/R/query-engine.R +++ b/r/R/query-engine.R @@ -361,6 +361,14 @@ ExecPlanReader <- R6Class("ExecPlanReader", ) ) +#' @export +head.ExecPlanReader <- function(x, n = 6L, ...) { + # We need to make sure that the head() of an ExecPlanReader + # is also an ExecPlanReader so that the evaluation takes place + # in a way that supports calls into R. + as_record_batch_reader(as_adq(RecordBatchReader__Head(x, n))) +} + do_exec_plan_substrait <- function(substrait_plan) { if (is.string(substrait_plan)) { substrait_plan <- substrait__internal__SubstraitFromJSON(substrait_plan) diff --git a/r/tests/testthat/test-query-engine.R b/r/tests/testthat/test-query-engine.R index 86d89898083..1d8d876bf53 100644 --- a/r/tests/testthat/test-query-engine.R +++ b/r/tests/testthat/test-query-engine.R @@ -79,6 +79,17 @@ test_that("ExecPlanReader evaluates head() lazily", { # evaluate to TRUE (i.e., the reader may or may not be completely drained). }) +test_that("head() of an ExecPlanReader is an ExecPlanReader", { + reader <- as_record_batch_reader(as_adq(arrow_table(x = 1:10))) + expect_r6_class(reader, "ExecPlanReader") + reader_head <- head(reader, 6) + expect_r6_class(reader_head, "ExecPlanReader") + expect_equal( + as_arrow_table(reader_head), + arrow_table(x = 1:6) + ) +}) + test_that("do_exec_plan_substrait can evaluate a simple plan", { skip_if_not_available("substrait")