attempt RecordBatch$to_file

romainfrancois · romainfrancois · commit f5e3eff490b8 · 2018-09-24T20:15:02.000+02:00
diff --git a/r/DESCRIPTION b/r/DESCRIPTION
@@ -19,7 +19,9 @@ Imports:
     purrr,
     assertthat,
     glue,
-    R6
+    R6,
+    vctrs,
+    fs
 Roxygen: list(markdown = TRUE)
 RoxygenNote: 6.0.1.9000
 Suggests:
diff --git a/r/R/RcppExports.R b/r/R/RcppExports.R
@@ -101,6 +101,10 @@ Table_schema <- function(x) {
     .Call(`_arrow_Table_schema`, x)
 }
 
+RecordBatch_to_file <- function(batch, path) {
+    .Call(`_arrow_RecordBatch_to_file`, batch, path)
+}
+
 Field_initialize <- function(name, type, nullable = TRUE) {
     .Call(`_arrow_Field_initialize`, name, type, nullable)
 }
diff --git a/r/R/array.R b/r/R/array.R
@@ -97,7 +97,8 @@ array <- function(...){
     },
     num_columns = function() RecordBatch_num_columns(self),
     num_rows = function() RecordBatch_num_rows(self),
-    schema = function() schema(.xp = RecordBatch_schema(self))
+    schema = function() schema(.xp = RecordBatch_schema(self)),
+    to_file = function(path) RecordBatch_to_file(self, fs::path_abs(path))
   )
 )
 
diff --git a/r/src/RcppExports.cpp b/r/src/RcppExports.cpp
@@ -288,6 +288,18 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// RecordBatch_to_file
+int RecordBatch_to_file(const std::shared_ptr<arrow::RecordBatch>& batch, std::string path);
+RcppExport SEXP _arrow_RecordBatch_to_file(SEXP batchSEXP, SEXP pathSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const std::shared_ptr<arrow::RecordBatch>& >::type batch(batchSEXP);
+    Rcpp::traits::input_parameter< std::string >::type path(pathSEXP);
+    rcpp_result_gen = Rcpp::wrap(RecordBatch_to_file(batch, path));
+    return rcpp_result_gen;
+END_RCPP
+}
 // Field_initialize
 std::shared_ptr<arrow::Field> Field_initialize(const std::string& name, const std::shared_ptr<arrow::DataType>& type, bool nullable);
 RcppExport SEXP _arrow_Field_initialize(SEXP nameSEXP, SEXP typeSEXP, SEXP nullableSEXP) {
@@ -1598,6 +1610,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_arrow_Table_num_columns", (DL_FUNC) &_arrow_Table_num_columns, 1},
     {"_arrow_Table_num_rows", (DL_FUNC) &_arrow_Table_num_rows, 1},
     {"_arrow_Table_schema", (DL_FUNC) &_arrow_Table_schema, 1},
+    {"_arrow_RecordBatch_to_file", (DL_FUNC) &_arrow_RecordBatch_to_file, 2},
     {"_arrow_Field_initialize", (DL_FUNC) &_arrow_Field_initialize, 3},
     {"_arrow_Field_ToString", (DL_FUNC) &_arrow_Field_ToString, 1},
     {"_arrow_Field_name", (DL_FUNC) &_arrow_Field_name, 1},
diff --git a/r/src/buffer.cpp b/r/src/buffer.cpp
@@ -16,6 +16,8 @@
 // under the License.
 
 #include "arrow_types.h"
+#include <arrow/io/file.h>
+#include <arrow/ipc/writer.h>
 
 using namespace Rcpp;
 using namespace arrow;
@@ -135,3 +137,22 @@ int Table_num_rows(const std::shared_ptr<arrow::Table>& x){
 std::shared_ptr<arrow::Schema> Table_schema(const std::shared_ptr<arrow::Table>& x){
   return x->schema();
 }
+
+// [[Rcpp::export]]
+int RecordBatch_to_file(const std::shared_ptr<arrow::RecordBatch>& batch, std::string path) {
+  std::shared_ptr<arrow::io::OutputStream> stream;
+  auto s = arrow::io::FileOutputStream::Open(path, &stream);
+
+  std::shared_ptr<arrow::ipc::RecordBatchWriter> file_writer;
+  s = arrow::ipc::RecordBatchFileWriter::Open(stream.get(), batch->schema(), &file_writer);
+  s = file_writer->WriteRecordBatch(*batch, true);
+  s = file_writer->Close();
+  int64_t offset;
+  s = stream->Tell(&offset);
+  s = stream->Close();
+  return offset;
+}
+
+
+
+

Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,10 @@ Table_schema <- function(x) {`
`101`	`101`	.Call(`_arrow_Table_schema`, x)
`102`	`102`	`}`
`103`	`103`
	`104`	`+RecordBatch_to_file <- function(batch, path) {`
	`105`	+ .Call(`_arrow_RecordBatch_to_file`, batch, path)
	`106`	`+}`
	`107`	`+`
`104`	`108`	`Field_initialize <- function(name, type, nullable = TRUE) {`
`105`	`109`	.Call(`_arrow_Field_initialize`, name, type, nullable)
`106`	`110`	`}`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,8 @@ array <- function(...){`
`97`	`97`	`},`
`98`	`98`	`num_columns = function() RecordBatch_num_columns(self),`
`99`	`99`	`num_rows = function() RecordBatch_num_rows(self),`
`100`		`- schema = function() schema(.xp = RecordBatch_schema(self))`
	`100`	`+ schema = function() schema(.xp = RecordBatch_schema(self)),`
	`101`	`+ to_file = function(path) RecordBatch_to_file(self, fs::path_abs(path))`
`101`	`102`	`)`
`102`	`103`	`)`
`103`	`104`