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..d09ddd433 100644 --- a/r/R/array-stream.R +++ b/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] 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 new file mode 100644 index 000000000..290a608bf --- /dev/null +++ b/r/man/function_array_stream.Rd @@ -0,0 +1,54 @@ +% 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}{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} +} +\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{ +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}. +\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 +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..72240b4ae 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,155 @@ #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. diff --git a/r/src/init.c b/r/src/init.c index 97f8f4c9e..334e08f3d 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 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}, diff --git a/r/src/nanoarrow_cpp.cc b/r/src/nanoarrow_cpp.cc index 83ebe07ac..44888f338 100644 --- a/r/src/nanoarrow_cpp.cc +++ b/r/src/nanoarrow_cpp.cc @@ -205,6 +205,31 @@ 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); + 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..4eab65829 100644 --- a/r/tests/testthat/test-array-stream.R +++ b/r/tests/testthat/test-array-stream.R @@ -26,6 +26,343 @@ 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() 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()) + 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( + 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( + 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 + ) + # 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( + list(message = "original error"), + class = c("nanoarrow_test_broken", "error", "condition") + )), + na_int32() + ) + expect_error( + stream$get_next(), + "\\[5\\].*error message could not be retrieved" + ) + expect_error(stream$get_next(), "\\[5\\].*error message could not be retrieved") + 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( + 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( + 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() 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("basic_array_stream() can create streams from batches", { stream <- basic_array_stream( list(