ARROW-3818: [R] Table$from_batches

r/NAMESPACE

@@ -145,6 +145,7 @@ export(schema)
 export(str.integer64)
 export(struct)
 export(table)
+export(table_from_batches)
 export(time32)
 export(time64)
 export(timestamp)
@@ -165,6 +166,7 @@ importFrom(glue,glue)
 importFrom(purrr,map)
 importFrom(purrr,map2)
 importFrom(purrr,map_int)
+importFrom(purrr,walk)
 importFrom(rlang,abort)
 importFrom(rlang,dots_n)
 importFrom(rlang,list2)

r/R/R6.R

@@ -18,7 +18,7 @@
 #' @include enums.R
 #' @importFrom R6 R6Class
 #' @importFrom glue glue
-#' @importFrom purrr map map_int map2
+#' @importFrom purrr map map_int map2 walk
 #' @importFrom rlang dots_n
 #' @importFrom assertthat assert_that
 

r/R/Table.R

@@ -60,6 +60,27 @@ table <- function(.data){
   shared_ptr(`arrow::Table`, Table__from_dataframe(.data))
 }
 
+#' Create a arrow::Table from record batches
+#'
+#' @param ... variable number of arrow::RecordBatch. This supports tidy dots splicing
+#' @param schema either NULL or a arrow::Schema
+#'
+#' @return a arrow::Table
+#'
+#' @export
+table_from_batches <- function(..., schema = NULL) {
+  batches <- list2(...)
+  walk(batches, ~stopifnot(inherits(., "arrow::RecordBatch")))
+
+  if(is.null(schema)) {
+    shared_ptr(`arrow::Table`, Table__FromRecordBatches(batches))
+  } else if(inherits(schema, "arrow::Schema")) {
+    shared_ptr(`arrow::Table`, Table__FromRecordBatches_Schema(batches, schema))
+  } else {
+    abort("schema should be NULL or a `arrow::Scherma")
+  }
+}
+
 #' @export
 `as_tibble.arrow::Table` <- function(x, use_threads = TRUE, ...){
   Table__to_dataframe(x, use_threads = use_threads)

r/src/arrow_types.h

@@ -198,5 +198,16 @@ class RBuffer : public MutableBuffer {
   Vec vec_;
 };
 
+template <typename T>
+std::vector<std::shared_ptr<T>> list_to_shared_ptr_vector(SEXP lst) {
+  R_xlen_t n = XLENGTH(lst);
+  std::vector<std::shared_ptr<T>> res(n);
+  for (R_xlen_t i = 0; i < n; i++) {
+    res[i] = Rcpp::ConstReferenceSmartPtrInputParameter<std::shared_ptr<T>>(
+        VECTOR_ELT(lst, i));
+  }
+  return res;
+}
+
 }  // namespace r
 }  // namespace arrow

r/src/table.cpp

@@ -32,6 +32,22 @@ std::shared_ptr<arrow::Table> Table__from_dataframe(DataFrame tbl) {
   return out;
 }
 
+// [[Rcpp::export]]
+std::shared_ptr<arrow::Table> Table__FromRecordBatches(List_ lst_batches) {
+  auto batches = arrow::r::list_to_shared_ptr_vector<arrow::RecordBatch>(lst_batches);
+  std::shared_ptr<arrow::Table> out;
+  STOP_IF_NOT_OK(arrow::Table::FromRecordBatches(std::move(batches), &out));
+  return out;
+}
+
+// [[Rcpp::export]]
+std::shared_ptr<arrow::Table> Table__FromRecordBatches_Schema(List_ lst_batches, const std::shared_ptr<arrow::Schema>& schema) {
+  auto batches = arrow::r::list_to_shared_ptr_vector<arrow::RecordBatch>(lst_batches);
+  std::shared_ptr<arrow::Table> out;
+  STOP_IF_NOT_OK(arrow::Table::FromRecordBatches(schema, std::move(batches), &out));
+  return out;
+}
+
 // [[Rcpp::export]]
 int Table__num_columns(const std::shared_ptr<arrow::Table>& x) {
   return x->num_columns();

r/tests/testthat/test-Table.R

@@ -68,3 +68,15 @@ test_that("Table cast (ARROW-3741)", {
   expect_equal(tab2$column(0L)$type, int16())
   expect_equal(tab2$column(1L)$type, int64())
 })
+
+test_that("table_from_batches() cast (ARROW-3818)", {
+  d <- tibble::tibble(int = 1:10, dbl = rnorm(10))
+  batches <- purrr::rerun(10, record_batch(d))
+  table <- table_from_batches(!!!batches)
+  expect_is(table, "arrow::Table")
+  expect_equal(as_tibble(table), vctrs::vec_rbind(!!!purrr::rerun(10, d)))
+
+  table <- table_from_batches(!!!batches, schema = schema(int = int32(), dbl = float64()))
+  expect_is(table, "arrow::Table")
+  expect_equal(as_tibble(table), vctrs::vec_rbind(!!!purrr::rerun(10, d)))
+})