Enable writeCSV from RecordBatchReader

apache · Dec 16, 2021 · b45f9e6 · b45f9e6
1 parent b9ad11f
commit b45f9e6
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 1 deletion.
diff --git a/r/R/csv.R b/r/R/csv.R
@@ -692,7 +692,7 @@ write_csv_arrow <- function(x,
     x <- Table$create(x)
   }
 
-  assert_that(is_writable_table(x))
+  assert_that(is_writable_table(x) | inherits(x, "FileSystemDataset"))
 
   if (!inherits(sink, "OutputStream")) {
     sink <- make_output_stream(sink)
@@ -703,6 +703,8 @@ write_csv_arrow <- function(x,
     csv___WriteCSV__RecordBatch(x, write_options, sink)
   } else if (inherits(x, "Table")) {
     csv___WriteCSV__Table(x, write_options, sink)
+  } else if (inherits(x, "FileSystemDataset")) {
+    csv___WriteCSV__RecordBatchReader(x, write_options, sink)
   }
 
   invisible(x_out)

diff --git a/r/src/arrowExports.cpp b/r/src/arrowExports.cpp
diff --git a/r/tests/testthat/test-csv.R b/r/tests/testthat/test-csv.R
@@ -457,3 +457,19 @@ test_that("read_csv_arrow() deals with BOMs (byte-order-marks) correctly", {
     tibble(a = 1, b = 2)
   )
 })
+
+test_that("write_csv_arrow can write from FileSystemDataset objects", {
+
+  skip_if_not_available("dataset")
+  data_dir <- make_temp_dir()
+  write_dataset(tbl_no_dates, data_dir, partitioning = "lgl")
+  data_in <- open_dataset(data_dir)
+
+  csv_file <- tempfile()
+  tbl_out <- write_csv_arrow(data_in, csv_file)
+
+
+  expect_true(file.exists(csv_file))
+  expect_identical(tbl_out, tbl_no_dates)
+
+})