diff --git a/r/R/record-batch-writer.R b/r/R/record-batch-writer.R index 6f43c52d96a..a1cfd8a0040 100644 --- a/r/R/record-batch-writer.R +++ b/r/R/record-batch-writer.R @@ -58,12 +58,11 @@ RecordBatchWriter <- R6Class("RecordBatchWriter", inherit = Object, write_table = function(table) ipc___RecordBatchWriter__WriteTable(self, table), write = function(x) { + x <- to_arrow(x) if (inherits(x, "RecordBatch")) { self$write_batch(x) } else if (inherits(x, "Table")) { self$write_table(x) - } else if (inherits(x, "data.frame")) { - self$write_table(table(x)) } else { abort("unexpected type for RecordBatchWriter$write(), must be an arrow::RecordBatch or an arrow::Table") } diff --git a/r/tests/testthat/test-record-batch-reader.R b/r/tests/testthat/test-record-batch-reader.R index b557f0669ac..2851538010b 100644 --- a/r/tests/testthat/test-record-batch-reader.R +++ b/r/tests/testthat/test-record-batch-reader.R @@ -18,16 +18,20 @@ context("RecordBatch.*(Reader|Writer)") test_that("RecordBatchStreamReader / Writer", { - batch <- record_batch( + tbl <- tibble::tibble( x = 1:10, y = letters[1:10] ) + batch <- record_batch(tbl) + tab <- Table$create(tbl) sink <- BufferOutputStream$create() expect_equal(sink$tell(), 0) writer <- RecordBatchStreamWriter$create(sink, batch$schema) expect_is(writer, "RecordBatchStreamWriter") - writer$write_batch(batch) + writer$write(batch) + writer$write(tab) + writer$write(tbl) expect_true(sink$tell() > 0) writer$close() @@ -40,20 +44,22 @@ test_that("RecordBatchStreamReader / Writer", { batch1 <- reader$read_next_batch() expect_is(batch1, "RecordBatch") expect_equal(batch, batch1) - + batch2 <- reader$read_next_batch() + expect_is(batch2, "RecordBatch") + expect_equal(batch, batch2) + batch3 <- reader$read_next_batch() + expect_is(batch3, "RecordBatch") + expect_equal(batch, batch3) expect_null(reader$read_next_batch()) }) test_that("RecordBatchFileReader / Writer", { - batch <- record_batch( - x = 1:10, - y = letters[1:10] - ) - sink <- BufferOutputStream$create() writer <- RecordBatchFileWriter$create(sink, batch$schema) expect_is(writer, "RecordBatchFileWriter") - writer$write_batch(batch) + writer$write(batch) + writer$write(tab) + writer$write(tbl) writer$close() buf <- sink$getvalue() @@ -66,5 +72,5 @@ test_that("RecordBatchFileReader / Writer", { expect_is(batch1, "RecordBatch") expect_equal(batch, batch1) - expect_equal(reader$num_record_batches, 1) + expect_equal(reader$num_record_batches, 3) })