feat: import_stream internal method for Series to support Arrow C stream interface

NAMESPACE

@@ -157,6 +157,7 @@ S3method(as_polars_df,RPolarsLazyFrame)
 S3method(as_polars_df,RPolarsLazyGroupBy)
 S3method(as_polars_df,RPolarsRollingGroupBy)
 S3method(as_polars_df,RPolarsSeries)
+S3method(as_polars_df,RecordBatchReader)
 S3method(as_polars_df,data.frame)
 S3method(as_polars_df,default)
 S3method(as_polars_df,nanoarrow_array)
@@ -171,6 +172,7 @@ S3method(as_polars_series,RPolarsChainedThen)
 S3method(as_polars_series,RPolarsExpr)
 S3method(as_polars_series,RPolarsSeries)
 S3method(as_polars_series,RPolarsThen)
+S3method(as_polars_series,RecordBatchReader)
 S3method(as_polars_series,clock_sys_time)
 S3method(as_polars_series,clock_time_point)
 S3method(as_polars_series,clock_zoned_time)

NEWS.md

@@ -15,6 +15,9 @@
 - New S3 methods `nanoarrow::as_nanoarrow_array_stream()` and `nanoarrow::infer_nanoarrow_schema()`
   for `RPolarsSeries` (#1076).
 - New method `$dt$is_leap_year()` (#1077).
+- `as_polars_df()` and `as_polars_series()` supports `arrow::RecordBatchReader` (#1078).
+- The new `experimental` argument for `as_polars_df(<ArrowTabular>)`, `as_polars_df(<RecordBatchReader>)`,
+  `as_polars_series(<nanoarrow_array_stream>)` (#1078).
 
 ## Polars R Package 0.16.3
 

R/as_polars.R

@@ -7,35 +7,36 @@
 #' [$collect()][LazyFrame_collect] or [$fetch()][LazyFrame_fetch], depending on
 #' whether the number of rows to fetch is infinite or not.
 #' @rdname as_polars_df
+#' @inheritParams as_polars_series
 #' @param x Object to convert to a polars DataFrame.
 #' @param ... Additional arguments passed to methods.
 #' @return a [DataFrame][DataFrame_class]
 #' @examplesIf requireNamespace("arrow", quietly = TRUE)
 #' # Convert the row names of a data frame to a column
 #' as_polars_df(mtcars, rownames = "car")
 #'
-#' # Convert an arrow Table to a polars DataFrame
-#' at = arrow::arrow_table(x = 1:5, y = 6:10)
-#' as_polars_df(at)
-#'
-#' # Convert an arrow Table, with renaming all columns
+#' # Convert a data frame, with renaming all columns
 #' as_polars_df(
-#'   at,
+#'   data.frame(x = 1, y = 2),
 #'   schema = c("a", "b")
 #' )
 #'
-#' # Convert an arrow Table, with renaming and casting all columns
+#' # Convert a data frame, with renaming and casting all columns
 #' as_polars_df(
-#'   at,
+#'   data.frame(x = 1, y = 2),
 #'   schema = list(b = pl$Int64, a = pl$String)
 #' )
 #'
-#' # Convert an arrow Table, with casting some columns
+#' # Convert a data frame, with casting some columns
 #' as_polars_df(
-#'   at,
+#'   data.frame(x = 1, y = 2),
 #'   schema_overrides = list(y = pl$String) # cast some columns
 #' )
 #'
+#' # Convert an arrow Table to a polars DataFrame
+#' at = arrow::arrow_table(x = 1:5, y = 6:10)
+#' as_polars_df(at)
+#'
 #' # Create a polars DataFrame from a data.frame
 #' lf = as_polars_df(mtcars)$lazy()
 #'
@@ -212,13 +213,33 @@ as_polars_df.ArrowTabular = function(
     ...,
     rechunk = TRUE,
     schema = NULL,
-    schema_overrides = NULL) {
+    schema_overrides = NULL,
+    experimental = FALSE) {
   arrow_to_rpldf(
     x,
     rechunk = rechunk,
     schema = schema,
-    schema_overrides = schema_overrides
-  )
+    schema_overrides = schema_overrides,
+    experimental = experimental
+  ) |>
+    result() |>
+    unwrap("in as_polars_df():")
+}
+
+
+#' @rdname as_polars_df
+#' @export
+as_polars_df.RecordBatchReader = function(x, ..., experimental = FALSE) {
+  uw = \(res) unwrap(res, "in as_polars_df(<RecordBatchReader>):")
+
+  if (isTRUE(experimental)) {
+    as_polars_series(x, name = "")$to_frame()$unnest("") |>
+      result() |>
+      uw()
+  } else {
+    .pr$DataFrame$from_arrow_record_batches(x$batches()) |>
+      uw()
+  }
 }
 
 
@@ -248,19 +269,12 @@ as_polars_df.nanoarrow_array = function(x, ...) {
 #' @rdname as_polars_df
 #' @export
 as_polars_df.nanoarrow_array_stream = function(x, ...) {
-  if (!inherits(nanoarrow::infer_nanoarrow_ptype(x$get_schema()), "data.frame")) {
+  if (!identical(nanoarrow::nanoarrow_schema_parse(x$get_schema())$type, "struct")) {
     Err_plain("Can't convert non-struct array stream to RPolarsDataFrame") |>
       unwrap("in as_polars_df(<nanoarrow_array_stream>):")
   }
 
-  series = as_polars_series.nanoarrow_array_stream(x, name = NULL)
-
-  if (length(series)) {
-    series$to_frame()$unnest("")
-  } else {
-    # TODO: support 0-length array stream
-    pl$DataFrame()
-  }
+  as_polars_series.nanoarrow_array_stream(x, name = "")$to_frame()$unnest("")
 }
 
 
@@ -397,6 +411,20 @@ as_polars_series.Array = function(x, name = NULL, ..., rechunk = TRUE) {
 as_polars_series.ChunkedArray = as_polars_series.Array
 
 
+#' @rdname as_polars_series
+#' @export
+as_polars_series.RecordBatchReader = function(x, name = NULL, ...) {
+  stream_out = polars_allocate_array_stream()
+  x$export_to_c(stream_out)
+
+  .pr$Series$import_stream(
+    name %||% "",
+    stream_out
+  ) |>
+    unwrap("in as_polars_series(<RecordBatchReader>):")
+}
+
+
 #' @rdname as_polars_series
 #' @export
 as_polars_series.nanoarrow_array = function(x, name = NULL, ...) {
@@ -406,26 +434,39 @@ as_polars_series.nanoarrow_array = function(x, name = NULL, ...) {
 }
 
 
+#' @param experimental If `TRUE`, use experimental Arrow C stream interface inside the function.
+#' This argument is experimental and may be removed in the future.
 #' @rdname as_polars_series
 #' @export
-as_polars_series.nanoarrow_array_stream = function(x, name = NULL, ...) {
+as_polars_series.nanoarrow_array_stream = function(x, name = NULL, ..., experimental = FALSE) {
   on.exit(x$release())
 
-  list_of_arrays = nanoarrow::collect_array_stream(x, validate = FALSE)
+  if (isTRUE(experimental)) {
+    stream_out = polars_allocate_array_stream()
+    nanoarrow::nanoarrow_pointer_export(x, stream_out)
 
-  if (length(list_of_arrays) < 1L) {
-    # TODO: support 0-length array stream
-    out = pl$Series(name = name)
-  } else {
-    out = as_polars_series.nanoarrow_array(list_of_arrays[[1L]], name = name)
-    lapply(
-      list_of_arrays[-1L],
-      \(array) .pr$Series$append_mut(out, as_polars_series.nanoarrow_array(array))
+    .pr$Series$import_stream(
+      name %||% "",
+      stream_out
     ) |>
-      invisible()
-  }
+      unwrap("in as_polars_series(<nanoarrow_array_stream>):")
+  } else {
+    list_of_arrays = nanoarrow::collect_array_stream(x, validate = FALSE)
 
-  out
+    if (length(list_of_arrays) < 1L) {
+      # TODO: support 0-length array stream
+      out = pl$Series(name = name)
+    } else {
+      out = as_polars_series.nanoarrow_array(list_of_arrays[[1L]], name = name)
+      lapply(
+        list_of_arrays[-1L],
+        \(array) .pr$Series$append_mut(out, as_polars_series.nanoarrow_array(array))
+      ) |>
+        invisible()
+    }
+
+    out
+  }
 }
 
 

R/construction.R

@@ -9,9 +9,11 @@
 #' If schema names or types do not match `x`, the columns will be renamed/recast.
 #' If `NULL` (default), convert columns as is.
 #' @param schema_overrides named list of DataTypes. Cast some columns to the DataType.
+#' @param experimental If `TRUE`, use the Arrow C stream interface.
 #' @noRd
 #' @return RPolarsDataFrame
-arrow_to_rpldf = function(at, schema = NULL, schema_overrides = NULL, rechunk = TRUE) {
+arrow_to_rpldf = function(
+    at, schema = NULL, schema_overrides = NULL, rechunk = TRUE, ..., experimental = FALSE) {
   # new column names by schema, #todo get names if schema not NULL
   n_cols = at$num_columns
 
@@ -53,9 +55,12 @@ arrow_to_rpldf = function(at, schema = NULL, schema_overrides = NULL, rechunk =
     if (tbl$num_rows == 0L) {
       rdf = pl$DataFrame() # TODO: support creating 0-row DataFrame
     } else {
-      rdf = unwrap(
-        .pr$DataFrame$from_arrow_record_batches(arrow::as_record_batch_reader(tbl)$batches())
-      )
+      if (isTRUE(experimental)) {
+        rdf = as_polars_series(arrow::as_record_batch_reader(tbl))$to_frame()$unnest("")
+      } else {
+        rdf = .pr$DataFrame$from_arrow_record_batches(arrow::as_record_batch_reader(tbl)$batches()) |>
+          unwrap()
+      }
     }
   } else {
     rdf = pl$DataFrame()

R/extendr-wrappers.R

@@ -10,6 +10,8 @@
 #' @useDynLib polars, .registration = TRUE
 NULL
 
+polars_allocate_array_stream <- function() .Call(wrap__polars_allocate_array_stream)
+
 all_horizontal <- function(dotdotdot) .Call(wrap__all_horizontal, dotdotdot)
 
 any_horizontal <- function(dotdotdot) .Call(wrap__any_horizontal, dotdotdot)
@@ -58,12 +60,6 @@ struct_ <- function(exprs, eager, schema) .Call(wrap__struct_, exprs, eager, sch
 
 dtype_str_repr <- function(dtype) .Call(wrap__dtype_str_repr, dtype)
 
-new_arrow_stream <- function() .Call(wrap__new_arrow_stream)
-
-arrow_stream_to_df <- function(robj_str) .Call(wrap__arrow_stream_to_df, robj_str)
-
-arrow_stream_to_series <- function(robj_str) .Call(wrap__arrow_stream_to_series, robj_str)
-
 mem_address <- function(robj) .Call(wrap__mem_address, robj)
 
 clone_robj <- function(robj) .Call(wrap__clone_robj, robj)
@@ -1376,7 +1372,7 @@ RPolarsSeries$struct_fields <- function() .Call(wrap__RPolarsSeries__struct_fiel
 
 RPolarsSeries$export_stream <- function(stream_ptr, pl_flavor) invisible(.Call(wrap__RPolarsSeries__export_stream, self, stream_ptr, pl_flavor))
 
-RPolarsSeries$from_arrow_array_stream_str <- function(name, robj_str) .Call(wrap__RPolarsSeries__from_arrow_array_stream_str, name, robj_str)
+RPolarsSeries$import_stream <- function(name, stream_ptr) .Call(wrap__RPolarsSeries__import_stream, name, stream_ptr)
 
 RPolarsSeries$from_arrow_array_robj <- function(name, array) .Call(wrap__RPolarsSeries__from_arrow_array_robj, name, array)
 

src/rust/src/arrow_interop/mod.rs

@@ -1,5 +1,7 @@
 pub mod to_rust;
 
+use polars_core::utils::arrow;
+
 use extendr_api::prelude::*;
 use std::result::Result;
 
@@ -61,3 +63,16 @@ impl RPackage for NanoArrowRPackage {
         "#)
     }
 }
+
+#[extendr]
+pub fn polars_allocate_array_stream() -> Robj {
+    let aas = Box::new(arrow::ffi::ArrowArrayStream::empty());
+    let x = Box::leak(aas); // leak box to make lifetime static
+    let x = x as *mut arrow::ffi::ArrowArrayStream;
+    format!("{:?}", x as usize).into()
+}
+
+extendr_module! {
+    mod arrow_interop;
+    fn polars_allocate_array_stream;
+}