diff --git a/r/R/flight.R b/r/R/flight.R index f56308f958495..0bd661e58d565 100644 --- a/r/R/flight.R +++ b/r/R/flight.R @@ -56,9 +56,11 @@ flight_disconnect <- function(client) { #' @param overwrite logical: if `path` exists on `client` already, should we #' replace it with the contents of `data`? Default is `TRUE`; if `FALSE` and #' `path` exists, the function will error. +#' @param max_chunksize integer: Maximum size for RecordBatch chunks when a `data.frame` is sent. +#' Individual chunks may be smaller depending on the chunk layout of individual columns. #' @return `client`, invisibly. #' @export -flight_put <- function(client, data, path, overwrite = TRUE) { +flight_put <- function(client, data, path, overwrite = TRUE, max_chunksize = NULL) { assert_is(data, c("data.frame", "Table", "RecordBatch")) if (!overwrite && flight_path_exists(client, path)) { @@ -70,8 +72,13 @@ flight_put <- function(client, data, path, overwrite = TRUE) { py_data <- reticulate::r_to_py(data) writer <- client$do_put(descriptor_for_path(path), py_data$schema)[[1]] - if (inherits(data, "RecordBatch")) { + if (inherits(data, "RecordBatch") && !is.null(max_chunksize)) { + warning("`max_chunksize` is not supported for flight_put with RecordBatch") writer$write_batch(py_data) + } else if (inherits(data, "RecordBatch")) { + writer$write_batch(py_data) + } else if (!is.null(max_chunksize)) { + writer$write_table(py_data, max_chunksize) } else { writer$write_table(py_data) } diff --git a/r/man/flight_put.Rd b/r/man/flight_put.Rd index 13a8da16fead5..c306b0f7bb9e0 100644 --- a/r/man/flight_put.Rd +++ b/r/man/flight_put.Rd @@ -4,7 +4,7 @@ \alias{flight_put} \title{Send data to a Flight server} \usage{ -flight_put(client, data, path, overwrite = TRUE) +flight_put(client, data, path, overwrite = TRUE, max_chunksize = NULL) } \arguments{ \item{client}{\code{pyarrow.flight.FlightClient}, as returned by \code{\link[=flight_connect]{flight_connect()}}} @@ -16,6 +16,9 @@ flight_put(client, data, path, overwrite = TRUE) \item{overwrite}{logical: if \code{path} exists on \code{client} already, should we replace it with the contents of \code{data}? Default is \code{TRUE}; if \code{FALSE} and \code{path} exists, the function will error.} + +\item{max_chunksize}{integer: Maximum size for RecordBatch chunks when a \code{data.frame} is sent. +Individual chunks may be smaller depending on the chunk layout of individual columns.} } \value{ \code{client}, invisibly. diff --git a/r/tests/testthat/test-python-flight.R b/r/tests/testthat/test-python-flight.R index 6fdf38f815ba6..f41abed1ffa80 100644 --- a/r/tests/testthat/test-python-flight.R +++ b/r/tests/testthat/test-python-flight.R @@ -37,6 +37,20 @@ if (process_is_running("demo_flight_server")) { regexp = 'data must be a "data.frame", "Table", or "RecordBatch"' ) }) + + test_that("flight_put with max_chunksize", { + flight_put(client, example_data, path = flight_obj, max_chunksize = 1) + expect_true(flight_path_exists(client, flight_obj)) + expect_true(flight_obj %in% list_flights(client)) + expect_warning( + flight_put(client, record_batch(example_data), path = flight_obj, max_chunksize = 123), + regexp = "`max_chunksize` is not supported for flight_put with RecordBatch" + ) + expect_error( + flight_put(client, Array$create(c(1:3)), path = flight_obj), + regexp = 'data must be a "data.frame", "Table", or "RecordBatch"' + ) + }) test_that("flight_get", { expect_identical(as.data.frame(flight_get(client, flight_obj)), example_data)