ARROW-15271: [R] Refactor do_exec_plan to return a RecordBatchReader

r/NAMESPACE

@@ -61,6 +61,7 @@ S3method(as_data_type,pyarrow.lib.DataType)
 S3method(as_data_type,pyarrow.lib.Field)
 S3method(as_record_batch,RecordBatch)
 S3method(as_record_batch,Table)
+S3method(as_record_batch,arrow_dplyr_query)
 S3method(as_record_batch,data.frame)
 S3method(as_record_batch,pyarrow.lib.RecordBatch)
 S3method(as_record_batch,pyarrow.lib.Table)
@@ -227,6 +228,7 @@ export(ReadableFile)
 export(RecordBatch)
 export(RecordBatchFileReader)
 export(RecordBatchFileWriter)
+export(RecordBatchReader)
 export(RecordBatchStreamReader)
 export(RecordBatchStreamWriter)
 export(RoundMode)

r/R/dataset-scan.R

@@ -153,7 +153,7 @@ head.Scanner <- function(x, n = 6L, ...) {
 
 #' @export
 tail.Scanner <- function(x, n = 6L, ...) {
-  tail_from_batches(dataset___Scanner__ScanBatches(x), n)
+  tail_from_batches(dataset___Scanner__ScanBatches(x), n)$read_table()
 }
 
 tail_from_batches <- function(batches, n) {
@@ -169,43 +169,57 @@ tail_from_batches <- function(batches, n) {
     if (n <= 0) break
   }
   # rev() the result to put the batches back in the right order
-  Table$create(!!!rev(result))
+  RecordBatchReader$create(batches = rev(result))
 }
 
 #' Apply a function to a stream of RecordBatches
 #'
 #' As an alternative to calling `collect()` on a `Dataset` query, you can
 #' use this function to access the stream of `RecordBatch`es in the `Dataset`.
-#' This lets you aggregate on each chunk and pull the intermediate results into
-#' a `data.frame` for further aggregation, even if you couldn't fit the whole
-#' `Dataset` result in memory.
+#' This lets you do more complex operations in R that operate on chunks of data
+#' without having to hold the entire Dataset in memory at once. You can include
+#' `map_batches()` in a dplyr pipeline and do additional dplyr methods on the
+#' stream of data in Arrow after it.
 #'
-#' This is experimental and not recommended for production use.
+#' Note that, unlike the core dplyr methods that are implemented in the Arrow
+#' query engine, `map_batches()` is not lazy: it starts evaluating on the data
+#' when you call it, even if you send its result to another pipeline function.
+#'
+#' This is experimental and not recommended for production use. It is also
+#' single-threaded and runs in R not C++, so it won't be as fast as core
+#' Arrow methods.
 #'
 #' @param X A `Dataset` or `arrow_dplyr_query` object, as returned by the
 #' `dplyr` methods on `Dataset`.
 #' @param FUN A function or `purrr`-style lambda expression to apply to each
-#' batch
+#' batch. It must return a RecordBatch or something coercible to one via
+#' `as_record_batch()'.
 #' @param ... Additional arguments passed to `FUN`
-#' @param .data.frame logical: collect the resulting chunks into a single
-#' `data.frame`? Default `TRUE`
+#' @param .data.frame Deprecated argument, ignored
+#' @return An `arrow_dplyr_query`.
 #' @export
-map_batches <- function(X, FUN, ..., .data.frame = TRUE) {
-  # TODO: ARROW-15271 possibly refactor do_exec_plan to return a RecordBatchReader
+map_batches <- function(X, FUN, ..., .data.frame = NULL) {
+  if (!is.null(.data.frame)) {
+    warning(
+      "The .data.frame argument is deprecated. ",
+      "Call collect() on the result to get a data.frame.",
+      call. = FALSE
+    )
+  }
   plan <- ExecPlan$create()
   final_node <- plan$Build(as_adq(X))
   reader <- plan$Run(final_node)
   FUN <- as_mapper(FUN)
 
-  # TODO: wrap batch in arrow_dplyr_query with X$selected_columns,
-  # X$temp_columns, and X$group_by_vars
-  # if X is arrow_dplyr_query, if some other arg (.dplyr?) == TRUE
+  # TODO: for future consideration
+  # * Move eval to C++ and make it a generator so it can stream, not block
+  # * Accept an output schema argument: with that, we could make this lazy (via collapse)
   batch <- reader$read_next_batch()
   res <- vector("list", 1024)
   i <- 0L
   while (!is.null(batch)) {
     i <- i + 1L
-    res[[i]] <- FUN(batch, ...)
+    res[[i]] <- as_record_batch(FUN(batch, ...))
     batch <- reader$read_next_batch()
   }
 
@@ -214,13 +228,7 @@ map_batches <- function(X, FUN, ..., .data.frame = TRUE) {
     res <- res[seq_len(i)]
   }
 
-  if (.data.frame & inherits(res[[1]], "arrow_dplyr_query")) {
-    res <- dplyr::bind_rows(map(res, dplyr::collect))
-  } else if (.data.frame) {
-    res <- dplyr::bind_rows(map(res, as.data.frame))
-  }
-
-  res
+  RecordBatchReader$create(batches = res)
 }
 
 #' @usage NULL

r/R/dplyr-collect.R

@@ -27,15 +27,30 @@ collect.arrow_dplyr_query <- function(x, as_data_frame = TRUE, ...) {
   }
 
   # See query-engine.R for ExecPlan/Nodes
+  plan <- ExecPlan$create()
+  final_node <- plan$Build(x)
   tryCatch(
-    tab <- do_exec_plan(x),
+    tab <- plan$Run(final_node)$read_table(),
     # n = 4 because we want the error to show up as being from collect()
     # and not handle_csv_read_error()
     error = function(e, call = caller_env(n = 4)) {
       handle_csv_read_error(e, x$.data$schema, call)
     }
   )
 
+  # TODO(ARROW-16607): move KVM handling into ExecPlan
+  if (ncol(tab)) {
+    # Apply any column metadata from the original schema, where appropriate
+    new_r_metadata <- get_r_metadata_from_old_schema(
+      tab$schema,
+      source_data(x)$schema,
+      drop_attributes = has_aggregation(x)
+    )
+    if (!is.null(new_r_metadata)) {
+      tab$r_metadata <- new_r_metadata
+    }
+  }
+
   if (as_data_frame) {
     df <- as.data.frame(tab)
     restore_dplyr_features(df, x)

r/R/duckdb.R

@@ -123,11 +123,7 @@ duckdb_disconnector <- function(con, tbl_name) {
 #' other processes (like DuckDB).
 #'
 #' @param .data the object to be converted
-#' @param as_arrow_query should the returned object be wrapped as an
-#' `arrow_dplyr_query`? (logical, default: `TRUE`)
-#'
-#' @return a `RecordBatchReader` object, wrapped as an arrow dplyr query which
-#'  can be used in dplyr pipelines.
+#' @return A `RecordBatchReader`.
 #' @export
 #'
 #' @examplesIf getFromNamespace("run_duckdb_examples", "arrow")()
@@ -142,7 +138,7 @@ duckdb_disconnector <- function(con, tbl_name) {
 #'   summarize(mean_mpg = mean(mpg, na.rm = TRUE)) %>%
 #'   to_arrow() %>%
 #'   collect()
-to_arrow <- function(.data, as_arrow_query = TRUE) {
+to_arrow <- function(.data) {
   # If this is an Arrow object already, return quickly since we're already Arrow
   if (inherits(.data, c("arrow_dplyr_query", "ArrowObject"))) {
     return(.data)
@@ -161,9 +157,5 @@ to_arrow <- function(.data, as_arrow_query = TRUE) {
   # Run the query
   res <- DBI::dbSendQuery(dbplyr::remote_con(.data), dbplyr::remote_query(.data), arrow = TRUE)
 
-  if (as_arrow_query) {
-    arrow_dplyr_query(duckdb::duckdb_fetch_record_batch(res))
-  } else {
-    duckdb::duckdb_fetch_record_batch(res)
-  }
+  duckdb::duckdb_fetch_record_batch(res)
 }

r/R/query-engine.R

@@ -15,36 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-do_exec_plan <- function(.data) {
-  plan <- ExecPlan$create()
-  final_node <- plan$Build(.data)
-  tab <- plan$Run(final_node)
-  # TODO (ARROW-14289): make the head/tail methods return RBR not Table
-  if (inherits(tab, "RecordBatchReader")) {
-    tab <- tab$read_table()
-  }
-
-  # If arrange() created $temp_columns, make sure to omit them from the result
-  # We can't currently handle this in the ExecPlan itself because sorting
-  # happens in the end (SinkNode) so nothing comes after it.
-  if (length(final_node$sort$temp_columns) > 0) {
-    tab <- tab[, setdiff(names(tab), final_node$sort$temp_columns), drop = FALSE]
-  }
-
-  if (ncol(tab)) {
-    # Apply any column metadata from the original schema, where appropriate
-    new_r_metadata <- get_r_metadata_from_old_schema(
-      tab$schema,
-      source_data(.data)$schema,
-      drop_attributes = has_aggregation(.data)
-    )
-    if (!is.null(new_r_metadata)) {
-      tab$r_metadata <- new_r_metadata
-    }
-  }
-  tab
-}
-
 ExecPlan <- R6Class("ExecPlan",
   inherit = ArrowObject,
   public = list(
@@ -220,17 +190,27 @@ ExecPlan <- R6Class("ExecPlan",
         # just use it to take the random slice
         slice_size <- node$head %||% node$tail
         if (!is.null(slice_size)) {
-          # TODO (ARROW-14289): make the head methods return RBR not Table
           out <- head(out, slice_size)
         }
         # Can we now tell `self$Stop()` to StopProducing? We already have
         # everything we need for the head (but it seems to segfault: ARROW-14329)
       } else if (!is.null(node$tail)) {
         # Reverse the row order to get back what we expect
-        # TODO: don't return Table, return RecordBatchReader
         out <- out$read_table()
         out <- out[rev(seq_len(nrow(out))), , drop = FALSE]
+        # Put back into RBR
+        out <- as_record_batch_reader(out)
       }
+
+      # If arrange() created $temp_columns, make sure to omit them from the result
+      # We can't currently handle this in ExecPlan_run itself because sorting
+      # happens in the end (SinkNode) so nothing comes after it.
+      if (length(node$sort$temp_columns) > 0) {
+        tab <- out$read_table()
+        tab <- tab[, setdiff(names(tab), node$sort$temp_columns), drop = FALSE]
+        out <- as_record_batch_reader(tab)
+      }
+
       out
     },
     Write = function(node, ...) {

r/R/record-batch-reader.R

@@ -56,6 +56,7 @@
 #'
 #' @rdname RecordBatchReader
 #' @name RecordBatchReader
+#' @export
 #' @include arrow-object.R
 #' @examples
 #' tf <- tempfile()
@@ -104,6 +105,16 @@ RecordBatchReader <- R6Class("RecordBatchReader",
     schema = function() RecordBatchReader__schema(self)
   )
 )
+RecordBatchReader$create <- function(..., batches = list(...), schema = NULL) {
+  are_batches <- map_lgl(batches, ~ inherits(., "RecordBatch"))
+  if (!all(are_batches)) {
+    stop(
+      "All inputs to RecordBatchReader$create must be RecordBatches",
+      call. = FALSE
+    )
+  }
+  RecordBatchReader__from_batches(batches, schema)
+}
 
 #' @export
 names.RecordBatchReader <- function(x) names(x$schema)
@@ -208,13 +219,13 @@ as_record_batch_reader.Table <- function(x, ...) {
 #' @rdname as_record_batch_reader
 #' @export
 as_record_batch_reader.RecordBatch <- function(x, ...) {
-  RecordBatchReader__from_batches(list(x), NULL)
+  RecordBatchReader$create(x, schema = x$schema)
 }
 
 #' @rdname as_record_batch_reader
 #' @export
 as_record_batch_reader.data.frame <- function(x, ...) {
-  as_record_batch_reader(as_record_batch(x))
+  RecordBatchReader$create(as_record_batch(x))
 }
 
 #' @rdname as_record_batch_reader
@@ -226,8 +237,8 @@ as_record_batch_reader.Dataset <- function(x, ...) {
 #' @rdname as_record_batch_reader
 #' @export
 as_record_batch_reader.arrow_dplyr_query <- function(x, ...) {
-  # TODO(ARROW-15271): make ExecPlan return RBR
-  as_record_batch_reader(collect.arrow_dplyr_query(x, as_data_frame = FALSE))
+  # TODO(ARROW-16607): use ExecPlan directly when it handles metadata
+  as_record_batch_reader(compute.arrow_dplyr_query(x))
 }
 
 #' @rdname as_record_batch_reader

r/R/record-batch.R

@@ -297,6 +297,12 @@ as_record_batch.Table <- function(x, ..., schema = NULL) {
   out
 }
 
+#' @rdname as_record_batch
+#' @export
+as_record_batch.arrow_dplyr_query <- function(x, ...) {
+  as_record_batch(compute.arrow_dplyr_query(x), ...)
+}
+
 #' @rdname as_record_batch
 #' @export
 as_record_batch.data.frame <- function(x, ..., schema = NULL) {

r/R/table.R

@@ -139,6 +139,13 @@ Table$create <- function(..., schema = NULL) {
   if (all_record_batches(dots)) {
     return(Table__from_record_batches(dots, schema))
   }
+  if (length(dots) == 1 && inherits(dots[[1]], c("RecordBatchReader", "RecordBatchFileReader"))) {
+    tab <- dots[[1]]$read_table()
+    if (!is.null(schema)) {
+      tab <- tab$cast(schema)
+    }
+    return(tab)
+  }
 
   # If any arrays are length 1, recycle them
   dots <- recycle_scalars(dots)