Skip to content

Commit 906331c

Browse files
ARROW-11704: [R] Wire up dplyr::mutate() for datasets
Closes apache#9586 from nealrichardson/r-dataset-projection Authored-by: Neal Richardson <neal.p.richardson@gmail.com> Signed-off-by: Neal Richardson <neal.p.richardson@gmail.com>
1 parent e2c7d95 commit 906331c

8 files changed

Lines changed: 157 additions & 23 deletions

File tree

‎r/NEWS.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
## dplyr methods
2323

24-
* `dplyr::mutate()` on Arrow `Table` and `RecordBatch` is now supported in Arrow for many applications. Where not yet supported, the implementation falls back to pulling data into an R `data.frame` first.
24+
* `dplyr::mutate()` is now supported in Arrow for many applications. For queries on `Table` and `RecordBatch` that are not yet supported in Arrow, the implementation falls back to pulling data into an R `data.frame` first, as in the previous release. For queries on `Dataset`, it raises an error if the feature is not implemented.
2525
* String functions `nchar()`, `tolower()`, and `toupper()`, along with their `stringr` spellings `str_length()`, `str_to_lower()`, and `str_to_upper()`, are supported in Arrow `dplyr` calls. `str_trim()` is also supported.
2626

2727
## Other improvements

‎r/R/arrowExports.R‎

Lines changed: 6 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎r/R/dataset-scan.R‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -157,13 +157,19 @@ ScannerBuilder <- R6Class("ScannerBuilder", inherit = ArrowObject,
157157
public = list(
158158
Project = function(cols) {
159159
# cols is either a character vector or a named list of Expressions
160-
if (!is.character(cols)) {
161-
# We don't yet support mutate() on datasets, so this is just a list
162-
# of FieldRefs, and we need to back out the field names
163-
cols <- get_field_names(cols)
160+
if (is.character(cols)) {
161+
dataset___ScannerBuilder__ProjectNames(self, cols)
162+
} else {
163+
# If we have expressions, but they all turn out to be field_refs,
164+
# we can still call the simple method
165+
field_names <- get_field_names(cols)
166+
if (all(nzchar(field_names))) {
167+
dataset___ScannerBuilder__ProjectNames(self, field_names)
168+
} else {
169+
# Else, we are projecting/mutating
170+
dataset___ScannerBuilder__ProjectExprs(self, cols, names(cols))
171+
}
164172
}
165-
assert_is(cols, "character")
166-
dataset___ScannerBuilder__Project(self, cols)
167173
self
168174
},
169175
Filter = function(expr) {

‎r/R/dplyr.R‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -506,9 +506,6 @@ mutate.arrow_dplyr_query <- function(.data,
506506
}
507507

508508
.data <- arrow_dplyr_query(.data)
509-
if (query_on_dataset(.data)) {
510-
not_implemented_for_dataset("mutate()")
511-
}
512509

513510
.keep <- match.arg(.keep)
514511
.before <- enquo(.before)
@@ -529,6 +526,7 @@ mutate.arrow_dplyr_query <- function(.data,
529526
# Deparse and take the first element in case they're long expressions
530527
names(exprs)[unnamed] <- map_chr(exprs[unnamed], as_label)
531528

529+
is_dataset <- query_on_dataset(.data)
532530
mask <- arrow_mask(.data)
533531
results <- list()
534532
for (i in seq_along(exprs)) {
@@ -539,6 +537,15 @@ mutate.arrow_dplyr_query <- function(.data,
539537
if (inherits(results[[new_var]], "try-error")) {
540538
msg <- paste('Expression', as_label(exprs[[i]]), 'not supported in Arrow')
541539
return(abandon_ship(call, .data, msg))
540+
} else if (is_dataset &&
541+
!inherits(results[[new_var]], "Expression") &&
542+
!is.null(results[[new_var]])) {
543+
# We need some wrapping to handle literal values
544+
if (length(results[[new_var]]) != 1) {
545+
msg <- paste0('In ', new_var, " = ", as_label(exprs[[i]]), ", only values of size one are recycled")
546+
return(abandon_ship(call, .data, msg))
547+
}
548+
results[[new_var]] <- Expression$scalar(results[[new_var]])
542549
}
543550
# Put it in the data mask too
544551
mask[[new_var]] <- mask$.data[[new_var]] <- results[[new_var]]
@@ -583,7 +590,7 @@ abandon_ship <- function(call, .data, msg = NULL) {
583590
# Default message: function not implemented
584591
not_implemented_for_dataset(paste0(dplyr_fun_name, "()"))
585592
} else {
586-
stop(msg, call. = FALSE)
593+
stop(msg, "\nCall collect() first to pull data into R.", call. = FALSE)
587594
}
588595
}
589596

‎r/src/arrowExports.cpp‎

Lines changed: 16 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎r/src/dataset.cpp‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -308,11 +308,24 @@ std::shared_ptr<ds::PartitioningFactory> dataset___HivePartitioning__MakeFactory
308308
// ScannerBuilder, Scanner
309309

310310
// [[arrow::export]]
311-
void dataset___ScannerBuilder__Project(const std::shared_ptr<ds::ScannerBuilder>& sb,
312-
const std::vector<std::string>& cols) {
311+
void dataset___ScannerBuilder__ProjectNames(const std::shared_ptr<ds::ScannerBuilder>& sb,
312+
const std::vector<std::string>& cols) {
313313
StopIfNotOk(sb->Project(cols));
314314
}
315315

316+
// [[arrow::export]]
317+
void dataset___ScannerBuilder__ProjectExprs(
318+
const std::shared_ptr<ds::ScannerBuilder>& sb,
319+
const std::vector<std::shared_ptr<ds::Expression>>& exprs,
320+
const std::vector<std::string>& names) {
321+
// We have shared_ptrs of expressions but need the Expressions
322+
std::vector<ds::Expression> expressions;
323+
for (auto expr : exprs) {
324+
expressions.push_back(*expr);
325+
}
326+
StopIfNotOk(sb->Project(expressions, names));
327+
}
328+
316329
// [[arrow::export]]
317330
void dataset___ScannerBuilder__Filter(const std::shared_ptr<ds::ScannerBuilder>& sb,
318331
const std::shared_ptr<ds::Expression>& expr) {

‎r/src/expression.cpp‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,11 @@ std::shared_ptr<ds::Expression> dataset___expr__field_ref(std::string name) {
5050
// [[arrow::export]]
5151
std::string dataset___expr__get_field_ref_name(
5252
const std::shared_ptr<ds::Expression>& ref) {
53-
auto refname = ref->field_ref()->name();
54-
return *refname;
53+
auto field_ref = ref->field_ref();
54+
if (field_ref == nullptr) {
55+
return "";
56+
}
57+
return *field_ref->name();
5558
}
5659

5760
// [[arrow::export]]

‎r/tests/testthat/test-dataset.R‎

Lines changed: 91 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -717,6 +717,96 @@ test_that("filter() with expressions", {
717717
)
718718
})
719719

720+
test_that("mutate()", {
721+
ds <- open_dataset(dataset_dir, partitioning = schema(part = uint8()))
722+
mutated <- ds %>%
723+
select(chr, dbl, int) %>%
724+
filter(dbl * 2 > 14 & dbl - 50 < 3L) %>%
725+
mutate(twice = int * 2)
726+
expect_output(
727+
print(mutated),
728+
"FileSystemDataset (query)
729+
chr: string
730+
dbl: double
731+
int: int32
732+
twice: expr
733+
734+
* Filter: ((multiply_checked(dbl, 2) > 14) and (subtract_checked(dbl, 50) < 3))
735+
See $.data for the source Arrow object",
736+
fixed = TRUE
737+
)
738+
expect_equivalent(
739+
mutated %>%
740+
collect() %>%
741+
arrange(dbl),
742+
rbind(
743+
df1[8:10, c("chr", "dbl", "int")],
744+
df2[1:2, c("chr", "dbl", "int")]
745+
) %>%
746+
mutate(
747+
twice = int * 2
748+
)
749+
)
750+
})
751+
752+
test_that("transmute()", {
753+
ds <- open_dataset(dataset_dir, partitioning = schema(part = uint8()))
754+
mutated <-
755+
expect_equivalent(
756+
ds %>%
757+
select(chr, dbl, int) %>%
758+
filter(dbl * 2 > 14 & dbl - 50 < 3L) %>%
759+
transmute(twice = int * 2) %>%
760+
collect() %>%
761+
arrange(twice),
762+
rbind(
763+
df1[8:10, "int", drop = FALSE],
764+
df2[1:2, "int", drop = FALSE]
765+
) %>%
766+
transmute(
767+
twice = int * 2
768+
)
769+
)
770+
})
771+
772+
test_that("mutate() features not yet implemented", {
773+
expect_error(
774+
ds %>%
775+
group_by(int) %>%
776+
mutate(avg = mean(int)),
777+
"mutate() on grouped data not supported in Arrow\nCall collect() first to pull data into R.",
778+
fixed = TRUE
779+
)
780+
})
781+
782+
783+
test_that("mutate() with scalar (length 1) literal inputs", {
784+
expect_equal(
785+
ds %>%
786+
mutate(the_answer = 42) %>%
787+
collect() %>%
788+
pull(the_answer),
789+
rep(42, nrow(ds))
790+
)
791+
792+
expect_error(
793+
ds %>% mutate(the_answer = c(42, 42)),
794+
"In the_answer = c(42, 42), only values of size one are recycled\nCall collect() first to pull data into R.",
795+
fixed = TRUE
796+
)
797+
})
798+
799+
test_that("mutate() with NULL inputs", {
800+
expect_equal(
801+
ds %>%
802+
mutate(int = NULL) %>%
803+
collect(),
804+
ds %>%
805+
select(-int) %>%
806+
collect()
807+
)
808+
})
809+
720810
test_that("filter scalar validation doesn't crash (ARROW-7772)", {
721811
expect_error(
722812
ds %>%
@@ -832,7 +922,6 @@ test_that("dplyr method not implemented messages", {
832922
expect_error(x, "is not currently implemented for Arrow Datasets")
833923
}
834924
expect_not_implemented(ds %>% arrange(int))
835-
expect_not_implemented(ds %>% mutate(int = int + 2))
836925
expect_not_implemented(ds %>% filter(int == 1) %>% summarize(n()))
837926
})
838927

@@ -1137,7 +1226,7 @@ test_that("Dataset writing: no partitioning", {
11371226
test_that("Dataset writing: partition on null", {
11381227
skip_on_os("windows") # https://issues.apache.org/jira/browse/ARROW-9651
11391228
ds <- open_dataset(hive_dir)
1140-
1229+
11411230
dst_dir <- tempfile()
11421231
partitioning = hive_partition(lgl = boolean())
11431232
write_dataset(ds, dst_dir, partitioning = partitioning)

0 commit comments

Comments
 (0)