From f794df9f54227dfd7b378a03714847f532532cec Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 18 Oct 2026 01:12:07 +0000 Subject: [PATCH 1/3] feat(r): Add function_array_stream() for lazy R callbacks --- r/NAMESPACE | 1 + r/R/array-stream.R | 67 ++++++ r/man/function_array_stream.Rd | 51 ++++ r/src/array_stream.c | 223 ++++++++++++++++++ r/src/init.c | 6 + r/src/nanoarrow_cpp.cc | 13 + r/tests/testthat/test-array-stream.R | 341 +++++++++++++++++++++++++++ 7 files changed, 702 insertions(+) create mode 100644 r/man/function_array_stream.Rd diff --git a/r/NAMESPACE b/r/NAMESPACE index dbd9d07d4..1ce6c21ef 100644 --- a/r/NAMESPACE +++ b/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) diff --git a/r/R/array-stream.R b/r/R/array-stream.R index eefa6ae92..e27d9b235 100644 --- a/r/R/array-stream.R +++ b/r/R/array-stream.R @@ -50,6 +50,73 @@ 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. +#' +#' Evaluation is isolated in a top-level R context so 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`. +#' 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 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)) + schema <- as_nanoarrow_schema(schema) + validate <- as.logical(validate)[1] + if (is.na(validate)) { + stop("`validate` must not be NA") + } + + .Call(nanoarrow_c_function_array_stream, fun, schema, validate) +} + +.nanoarrow_function_array_stream_get_next <- function(fun, schema, validate) { + 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 + } +} + +.nanoarrow_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] diff --git a/r/man/function_array_stream.Rd b/r/man/function_array_stream.Rd new file mode 100644 index 000000000..4616aad42 --- /dev/null +++ b/r/man/function_array_stream.Rd @@ -0,0 +1,51 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/array-stream.R +\name{function_array_stream} +\alias{function_array_stream} +\title{Create an array stream from an R function} +\usage{ +function_array_stream(fun, schema, validate = TRUE) +} +\arguments{ +\item{fun}{A function called with no arguments to produce each array.} + +\item{schema}{A schema coercible by \code{\link[=as_nanoarrow_schema]{as_nanoarrow_schema()}}, used to coerce +and validate each result.} + +\item{validate}{Use \code{FALSE} to skip validation within the callback. The +stream's \code{get_next()} method can validate the returned array independently.} +} +\value{ +An \link[=as_nanoarrow_array_stream]{nanoarrow_array_stream} +} +\description{ +The function is called once for each requested array. It should return an +object coercible by \code{\link[=as_nanoarrow_array]{as_nanoarrow_array()}}, or \code{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 \code{fun}. +After \code{NULL} or an error is returned, the callback is not called again. +} +\details{ +Evaluation is isolated in a top-level R context so non-local exits (for +example, invoking an \code{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 \code{fun}. +The stream retains \code{fun} until it is released; explicitly release streams +when \code{fun} or its environment refers to the stream itself. +} +\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() + +} diff --git a/r/src/array_stream.c b/r/src/array_stream.c index 603645d2f..1f08613e4 100644 --- a/r/src/array_stream.c +++ b/r/src/array_stream.c @@ -18,6 +18,8 @@ #define R_NO_REMAP #include #include +#include +#include #include "array.h" #include "array_stream.h" @@ -25,6 +27,227 @@ #include "schema.h" #include "util.h" +struct FunctionArrayStreamData { + struct ArrowSchema schema; + SEXP fun; + SEXP schema_xptr; + SEXP validate; + struct ArrowError error; + int error_code; + int finished; + int errored; + int evaluating; + int release_pending; +}; + +struct FunctionArrayStreamCall { + struct FunctionArrayStreamData* data; + struct ArrowArray* out; + int code; + int caught; +}; + +static void release_function_array_stream_data(struct FunctionArrayStreamData* data) { + if (data == NULL) { + return; + } + if (data->schema.release != NULL) { + data->schema.release(&data->schema); + } + nanoarrow_release_sexp(data->fun); + nanoarrow_release_sexp(data->schema_xptr); + nanoarrow_release_sexp(data->validate); + ArrowFree(data); +} + +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 == NULL) { + return; + } + if (data->evaluating) { + // The callback may release its own stream. Keep its state alive until it returns. + data->release_pending = 1; + return; + } + release_function_array_stream_data(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; + if (data == NULL) { + return EINVAL; + } + int result = ArrowSchemaDeepCopy(&data->schema, out); + if (result != NANOARROW_OK) { + if (!data->errored) { + ArrowErrorSet(&data->error, "ArrowSchemaDeepCopy() failed"); + } + return result; + } + return NANOARROW_OK; +} + +static SEXP function_array_stream_condition_handler(SEXP condition, void* opaque) { + struct FunctionArrayStreamCall* call = (struct FunctionArrayStreamCall*)opaque; + call->caught = 1; + if (Rf_inherits(condition, "interrupt")) { + call->caught = 2; + } + SEXP message_call = PROTECT(Rf_lang2(Rf_install("conditionMessage"), condition)); + SEXP message = Rf_eval(message_call, R_BaseEnv); + UNPROTECT(1); + return message; +} + +static SEXP function_array_stream_eval_call(void* opaque) { + return Rf_eval((SEXP)opaque, nanoarrow_ns_pkg); +} + +static void function_array_stream_eval(void* opaque) { + struct FunctionArrayStreamCall* call = (struct FunctionArrayStreamCall*)opaque; + struct FunctionArrayStreamData* data = call->data; + SEXP classes = PROTECT(Rf_allocVector(STRSXP, 2)); + SET_STRING_ELT(classes, 0, Rf_mkChar("error")); + SET_STRING_ELT(classes, 1, Rf_mkChar("interrupt")); + SEXP callback = + PROTECT(Rf_lang4(Rf_install(".nanoarrow_function_array_stream_get_next"), data->fun, + data->schema_xptr, data->validate)); + SEXP result = + PROTECT(R_tryCatch(&function_array_stream_eval_call, callback, classes, + &function_array_stream_condition_handler, call, NULL, NULL)); + if (data->errored) { + call->code = data->error_code; + } else if (call->caught) { + const char* message = Rf_translateCharUTF8(STRING_ELT(result, 0)); + if (call->caught == 2) { + struct ArrowError error; + ArrowErrorSet(&error, "R callback interrupted: %s", message); + function_array_stream_set_error(data, EIO, ArrowErrorMessage(&error)); + } else { + function_array_stream_set_error(data, EIO, message); + } + call->code = data->error_code; + } else if (result == R_NilValue) { + data->finished = 1; + call->code = NANOARROW_OK; + } else { + // 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; + } + UNPROTECT(3); +} + +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, 0}; + // R_tryCatch() does not catch restart jumps or failures in conditionMessage(). + // All R API calls must stay inside this outer boundary. + 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; + + if (call.code != NANOARROW_OK) { + if (out->release != NULL) { + out->release(out); + } + out->release = NULL; + } + if (data->release_pending) { + if (out->release != NULL) { + out->release(out); + } + out->release = NULL; + // The stream itself may already have been freed by a finalizer during evaluation. + release_function_array_stream_data(data); + return NANOARROW_OK; + } + return call.code; +} + +SEXP nanoarrow_c_function_array_stream(SEXP fun, SEXP schema_xptr, SEXP validate) { + SEXP stream_xptr = PROTECT(nanoarrow_array_stream_owning_xptr()); + struct ArrowArrayStream* stream = nanoarrow_output_array_stream_from_xptr(stream_xptr); + SEXP schema_copy_xptr = PROTECT(nanoarrow_schema_owning_xptr()); + struct ArrowSchema* schema_copy = nanoarrow_output_schema_from_xptr(schema_copy_xptr); + schema_export(schema_xptr, schema_copy); + struct FunctionArrayStreamData* data = + (struct FunctionArrayStreamData*)ArrowMalloc(sizeof(*data)); + check_trivial_alloc(data, "struct FunctionArrayStreamData"); + memset(data, 0, sizeof(*data)); + ArrowErrorInit(&data->error); + data->fun = R_NilValue; + data->schema_xptr = R_NilValue; + data->validate = 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(fun); + data->fun = fun; + nanoarrow_preserve_sexp(schema_copy_xptr); + data->schema_xptr = schema_copy_xptr; + nanoarrow_preserve_sexp(validate); + data->validate = validate; + if (ArrowSchemaDeepCopy(schema_copy, &data->schema) != NANOARROW_OK) { + Rf_error("ArrowSchemaDeepCopy() failed"); + } + UNPROTECT(2); + 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. diff --git a/r/src/init.c b/r/src/init.c index 97f8f4c9e..b40dcdc0c 100644 --- a/r/src/init.c +++ b/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 fun, SEXP schema_xptr, SEXP validate); +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); @@ -114,6 +117,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}, diff --git a/r/src/nanoarrow_cpp.cc b/r/src/nanoarrow_cpp.cc index 83ebe07ac..01030051a 100644 --- a/r/src/nanoarrow_cpp.cc +++ b/r/src/nanoarrow_cpp.cc @@ -205,6 +205,19 @@ extern "C" void nanoarrow_preserve_and_release_on_other_thread(SEXP obj) { worker.join(); } +extern "C" SEXP nanoarrow_c_function_array_stream_get_next_on_other_thread( + SEXP array_stream_xptr) { + struct ArrowArrayStream* stream = nanoarrow_array_stream_from_xptr(array_stream_xptr); + int result = 0; + std::thread worker([&] { + struct ArrowArray out {}; + result = stream->get_next(stream, &out); + if (out.release != nullptr) out.release(&out); + }); + worker.join(); + return Rf_ScalarInteger(result); +} + // Collector utility for iterating over and collecting batches // Keeping this all in a single object reduces the amount of C++ deletion // we need to keep track of. diff --git a/r/tests/testthat/test-array-stream.R b/r/tests/testthat/test-array-stream.R index 36859c05a..cb1181c4e 100644 --- a/r/tests/testthat/test-array-stream.R +++ b/r/tests/testthat/test-array-stream.R @@ -26,6 +26,347 @@ test_that("basic_array_stream() can create empty streams", { ) }) +test_that("function_array_stream() lazily creates arrays and reaches EOF", { + calls <- 0L + stream <- function_array_stream( + function() { + calls <<- calls + 1L + if (calls > 2L) NULL else data.frame(x = calls) + }, + na_struct(list(x = na_int32())) + ) + + expect_identical(calls, 0L) + expect_identical(stream$get_schema()$format, "+s") + expect_identical(calls, 0L) + expect_identical(as.data.frame(stream$get_next()), data.frame(x = 1L)) + expect_identical(as.data.frame(stream$get_next()), data.frame(x = 2L)) + expect_null(stream$get_next()) + expect_null(stream$get_next()) + expect_identical(calls, 3L) + stream$release() +}) + +test_that("function_array_stream() owns a schema copy", { + schema <- na_int32() + stream <- function_array_stream(function() 42L, schema) + nanoarrow_pointer_release(schema) + + expect_identical(stream$get_schema()$format, "i") + expect_identical(stream$get_next()$length, 1L) + stream$release() +}) + +test_that("function_array_stream() accepts NULL immediately and zero-row arrays", { + calls <- 0L + empty <- function_array_stream( + function() { + calls <<- calls + 1L + NULL + }, na_int32() + ) + expect_null(empty$get_next()) + expect_null(empty$get_next()) + expect_identical(calls, 1L) + empty$release() + + stream <- function_array_stream( + function() data.frame(x = integer()), na_struct(list(x = na_int32())) + ) + expect_identical(stream$get_next()$length, 0L) + stream$release() +}) + +test_that("function_array_stream() coerces common array inputs", { + for (value in list(1:3, as_nanoarrow_array(1:3))) { + stream <- function_array_stream(function() value, na_int32()) + expect_identical(stream$get_next()$length, 3L) + stream$release() + } +}) + +test_that("function_array_stream() validates schemas and preserves unchecked behavior", { + malformed <- as_nanoarrow_array(data.frame(x = 1L)) + nanoarrow_array_set_schema(malformed, na_int32(), validate = FALSE) + + stream <- function_array_stream(function() malformed, na_int32()) + expect_error(stream$get_next(), "Expected array with 2 buffer") + expect_error(stream$get_next(), "Expected array with 2 buffer") + stream$release() + + stream <- function_array_stream( + function() malformed, na_int32(), validate = FALSE + ) + expect_s3_class(stream$get_next(validate = FALSE), "nanoarrow_array") + stream$release() +}) + +test_that("function_array_stream() errors become stable stream errors", { + stream <- function_array_stream(function() stop("callback failed"), na_int32()) + expect_error(stream$get_next(), "callback failed") + expect_error(stream$get_next(), "callback failed") + stream$release() +}) + +test_that("function_array_stream() contains interrupts and re-entrant calls", { + interrupted <- function() { + condition <- structure( + list(message = "test interrupt"), + class = c("interrupt", "condition") + ) + signalCondition(condition) + 1L + } + stream <- function_array_stream(interrupted, na_int32()) + expect_error(stream$get_next(), "R callback interrupted") + expect_error(stream$get_next(), "R callback interrupted") + stream$release() + + stream <- NULL + stream <- function_array_stream(function() stream$get_next(), na_int32()) + expect_error(stream$get_next(), "Re-entrant call to function array stream") + expect_error(stream$get_next(), "Re-entrant call to function array stream") + expect_identical( + .nanoarrow_function_array_stream_get_next_on_other_thread(stream), 22L + ) + stream$release() + + stream <- NULL + stream <- function_array_stream( + function() { + tryCatch(stream$get_next(), error = function(e) NULL) + 1L + }, + na_int32() + ) + expect_error(stream$get_next(), "Re-entrant call to function array stream") + expect_identical( + .nanoarrow_function_array_stream_get_next_on_other_thread(stream), 22L + ) + stream$release() +}) + +test_that("function_array_stream() contains abort restarts and allows warnings", { + stream <- function_array_stream( + function() invokeRestart("abort"), na_int32() + ) + expect_error( + stream$get_next(), + "R callback exited non-locally .*abort restart" + ) + expect_error(stream$get_next(), "R callback exited non-locally") + stream$release() + + stream <- function_array_stream(function() { + warning("callback warning") + 1L + }, na_int32()) + withr::local_options(warn = 1) + caller_handled_warning <- FALSE + warning_output <- capture.output( + result <- withCallingHandlers( + stream$get_next(), + warning = function(w) { + caller_handled_warning <<- TRUE + invokeRestart("muffleWarning") + } + ), + type = "message" + ) + expect_match(paste(warning_output, collapse = "\n"), "callback warning") + expect_false(caller_handled_warning) + expect_identical(result$length, 1L) + stream$release() +}) + +test_that("function_array_stream() reports conditionMessage failures safely", { + method <- "conditionMessage.nanoarrow_test_broken" + had_method <- exists(method, envir = globalenv(), inherits = FALSE) + if (had_method) { + old_method <- get(method, envir = globalenv(), inherits = FALSE) + } + on.exit( + if (had_method) { + assign(method, old_method, envir = globalenv()) + } else { + rm(list = method, envir = globalenv()) + }, + add = TRUE + ) + assign(method, function(condition) stop("conditionMessage failed"), globalenv()) + + stream <- function_array_stream( + function() stop(structure( + list(message = "original error"), + class = c("nanoarrow_test_broken", "error", "condition") + )), + na_int32() + ) + expect_error( + stream$get_next(), + "\\[5\\].*conditionMessage failed" + ) + expect_error(stream$get_next(), "\\[5\\].*conditionMessage failed") + stream$release() +}) + +test_that("function_array_stream() protects callback state across export and GC", { + value <- 1:3 + stream <- function_array_stream(function() value, na_int32()) + stream_copy <- nanoarrow_allocate_array_stream() + nanoarrow_pointer_export(stream, stream_copy) + rm(stream) + gc() + expect_identical(stream_copy$get_next()$length, 3L) + stream_copy$release() +}) + +test_that("function_array_stream() keeps callbacks alive through finalizer export", { + calls <- 0L + finalized <- 0L + stream <- function_array_stream(function() { + calls <<- calls + 1L + calls + }, na_int32()) + stream <- array_stream_set_finalizer(stream, function() { + finalized <<- finalized + 1L + }) + stream_copy <- nanoarrow_allocate_array_stream() + nanoarrow_pointer_export(stream, stream_copy) + rm(stream) + gc() + + expect_identical(stream_copy$get_next()$length, 1L) + stream_copy$release() + stream_copy$release() + expect_identical(finalized, 1L) +}) + +test_that("function_array_stream() rejects evaluation on a non-R thread", { + calls <- 0L + stream <- function_array_stream(function() { + calls <<- calls + 1L + 1L + }, na_int32()) + expect_identical( + .nanoarrow_function_array_stream_get_next_on_other_thread(stream), + 5L + ) + expect_error(stream$get_next(), "Can't evaluate R callback on a non-R thread") + expect_identical(calls, 0L) + stream$release() + + stream <- function_array_stream(function() NULL, na_int32()) + expect_null(stream$get_next()) + expect_identical( + .nanoarrow_function_array_stream_get_next_on_other_thread(stream), 0L + ) + stream$release() +}) + +test_that("function_array_stream() releases preserved state exactly once", { + gc() + preserved_empty() + baseline <- preserved_count() + stream <- function_array_stream(function() 1L, na_int32()) + after_create <- preserved_count() + expect_gt(after_create, baseline) + stream$release() + expect_lte(preserved_count(), baseline) + stream$release() + expect_lte(preserved_count(), baseline) + + stream <- function_array_stream(function() 1L, na_int32()) + rm(stream) + gc() + preserved_empty() + expect_lte(preserved_count(), baseline) +}) + +test_that("function_array_stream() integrates with stream conversion", { + stream <- function_array_stream( + local({ + i <- 0L + function() { + i <<- i + 1L + if (i > 2L) NULL else data.frame(x = i) + } + }), + na_struct(list(x = na_int32())) + ) + expect_identical(as.data.frame(stream), data.frame(x = 1:2)) + + stream <- function_array_stream( + local({ + i <- 0L + function() { + i <<- i + 1L + if (i > 2L) NULL else data.frame(x = i) + } + }), + na_struct(list(x = na_int32())) + ) + expect_identical( + lapply(collect_array_stream(stream), as.data.frame), + list(data.frame(x = 1L), data.frame(x = 2L)) + ) + stream$release() + + stream <- function_array_stream( + local({ + i <- 0L + function() { + i <<- i + 1L + if (i > 2L) NULL else data.frame(x = i) + } + }), + na_struct(list(x = na_int32())) + ) + expect_identical(convert_array_stream(stream), data.frame(x = 1:2)) + stream$release() +}) + +test_that("function_array_stream() methods are safe after release", { + stream <- function_array_stream(function() NULL, na_int32()) + stream$release() + expect_error(stream$get_next()) + expect_error(stream$get_schema()) +}) + +test_that("function_array_stream() safely defers release from its callback", { + stream <- NULL + stream <- function_array_stream( + function() { + stream$release() + 1L + }, na_int32() + ) + expect_null(stream$get_next()) + expect_false(nanoarrow_pointer_is_valid(stream)) + expect_silent(stream$release()) +}) + +test_that("function_array_stream() defers wrapped release across callback GC", { + released <- FALSE + holder <- new.env(parent = emptyenv()) + holder$stream <- NULL + parent <- function_array_stream( + function() { + holder$stream$release() + gc() + 1L + }, + na_int32() + ) + holder$stream <- array_stream_set_finalizer(parent, function() { + released <<- TRUE + }) + + expect_null(holder$stream$get_next()) + expect_true(released) + expect_false(nanoarrow_pointer_is_valid(holder$stream)) +}) + test_that("basic_array_stream() can create streams from batches", { stream <- basic_array_stream( list( From 2fad88bf3e9984a23b63ca7a06b39c125f103eca Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 18 Oct 2026 01:39:43 +0000 Subject: [PATCH 2/3] fix(r): Tighten function array stream validation and release handling --- r/R/array-stream.R | 14 ++--- r/R/util.R | 4 ++ r/man/function_array_stream.Rd | 5 +- r/src/array_stream.c | 85 +++++++++++----------------- r/src/init.c | 3 + r/src/nanoarrow_cpp.cc | 12 ++++ r/tests/testthat/test-array-stream.R | 57 ++++++++++++++++--- 7 files changed, 111 insertions(+), 69 deletions(-) diff --git a/r/R/array-stream.R b/r/R/array-stream.R index e27d9b235..5d63c0106 100644 --- a/r/R/array-stream.R +++ b/r/R/array-stream.R @@ -69,8 +69,9 @@ basic_array_stream <- function(batches, schema = NULL, validate = TRUE) { #' @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 Use `FALSE` to skip validation within the callback. The -#' stream's `get_next()` method can validate the returned array independently. +#' @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 @@ -92,15 +93,14 @@ basic_array_stream <- function(batches, schema = NULL, validate = TRUE) { function_array_stream <- function(fun, schema, validate = TRUE) { stopifnot(is.function(fun)) schema <- as_nanoarrow_schema(schema) - validate <- as.logical(validate)[1] - if (is.na(validate)) { - stop("`validate` must not be NA") + if (!is.logical(validate) || length(validate) != 1L || is.na(validate)) { + stop("`validate` must be a single non-NA logical value") } .Call(nanoarrow_c_function_array_stream, fun, schema, validate) } -.nanoarrow_function_array_stream_get_next <- function(fun, schema, validate) { +function_array_stream_get_next_wrapper <- function(fun, schema, validate) { result <- fun() if (is.null(result)) { NULL @@ -113,7 +113,7 @@ function_array_stream <- function(fun, schema, validate = TRUE) { } } -.nanoarrow_function_array_stream_get_next_on_other_thread <- function(x) { +function_array_stream_get_next_on_other_thread <- function(x) { .Call(nanoarrow_c_function_array_stream_get_next_on_other_thread, x) } diff --git a/r/R/util.R b/r/R/util.R index 15f029210..40b29baea 100644 --- a/r/R/util.R +++ b/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() { diff --git a/r/man/function_array_stream.Rd b/r/man/function_array_stream.Rd index 4616aad42..cf4294216 100644 --- a/r/man/function_array_stream.Rd +++ b/r/man/function_array_stream.Rd @@ -12,8 +12,9 @@ function_array_stream(fun, schema, validate = TRUE) \item{schema}{A schema coercible by \code{\link[=as_nanoarrow_schema]{as_nanoarrow_schema()}}, used to coerce and validate each result.} -\item{validate}{Use \code{FALSE} to skip validation within the callback. The -stream's \code{get_next()} method can validate the returned array independently.} +\item{validate}{A single \code{TRUE} or \code{FALSE}. Use \code{FALSE} to skip validation +within the callback. The stream's \code{get_next()} method can validate the +returned array independently.} } \value{ An \link[=as_nanoarrow_array_stream]{nanoarrow_array_stream} diff --git a/r/src/array_stream.c b/r/src/array_stream.c index 1f08613e4..c040777af 100644 --- a/r/src/array_stream.c +++ b/r/src/array_stream.c @@ -31,7 +31,7 @@ struct FunctionArrayStreamData { struct ArrowSchema schema; SEXP fun; SEXP schema_xptr; - SEXP validate; + int validate; struct ArrowError error; int error_code; int finished; @@ -48,15 +48,11 @@ struct FunctionArrayStreamCall { }; static void release_function_array_stream_data(struct FunctionArrayStreamData* data) { - if (data == NULL) { - return; - } if (data->schema.release != NULL) { data->schema.release(&data->schema); } nanoarrow_release_sexp(data->fun); nanoarrow_release_sexp(data->schema_xptr); - nanoarrow_release_sexp(data->validate); ArrowFree(data); } @@ -74,9 +70,6 @@ static void release_function_array_stream(struct ArrowArrayStream* stream) { (struct FunctionArrayStreamData*)stream->private_data; stream->release = NULL; stream->private_data = NULL; - if (data == NULL) { - return; - } if (data->evaluating) { // The callback may release its own stream. Keep its state alive until it returns. data->release_pending = 1; @@ -96,17 +89,7 @@ static int function_array_stream_get_schema(struct ArrowArrayStream* stream, struct ArrowSchema* out) { struct FunctionArrayStreamData* data = (struct FunctionArrayStreamData*)stream->private_data; - if (data == NULL) { - return EINVAL; - } - int result = ArrowSchemaDeepCopy(&data->schema, out); - if (result != NANOARROW_OK) { - if (!data->errored) { - ArrowErrorSet(&data->error, "ArrowSchemaDeepCopy() failed"); - } - return result; - } - return NANOARROW_OK; + return ArrowSchemaDeepCopy(&data->schema, out); } static SEXP function_array_stream_condition_handler(SEXP condition, void* opaque) { @@ -131,9 +114,9 @@ static void function_array_stream_eval(void* opaque) { SEXP classes = PROTECT(Rf_allocVector(STRSXP, 2)); SET_STRING_ELT(classes, 0, Rf_mkChar("error")); SET_STRING_ELT(classes, 1, Rf_mkChar("interrupt")); - SEXP callback = - PROTECT(Rf_lang4(Rf_install(".nanoarrow_function_array_stream_get_next"), data->fun, - data->schema_xptr, data->validate)); + SEXP validate = PROTECT(Rf_ScalarLogical(data->validate)); + SEXP callback = PROTECT(Rf_lang4(Rf_install("function_array_stream_get_next_wrapper"), + data->fun, data->schema_xptr, validate)); SEXP result = PROTECT(R_tryCatch(&function_array_stream_eval_call, callback, classes, &function_array_stream_condition_handler, call, NULL, NULL)); @@ -161,7 +144,7 @@ static void function_array_stream_eval(void* opaque) { UNPROTECT(1); call->code = NANOARROW_OK; } - UNPROTECT(3); + UNPROTECT(4); } static int function_array_stream_get_next(struct ArrowArrayStream* stream, @@ -197,12 +180,6 @@ static int function_array_stream_get_next(struct ArrowArrayStream* stream, } data->evaluating = 0; - if (call.code != NANOARROW_OK) { - if (out->release != NULL) { - out->release(out); - } - out->release = NULL; - } if (data->release_pending) { if (out->release != NULL) { out->release(out); @@ -210,7 +187,7 @@ static int function_array_stream_get_next(struct ArrowArrayStream* stream, out->release = NULL; // The stream itself may already have been freed by a finalizer during evaluation. release_function_array_stream_data(data); - return NANOARROW_OK; + return EINVAL; } return call.code; } @@ -228,7 +205,7 @@ SEXP nanoarrow_c_function_array_stream(SEXP fun, SEXP schema_xptr, SEXP validate ArrowErrorInit(&data->error); data->fun = R_NilValue; data->schema_xptr = R_NilValue; - data->validate = R_NilValue; + data->validate = LOGICAL(validate)[0]; stream->private_data = data; stream->get_last_error = &function_array_stream_get_last_error; stream->get_schema = &function_array_stream_get_schema; @@ -239,11 +216,7 @@ SEXP nanoarrow_c_function_array_stream(SEXP fun, SEXP schema_xptr, SEXP validate data->fun = fun; nanoarrow_preserve_sexp(schema_copy_xptr); data->schema_xptr = schema_copy_xptr; - nanoarrow_preserve_sexp(validate); - data->validate = validate; - if (ArrowSchemaDeepCopy(schema_copy, &data->schema) != NANOARROW_OK) { - Rf_error("ArrowSchemaDeepCopy() failed"); - } + schema_export(schema_copy_xptr, &data->schema); UNPROTECT(2); return stream_xptr; } @@ -364,30 +337,36 @@ struct WrapperArrayStreamData { }; static void finalize_wrapper_array_stream(struct ArrowArrayStream* array_stream) { - if (array_stream->private_data != NULL) { - struct WrapperArrayStreamData* data = - (struct WrapperArrayStreamData*)array_stream->private_data; - - // Run the parent array stream release callback - data->parent_array_stream->release(data->parent_array_stream); - - // If safe to do so, attempt to do an eager evaluation of a release - // callback that may have been registered. If it is not safe to do so, - // garbage collection will run any finalizers that have been set - // on the chain of environments leading up to the finalizer. - if (nanoarrow_is_main_thread()) { - run_user_array_stream_finalizer(data->parent_array_stream_xptr); - } + struct WrapperArrayStreamData* data = + (struct WrapperArrayStreamData*)array_stream->private_data; + array_stream->release = NULL; + array_stream->private_data = NULL; - nanoarrow_release_sexp(data->parent_array_stream_xptr); - ArrowFree(array_stream->private_data); + // Export may fail before allocating the wrapper's private data. + if (data == NULL) { + return; } - array_stream->release = NULL; + // Run the parent array stream release callback + data->parent_array_stream->release(data->parent_array_stream); + + // If safe to do so, attempt to do an eager evaluation of a release + // callback that may have been registered. If it is not safe to do so, + // garbage collection will run any finalizers that have been set + // on the chain of environments leading up to the finalizer. + if (nanoarrow_is_main_thread()) { + run_user_array_stream_finalizer(data->parent_array_stream_xptr); + } + + nanoarrow_release_sexp(data->parent_array_stream_xptr); + ArrowFree(data); } static const char* wrapper_array_stream_get_last_error( struct ArrowArrayStream* array_stream) { + if (array_stream->release == NULL) { + return "Array stream has been released"; + } struct WrapperArrayStreamData* data = (struct WrapperArrayStreamData*)array_stream->private_data; return data->parent_array_stream->get_last_error(data->parent_array_stream); diff --git a/r/src/init.c b/r/src/init.c index b40dcdc0c..a52b3979f 100644 --- a/r/src/init.c +++ b/r/src/init.c @@ -98,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); @@ -185,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}, diff --git a/r/src/nanoarrow_cpp.cc b/r/src/nanoarrow_cpp.cc index 01030051a..44888f338 100644 --- a/r/src/nanoarrow_cpp.cc +++ b/r/src/nanoarrow_cpp.cc @@ -205,6 +205,18 @@ extern "C" void nanoarrow_preserve_and_release_on_other_thread(SEXP obj) { worker.join(); } +extern "C" void nanoarrow_release_array_stream_on_other_thread( + struct ArrowArrayStream* stream) { + std::thread worker([stream] { stream->release(stream); }); + worker.join(); +} + +extern "C" SEXP nanoarrow_c_release_array_stream_on_other_thread(SEXP array_stream_xptr) { + struct ArrowArrayStream* stream = nanoarrow_array_stream_from_xptr(array_stream_xptr); + nanoarrow_release_array_stream_on_other_thread(stream); + return R_NilValue; +} + extern "C" SEXP nanoarrow_c_function_array_stream_get_next_on_other_thread( SEXP array_stream_xptr) { struct ArrowArrayStream* stream = nanoarrow_array_stream_from_xptr(array_stream_xptr); diff --git a/r/tests/testthat/test-array-stream.R b/r/tests/testthat/test-array-stream.R index cb1181c4e..8b1e4b20c 100644 --- a/r/tests/testthat/test-array-stream.R +++ b/r/tests/testthat/test-array-stream.R @@ -77,6 +77,15 @@ test_that("function_array_stream() accepts NULL immediately and zero-row arrays" stream$release() }) +test_that("function_array_stream() validates its validate argument", { + for (validate in list(NA, logical(), c(TRUE, FALSE), 1L, "TRUE")) { + expect_error( + function_array_stream(function() NULL, na_int32(), validate = validate), + "`validate` must be a single non-NA logical value" + ) + } +}) + test_that("function_array_stream() coerces common array inputs", { for (value in list(1:3, as_nanoarrow_array(1:3))) { stream <- function_array_stream(function() value, na_int32()) @@ -127,7 +136,7 @@ test_that("function_array_stream() contains interrupts and re-entrant calls", { expect_error(stream$get_next(), "Re-entrant call to function array stream") expect_error(stream$get_next(), "Re-entrant call to function array stream") expect_identical( - .nanoarrow_function_array_stream_get_next_on_other_thread(stream), 22L + function_array_stream_get_next_on_other_thread(stream), 22L ) stream$release() @@ -141,7 +150,7 @@ test_that("function_array_stream() contains interrupts and re-entrant calls", { ) expect_error(stream$get_next(), "Re-entrant call to function array stream") expect_identical( - .nanoarrow_function_array_stream_get_next_on_other_thread(stream), 22L + function_array_stream_get_next_on_other_thread(stream), 22L ) stream$release() }) @@ -249,7 +258,7 @@ test_that("function_array_stream() rejects evaluation on a non-R thread", { 1L }, na_int32()) expect_identical( - .nanoarrow_function_array_stream_get_next_on_other_thread(stream), + function_array_stream_get_next_on_other_thread(stream), 5L ) expect_error(stream$get_next(), "Can't evaluate R callback on a non-R thread") @@ -259,7 +268,7 @@ test_that("function_array_stream() rejects evaluation on a non-R thread", { stream <- function_array_stream(function() NULL, na_int32()) expect_null(stream$get_next()) expect_identical( - .nanoarrow_function_array_stream_get_next_on_other_thread(stream), 0L + function_array_stream_get_next_on_other_thread(stream), 0L ) stream$release() }) @@ -333,7 +342,7 @@ test_that("function_array_stream() methods are safe after release", { expect_error(stream$get_schema()) }) -test_that("function_array_stream() safely defers release from its callback", { +test_that("function_array_stream() reports release from its callback", { stream <- NULL stream <- function_array_stream( function() { @@ -341,7 +350,7 @@ test_that("function_array_stream() safely defers release from its callback", { 1L }, na_int32() ) - expect_null(stream$get_next()) + expect_error(stream$get_next(), "\\[22\\].*Array stream has been released") expect_false(nanoarrow_pointer_is_valid(stream)) expect_silent(stream$release()) }) @@ -362,11 +371,45 @@ test_that("function_array_stream() defers wrapped release across callback GC", { released <<- TRUE }) - expect_null(holder$stream$get_next()) + expect_error( + holder$stream$get_next(), + "\\[22\\].*Array stream has been released" + ) expect_true(released) expect_false(nanoarrow_pointer_is_valid(holder$stream)) }) +test_that("function_array_stream() releases preserved state on another thread", { + gc() + preserved_empty() + baseline <- preserved_count() + stream <- function_array_stream(function() NULL, na_int32()) + expect_gt(preserved_count(), baseline) + + release_array_stream_on_other_thread(stream) + expect_false(nanoarrow_pointer_is_valid(stream)) + # Both R references must remain queued until the main thread drains them. + expect_identical(preserved_empty(), 2) + expect_lte(preserved_count(), baseline) +}) + +test_that("wrapped stream release is safe during its finalizer", { + stream <- NULL + finalizer_results <- new.env(parent = emptyenv()) + stream <- function_array_stream(function() NULL, na_int32()) + stream <- array_stream_set_finalizer(stream, function() { + finalizer_results$schema <- tryCatch(stream$get_schema(), error = conditionMessage) + finalizer_results$next_result <- tryCatch(stream$get_next(), error = conditionMessage) + finalizer_results$release <- tryCatch(stream$release(), error = conditionMessage) + }) + + expect_silent(stream$release()) + expect_match(finalizer_results$schema, "already been released") + expect_match(finalizer_results$next_result, "already been released") + expect_null(finalizer_results$release) + expect_false(nanoarrow_pointer_is_valid(stream)) +}) + test_that("basic_array_stream() can create streams from batches", { stream <- basic_array_stream( list( From 7ed4697c265009480be611603e63d9632d1f217e Mon Sep 17 00:00:00 2001 From: eitsupi Date: Mon, 19 Oct 2026 15:49:15 +0000 Subject: [PATCH 3/3] refactor(r): Move function_array_stream() callback handling to R Build the callback environment and its tryCatch() wrapper in R so the C stream only preserves and evaluates an expression. Keep R_ToplevelExec() as the boundary that stops other non-local exits, and drop support for releasing the stream from its own callback. --- r/R/array-stream.R | 63 ++++++++--- r/man/function_array_stream.Rd | 10 +- r/src/array_stream.c | 159 +++++++++------------------ r/src/init.c | 2 +- r/tests/testthat/test-array-stream.R | 67 ++--------- 5 files changed, 116 insertions(+), 185 deletions(-) diff --git a/r/R/array-stream.R b/r/R/array-stream.R index 5d63c0106..d09ddd433 100644 --- a/r/R/array-stream.R +++ b/r/R/array-stream.R @@ -58,13 +58,15 @@ basic_array_stream <- function(batches, schema = NULL, validate = TRUE) { #' array is requested. Neither construction nor schema access calls `fun`. #' After `NULL` or an error is returned, the callback is not called again. #' -#' Evaluation is isolated in a top-level R context so non-local exits (for -#' example, invoking an `abort` restart) cannot escape through the C stream +#' 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`. -#' The stream retains `fun` until it is released; explicitly release streams -#' when `fun` or its environment refers to the stream itself. +#' `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 @@ -92,25 +94,50 @@ basic_array_stream <- function(batches, schema = NULL, validate = TRUE) { #' function_array_stream <- function(fun, schema, validate = TRUE) { stopifnot(is.function(fun)) - schema <- as_nanoarrow_schema(schema) + # 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") } - .Call(nanoarrow_c_function_array_stream, fun, schema, validate) -} - -function_array_stream_get_next_wrapper <- function(fun, schema, validate) { - 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 + # 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) { diff --git a/r/man/function_array_stream.Rd b/r/man/function_array_stream.Rd index cf4294216..290a608bf 100644 --- a/r/man/function_array_stream.Rd +++ b/r/man/function_array_stream.Rd @@ -27,13 +27,15 @@ array is requested. Neither construction nor schema access calls \code{fun}. After \code{NULL} or an error is returned, the callback is not called again. } \details{ -Evaluation is isolated in a top-level R context so non-local exits (for -example, invoking an \code{abort} restart) cannot escape through the C stream +Errors and interrupts raised by \code{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 \code{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 \code{fun}. -The stream retains \code{fun} until it is released; explicitly release streams -when \code{fun} or its environment refers to the stream itself. +\code{fun} must not release the stream (or call its methods) while it is being +evaluated. The stream retains \code{fun} until it is released; explicitly release +streams when \code{fun} or its environment refers to the stream itself. } \examples{ i <- 0L diff --git a/r/src/array_stream.c b/r/src/array_stream.c index c040777af..72240b4ae 100644 --- a/r/src/array_stream.c +++ b/r/src/array_stream.c @@ -29,33 +29,21 @@ struct FunctionArrayStreamData { struct ArrowSchema schema; - SEXP fun; - SEXP schema_xptr; - int validate; + SEXP expr; + SEXP env; struct ArrowError error; int error_code; int finished; int errored; int evaluating; - int release_pending; }; struct FunctionArrayStreamCall { struct FunctionArrayStreamData* data; struct ArrowArray* out; int code; - int caught; }; -static void release_function_array_stream_data(struct FunctionArrayStreamData* data) { - if (data->schema.release != NULL) { - data->schema.release(&data->schema); - } - nanoarrow_release_sexp(data->fun); - nanoarrow_release_sexp(data->schema_xptr); - ArrowFree(data); -} - static void function_array_stream_set_error(struct FunctionArrayStreamData* data, int code, const char* message) { if (!data->errored) { @@ -70,12 +58,12 @@ static void release_function_array_stream(struct ArrowArrayStream* stream) { (struct FunctionArrayStreamData*)stream->private_data; stream->release = NULL; stream->private_data = NULL; - if (data->evaluating) { - // The callback may release its own stream. Keep its state alive until it returns. - data->release_pending = 1; - return; + if (data->schema.release != NULL) { + data->schema.release(&data->schema); } - release_function_array_stream_data(data); + nanoarrow_release_sexp(data->expr); + nanoarrow_release_sexp(data->env); + ArrowFree(data); } static const char* function_array_stream_get_last_error(struct ArrowArrayStream* stream) { @@ -92,50 +80,23 @@ static int function_array_stream_get_schema(struct ArrowArrayStream* stream, return ArrowSchemaDeepCopy(&data->schema, out); } -static SEXP function_array_stream_condition_handler(SEXP condition, void* opaque) { - struct FunctionArrayStreamCall* call = (struct FunctionArrayStreamCall*)opaque; - call->caught = 1; - if (Rf_inherits(condition, "interrupt")) { - call->caught = 2; - } - SEXP message_call = PROTECT(Rf_lang2(Rf_install("conditionMessage"), condition)); - SEXP message = Rf_eval(message_call, R_BaseEnv); - UNPROTECT(1); - return message; -} - -static SEXP function_array_stream_eval_call(void* opaque) { - return Rf_eval((SEXP)opaque, nanoarrow_ns_pkg); -} - +// 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 classes = PROTECT(Rf_allocVector(STRSXP, 2)); - SET_STRING_ELT(classes, 0, Rf_mkChar("error")); - SET_STRING_ELT(classes, 1, Rf_mkChar("interrupt")); - SEXP validate = PROTECT(Rf_ScalarLogical(data->validate)); - SEXP callback = PROTECT(Rf_lang4(Rf_install("function_array_stream_get_next_wrapper"), - data->fun, data->schema_xptr, validate)); - SEXP result = - PROTECT(R_tryCatch(&function_array_stream_eval_call, callback, classes, - &function_array_stream_condition_handler, call, NULL, NULL)); + SEXP result = PROTECT(Rf_eval(data->expr, data->env)); if (data->errored) { - call->code = data->error_code; - } else if (call->caught) { - const char* message = Rf_translateCharUTF8(STRING_ELT(result, 0)); - if (call->caught == 2) { - struct ArrowError error; - ArrowErrorSet(&error, "R callback interrupted: %s", message); - function_array_stream_set_error(data, EIO, ArrowErrorMessage(&error)); - } else { - function_array_stream_set_error(data, EIO, message); - } + // 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 { + } 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); @@ -143,8 +104,12 @@ static void function_array_stream_eval(void* opaque) { 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(4); + UNPROTECT(1); } static int function_array_stream_get_next(struct ArrowArrayStream* stream, @@ -170,54 +135,44 @@ static int function_array_stream_get_next(struct ArrowArrayStream* stream, } data->evaluating = 1; - struct FunctionArrayStreamCall call = {data, out, NANOARROW_OK, 0}; - // R_tryCatch() does not catch restart jumps or failures in conditionMessage(). - // All R API calls must stay inside this outer boundary. + 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; - - if (data->release_pending) { - if (out->release != NULL) { - out->release(out); - } - out->release = NULL; - // The stream itself may already have been freed by a finalizer during evaluation. - release_function_array_stream_data(data); - return EINVAL; - } return call.code; } -SEXP nanoarrow_c_function_array_stream(SEXP fun, SEXP schema_xptr, SEXP validate) { +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); - SEXP schema_copy_xptr = PROTECT(nanoarrow_schema_owning_xptr()); - struct ArrowSchema* schema_copy = nanoarrow_output_schema_from_xptr(schema_copy_xptr); - schema_export(schema_xptr, schema_copy); struct FunctionArrayStreamData* data = (struct FunctionArrayStreamData*)ArrowMalloc(sizeof(*data)); check_trivial_alloc(data, "struct FunctionArrayStreamData"); memset(data, 0, sizeof(*data)); ArrowErrorInit(&data->error); - data->fun = R_NilValue; - data->schema_xptr = R_NilValue; - data->validate = LOGICAL(validate)[0]; + 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(fun); - data->fun = fun; - nanoarrow_preserve_sexp(schema_copy_xptr); - data->schema_xptr = schema_copy_xptr; - schema_export(schema_copy_xptr, &data->schema); - UNPROTECT(2); + 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; } @@ -337,36 +292,30 @@ struct WrapperArrayStreamData { }; static void finalize_wrapper_array_stream(struct ArrowArrayStream* array_stream) { - struct WrapperArrayStreamData* data = - (struct WrapperArrayStreamData*)array_stream->private_data; - array_stream->release = NULL; - array_stream->private_data = NULL; - - // Export may fail before allocating the wrapper's private data. - if (data == NULL) { - return; - } - - // Run the parent array stream release callback - data->parent_array_stream->release(data->parent_array_stream); + if (array_stream->private_data != NULL) { + struct WrapperArrayStreamData* data = + (struct WrapperArrayStreamData*)array_stream->private_data; + + // Run the parent array stream release callback + data->parent_array_stream->release(data->parent_array_stream); + + // If safe to do so, attempt to do an eager evaluation of a release + // callback that may have been registered. If it is not safe to do so, + // garbage collection will run any finalizers that have been set + // on the chain of environments leading up to the finalizer. + if (nanoarrow_is_main_thread()) { + run_user_array_stream_finalizer(data->parent_array_stream_xptr); + } - // If safe to do so, attempt to do an eager evaluation of a release - // callback that may have been registered. If it is not safe to do so, - // garbage collection will run any finalizers that have been set - // on the chain of environments leading up to the finalizer. - if (nanoarrow_is_main_thread()) { - run_user_array_stream_finalizer(data->parent_array_stream_xptr); + nanoarrow_release_sexp(data->parent_array_stream_xptr); + ArrowFree(array_stream->private_data); } - nanoarrow_release_sexp(data->parent_array_stream_xptr); - ArrowFree(data); + array_stream->release = NULL; } static const char* wrapper_array_stream_get_last_error( struct ArrowArrayStream* array_stream) { - if (array_stream->release == NULL) { - return "Array stream has been released"; - } struct WrapperArrayStreamData* data = (struct WrapperArrayStreamData*)array_stream->private_data; return data->parent_array_stream->get_last_error(data->parent_array_stream); diff --git a/r/src/init.c b/r/src/init.c index a52b3979f..334e08f3d 100644 --- a/r/src/init.c +++ b/r/src/init.c @@ -31,7 +31,7 @@ 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 fun, SEXP schema_xptr, SEXP validate); +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); diff --git a/r/tests/testthat/test-array-stream.R b/r/tests/testthat/test-array-stream.R index 8b1e4b20c..4eab65829 100644 --- a/r/tests/testthat/test-array-stream.R +++ b/r/tests/testthat/test-array-stream.R @@ -202,7 +202,14 @@ test_that("function_array_stream() reports conditionMessage failures safely", { }, add = TRUE ) - assign(method, function(condition) stop("conditionMessage failed"), globalenv()) + # stop() calls conditionMessage() once itself, so only fail on later calls to + # exercise the handler in the R-level tryCatch(). + n_calls <- 0L + assign(method, function(condition) { + n_calls <<- n_calls + 1L + if (n_calls > 1L) stop("conditionMessage failed") + "original error" + }, globalenv()) stream <- function_array_stream( function() stop(structure( @@ -213,9 +220,9 @@ test_that("function_array_stream() reports conditionMessage failures safely", { ) expect_error( stream$get_next(), - "\\[5\\].*conditionMessage failed" + "\\[5\\].*error message could not be retrieved" ) - expect_error(stream$get_next(), "\\[5\\].*conditionMessage failed") + expect_error(stream$get_next(), "\\[5\\].*error message could not be retrieved") stream$release() }) @@ -342,43 +349,6 @@ test_that("function_array_stream() methods are safe after release", { expect_error(stream$get_schema()) }) -test_that("function_array_stream() reports release from its callback", { - stream <- NULL - stream <- function_array_stream( - function() { - stream$release() - 1L - }, na_int32() - ) - expect_error(stream$get_next(), "\\[22\\].*Array stream has been released") - expect_false(nanoarrow_pointer_is_valid(stream)) - expect_silent(stream$release()) -}) - -test_that("function_array_stream() defers wrapped release across callback GC", { - released <- FALSE - holder <- new.env(parent = emptyenv()) - holder$stream <- NULL - parent <- function_array_stream( - function() { - holder$stream$release() - gc() - 1L - }, - na_int32() - ) - holder$stream <- array_stream_set_finalizer(parent, function() { - released <<- TRUE - }) - - expect_error( - holder$stream$get_next(), - "\\[22\\].*Array stream has been released" - ) - expect_true(released) - expect_false(nanoarrow_pointer_is_valid(holder$stream)) -}) - test_that("function_array_stream() releases preserved state on another thread", { gc() preserved_empty() @@ -393,23 +363,6 @@ test_that("function_array_stream() releases preserved state on another thread", expect_lte(preserved_count(), baseline) }) -test_that("wrapped stream release is safe during its finalizer", { - stream <- NULL - finalizer_results <- new.env(parent = emptyenv()) - stream <- function_array_stream(function() NULL, na_int32()) - stream <- array_stream_set_finalizer(stream, function() { - finalizer_results$schema <- tryCatch(stream$get_schema(), error = conditionMessage) - finalizer_results$next_result <- tryCatch(stream$get_next(), error = conditionMessage) - finalizer_results$release <- tryCatch(stream$release(), error = conditionMessage) - }) - - expect_silent(stream$release()) - expect_match(finalizer_results$schema, "already been released") - expect_match(finalizer_results$next_result, "already been released") - expect_null(finalizer_results$release) - expect_false(nanoarrow_pointer_is_valid(stream)) -}) - test_that("basic_array_stream() can create streams from batches", { stream <- basic_array_stream( list(