feat(r): add a function for lazy R callbacks

r/NAMESPACE

@@ -148,6 +148,7 @@ export(convert_array_extension)
 export(convert_array_stream)
 export(convert_buffer)
 export(example_ipc_stream)
+export(function_array_stream)
 export(infer_nanoarrow_ptype)
 export(infer_nanoarrow_ptype_extension)
 export(infer_nanoarrow_schema)

r/R/array-stream.R

@@ -50,6 +50,100 @@ basic_array_stream <- function(batches, schema = NULL, validate = TRUE) {
   .Call(nanoarrow_c_basic_array_stream, batches, schema, validate)
 }
 
+#' Create an array stream from an R function
+#'
+#' The function is called once for each requested array. It should return an
+#' object coercible by [as_nanoarrow_array()], or `NULL` to indicate end of
+#' stream. The schema is required because it is available before the first
+#' array is requested. Neither construction nor schema access calls `fun`.
+#' After `NULL` or an error is returned, the callback is not called again.
+#'
+#' Errors and interrupts raised by `fun` are converted to stream errors.
+#' Evaluation is also isolated in a top-level R context so other non-local exits
+#' (for example, invoking an `abort` restart) cannot escape through the C stream
+#' interface. Calling handlers in the caller are therefore not applied to
+#' callback evaluation. The callback is evaluated only on R's main thread.
+#' Warnings are handled in this isolated context, unless handled within `fun`.
+#' `fun` must not release the stream (or call its methods) while it is being
+#' evaluated. The stream retains `fun` until it is released; explicitly release
+#' streams when `fun` or its environment refers to the stream itself.
+#'
+#' @param fun A function called with no arguments to produce each array.
+#' @param schema A schema coercible by [as_nanoarrow_schema()], used to coerce
+#'   and validate each result.
+#' @param validate A single `TRUE` or `FALSE`. Use `FALSE` to skip validation
+#'   within the callback. The stream's `get_next()` method can validate the
+#'   returned array independently.
+#'
+#' @return An [nanoarrow_array_stream][as_nanoarrow_array_stream]
+#' @export
+#'
+#' @examples
+#' i <- 0L
+#' stream <- function_array_stream(
+#'   function() {
+#'     i <<- i + 1L
+#'     if (i > 2L) NULL else mtcars[i, , drop = FALSE]
+#'   },
+#'   infer_nanoarrow_schema(mtcars)
+#' )
+#' as.data.frame(stream$get_next())
+#' as.data.frame(stream$get_next())
+#' stream$get_next()
+#' stream$release()
+#'
+function_array_stream <- function(fun, schema, validate = TRUE) {
+  stopifnot(is.function(fun))
+  # Copy the schema so the callback does not depend on the caller's object.
+  schema_copy <- nanoarrow_allocate_schema()
+  nanoarrow_pointer_export(as_nanoarrow_schema(schema), schema_copy)
+  schema <- schema_copy
+  if (!is.logical(validate) || length(validate) != 1L || is.na(validate)) {
+    stop("`validate` must be a single non-NA logical value")
+  }
+
+  # The callback is evaluated from C as `next_array()` in this environment.
+  # It returns NULL (end of stream), an array, or a character(1) error message.
+  # Only `next_array` needs to be visible there, so the environment is empty.
+  env <- new.env(parent = emptyenv())
+  env$next_array <- function() {
+    tryCatch(
+      {
+        result <- fun()
+        if (is.null(result)) {
+          NULL
+        } else {
+          array <- as_nanoarrow_array(result, schema = schema)
+          if (validate) {
+            nanoarrow_array_set_schema(array, schema, validate = TRUE)
+          }
+          array
+        }
+      },
+      error = function(e) {
+        tryCatch(
+          as.character(conditionMessage(e))[1],
+          error = function(e2) {
+            "R callback failed and its error message could not be retrieved"
+          }
+        )
+      },
+      interrupt = function(e) {
+        tryCatch(
+          paste0("R callback interrupted: ", conditionMessage(e)),
+          error = function(e2) "R callback interrupted"
+        )
+      }
+    )
+  }
+
+  .Call(nanoarrow_c_function_array_stream, quote(next_array()), env, schema)
+}
+
+function_array_stream_get_next_on_other_thread <- function(x) {
+  .Call(nanoarrow_c_function_array_stream_get_next_on_other_thread, x)
+}
+
 #' Register an array stream finalizer
 #'
 #' In some cases, R functions that return a [nanoarrow_array_stream][as_nanoarrow_array_stream]

r/R/util.R

@@ -74,6 +74,10 @@ preserve_and_release_on_other_thread <- function(obj) {
   invisible(.Call(nanoarrow_c_preserve_and_release_on_other_thread, obj))
 }
 
+release_array_stream_on_other_thread <- function(array_stream) {
+  invisible(.Call(nanoarrow_c_release_array_stream_on_other_thread, array_stream))
+}
+
 # This is used by bookkeeping infrastructure when debugging an imbalance in
 # preserved/released SEXPs.
 current_stack_trace_chr <- function() {

r/src/array_stream.c

@@ -18,13 +18,164 @@
 #define R_NO_REMAP
 #include <R.h>
 #include <Rinternals.h>
+#include <errno.h>
+#include <string.h>
 
 #include "array.h"
 #include "array_stream.h"
 #include "nanoarrow.h"
 #include "schema.h"
 #include "util.h"
 
+struct FunctionArrayStreamData {
+  struct ArrowSchema schema;
+  SEXP expr;
+  SEXP env;
+  struct ArrowError error;
+  int error_code;
+  int finished;
+  int errored;
+  int evaluating;
+};
+
+struct FunctionArrayStreamCall {
+  struct FunctionArrayStreamData* data;
+  struct ArrowArray* out;
+  int code;
+};
+
+static void function_array_stream_set_error(struct FunctionArrayStreamData* data,
+                                            int code, const char* message) {
+  if (!data->errored) {
+    ArrowErrorSet(&data->error, "%s", message);
+    data->error_code = code;
+    data->errored = 1;
+  }
+}
+
+static void release_function_array_stream(struct ArrowArrayStream* stream) {
+  struct FunctionArrayStreamData* data =
+      (struct FunctionArrayStreamData*)stream->private_data;
+  stream->release = NULL;
+  stream->private_data = NULL;
+  if (data->schema.release != NULL) {
+    data->schema.release(&data->schema);
+  }
+  nanoarrow_release_sexp(data->expr);
+  nanoarrow_release_sexp(data->env);
+  ArrowFree(data);
+}
+
+static const char* function_array_stream_get_last_error(struct ArrowArrayStream* stream) {
+  struct FunctionArrayStreamData* data =
+      (struct FunctionArrayStreamData*)stream->private_data;
+  return data == NULL ? "Array stream has been released"
+                      : ArrowErrorMessage(&data->error);
+}
+
+static int function_array_stream_get_schema(struct ArrowArrayStream* stream,
+                                            struct ArrowSchema* out) {
+  struct FunctionArrayStreamData* data =
+      (struct FunctionArrayStreamData*)stream->private_data;
+  return ArrowSchemaDeepCopy(&data->schema, out);
+}
+
+// Runs inside R_ToplevelExec(). The R-level expression returns NULL (end of
+// stream), a nanoarrow_array, or a character(1) error message.
+static void function_array_stream_eval(void* opaque) {
+  struct FunctionArrayStreamCall* call = (struct FunctionArrayStreamCall*)opaque;
+  struct FunctionArrayStreamData* data = call->data;
+  SEXP result = PROTECT(Rf_eval(data->expr, data->env));
+  if (data->errored) {
+    // A re-entrant call was rejected while the expression was being evaluated.
+    call->code = data->error_code;
+  } else if (result == R_NilValue) {
+    data->finished = 1;
+    call->code = NANOARROW_OK;
+  } else if (TYPEOF(result) == STRSXP && Rf_length(result) == 1) {
+    function_array_stream_set_error(data, EIO,
+                                    Rf_translateCharUTF8(STRING_ELT(result, 0)));
+    call->code = data->error_code;
+  } else if (Rf_inherits(result, "nanoarrow_array")) {
+    // Let the external pointer clean up a partial export if R jumps during allocation.
+    SEXP array_xptr = PROTECT(nanoarrow_array_owning_xptr());
+    struct ArrowArray* array = nanoarrow_output_array_from_xptr(array_xptr);
+    array_export(result, array);
+    ArrowArrayMove(array, call->out);
+    UNPROTECT(1);
+    call->code = NANOARROW_OK;
+  } else {
+    function_array_stream_set_error(data, EIO,
+                                    "R callback returned an unexpected result");
+    call->code = data->error_code;
+  }
+  UNPROTECT(1);
+}
+
+static int function_array_stream_get_next(struct ArrowArrayStream* stream,
+                                          struct ArrowArray* out) {
+  struct FunctionArrayStreamData* data =
+      (struct FunctionArrayStreamData*)stream->private_data;
+  out->release = NULL;
+  if (data == NULL || data->finished) {
+    return NANOARROW_OK;
+  }
+  if (data->errored) {
+    return data->error_code;
+  }
+  if (!nanoarrow_is_main_thread()) {
+    function_array_stream_set_error(data, EIO,
+                                    "Can't evaluate R callback on a non-R thread");
+    return data->error_code;
+  }
+  if (data->evaluating) {
+    function_array_stream_set_error(data, EINVAL,
+                                    "Re-entrant call to function array stream");
+    return data->error_code;
+  }
+
+  data->evaluating = 1;
+  struct FunctionArrayStreamCall call = {data, out, NANOARROW_OK};
+  // Rf_eval() alone would let R non-local exits (errors that escape, an abort
+  // restart, etc.) longjmp through this Arrow C stream callback and its caller,
+  // skipping the caller's cleanup. R_UnwindProtect() only runs cleanup before
+  // continuing the jump, so it cannot stop it. R_ToplevelExec() contains the
+  // jump and lets us report it as an error code. Errors and interrupts are
+  // normally handled by the R-level tryCatch(); this is the last line of defense.
+  if (!R_ToplevelExec(&function_array_stream_eval, &call)) {
+    function_array_stream_set_error(
+        data, EIO, "R callback exited non-locally (e.g., via an abort restart)");
+    call.code = data->error_code;
+  }
+  data->evaluating = 0;
+  return call.code;
+}
+
+SEXP nanoarrow_c_function_array_stream(SEXP expr, SEXP env, SEXP schema_xptr) {
+  SEXP stream_xptr = PROTECT(nanoarrow_array_stream_owning_xptr());
+  struct ArrowArrayStream* stream = nanoarrow_output_array_stream_from_xptr(stream_xptr);
+  struct FunctionArrayStreamData* data =
+      (struct FunctionArrayStreamData*)ArrowMalloc(sizeof(*data));
+  check_trivial_alloc(data, "struct FunctionArrayStreamData");
+  memset(data, 0, sizeof(*data));
+  ArrowErrorInit(&data->error);
+  data->expr = R_NilValue;
+  data->env = R_NilValue;
+  stream->private_data = data;
+  stream->get_last_error = &function_array_stream_get_last_error;
+  stream->get_schema = &function_array_stream_get_schema;
+  stream->get_next = &function_array_stream_get_next;
+  stream->release = &release_function_array_stream;
+  // Attach the finalizer before any operation that can allocate or fail.
+  nanoarrow_preserve_sexp(expr);
+  data->expr = expr;
+  nanoarrow_preserve_sexp(env);
+  data->env = env;
+  schema_export(schema_xptr, &data->schema);
+  UNPROTECT(1);
+  return stream_xptr;
+}
+
 // Ideally user-supplied finalizers are written in such a way that they don't jump;
 // however if they do it is likely that memory will leak. Here, we use
 // R_tryCatchError to minimize the chances of that happening.

r/src/init.c

@@ -31,6 +31,9 @@ extern SEXP nanoarrow_c_array_stream_get_schema(SEXP array_stream_xptr);
 extern SEXP nanoarrow_c_array_stream_get_next(SEXP array_stream_xptr);
 extern SEXP nanoarrow_c_basic_array_stream(SEXP batches_sexp, SEXP schema_xptr,
                                            SEXP validate_sexp);
+extern SEXP nanoarrow_c_function_array_stream(SEXP expr, SEXP env, SEXP schema_xptr);
+extern SEXP nanoarrow_c_function_array_stream_get_next_on_other_thread(
+    SEXP array_stream_xptr);
 extern SEXP nanoarrow_c_array_list_total_length(SEXP list_of_array_xptr);
 extern SEXP nanoarrow_c_array_view(SEXP array_xptr, SEXP schema_xptr);
 extern SEXP nanoarrow_c_array_init(SEXP schema_xptr);
@@ -95,6 +98,7 @@ extern SEXP nanoarrow_c_schema_set_dictionary(SEXP schema_mut_xptr, SEXP diction
 extern SEXP nanoarrow_c_preserved_count(void);
 extern SEXP nanoarrow_c_preserved_empty(void);
 extern SEXP nanoarrow_c_preserve_and_release_on_other_thread(SEXP obj);
+extern SEXP nanoarrow_c_release_array_stream_on_other_thread(SEXP array_stream_xptr);
 extern SEXP nanoarrow_c_vctr_chunk_offsets(SEXP array_list);
 extern SEXP nanoarrow_c_vctr_chunk_resolve(SEXP indices_sexp, SEXP offsets_sexp);
 extern SEXP nanoarrow_c_vctr_as_slice(SEXP indices_sexp);
@@ -114,6 +118,9 @@ static const R_CallMethodDef CallEntries[] = {
      1},
     {"nanoarrow_c_array_stream_get_next", (DL_FUNC)&nanoarrow_c_array_stream_get_next, 1},
     {"nanoarrow_c_basic_array_stream", (DL_FUNC)&nanoarrow_c_basic_array_stream, 3},
+    {"nanoarrow_c_function_array_stream", (DL_FUNC)&nanoarrow_c_function_array_stream, 3},
+    {"nanoarrow_c_function_array_stream_get_next_on_other_thread",
+     (DL_FUNC)&nanoarrow_c_function_array_stream_get_next_on_other_thread, 1},
     {"nanoarrow_c_array_list_total_length", (DL_FUNC)&nanoarrow_c_array_list_total_length,
      1},
     {"nanoarrow_c_array_view", (DL_FUNC)&nanoarrow_c_array_view, 2},
@@ -179,6 +186,8 @@ static const R_CallMethodDef CallEntries[] = {
     {"nanoarrow_c_preserved_empty", (DL_FUNC)&nanoarrow_c_preserved_empty, 0},
     {"nanoarrow_c_preserve_and_release_on_other_thread",
      (DL_FUNC)&nanoarrow_c_preserve_and_release_on_other_thread, 1},
+    {"nanoarrow_c_release_array_stream_on_other_thread",
+     (DL_FUNC)&nanoarrow_c_release_array_stream_on_other_thread, 1},
     {"nanoarrow_c_vctr_chunk_offsets", (DL_FUNC)&nanoarrow_c_vctr_chunk_offsets, 1},
     {"nanoarrow_c_vctr_chunk_resolve", (DL_FUNC)&nanoarrow_c_vctr_chunk_resolve, 2},
     {"nanoarrow_c_vctr_as_slice", (DL_FUNC)&nanoarrow_c_vctr_as_slice, 1},