ARROW-16444: [R] Implement user-defined scalar functions in R bindings

r/NAMESPACE

@@ -45,7 +45,9 @@ S3method(as_arrow_array,data.frame)
 S3method(as_arrow_array,default)
 S3method(as_arrow_array,pyarrow.lib.Array)
 S3method(as_arrow_table,RecordBatch)
+S3method(as_arrow_table,RecordBatchReader)
 S3method(as_arrow_table,Table)
+S3method(as_arrow_table,arrow_dplyr_query)
 S3method(as_arrow_table,data.frame)
 S3method(as_arrow_table,default)
 S3method(as_arrow_table,pyarrow.lib.RecordBatch)
@@ -343,6 +345,7 @@ export(read_schema)
 export(read_tsv_arrow)
 export(record_batch)
 export(register_extension_type)
+export(register_scalar_function)
 export(reregister_extension_type)
 export(s3_bucket)
 export(schema)

r/R/compute.R

@@ -306,3 +306,179 @@ cast_options <- function(safe = TRUE, ...) {
   )
   modifyList(opts, list(...))
 }
+
+#' Register user-defined functions
+#'
+#' These functions support calling R code from query engine execution
+#' (i.e., a [dplyr::mutate()] or [dplyr::filter()] on a [Table] or [Dataset]).
+#' Use [register_scalar_function()] attach Arrow input and output types to an
+#' R function and make it available for use in the dplyr interface and/or
+#' [call_function()]. Scalar functions are currently the only type of
+#' user-defined function supported. In Arrow, scalar functions must be
+#' stateless and return output with the same shape (i.e., the same number
+#' of rows) as the input.
+#'
+#' @param name The function name to be used in the dplyr bindings
+#' @param in_type A [DataType] of the input type or a [schema()]
+#'   for functions with more than one argument. This signature will be used
+#'   to determine if this function is appropriate for a given set of arguments.
+#'   If this function is appropriate for more than one signature, pass a
+#'   `list()` of the above.
+#' @param out_type A [DataType] of the output type or a function accepting
+#'   a single argument (`types`), which is a `list()` of [DataType]s. If a
+#'   function it must return a [DataType].
+#' @param fun An R function or rlang-style lambda expression. The function
+#'   will be called with a first argument `context` which is a `list()`
+#'   with elements `batch_size` (the expected length of the output) and
+#'   `output_type` (the required [DataType] of the output) that may be used
+#'   to ensure that the output has the correct type and length. Subsequent
+#'   arguments are passed by position as specified by `in_types`. If
+#'   `auto_convert` is `TRUE`, subsequent arguments are converted to
+#'   R vectors before being passed to `fun` and the output is automatically
+#'   constructed with the expected output type via [as_arrow_array()].
+#' @param auto_convert Use `TRUE` to convert inputs before passing to `fun`
+#'   and construct an Array of the correct type from the output. Use this
+#'   option to write functions of R objects as opposed to functions of
+#'   Arrow R6 objects.
+#'
+#' @return `NULL`, invisibly
+#' @export
+#'
+#' @examplesIf arrow_with_dataset()
+#' library(dplyr, warn.conflicts = FALSE)
+#'
+#' some_model <- lm(mpg ~ disp + cyl, data = mtcars)
+#' register_scalar_function(
+#'   "mtcars_predict_mpg",
+#'   function(context, disp, cyl) {
+#'     predict(some_model, newdata = data.frame(disp, cyl))
+#'   },
+#'   in_type = schema(disp = float64(), cyl = float64()),
+#'   out_type = float64(),
+#'   auto_convert = TRUE
+#' )
+#'
+#' as_arrow_table(mtcars) %>%
+#'   transmute(mpg, mpg_predicted = mtcars_predict_mpg(disp, cyl)) %>%
+#'   collect() %>%
+#'   head()
+#'
+register_scalar_function <- function(name, fun, in_type, out_type,
+                                     auto_convert = FALSE) {
+  assert_that(is.string(name))
+
+  scalar_function <- arrow_scalar_function(
+    fun,
+    in_type,
+    out_type,
+    auto_convert = auto_convert
+  )
+
+  # register with Arrow C++ function registry (enables its use in
+  # call_function() and Expression$create())
+  RegisterScalarUDF(name, scalar_function)
+
+  # register with dplyr binding (enables its use in mutate(), filter(), etc.)
+  register_binding(
+    name,
+    function(...) build_expr(name, ...),
+    update_cache = TRUE
+  )
+
+  invisible(NULL)
+}
+
+arrow_scalar_function <- function(fun, in_type, out_type, auto_convert = FALSE) {
+  assert_that(is.function(fun))
+
+  # Create a small wrapper function that is easier to call from C++.
+  # TODO(ARROW-17148): This wrapper could be implemented in C/C++ to
+  # reduce evaluation overhead and generate prettier backtraces when
+  # errors occur (probably using a similar approach to purrr).
+  if (auto_convert) {
+    wrapper_fun <- function(context, args) {
+      args <- lapply(args, as.vector)
+      result <- do.call(fun, c(list(context), args))
+      as_arrow_array(result, type = context$output_type)
+    }
+  } else {
+    wrapper_fun <- function(context, args) {
+      do.call(fun, c(list(context), args))
+    }
+  }
+
+  # in_type can be a list() if registering multiple kernels at once
+  if (is.list(in_type)) {
+    in_type <- lapply(in_type, in_type_as_schema)
+  } else {
+    in_type <- list(in_type_as_schema(in_type))
+  }
+
+  # out_type can be a list() if registering multiple kernels at once
+  if (is.list(out_type)) {
+    out_type <- lapply(out_type, out_type_as_function)
+  } else {
+    out_type <- list(out_type_as_function(out_type))
+  }
+
+  # recycle out_type (which is frequently length 1 even if multiple kernels
+  # are being registered at once)
+  out_type <- rep_len(out_type, length(in_type))
+
+  # check n_kernels and number of args in fun
+  n_kernels <- length(in_type)
+  if (n_kernels == 0) {
+    abort("Can't register user-defined scalar function with 0 kernels")
+  }
+
+  expected_n_args <- in_type[[1]]$num_fields + 1L
+  fun_formals_have_dots <- any(names(formals(fun)) == "...")
+  if (!fun_formals_have_dots && length(formals(fun)) != expected_n_args) {
+    abort(
+      sprintf(
+        paste0(
+          "Expected `fun` to accept %d argument(s)\n",
+          "but found a function that acccepts %d argument(s)\n",
+          "Did you forget to include `context` as the first argument?"
+        ),
+        expected_n_args,
+        length(formals(fun))
+      )
+    )
+  }
+
+  structure(
+    list(
+      wrapper_fun = wrapper_fun,
+      in_type = in_type,
+      out_type = out_type
+    ),
+    class = "arrow_scalar_function"
+  )
+}
+
+# This function sanitizes the in_type argument for arrow_scalar_function(),
+# which can be a data type (e.g., int32()), a field for a unary function
+# or a schema() for functions accepting more than one argument. C++ expects
+# a schema().
+in_type_as_schema <- function(x) {
+  if (inherits(x, "Field")) {
+    schema(x)
+  } else if (inherits(x, "DataType")) {
+    schema(field("", x))
+  } else {
+    as_schema(x)
+  }
+}
+
+# This function sanitizes the out_type argument for arrow_scalar_function(),
+# which can be a data type (e.g., int32()) or a function of the input types.
+# C++ currently expects a function.
+out_type_as_function <- function(x) {
+  if (is.function(x)) {
+    x
+  } else {
+    x <- as_data_type(x)
+    function(types) x
+  }
+}

r/R/dplyr-collect.R

@@ -20,7 +20,7 @@
 
 collect.arrow_dplyr_query <- function(x, as_data_frame = TRUE, ...) {
   tryCatch(
-    out <- as_record_batch_reader(x)$read_table(),
+    out <- as_arrow_table(x),
     # 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)) {

r/R/dplyr-funcs.R

@@ -50,20 +50,28 @@ NULL
 #'   - `fun`: string function name
 #'   - `data`: `Expression` (these are all currently a single field)
 #'   - `options`: list of function options, as passed to call_function
+#' @param update_cache Update .cache$functions at the time of registration.
+#'   the default is FALSE because the majority of usage is to register
+#'   bindings at package load, after which we create the cache once. The
+#'   reason why .cache$functions is needed in addition to nse_funcs for
+#'   non-aggregate functions could be revisited...it is currently used
+#'   as the data mask in mutate, filter, and aggregate (but not
+#'   summarise) because the data mask has to be a list.
 #' @param registry An environment in which the functions should be
 #'   assigned.
 #'
 #' @return The previously registered binding or `NULL` if no previously
 #'   registered function existed.
 #' @keywords internal
 #'
-register_binding <- function(fun_name, fun, registry = nse_funcs) {
+register_binding <- function(fun_name, fun, registry = nse_funcs,
+                             update_cache = FALSE) {
   unqualified_name <- sub("^.*?:{+}", "", fun_name)
 
   previous_fun <- registry[[unqualified_name]]
 
   # if the unqualified name exists in the registry, warn
-  if (!is.null(fun) && !is.null(previous_fun)) {
+  if (!is.null(previous_fun)) {
     warn(
       paste0(
         "A \"",
@@ -73,11 +81,36 @@ register_binding <- function(fun_name, fun, registry = nse_funcs) {
   }
 
   # register both as `pkg::fun` and as `fun` if `qualified_name` is prefixed
-  if (grepl("::", fun_name)) {
-    registry[[unqualified_name]] <- fun
-    registry[[fun_name]] <- fun
-  } else {
-    registry[[unqualified_name]] <- fun
+  # unqualified_name and fun_name will be the same if not prefixed
+  registry[[unqualified_name]] <- fun
+  registry[[fun_name]] <- fun
+
+  if (update_cache) {
+    fun_cache <- .cache$functions
+    fun_cache[[unqualified_name]] <- fun
+    fun_cache[[fun_name]] <- fun
+    .cache$functions <- fun_cache
+  }
+
+  invisible(previous_fun)
+}
+
+unregister_binding <- function(fun_name, registry = nse_funcs,
+                               update_cache = FALSE) {
+  unqualified_name <- sub("^.*?:{+}", "", fun_name)
+  previous_fun <- registry[[unqualified_name]]
+
+  rm(
+    list = unique(c(fun_name, unqualified_name)),
+    envir = registry,
+    inherits = FALSE
+  )
+
+  if (update_cache) {
+    fun_cache <- .cache$functions
+    fun_cache[[unqualified_name]] <- NULL
+    fun_cache[[fun_name]] <- NULL
+    .cache$functions <- fun_cache
   }
 
   invisible(previous_fun)

r/R/feather.R

@@ -190,7 +190,7 @@ FeatherReader <- R6Class("FeatherReader",
   inherit = ArrowObject,
   public = list(
     Read = function(columns) {
-      ipc___feather___Reader__Read(self, columns, on_old_windows())
+      ipc___feather___Reader__Read(self, columns)
     },
     print = function(...) {
       cat("FeatherReader:\n")
@@ -211,5 +211,5 @@ names.FeatherReader <- function(x) x$column_names
 
 FeatherReader$create <- function(file) {
   assert_is(file, "RandomAccessFile")
-  ipc___feather___Reader__Open(file, on_old_windows())
+  ipc___feather___Reader__Open(file)
 }