From c1cc8add89354ca94c2ab059cf95857f25f8538c Mon Sep 17 00:00:00 2001
From: Will Jones <willjones127@gmail.com>
Date: Thu, 21 Apr 2022 12:10:12 -0700
Subject: [PATCH] Convert to adq first

---
 r/R/dataset-scan.R              | 2 +-
 r/tests/testthat/test-dataset.R | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/r/R/dataset-scan.R b/r/R/dataset-scan.R
index 870b1ac100b..1675bce7b85 100644
--- a/r/R/dataset-scan.R
+++ b/r/R/dataset-scan.R
@@ -193,7 +193,7 @@ tail_from_batches <- function(batches, n) {
 map_batches <- function(X, FUN, ..., .data.frame = TRUE) {
   # TODO: ARROW-15271 possibly refactor do_exec_plan to return a RecordBatchReader
   plan <- ExecPlan$create()
-  final_node <- plan$Build(X)
+  final_node <- plan$Build(as_adq(X))
   reader <- plan$Run(final_node)
   FUN <- as_mapper(FUN)
 
diff --git a/r/tests/testthat/test-dataset.R b/r/tests/testthat/test-dataset.R
index 3a62c466c21..6be88be0585 100644
--- a/r/tests/testthat/test-dataset.R
+++ b/r/tests/testthat/test-dataset.R
@@ -619,6 +619,14 @@ test_that("map_batches", {
     c(5, 10)
   )
 
+  # Can take a raw dataset as X argument
+  expect_equal(
+    ds %>%
+      map_batches(~ count(., part)) %>%
+      arrange(part),
+    tibble(part = c(1, 2), n = c(10, 10))
+  )
+
   # $Take returns RecordBatch, which gets binded into a tibble
   expect_equal(
     ds %>%