You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by np...@apache.org on 2022/05/26 17:14:45 UTC
[arrow] branch master updated: ARROW-16607: [R] Improve KeyValueMetadata handling
This is an automated email from the ASF dual-hosted git repository.
npr pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new a6025f1571 ARROW-16607: [R] Improve KeyValueMetadata handling
a6025f1571 is described below
commit a6025f15712aa0829aab748a8d3e776f335265cc
Author: Neal Richardson <ne...@gmail.com>
AuthorDate: Thu May 26 13:14:31 2022 -0400
ARROW-16607: [R] Improve KeyValueMetadata handling
* Pushes KVM handling into ExecPlan so that Run() preserves the R metadata we want.
* Also pushes special handling for a kind of collapsed query from collect() into Build().
* Better encapsulate KVM for the the $metadata and $r_metadata so that as a user/developer, you never have to touch the serialize/deserialize functions, you just have a list to work with. This is a slight API change, most noticeable if you were to `print(tab$metadata)`; better is to `print(str(tab$metdata))`.
* Factor out a common utility in r/src for taking cpp11::strings (named character vector) and producing arrow::KeyValueMetadata
The upshot of all of this is that we can push the ExecPlan evaluation into `as_record_batch_reader()`, and all that `collect()` does on top is read the RBR into a Table/data.frame. This means that we can plug dplyr queries into anything else that expects a RecordBatchReader, and it will be (to the maximum extent possible, given the limitations of ExecPlan) streaming, not requiring you to `compute()` and materialize things first.
Closes #13210 from nealrichardson/kvm
Authored-by: Neal Richardson <ne...@gmail.com>
Signed-off-by: Neal Richardson <ne...@gmail.com>
---
r/R/arrow-tabular.R | 17 +----
r/R/arrowExports.R | 4 +-
r/R/dataset-scan.R | 4 +-
r/R/dataset-write.R | 16 +---
r/R/dplyr-collect.R | 31 +-------
r/R/dplyr-group-by.R | 4 +-
r/R/metadata.R | 15 ++--
r/R/query-engine.R | 159 ++++++++++++++++++++++++++-------------
r/R/record-batch-reader.R | 6 +-
r/R/record-batch.R | 2 +-
r/R/schema.R | 26 ++++---
r/R/table.R | 2 +-
r/src/arrowExports.cpp | 9 ++-
r/src/compute-exec.cpp | 22 +++---
r/src/schema.cpp | 13 ++--
r/tests/testthat/test-metadata.R | 7 +-
16 files changed, 177 insertions(+), 160 deletions(-)
diff --git a/r/R/arrow-tabular.R b/r/R/arrow-tabular.R
index 43110ccf24..58a604ba61 100644
--- a/r/R/arrow-tabular.R
+++ b/r/R/arrow-tabular.R
@@ -70,7 +70,6 @@ ArrowTabular <- R6Class("ArrowTabular",
self$schema$metadata
} else {
# Set the metadata
- new <- prepare_key_value_metadata(new)
out <- self$ReplaceSchemaMetadata(new)
# ReplaceSchemaMetadata returns a new object but we're modifying in place,
# so swap in that new C++ object pointer into our R6 object
@@ -82,16 +81,10 @@ ArrowTabular <- R6Class("ArrowTabular",
# Helper for the R metadata that handles the serialization
# See also method on Schema
if (missing(new)) {
- out <- self$metadata$r
- if (!is.null(out)) {
- # Can't unserialize NULL
- out <- .unserialize_arrow_r_metadata(out)
- }
- # Returns either NULL or a named list
- out
+ self$metadata$r
} else {
# Set the R metadata
- self$metadata$r <- .serialize_arrow_r_metadata(new)
+ self$metadata$r <- new
self
}
}
@@ -101,11 +94,7 @@ ArrowTabular <- R6Class("ArrowTabular",
#' @export
as.data.frame.ArrowTabular <- function(x, row.names = NULL, optional = FALSE, ...) {
df <- x$to_data_frame()
-
- if (!is.null(r_metadata <- x$metadata$r)) {
- df <- apply_arrow_r_metadata(df, .unserialize_arrow_r_metadata(r_metadata))
- }
- df
+ apply_arrow_r_metadata(df, x$metadata$r)
}
#' @export
diff --git a/r/R/arrowExports.R b/r/R/arrowExports.R
index 3414c9b21c..8ad56f227f 100644
--- a/r/R/arrowExports.R
+++ b/r/R/arrowExports.R
@@ -404,8 +404,8 @@ ExecPlan_create <- function(use_threads) {
.Call(`_arrow_ExecPlan_create`, use_threads)
}
-ExecPlan_run <- function(plan, final_node, sort_options, head) {
- .Call(`_arrow_ExecPlan_run`, plan, final_node, sort_options, head)
+ExecPlan_run <- function(plan, final_node, sort_options, metadata, head) {
+ .Call(`_arrow_ExecPlan_run`, plan, final_node, sort_options, metadata, head)
}
ExecPlan_StopProducing <- function(plan) {
diff --git a/r/R/dataset-scan.R b/r/R/dataset-scan.R
index 72f9dec276..a8da1fb60d 100644
--- a/r/R/dataset-scan.R
+++ b/r/R/dataset-scan.R
@@ -206,10 +206,8 @@ map_batches <- function(X, FUN, ..., .data.frame = NULL) {
call. = FALSE
)
}
- plan <- ExecPlan$create()
- final_node <- plan$Build(as_adq(X))
- reader <- plan$Run(final_node)
FUN <- as_mapper(FUN)
+ reader <- as_record_batch_reader(X)
# TODO: for future consideration
# * Move eval to C++ and make it a generator so it can stream, not block
diff --git a/r/R/dataset-write.R b/r/R/dataset-write.R
index bcb90a7ffc..041adc7afd 100644
--- a/r/R/dataset-write.R
+++ b/r/R/dataset-write.R
@@ -162,7 +162,7 @@ write_dataset <- function(dataset,
plan <- ExecPlan$create()
final_node <- plan$Build(dataset)
- if (!is.null(final_node$sort %||% final_node$head %||% final_node$tail)) {
+ if (!is.null(final_node$extras$sort %||% final_node$extras$head %||% final_node$extras$tail)) {
# Because sorting and topK are only handled in the SinkNode (or in R!),
# they wouldn't get picked up in the WriteNode. So let's Run this ExecPlan
# to capture those, and then create a new plan for writing
@@ -206,20 +206,8 @@ write_dataset <- function(dataset,
validate_positive_int_value(min_rows_per_group)
validate_positive_int_value(max_rows_per_group)
- source_schema <- source_data(dataset)$schema
- # For backwards compatibility with Scanner-based writer (arrow <= 7.0.0):
- # retain metadata from source dataset
- output_schema$metadata <- source_schema$metadata
- new_r_meta <- get_r_metadata_from_old_schema(
- output_schema,
- source_schema,
- drop_attributes = has_aggregation(dataset)
- )
- if (!is.null(new_r_meta)) {
- output_schema$r_metadata <- new_r_meta
- }
plan$Write(
- final_node, prepare_key_value_metadata(output_schema$metadata),
+ final_node,
options, path_and_fs$fs, path_and_fs$path,
partitioning, basename_template,
existing_data_behavior, max_partitions,
diff --git a/r/R/dplyr-collect.R b/r/R/dplyr-collect.R
index 1c875b15b0..7f10ed307e 100644
--- a/r/R/dplyr-collect.R
+++ b/r/R/dplyr-collect.R
@@ -19,18 +19,8 @@
# The following S3 methods are registered on load if dplyr is present
collect.arrow_dplyr_query <- function(x, as_data_frame = TRUE, ...) {
- # head and tail are not ExecNodes, at best we can handle them via sink node
- # so if there are any steps done after head/tail, we need to
- # evaluate the query up to then and then do a new query for the rest
- if (is_collapsed(x) && has_head_tail(x$.data)) {
- x$.data <- as_adq(dplyr::compute(x$.data))$.data
- }
-
- # See query-engine.R for ExecPlan/Nodes
- plan <- ExecPlan$create()
- final_node <- plan$Build(x)
tryCatch(
- tab <- plan$Run(final_node)$read_table(),
+ out <- as_record_batch_reader(x)$read_table(),
# n = 4 because we want the error to show up as being from collect()
# and not handle_csv_read_error()
error = function(e, call = caller_env(n = 4)) {
@@ -38,25 +28,10 @@ collect.arrow_dplyr_query <- function(x, as_data_frame = TRUE, ...) {
}
)
- # TODO(ARROW-16607): move KVM handling into ExecPlan
- if (ncol(tab)) {
- # Apply any column metadata from the original schema, where appropriate
- new_r_metadata <- get_r_metadata_from_old_schema(
- tab$schema,
- source_data(x)$schema,
- drop_attributes = has_aggregation(x)
- )
- if (!is.null(new_r_metadata)) {
- tab$r_metadata <- new_r_metadata
- }
- }
-
if (as_data_frame) {
- df <- as.data.frame(tab)
- restore_dplyr_features(df, x)
- } else {
- restore_dplyr_features(tab, x)
+ out <- as.data.frame(out)
}
+ restore_dplyr_features(out, x)
}
collect.ArrowTabular <- function(x, as_data_frame = TRUE, ...) {
if (as_data_frame) {
diff --git a/r/R/dplyr-group-by.R b/r/R/dplyr-group-by.R
index 1014c20f4d..250dbedb18 100644
--- a/r/R/dplyr-group-by.R
+++ b/r/R/dplyr-group-by.R
@@ -64,7 +64,7 @@ group_vars.arrow_dplyr_query <- function(x) x$group_by_vars
group_vars.Dataset <- function(x) NULL
group_vars.RecordBatchReader <- function(x) NULL
group_vars.ArrowTabular <- function(x) {
- x$r_metadata$attributes$.group_vars
+ x$metadata$r$attributes$.group_vars
}
# the logical literal in the two functions below controls the default value of
@@ -81,6 +81,6 @@ ungroup.arrow_dplyr_query <- function(x, ...) {
}
ungroup.Dataset <- ungroup.RecordBatchReader <- force
ungroup.ArrowTabular <- function(x) {
- x$r_metadata$attributes$.group_vars <- NULL
+ x$metadata$r$attributes$.group_vars <- NULL
x
}
diff --git a/r/R/metadata.R b/r/R/metadata.R
index 147318f559..74080e8f48 100644
--- a/r/R/metadata.R
+++ b/r/R/metadata.R
@@ -35,7 +35,7 @@
rawToChar(out)
}
-.unserialize_arrow_r_metadata <- function(x) {
+.deserialize_arrow_r_metadata <- function(x) {
tryCatch(
expr = {
out <- unserialize(charToRaw(x))
@@ -55,6 +55,9 @@
#' @importFrom rlang trace_back
apply_arrow_r_metadata <- function(x, r_metadata) {
+ if (is.null(r_metadata)) {
+ return(x)
+ }
tryCatch(
expr = {
columns_metadata <- r_metadata$columns
@@ -196,12 +199,10 @@ arrow_attributes <- function(x, only_top_level = FALSE) {
}
}
-get_r_metadata_from_old_schema <- function(new_schema,
- old_schema,
- drop_attributes = FALSE) {
+get_r_metadata_from_old_schema <- function(new_schema, old_schema) {
# TODO: do we care about other (non-R) metadata preservation?
# How would we know if it were meaningful?
- r_meta <- old_schema$r_metadata
+ r_meta <- old_schema$metadata$r
if (!is.null(r_meta)) {
# Filter r_metadata$columns on columns with name _and_ type match
common_names <- intersect(names(r_meta$columns), names(new_schema))
@@ -209,10 +210,6 @@ get_r_metadata_from_old_schema <- function(new_schema,
map_lgl(common_names, ~ old_schema[[.]] == new_schema[[.]])
]
r_meta$columns <- r_meta$columns[keep]
- if (drop_attributes) {
- # dplyr drops top-level attributes if you do summarize
- r_meta$attributes <- NULL
- }
}
r_meta
}
diff --git a/r/R/query-engine.R b/r/R/query-engine.R
index 4dfb5d3e9b..c40c61e98a 100644
--- a/r/R/query-engine.R
+++ b/r/R/query-engine.R
@@ -19,17 +19,17 @@ ExecPlan <- R6Class("ExecPlan",
inherit = ArrowObject,
public = list(
Scan = function(dataset) {
- # Handle arrow_dplyr_query
- if (inherits(dataset, "arrow_dplyr_query")) {
- if (inherits(dataset$.data, "RecordBatchReader")) {
- return(ExecNode_SourceNode(self, dataset$.data))
- } else if (inherits(dataset$.data, "ArrowTabular")) {
- if (inherits(dataset$.data, "RecordBatch")) {
- dataset$.data <- Table$create(dataset$.data)
- }
- return(ExecNode_TableSourceNode(self, dataset$.data))
+ if (inherits(dataset, c("RecordBatchReader", "ArrowTabular"))) {
+ return(self$SourceNode(dataset))
+ } else if (inherits(dataset, "arrow_dplyr_query")) {
+ if (inherits(dataset$.data, c("RecordBatchReader", "ArrowTabular"))) {
+ # There's no predicate pushdown to do, so no need to deal with other
+ # arrow_dplyr_query attributes here. They'll be handled by other
+ # ExecNodes
+ return(self$SourceNode(dataset$.data))
}
-
+ # Else, we're scanning a Dataset, and we need to pull out the filter
+ # and projection (column selection) to push down into the scanner
filter <- dataset$filtered_rows
if (isTRUE(filter)) {
filter <- Expression$scalar(TRUE)
@@ -42,17 +42,30 @@ ExecPlan <- R6Class("ExecPlan",
dataset <- dataset$.data
assert_is(dataset, "Dataset")
} else {
- if (inherits(dataset, "ArrowTabular")) {
- dataset <- InMemoryDataset$create(dataset)
- }
assert_is(dataset, "Dataset")
- # Set some defaults
+ # Just a dataset, not a query, so there's no predicates to push down
+ # so set some defaults
filter <- Expression$scalar(TRUE)
colnames <- names(dataset)
}
- # ScanNode needs the filter to do predicate pushdown and skip partitions,
- # and it needs to know which fields to materialize (and which are unnecessary)
- ExecNode_Scan(self, dataset, filter, colnames %||% character(0))
+
+ out <- ExecNode_Scan(self, dataset, filter, colnames %||% character(0))
+ # Hold onto the source data's schema so we can preserve schema metadata
+ # in the resulting Scan/Write
+ out$extras$source_schema <- dataset$schema
+ out
+ },
+ SourceNode = function(.data) {
+ if (inherits(.data, "RecordBatchReader")) {
+ out <- ExecNode_SourceNode(self, .data)
+ } else {
+ assert_is(.data, "ArrowTabular")
+ out <- ExecNode_TableSourceNode(self, as_arrow_table(.data))
+ }
+ # Hold onto the source data's schema so we can preserve schema metadata
+ # in the resulting Scan/Write
+ out$extras$source_schema <- .data$schema
+ out
},
Build = function(.data) {
# This method takes an arrow_dplyr_query and chains together the
@@ -65,14 +78,24 @@ ExecPlan <- R6Class("ExecPlan",
.data <- ensure_group_vars(.data)
.data <- ensure_arrange_vars(.data) # this sets .data$temp_columns
- if (inherits(.data$.data, "arrow_dplyr_query")) {
- # We have a nested query. Recurse.
- node <- self$Build(.data$.data)
+ if (is_collapsed(.data)) {
+ # We have a nested query.
+ if (has_head_tail(.data$.data)) {
+ # head and tail are not ExecNodes; at best we can handle them via
+ # SinkNode, so if there are any steps done after head/tail, we need to
+ # evaluate the query up to then and then do a new query for the rest.
+ # as_record_batch_reader() will build and run an ExecPlan
+ node <- self$SourceNode(as_record_batch_reader(.data$.data))
+ } else {
+ # Recurse
+ node <- self$Build(.data$.data)
+ }
} else {
node <- self$Scan(.data)
}
- # ARROW-13498: Even though Scan takes the filter, apparently we have to do it again
+ # ARROW-13498: Even though Scan takes the filter (if you have a Dataset),
+ # we have to do it again
if (inherits(.data$filtered_rows, "Expression")) {
node <- node$Filter(.data$filtered_rows)
}
@@ -108,7 +131,7 @@ ExecPlan <- R6Class("ExecPlan",
if (getOption("arrow.summarise.sort", FALSE)) {
# Add sorting instructions for the rows too to match dplyr
# (see below about why sorting isn't itself a Node)
- node$sort <- list(
+ node$extras$sort <- list(
names = group_vars,
orders = rep(0L, length(group_vars))
)
@@ -140,7 +163,7 @@ ExecPlan <- R6Class("ExecPlan",
}
if (!is.null(.data$union_all)) {
- node <- node$UnionAll(self$Build(.data$union_all$right_data))
+ node <- node$Union(self$Build(.data$union_all$right_data))
}
}
@@ -150,7 +173,7 @@ ExecPlan <- R6Class("ExecPlan",
# (1) arrange > summarize > arrange
# (2) ARROW-13779: arrange then operation where order matters (e.g. cumsum)
if (length(.data$arrange_vars)) {
- node$sort <- list(
+ node$extras$sort <- list(
names = names(.data$arrange_vars),
orders = .data$arrange_desc,
temp_columns = names(.data$temp_columns)
@@ -160,10 +183,10 @@ ExecPlan <- R6Class("ExecPlan",
# This is only safe because we are going to evaluate queries that end
# with head/tail first, then evaluate any subsequent query as a new query
if (!is.null(.data$head)) {
- node$head <- .data$head
+ node$extras$head <- .data$head
}
if (!is.null(.data$tail)) {
- node$tail <- .data$tail
+ node$extras$tail <- .data$tail
}
node
},
@@ -172,33 +195,41 @@ ExecPlan <- R6Class("ExecPlan",
# Sorting and head/tail (if sorted) are handled in the SinkNode,
# created in ExecPlan_run
- sorting <- node$sort %||% list()
- select_k <- node$head %||% -1L
+ sorting <- node$extras$sort %||% list()
+ select_k <- node$extras$head %||% -1L
has_sorting <- length(sorting) > 0
if (has_sorting) {
- if (!is.null(node$tail)) {
+ if (!is.null(node$extras$tail)) {
# Reverse the sort order and take the top K, then after we'll reverse
# the resulting rows so that it is ordered as expected
sorting$orders <- !sorting$orders
- select_k <- node$tail
+ select_k <- node$extras$tail
}
sorting$orders <- as.integer(sorting$orders)
}
- out <- ExecPlan_run(self, node, sorting, select_k)
+ out <- ExecPlan_run(
+ self,
+ node,
+ sorting,
+ prepare_key_value_metadata(node$final_metadata()),
+ select_k
+ )
if (!has_sorting) {
# Since ExecPlans don't scan in deterministic order, head/tail are both
# essentially taking a random slice from somewhere in the dataset.
# And since the head() implementation is way more efficient than tail(),
# just use it to take the random slice
- slice_size <- node$head %||% node$tail
+ # TODO(ARROW-16628): handle limit in ExecNode
+ slice_size <- node$extras$head %||% node$extras$tail
if (!is.null(slice_size)) {
out <- head(out, slice_size)
}
# Can we now tell `self$Stop()` to StopProducing? We already have
# everything we need for the head (but it seems to segfault: ARROW-14329)
- } else if (!is.null(node$tail)) {
+ } else if (!is.null(node$extras$tail)) {
+ # TODO(ARROW-16630): proper BottomK support
# Reverse the row order to get back what we expect
out <- out$read_table()
out <- out[rev(seq_len(nrow(out))), , drop = FALSE]
@@ -209,9 +240,10 @@ ExecPlan <- R6Class("ExecPlan",
# If arrange() created $temp_columns, make sure to omit them from the result
# We can't currently handle this in ExecPlan_run itself because sorting
# happens in the end (SinkNode) so nothing comes after it.
- if (length(node$sort$temp_columns) > 0) {
+ # TODO(ARROW-16631): move into ExecPlan
+ if (length(node$extras$sort$temp_columns) > 0) {
tab <- out$read_table()
- tab <- tab[, setdiff(names(tab), node$sort$temp_columns), drop = FALSE]
+ tab <- tab[, setdiff(names(tab), node$extras$sort$temp_columns), drop = FALSE]
out <- as_record_batch_reader(tab)
}
@@ -219,7 +251,12 @@ ExecPlan <- R6Class("ExecPlan",
},
Write = function(node, ...) {
# TODO(ARROW-16200): take FileSystemDatasetWriteOptions not ...
- ExecPlan_Write(self, node, ...)
+ ExecPlan_Write(
+ self,
+ node,
+ prepare_key_value_metadata(node$final_metadata()),
+ ...
+ )
},
Stop = function() ExecPlan_StopProducing(self)
)
@@ -231,38 +268,52 @@ ExecPlan$create <- function(use_threads = option_use_threads()) {
ExecNode <- R6Class("ExecNode",
inherit = ArrowObject,
public = list(
- # `sort` is a slight hack to be able to keep around arrange() params,
- # which don't currently yield their own ExecNode but rather are consumed
- # in the SinkNode (in ExecPlan$run())
- sort = NULL,
- # Similar hacks for head and tail
- head = NULL,
- tail = NULL,
- preserve_sort = function(new_node) {
- new_node$sort <- self$sort
- new_node$head <- self$head
- new_node$tail <- self$tail
+ extras = list(
+ # `sort` is a slight hack to be able to keep around arrange() params,
+ # which don't currently yield their own ExecNode but rather are consumed
+ # in the SinkNode (in ExecPlan$run())
+ sort = NULL,
+ # Similar hacks for head and tail
+ head = NULL,
+ tail = NULL,
+ # `source_schema` is put here in Scan() so that at Run/Write, we can
+ # extract the relevant metadata and keep it in the result
+ source_schema = NULL
+ ),
+ preserve_extras = function(new_node) {
+ new_node$extras <- self$extras
new_node
},
+ final_metadata = function() {
+ # Copy metadata from source schema and trim R column metadata to match
+ # which columns are included in the result
+ old_schema <- self$extras$source_schema
+ old_meta <- old_schema$metadata
+ old_meta$r <- get_r_metadata_from_old_schema(self$schema, old_schema)
+ old_meta
+ },
Project = function(cols) {
if (length(cols)) {
assert_is_list_of(cols, "Expression")
- self$preserve_sort(ExecNode_Project(self, cols, names(cols)))
+ self$preserve_extras(ExecNode_Project(self, cols, names(cols)))
} else {
- self$preserve_sort(ExecNode_Project(self, character(0), character(0)))
+ self$preserve_extras(ExecNode_Project(self, character(0), character(0)))
}
},
Filter = function(expr) {
assert_is(expr, "Expression")
- self$preserve_sort(ExecNode_Filter(self, expr))
+ self$preserve_extras(ExecNode_Filter(self, expr))
},
Aggregate = function(options, target_names, out_field_names, key_names) {
- self$preserve_sort(
+ out <- self$preserve_extras(
ExecNode_Aggregate(self, options, target_names, out_field_names, key_names)
)
+ # dplyr drops top-level attributes when you call summarize()
+ out$extras$source_schema$metadata[["r"]]$attributes <- NULL
+ out
},
Join = function(type, right_node, by, left_output, right_output, left_suffix, right_suffix) {
- self$preserve_sort(
+ self$preserve_extras(
ExecNode_Join(
self,
type,
@@ -276,8 +327,8 @@ ExecNode <- R6Class("ExecNode",
)
)
},
- UnionAll = function(right_node) {
- self$preserve_sort(ExecNode_Union(self, right_node))
+ Union = function(right_node) {
+ self$preserve_extras(ExecNode_Union(self, right_node))
}
),
active = list(
diff --git a/r/R/record-batch-reader.R b/r/R/record-batch-reader.R
index 620be4c506..8f6a600dfb 100644
--- a/r/R/record-batch-reader.R
+++ b/r/R/record-batch-reader.R
@@ -237,8 +237,10 @@ as_record_batch_reader.Dataset <- function(x, ...) {
#' @rdname as_record_batch_reader
#' @export
as_record_batch_reader.arrow_dplyr_query <- function(x, ...) {
- # TODO(ARROW-16607): use ExecPlan directly when it handles metadata
- as_record_batch_reader(compute.arrow_dplyr_query(x))
+ # See query-engine.R for ExecPlan/Nodes
+ plan <- ExecPlan$create()
+ final_node <- plan$Build(x)
+ plan$Run(final_node)
}
#' @rdname as_record_batch_reader
diff --git a/r/R/record-batch.R b/r/R/record-batch.R
index c8579c1124..2a924afa56 100644
--- a/r/R/record-batch.R
+++ b/r/R/record-batch.R
@@ -102,7 +102,7 @@ RecordBatch <- R6Class("RecordBatch",
},
RemoveColumn = function(i) RecordBatch__RemoveColumn(self, i),
ReplaceSchemaMetadata = function(new) {
- RecordBatch__ReplaceSchemaMetadata(self, new)
+ RecordBatch__ReplaceSchemaMetadata(self, prepare_key_value_metadata(new))
},
Slice = function(offset, length = NULL) {
if (is.null(length)) {
diff --git a/r/R/schema.R b/r/R/schema.R
index 0a1245a34d..86a968b500 100644
--- a/r/R/schema.R
+++ b/r/R/schema.R
@@ -139,9 +139,18 @@ Schema <- R6Class("Schema",
num_fields = function() Schema__num_fields(self),
fields = function() Schema__fields(self),
HasMetadata = function() Schema__HasMetadata(self),
+ raw_metadata = function() {
+ # This is the named list of strings
+ Schema__metadata(self)
+ },
metadata = function(new_metadata) {
if (missing(new_metadata)) {
- Schema__metadata(self)
+ out <- self$raw_metadata
+ if (!is.null(out[["r"]])) {
+ # Can't unserialize NULL
+ out[["r"]] <- .deserialize_arrow_r_metadata(out[["r"]])
+ }
+ out
} else {
# Set the metadata
out <- self$WithMetadata(new_metadata)
@@ -155,16 +164,10 @@ Schema <- R6Class("Schema",
# Helper for the R metadata that handles the serialization
# See also method on ArrowTabular
if (missing(new)) {
- out <- self$metadata$r
- if (!is.null(out)) {
- # Can't unserialize NULL
- out <- .unserialize_arrow_r_metadata(out)
- }
- # Returns either NULL or a named list
- out
+ self$metadata$r
} else {
# Set the R metadata
- self$metadata$r <- .serialize_arrow_r_metadata(new)
+ self$metadata$r <- new
self
}
}
@@ -200,6 +203,9 @@ prepare_key_value_metadata <- function(metadata) {
call. = FALSE
)
}
+ if (is.list(metadata[["r"]])) {
+ metadata[["r"]] <- .serialize_arrow_r_metadata(metadata[["r"]])
+ }
map_chr(metadata, as.character)
}
@@ -347,7 +353,7 @@ unify_schemas <- function(..., schemas = list(...)) {
#' @export
print.arrow_r_metadata <- function(x, ...) {
utils::str(x)
- utils::str(.unserialize_arrow_r_metadata(x))
+ utils::str(.deserialize_arrow_r_metadata(x))
invisible(x)
}
diff --git a/r/R/table.R b/r/R/table.R
index c9789b8df8..305f305129 100644
--- a/r/R/table.R
+++ b/r/R/table.R
@@ -93,7 +93,7 @@ Table <- R6Class("Table",
AddColumn = function(i, new_field, value) Table__AddColumn(self, i, new_field, value),
SetColumn = function(i, new_field, value) Table__SetColumn(self, i, new_field, value),
ReplaceSchemaMetadata = function(new) {
- Table__ReplaceSchemaMetadata(self, new)
+ Table__ReplaceSchemaMetadata(self, prepare_key_value_metadata(new))
},
field = function(i) Table__field(self, i),
serialize = function(output_stream, ...) write_table(self, output_stream, ...),
diff --git a/r/src/arrowExports.cpp b/r/src/arrowExports.cpp
index 1cfc2ddf2c..2e3cf0f605 100644
--- a/r/src/arrowExports.cpp
+++ b/r/src/arrowExports.cpp
@@ -869,14 +869,15 @@ BEGIN_CPP11
END_CPP11
}
// compute-exec.cpp
-std::shared_ptr<arrow::RecordBatchReader> ExecPlan_run(const std::shared_ptr<compute::ExecPlan>& plan, const std::shared_ptr<compute::ExecNode>& final_node, cpp11::list sort_options, int64_t head);
-extern "C" SEXP _arrow_ExecPlan_run(SEXP plan_sexp, SEXP final_node_sexp, SEXP sort_options_sexp, SEXP head_sexp){
+std::shared_ptr<arrow::RecordBatchReader> ExecPlan_run(const std::shared_ptr<compute::ExecPlan>& plan, const std::shared_ptr<compute::ExecNode>& final_node, cpp11::list sort_options, cpp11::strings metadata, int64_t head);
+extern "C" SEXP _arrow_ExecPlan_run(SEXP plan_sexp, SEXP final_node_sexp, SEXP sort_options_sexp, SEXP metadata_sexp, SEXP head_sexp){
BEGIN_CPP11
arrow::r::Input<const std::shared_ptr<compute::ExecPlan>&>::type plan(plan_sexp);
arrow::r::Input<const std::shared_ptr<compute::ExecNode>&>::type final_node(final_node_sexp);
arrow::r::Input<cpp11::list>::type sort_options(sort_options_sexp);
+ arrow::r::Input<cpp11::strings>::type metadata(metadata_sexp);
arrow::r::Input<int64_t>::type head(head_sexp);
- return cpp11::as_sexp(ExecPlan_run(plan, final_node, sort_options, head));
+ return cpp11::as_sexp(ExecPlan_run(plan, final_node, sort_options, metadata, head));
END_CPP11
}
// compute-exec.cpp
@@ -5212,7 +5213,7 @@ static const R_CallMethodDef CallEntries[] = {
{ "_arrow_io___CompressedOutputStream__Make", (DL_FUNC) &_arrow_io___CompressedOutputStream__Make, 2},
{ "_arrow_io___CompressedInputStream__Make", (DL_FUNC) &_arrow_io___CompressedInputStream__Make, 2},
{ "_arrow_ExecPlan_create", (DL_FUNC) &_arrow_ExecPlan_create, 1},
- { "_arrow_ExecPlan_run", (DL_FUNC) &_arrow_ExecPlan_run, 4},
+ { "_arrow_ExecPlan_run", (DL_FUNC) &_arrow_ExecPlan_run, 5},
{ "_arrow_ExecPlan_StopProducing", (DL_FUNC) &_arrow_ExecPlan_StopProducing, 1},
{ "_arrow_ExecNode_output_schema", (DL_FUNC) &_arrow_ExecNode_output_schema, 1},
{ "_arrow_ExecNode_Scan", (DL_FUNC) &_arrow_ExecNode_Scan, 4},
diff --git a/r/src/compute-exec.cpp b/r/src/compute-exec.cpp
index d18c47260b..d9eaf50cd7 100644
--- a/r/src/compute-exec.cpp
+++ b/r/src/compute-exec.cpp
@@ -34,6 +34,8 @@ namespace compute = ::arrow::compute;
std::shared_ptr<compute::FunctionOptions> make_compute_options(std::string func_name,
cpp11::list options);
+std::shared_ptr<arrow::KeyValueMetadata> strings_to_kvm(cpp11::strings metadata);
+
// [[arrow::export]]
std::shared_ptr<compute::ExecPlan> ExecPlan_create(bool use_threads) {
static compute::ExecContext threaded_context{gc_memory_pool(),
@@ -57,7 +59,7 @@ std::shared_ptr<compute::ExecNode> MakeExecNodeOrStop(
std::shared_ptr<arrow::RecordBatchReader> ExecPlan_run(
const std::shared_ptr<compute::ExecPlan>& plan,
const std::shared_ptr<compute::ExecNode>& final_node, cpp11::list sort_options,
- int64_t head = -1) {
+ cpp11::strings metadata, int64_t head = -1) {
// For now, don't require R to construct SinkNodes.
// Instead, just pass the node we should collect as an argument.
arrow::AsyncGenerator<arrow::util::optional<compute::ExecBatch>> sink_gen;
@@ -101,9 +103,15 @@ std::shared_ptr<arrow::RecordBatchReader> ExecPlan_run(
}
}};
+ // Attach metadata to the schema
+ auto out_schema = final_node->output_schema();
+ if (metadata.size() > 0) {
+ auto kv = strings_to_kvm(metadata);
+ out_schema = out_schema->WithMetadata(kv);
+ }
return compute::MakeGeneratorReader(
- final_node->output_schema(),
- [stop_producing, plan, sink_gen] { return sink_gen(); }, gc_memory_pool());
+ out_schema, [stop_producing, plan, sink_gen] { return sink_gen(); },
+ gc_memory_pool());
}
// [[arrow::export]]
@@ -183,13 +191,7 @@ void ExecPlan_Write(
opts.min_rows_per_group = min_rows_per_group;
opts.max_rows_per_group = max_rows_per_group;
- // TODO: factor this out to a strings_to_KVM() helper
- auto values = cpp11::as_cpp<std::vector<std::string>>(metadata);
- auto names = cpp11::as_cpp<std::vector<std::string>>(metadata.attr("names"));
-
- auto kv =
- std::make_shared<arrow::KeyValueMetadata>(std::move(names), std::move(values));
-
+ auto kv = strings_to_kvm(metadata);
MakeExecNodeOrStop("write", final_node->plan(), {final_node.get()},
ds::WriteNodeOptions{std::move(opts), std::move(kv)});
diff --git a/r/src/schema.cpp b/r/src/schema.cpp
index 6a49431d8e..2bc58f0fa3 100644
--- a/r/src/schema.cpp
+++ b/r/src/schema.cpp
@@ -112,14 +112,17 @@ cpp11::writable::list Schema__metadata(const std::shared_ptr<arrow::Schema>& sch
return out;
}
-// [[arrow::export]]
-std::shared_ptr<arrow::Schema> Schema__WithMetadata(
- const std::shared_ptr<arrow::Schema>& schema, cpp11::strings metadata) {
+std::shared_ptr<arrow::KeyValueMetadata> strings_to_kvm(cpp11::strings metadata) {
auto values = cpp11::as_cpp<std::vector<std::string>>(metadata);
auto names = cpp11::as_cpp<std::vector<std::string>>(metadata.attr("names"));
- auto kv =
- std::make_shared<arrow::KeyValueMetadata>(std::move(names), std::move(values));
+ return std::make_shared<arrow::KeyValueMetadata>(std::move(names), std::move(values));
+}
+
+// [[arrow::export]]
+std::shared_ptr<arrow::Schema> Schema__WithMetadata(
+ const std::shared_ptr<arrow::Schema>& schema, cpp11::strings metadata) {
+ auto kv = strings_to_kvm(metadata);
return schema->WithMetadata(std::move(kv));
}
diff --git a/r/tests/testthat/test-metadata.R b/r/tests/testthat/test-metadata.R
index 986b1d802c..69475ccd39 100644
--- a/r/tests/testthat/test-metadata.R
+++ b/r/tests/testthat/test-metadata.R
@@ -52,7 +52,11 @@ test_that("Table metadata", {
test_that("Table R metadata", {
tab <- Table$create(example_with_metadata)
- expect_output(print(tab$metadata), "arrow_r_metadata")
+ expect_output(
+ print(tab$metadata),
+ "$r$columns$c$columns$c1$attributes$extra_attr",
+ fixed = TRUE
+ )
expect_identical(as.data.frame(tab), example_with_metadata)
})
@@ -96,6 +100,7 @@ test_that("Garbage R metadata doesn't break things", {
)
# serialize data like .serialize_arrow_r_metadata does, but don't call that
# directly since it checks to ensure that the data is a list
+ tab <- Table$create(example_data[1:6])
tab$metadata$r <- rawToChar(serialize("garbage", NULL, ascii = TRUE))
expect_warning(
expect_identical(as.data.frame(tab), example_data[1:6]),