From 663a44a447cdeb7a830b7e4b630a4390e0a48f1a Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Wed, 26 Oct 2022 10:40:23 -0300 Subject: [PATCH 1/3] make sure head() of an execplanreader is an execplanreader --- r/NAMESPACE | 1 + r/R/query-engine.R | 8 ++++++++ r/tests/testthat/test-query-engine.R | 11 +++++++++++ 3 files changed, 20 insertions(+) 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") From 8e098025c003c7f667ee719fd53b5056fb607019 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Thu, 3 Nov 2022 12:48:28 -0300 Subject: [PATCH 2/3] fix order --- r/src/recordbatchreader.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/r/src/recordbatchreader.cpp b/r/src/recordbatchreader.cpp index 9ea4d917016..8a4d4563bbe 100644 --- a/r/src/recordbatchreader.cpp +++ b/r/src/recordbatchreader.cpp @@ -166,6 +166,7 @@ class RecordBatchReaderHead : public arrow::RecordBatchReader { } else { done_ = true; arrow::Status result = reader_->Close(); + done_ = true; return result; } } From 5f5989042cd5ecd8a6194614fdec5a68e558048b Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Mon, 2 Jan 2023 14:11:59 -0400 Subject: [PATCH 3/3] fix merge --- r/src/recordbatchreader.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/r/src/recordbatchreader.cpp b/r/src/recordbatchreader.cpp index 8a4d4563bbe..9ea4d917016 100644 --- a/r/src/recordbatchreader.cpp +++ b/r/src/recordbatchreader.cpp @@ -166,7 +166,6 @@ class RecordBatchReaderHead : public arrow::RecordBatchReader { } else { done_ = true; arrow::Status result = reader_->Close(); - done_ = true; return result; } }