ARROW-10522: [R] Allow rename Table and RecordBatch columns with names()

r/NAMESPACE

@@ -18,6 +18,8 @@ S3method("[[",RecordBatch)
 S3method("[[",Schema)
 S3method("[[",Table)
 S3method("[[<-",Table)
+S3method("names<-",RecordBatch)
+S3method("names<-",Table)
 S3method(Ops,Array)
 S3method(Ops,ChunkedArray)
 S3method(Ops,Expression)

r/R/record-batch.R

@@ -75,6 +75,7 @@ RecordBatch <- R6Class("RecordBatch", inherit = ArrowObject,
     column = function(i) RecordBatch__column(self, i),
     column_name = function(i) RecordBatch__column_name(self, i),
     names = function() RecordBatch__names(self),
+    RenameColumns = function(value) RecordBatch__RenameColumns(self, value),
     Equals = function(other, check_metadata = FALSE, ...) {
       inherits(other, "RecordBatch") && RecordBatch__Equals(self, other, isTRUE(check_metadata))
     },
@@ -204,6 +205,9 @@ record_batch <- RecordBatch$create
 #' @export
 names.RecordBatch <- function(x) x$names()
 
+#' @export
+`names<-.RecordBatch` <- function(x, value) x$RenameColumns(value)
+
 #' @importFrom methods as
 #' @export
 `[.RecordBatch` <- function(x, i, j, ..., drop = FALSE) {

r/R/table.R

@@ -96,6 +96,7 @@ Table <- R6Class("Table", inherit = ArrowObject,
       Table__column(self, i)
     },
     ColumnNames = function() Table__ColumnNames(self),
+    RenameColumns = function(value) Table__RenameColumns(self, value),
     GetColumnByName = function(name) {
       assert_is(name, "character")
       assert_that(length(name) == 1)
@@ -257,6 +258,9 @@ dim.Table <- function(x) c(x$num_rows, x$num_columns)
 #' @export
 names.Table <- function(x) x$ColumnNames()
 
+#' @export
+`names<-.Table` <- function(x, value) x$RenameColumns(value)
+
 #' @export
 `[.Table` <- `[.RecordBatch`
 

r/src/recordbatch.cpp

@@ -42,6 +42,22 @@ std::shared_ptr<arrow::Schema> RecordBatch__schema(
   return x->schema();
 }
 
+// [[arrow::export]]
+std::shared_ptr<arrow::RecordBatch> RecordBatch__RenameColumns(
+    const std::shared_ptr<arrow::RecordBatch>& batch,
+    const std::vector<std::string>& names) {
+  int n = batch->num_columns();
+  if (names.size() != static_cast<size_t>(n)) {
+    cpp11::stop("RecordBatch has %d columns but %d names were provided", n, names.size());
+  }
+  std::vector<std::shared_ptr<arrow::Field>> fields(n);
+  for (int i = 0; i < n; i++) {
+    fields[i] = batch->schema()->field(i)->WithName(names[i]);
+  }
+  auto schema = std::make_shared<arrow::Schema>(std::move(fields));
+  return arrow::RecordBatch::Make(schema, batch->num_rows(), batch->columns());
+}
+
 // [[arrow::export]]
 std::shared_ptr<arrow::RecordBatch> RecordBatch__ReplaceSchemaMetadata(
     const std::shared_ptr<arrow::RecordBatch>& x, cpp11::strings metadata) {

r/src/table.cpp

@@ -74,6 +74,12 @@ std::vector<std::string> Table__ColumnNames(const std::shared_ptr<arrow::Table>&
   return table->ColumnNames();
 }
 
+// [[arrow::export]]
+std::shared_ptr<arrow::Table> Table__RenameColumns(
+    const std::shared_ptr<arrow::Table>& table, const std::vector<std::string>& names) {
+  return ValueOrStop(table->RenameColumns(names));
+}
+
 // [[arrow::export]]
 std::shared_ptr<arrow::Table> Table__Slice1(const std::shared_ptr<arrow::Table>& table,
                                             R_xlen_t offset) {

r/tests/testthat/test-RecordBatch.R

@@ -366,3 +366,15 @@ test_that("RecordBatch$Equals(check_metadata)", {
 
   expect_false(rb1$Equals(24)) # Not a RecordBatch
 })
+
+test_that("RecordBatch name assignment", {
+  rb <- record_batch(x = 1:10, y = 1:10)
+  expect_identical(names(rb), c("x", "y"))
+  names(rb) <- c("a", "b")
+  expect_identical(names(rb), c("a", "b"))
+  expect_error(names(rb) <- "f")
+  expect_error(names(rb) <- letters)
+  expect_error(names(rb) <- character(0))
+  expect_error(names(rb) <- NULL)
+  expect_error(names(rb) <- c(TRUE, FALSE))
+})

r/tests/testthat/test-Table.R

@@ -434,3 +434,15 @@ test_that("Table$SelectColumns()", {
   expect_error(tab$SelectColumns(2:4))
   expect_error(tab$SelectColumns(""))
 })
+
+test_that("Table name assignment", {
+  tab <- Table$create(x = 1:10, y = 1:10)
+  expect_identical(names(tab), c("x", "y"))
+  names(tab) <- c("a", "b")
+  expect_identical(names(tab), c("a", "b"))
+  expect_error(names(tab) <- "f")
+  expect_error(names(tab) <- letters)
+  expect_error(names(tab) <- character(0))
+  expect_error(names(tab) <- NULL)
+  expect_error(names(tab) <- c(TRUE, FALSE))
+})