Pass a function to be applied as function argument in dbplyr

user_3759418 0
Viewed 122

Let's say I want to create a function which can mutate a column using any function which is passed by the user. I need to know how to quote and unquote that function before it hits the dbplyr parser. Let's take a look at an example, say I have a function like this:

testFun <- function(data, fun, colName, colOut = "myAwesomeColumn") {
  dplyr::mutate(.data = data, !!colOut := fun(.data[[colName]]))
}

sc <- sparklyr::spark_connection(master = "local")
mtcars_spark <- dplyr::copy_to(sc, mtcars, "mtcars")
testFun(mtcars_spark, mean, "mpg")

So in the above example, I want to apply the mean() function to the "mpg" column and store it in a new column called "myAwesomeColumn".

When working with Spark, and specifically sparklyr, there will be an attempt by dbplyr to convert this code to SQL and send it to Spark. My understanding is that dbplyr applies the following rules:

  1. If it can find a Spark SQL equivalent, it will use that (e.g. mean() -> AVG)
  2. Otherwise it will pass the function as-is to look for Scala extensions or UDFs

The second option is what happens here since it cannot find the function fun and it therefore returns a Spark error

Error: org.apache.spark.sql.AnalysisException: Undefined function: 'fun'.
This function is neither a registered temporary function nor a permanent function
registered in the database 'default'.; line 1 pos 85
...

So we need another approach. The problem is getting rlang to convert fun to mean before this is interpreted by dbplyr. I know I can do this if I pass the function name as a string and use rlang::parse_expr(), for example:

testFun <- function(data, fun, colName, colOut = "myAwesomeColumn") {
  dplyr::mutate(data, !!colOut := rlang::parse_expr(paste0(fun, "(", colName, ")"))
}
testFun(mtcars_spark, "mean", "mpg")
# # Source: spark<?> [?? x 12]
#      mpg   cyl  disp    hp  drat    wt  qsec    vs    am  gear  carb myAwesomeColumn
#    <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl>           <dbl>
#  1  21       6  160    110  3.9   2.62  16.5     0     1     4     4            20.1
#  2  21       6  160    110  3.9   2.88  17.0     0     1     4     4            20.1
# # ... with more rows
dbplyr dplyr r
1 Answers

In order to get this to work, we must quote and unquote the fun argument. We also build the expression that we actually want to pass into our call to mutate(). See below for the solution.

testFun <- function(data, fun, colAmount, colOut = "output") { 
  fun <- rlang::enquo(fun) 
  dplyr::mutate(.data = data, !!colOut := rlang::call2(.fn = !!fun, rlang::sym(colAmount))) 
} 
     
testFun(mtcars_spark, mean, "mpg")                                                                                                                                 
# # Source: spark<?> [?? x 12]
#      mpg   cyl  disp    hp  drat    wt  qsec    vs    am  gear  carb output
#    <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl>  <dbl>
#  1  21       6  160    110  3.9   2.62  16.5     0     1     4     4   20.1
#  2  21       6  160    110  3.9   2.88  17.0     0     1     4     4   20.1
# # ... with more rows

Note that this is much simpler if you are using data.frames rather than tbl_sparks.

user_37594180
Related