GH-41323: [R] Redo how summarize() evaluates expressions

[nealrichardson]

Rationale for this change

Previously, the NSE expression handling in summarize() worked differently from filter/mutate/etc. Among the implications, it would not have been possible to define bindings for other aggregation functions that can be translated into some combination of supported aggregations, such as weighted.mean().

What changes are included in this PR?

• Expressions in summarize() can now be evaluated with "regular" arrow_eval(). Aggregation bindings stick the contents of the aggregation data they previously returned into an ..aggregations list that lives in an enclosing environment, and then return a FieldRef pointing to that. This makes the code in e.g. summarize_eval() a little harder to follow, since it's grabbing and pointing to objects out of its immediate scope, but I've tried to comment thoroughly and am happy to add more.
• arrow_eval() inspects the expression it receives for any functions that are not in the NSE mask and not in some other package's namespace (i.e. hopefully just user functions) and inserts them into the NSE mask, setting the enclosing environment for that copy of the function to be the mask, so that if the function calls other functions that we do have bindings for, the bindings get called. This is the approach I suggested back in [R] Try to arrow_eval user-defined functions #29667 (comment), and it is what fixes [R] Try to arrow_eval user-defined functions #29667 and [R] Improve user experience when mixing R code with Arrow dplyr pipelines #40938.

Are these changes tested?

Existing tests, which are pretty comprehensive, pass. But it would be good to try to be more evil in manual testing with the user-defined R function support.

Are there any user-facing changes?

Yes.

• GitHub Issue: [R] Redo how summarize() evaluates #41323

[jonkeane]

Thanks for this! I haven't pulled it locally (but should sometime soon). A few comments, mostly around comments

[jonkeane]

I'm curious to know more about going from the ctx object to storing these as ..aggregations. I'm not at all opposed, and think this looks more natural given some of our other machinery — but can't tell directly here if/why that's necessary

[nealrichardson]

I'll explain here, and then I'll add more in comments.

summarize() is complicated because you can do a mixture of scalar operations and aggregations, but that's not how acero works. So we have to pull out the aggregations, collect them in one list (that will become an Aggregate ExecNode), and in the expressions, replace them with FieldRefs so that further operations can happen (in what will become a ProjectNode that works on the result of the Aggregate).

In "normal" arrow_eval, like in mutate(), each expression/quosure results in a single Arrow Expression. But in summarize(), it could generate one or more aggregations that go into Aggregate, and then an Expression after that. Example from the comments in do_arrow_summarize():

  # For example,
  #   summarize(mean = sum(x) / n())
  # is effectively implemented as
  #   summarize(..temp0 = sum(x), ..temp1 = n()) %>%
  #   mutate(mean = ..temp0 / ..temp1) %>%
  #   select(-starts_with("..temp"))

So, each aggregation binding needs to push a ..tempN aggregation onto some list somewhere and return a FieldRef so that any projections happening after evaluate "normally".

ctx was useful previously because we weren't calling arrow_eval() on the expressions, we were walking them, looking for known aggregation functions, pulling them out and inserting into the expression the ..tempN symbols, and then once it was all scalar functions left, calling arrow_eval(). But that doesn't allow for defining a weighted.mean binding as function(x, w) sum(x * w) / sum(w) because you can't just substitute into the call to weighted.mean(x, w), you really just want to evaluate it and let the usual Arrow Expression logic work.

The challenge was: where is that "somewhere" to collect the aggregations? It needs to be somewhere where the binding functions can find it, we can't pass it in as an argument everywhere.

Generally, R looks up symbols in each parent frame of where the function is defined, like this:

> x <- 1
> f <- function() x
> f()
[1] 1

So you'd think that you could just have ..aggregations in this function, and when you evaluate the expressions, they would find them in the enclosing environment. But no. It's like this:

> g <- function() {
+   x <- 2
+   f()
+ }
> g()
[1] 1

f() has its own environment and looks up symbols from its parents.

So, I put ..aggregations in this environment, and then set this environment as the parent for each of the aggregation bindings in arrow_mask(). That way, they can find it and assign into it. This seemed better than the alternatives I almost gave up and fell back to, like something global/in the package namespace.

After doing that, I realized that I could do a similar thing with user functions: copy them into the mask and set their environment to be the mask, so they'd find the other bindings there.

Does that make sense?

[jonkeane]

Aaah yes, got it. I didn't totally put together that ..aggregations wasn't at the package scope, but that's really clever. And because it's transient within the call we don't have to worry about flushing it at the end / cleaning it up / managing state, yeah?

[nealrichardson]

Correct

[nealrichardson]

I added/reworked the comments in this function, and otherwise did some simplification so that I hope the code is more readable too. LMK what you think.

[jonkeane]

Sneaky!

[jonkeane]

Suggested change

	test_that("ARROW-14071 - user-defined R functions", {
	test_that("ARROW-14071 - R functions from a user's environment", {

Just to be super clear this isn't about UDFs

[jonkeane]

Nice, this is a helpful catch / test that honestly I could see being helpful when debugging this too — but makes this super clear that it's hitting this code and not some other path

[thisisnic]

This is cool. A few additional tests it might be worth chucking in to show how it works (or in the docs somewhere?) that are illustrative of a few things I wanted to check:

library(dplyr)
library(arrow)

single_transform <- function(x){
  str_remove_all(x, "[aeiou]")
}

multistep_transform <- function(x){
  y = stringr::str_replace_all(x, "B", "c")
  z = str_remove_all(y, "[aeiou]")
  z2 = str_to_upper(z)
  z2
}

multistep_transform_in_one <- function(x){
  str_to_upper(str_remove_all(stringr::str_replace_all(x, "B", "c"), "[aeiou]"))
}

tibble::tibble(x = c("Foo", "Bar", "Baz", "Qux")) %>%
  arrow_table() %>%
  mutate(y = single_transform(x)) %>%
  collect()
#> # A tibble: 4 × 2
#>   x     y    
#>   <chr> <chr>
#> 1 Foo   F    
#> 2 Bar   Br   
#> 3 Baz   Bz   
#> 4 Qux   Qx

tibble::tibble(x = c("Foo", "Bar", "Baz", "Qux")) %>%
  arrow_table() %>%
  mutate(y = multistep_transform(x)) %>%
  collect()
#> # A tibble: 4 × 2
#>   x     y    
#>   <chr> <chr>
#> 1 Foo   F    
#> 2 Bar   CR   
#> 3 Baz   CZ   
#> 4 Qux   QX

tibble::tibble(x = c("Foo", "Bar", "Baz", "Qux")) %>%
  arrow_table() %>%
  mutate(y = multistep_transform_in_one(x)) %>%
  collect()
#> # A tibble: 4 × 2
#>   x     y    
#>   <chr> <chr>
#> 1 Foo   F    
#> 2 Bar   CR   
#> 3 Baz   CZ   
#> 4 Qux   QX

[thisisnic]

This is what I was referring to with "not supported in arrow" @nealrichardson though I hadn't realised it already works really well when calling functions directly in terms of reporting the failed function, so things might be fine as-is, though printing out which function doesn't have bindings would be a nice-to-have.

library(arrow)
library(dplyr)

single_transform <- function(x){
  # this function does have bindings
  stringr::str_remove_all(x, "[aeiou]")
}

# succeeds
tibble::tibble(x = c("Foo", "Bar", "Baz", "Qux")) %>%
  arrow_table() %>%
  mutate(y = single_transform(x)) %>%
  collect()
#> # A tibble: 4 × 2
#>   x     y    
#>   <chr> <chr>
#> 1 Foo   F    
#> 2 Bar   Br   
#> 3 Baz   Bz   
#> 4 Qux   Qx

single_transform2 <- function(x){
  # this function doesn't have bindings
  stringr::str_to_sentence(x)
}

tibble::tibble(x = c("Foo", "Bar", "Baz", "Qux")) %>%
  arrow_table() %>%
  mutate(y = stringr::str_to_sentence(x)) %>%
  collect()
#> Warning: Expression stringr::str_to_sentence(x) not supported in Arrow; pulling
#> data into R
#> # A tibble: 4 × 2
#>   x     y    
#>   <chr> <chr>
#> 1 Foo   Foo  
#> 2 Bar   Bar  
#> 3 Baz   Baz  
#> 4 Qux   Qux

tibble::tibble(x = c("Foo", "Bar", "Baz", "Qux")) %>%
  arrow_table() %>%
  mutate(y = single_transform2(x)) %>%
  collect()
#> Warning: Expression single_transform2(x) not supported in Arrow; pulling data
#> into R
#> # A tibble: 4 × 2
#>   x     y    
#>   <chr> <chr>
#> 1 Foo   Foo  
#> 2 Bar   Bar  
#> 3 Baz   Baz  
#> 4 Qux   Qux

[github-actions]

⚠️ GitHub issue #41323 has been automatically assigned in GitHub to PR creator.

[nealrichardson]

This is ready to go AFAIK @jonkeane @thisisnic, LMK if you have any more feedback.

[thisisnic]

Looks good though probably worth addressing this comment in a follow-up ticket perhaps?

[nealrichardson]

Looks good though probably worth addressing this comment in a follow-up ticket perhaps?

I don't think I can do any better than what it does now. With single_transform2(x) failing, there's not anything in that expression that would give away what is not supported, you have to step through the function, and it could be doing anything so IDK that looking for unsupported functions named within that would do much. And unfortunately the error that stringr::str_to_sentence raises isn't obvious either: if you set options(arrow.debug=TRUE), the warning shown includes the actual error, saying: "Warning: In single_transform2(x), cannot coerce type 'environment' to vector of type 'character'; pulling data into R".

[conbench-apache-arrow]

After merging your PR, Conbench analyzed the 7 benchmarking runs that have been run so far on merge-commit 5865e96.

There was 1 benchmark result indicating a performance regression:

• Commit Run on ursa-i9-9960x at 2024-04-22 20:34:10Z
  • tpch (R) with engine=arrow, format=native, language=R, memory_map=False, query_id=TPCH-01, scale_factor=1

The full Conbench report has more details. It also includes information about 5 possible false positives for unstable benchmarks that are known to sometimes produce them.