feat(r/sedonadb): Add FFI support for ScalarUDF and TableProvider (#214)

Author: paleolimbot

Cargo.lock

@@ -5,6 +5,7 @@ S3method("[[<-",savvy_sedonadb__sealed)
 S3method(as.data.frame,sedonadb_dataframe)
 S3method(as_nanoarrow_array_stream,sedonadb_dataframe)
 S3method(as_sedonadb_dataframe,data.frame)
+S3method(as_sedonadb_dataframe,datafusion_table_provider)
 S3method(as_sedonadb_dataframe,nanoarrow_array)
 S3method(as_sedonadb_dataframe,nanoarrow_array_stream)
 S3method(as_sedonadb_dataframe,sedonadb_dataframe)
@@ -24,6 +25,7 @@ export(sd_count)
 export(sd_drop_view)
 export(sd_preview)
 export(sd_read_parquet)
+export(sd_register_udf)
 export(sd_sql)
 export(sd_to_view)
 export(sd_view)

r/sedonadb/NAMESPACE

@@ -77,6 +77,12 @@ NULL
   }
 }
 
+`InternalContext_data_frame_from_table_provider` <- function(self) {
+  function(`provider_xptr`) {
+    .savvy_wrap_InternalDataFrame(.Call(savvy_InternalContext_data_frame_from_table_provider__impl, `self`, `provider_xptr`))
+  }
+}
+
 `InternalContext_deregister_table` <- function(self) {
   function(`table_ref`) {
     invisible(.Call(savvy_InternalContext_deregister_table__impl, `self`, `table_ref`))
@@ -89,6 +95,18 @@ NULL
   }
 }
 
+`InternalContext_register_scalar_udf` <- function(self) {
+  function(`scalar_udf_xptr`) {
+    invisible(.Call(savvy_InternalContext_register_scalar_udf__impl, `self`, `scalar_udf_xptr`))
+  }
+}
+
+`InternalContext_scalar_udf_xptr` <- function(self) {
+  function(`name`) {
+    .Call(savvy_InternalContext_scalar_udf_xptr__impl, `self`, `name`)
+  }
+}
+
 `InternalContext_sql` <- function(self) {
   function(`query`) {
     .savvy_wrap_InternalDataFrame(.Call(savvy_InternalContext_sql__impl, `self`, `query`))
@@ -105,8 +123,11 @@ NULL
   e <- new.env(parent = emptyenv())
   e$.ptr <- ptr
   e$`data_frame_from_array_stream` <- `InternalContext_data_frame_from_array_stream`(ptr)
+  e$`data_frame_from_table_provider` <- `InternalContext_data_frame_from_table_provider`(ptr)
   e$`deregister_table` <- `InternalContext_deregister_table`(ptr)
   e$`read_parquet` <- `InternalContext_read_parquet`(ptr)
+  e$`register_scalar_udf` <- `InternalContext_register_scalar_udf`(ptr)
+  e$`scalar_udf_xptr` <- `InternalContext_scalar_udf_xptr`(ptr)
   e$`sql` <- `InternalContext_sql`(ptr)
   e$`view` <- `InternalContext_view`(ptr)
 
@@ -179,8 +200,8 @@ class(`InternalContext`) <- c("InternalContext__bundle", "savvy_sedonadb__sealed
 }
 
 `InternalDataFrame_to_arrow_stream` <- function(self) {
-  function(`out`) {
-    invisible(.Call(savvy_InternalDataFrame_to_arrow_stream__impl, `self`, `out`))
+  function(`out`, `requested_schema_xptr`) {
+    invisible(.Call(savvy_InternalDataFrame_to_arrow_stream__impl, `self`, `out`, `requested_schema_xptr`))
   }
 }
 
@@ -191,6 +212,12 @@ class(`InternalContext`) <- c("InternalContext__bundle", "savvy_sedonadb__sealed
   }
 }
 
+`InternalDataFrame_to_provider` <- function(self) {
+  function() {
+    .Call(savvy_InternalDataFrame_to_provider__impl, `self`)
+  }
+}
+
 `InternalDataFrame_to_view` <- function(self) {
   function(`ctx`, `table_ref`, `overwrite`) {
     `ctx` <- .savvy_extract_ptr(`ctx`, "InternalContext")
@@ -210,6 +237,7 @@ class(`InternalContext`) <- c("InternalContext__bundle", "savvy_sedonadb__sealed
   e$`to_arrow_schema` <- `InternalDataFrame_to_arrow_schema`(ptr)
   e$`to_arrow_stream` <- `InternalDataFrame_to_arrow_stream`(ptr)
   e$`to_parquet` <- `InternalDataFrame_to_parquet`(ptr)
+  e$`to_provider` <- `InternalDataFrame_to_provider`(ptr)
   e$`to_view` <- `InternalDataFrame_to_view`(ptr)
 
   class(e) <- c("InternalDataFrame", "savvy_sedonadb__sealed")

r/sedonadb/R/000-wrappers.R

@@ -80,6 +80,23 @@ sd_view <- function(table_ref) {
   new_sedonadb_dataframe(ctx, df)
 }
 
+#' Register a user-defined function
+#'
+#' Several types of user-defined functions can be registered into a session
+#' context. Currently, the only implemented variety is an external pointer
+#' to a Rust `FFI_ScalarUDF`, an example of which is available from the
+#' [DataFusion Python documentation](https://github.com/apache/datafusion-python/blob/6f3b1cab75cfaa0cdf914f9b6fa023cb9afccd7d/examples/datafusion-ffi-example/src/scalar_udf.rs).
+#'
+#' @param udf An object of class 'datafusion_scalar_udf'
+#'
+#' @returns NULL, invisibly
+#' @export
+#'
+sd_register_udf <- function(udf) {
+  ctx <- ctx()
+  ctx$register_scalar_udf(udf)
+}
+
 # We use just one context for now. In theory we could support multiple
 # contexts with a shared runtime, which would scope the registration
 # of various components more cleanly from the runtime.

r/sedonadb/R/context.R

@@ -67,6 +67,13 @@ as_sedonadb_dataframe.nanoarrow_array_stream <- function(x, ..., schema = NULL,
   as_sedonadb_dataframe(new_sedonadb_dataframe(ctx, df), schema = schema)
 }
 
+#' @export
+as_sedonadb_dataframe.datafusion_table_provider <- function(x, ..., schema = NULL) {
+  ctx <- ctx()
+  df <- ctx$data_frame_from_table_provider(x)
+  new_sedonadb_dataframe(ctx, df)
+}
+
 #' Count rows in a DataFrame
 #'
 #' @param .data A sedonadb_dataframe
@@ -297,9 +304,13 @@ infer_nanoarrow_schema.sedonadb_dataframe <- function(x, ...) {
 
 #' @importFrom nanoarrow as_nanoarrow_array_stream
 #' @export
-as_nanoarrow_array_stream.sedonadb_dataframe <- function(x, ...) {
+as_nanoarrow_array_stream.sedonadb_dataframe <- function(x, ..., schema = NULL) {
+  if (!is.null(schema)) {
+    schema <- nanoarrow::as_nanoarrow_schema(schema)
+  }
+
   stream <- nanoarrow::nanoarrow_allocate_array_stream()
-  x$df$to_arrow_stream(stream)
+  x$df$to_arrow_stream(stream, schema)
   stream
 }
 

r/sedonadb/R/dataframe.R

@@ -15,9 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include <Rinternals.h>
 #include <stdint.h>
 
+#include <Rinternals.h>
+
 #include <R_ext/Parse.h>
 
 #include "rust/api.h"
@@ -83,6 +84,13 @@ SEXP savvy_InternalContext_data_frame_from_array_stream__impl(
   return handle_result(res);
 }
 
+SEXP savvy_InternalContext_data_frame_from_table_provider__impl(
+    SEXP self__, SEXP c_arg__provider_xptr) {
+  SEXP res = savvy_InternalContext_data_frame_from_table_provider__ffi(
+      self__, c_arg__provider_xptr);
+  return handle_result(res);
+}
+
 SEXP savvy_InternalContext_deregister_table__impl(SEXP self__,
                                                   SEXP c_arg__table_ref) {
   SEXP res =
@@ -100,6 +108,19 @@ SEXP savvy_InternalContext_read_parquet__impl(SEXP self__, SEXP c_arg__paths) {
   return handle_result(res);
 }
 
+SEXP savvy_InternalContext_register_scalar_udf__impl(
+    SEXP self__, SEXP c_arg__scalar_udf_xptr) {
+  SEXP res = savvy_InternalContext_register_scalar_udf__ffi(
+      self__, c_arg__scalar_udf_xptr);
+  return handle_result(res);
+}
+
+SEXP savvy_InternalContext_scalar_udf_xptr__impl(SEXP self__,
+                                                 SEXP c_arg__name) {
+  SEXP res = savvy_InternalContext_scalar_udf_xptr__ffi(self__, c_arg__name);
+  return handle_result(res);
+}
+
 SEXP savvy_InternalContext_sql__impl(SEXP self__, SEXP c_arg__query) {
   SEXP res = savvy_InternalContext_sql__ffi(self__, c_arg__query);
   return handle_result(res);
@@ -149,9 +170,10 @@ SEXP savvy_InternalDataFrame_to_arrow_schema__impl(SEXP self__,
   return handle_result(res);
 }
 
-SEXP savvy_InternalDataFrame_to_arrow_stream__impl(SEXP self__,
-                                                   SEXP c_arg__out) {
-  SEXP res = savvy_InternalDataFrame_to_arrow_stream__ffi(self__, c_arg__out);
+SEXP savvy_InternalDataFrame_to_arrow_stream__impl(
+    SEXP self__, SEXP c_arg__out, SEXP c_arg__requested_schema_xptr) {
+  SEXP res = savvy_InternalDataFrame_to_arrow_stream__ffi(
+      self__, c_arg__out, c_arg__requested_schema_xptr);
   return handle_result(res);
 }
 
@@ -166,6 +188,11 @@ SEXP savvy_InternalDataFrame_to_parquet__impl(
   return handle_result(res);
 }
 
+SEXP savvy_InternalDataFrame_to_provider__impl(SEXP self__) {
+  SEXP res = savvy_InternalDataFrame_to_provider__ffi(self__);
+  return handle_result(res);
+}
+
 SEXP savvy_InternalDataFrame_to_view__impl(SEXP self__, SEXP c_arg__ctx,
                                            SEXP c_arg__table_ref,
                                            SEXP c_arg__overwrite) {
@@ -183,12 +210,18 @@ static const R_CallMethodDef CallEntries[] = {
      (DL_FUNC)&savvy_sedonadb_adbc_init_func__impl, 0},
     {"savvy_InternalContext_data_frame_from_array_stream__impl",
      (DL_FUNC)&savvy_InternalContext_data_frame_from_array_stream__impl, 3},
+    {"savvy_InternalContext_data_frame_from_table_provider__impl",
+     (DL_FUNC)&savvy_InternalContext_data_frame_from_table_provider__impl, 2},
     {"savvy_InternalContext_deregister_table__impl",
      (DL_FUNC)&savvy_InternalContext_deregister_table__impl, 2},
     {"savvy_InternalContext_new__impl",
      (DL_FUNC)&savvy_InternalContext_new__impl, 0},
     {"savvy_InternalContext_read_parquet__impl",
      (DL_FUNC)&savvy_InternalContext_read_parquet__impl, 2},
+    {"savvy_InternalContext_register_scalar_udf__impl",
+     (DL_FUNC)&savvy_InternalContext_register_scalar_udf__impl, 2},
+    {"savvy_InternalContext_scalar_udf_xptr__impl",
+     (DL_FUNC)&savvy_InternalContext_scalar_udf_xptr__impl, 2},
     {"savvy_InternalContext_sql__impl",
      (DL_FUNC)&savvy_InternalContext_sql__impl, 2},
     {"savvy_InternalContext_view__impl",
@@ -208,9 +241,11 @@ static const R_CallMethodDef CallEntries[] = {
     {"savvy_InternalDataFrame_to_arrow_schema__impl",
      (DL_FUNC)&savvy_InternalDataFrame_to_arrow_schema__impl, 2},
     {"savvy_InternalDataFrame_to_arrow_stream__impl",
-     (DL_FUNC)&savvy_InternalDataFrame_to_arrow_stream__impl, 2},
+     (DL_FUNC)&savvy_InternalDataFrame_to_arrow_stream__impl, 3},
     {"savvy_InternalDataFrame_to_parquet__impl",
      (DL_FUNC)&savvy_InternalDataFrame_to_parquet__impl, 8},
+    {"savvy_InternalDataFrame_to_provider__impl",
+     (DL_FUNC)&savvy_InternalDataFrame_to_provider__impl, 1},
     {"savvy_InternalDataFrame_to_view__impl",
      (DL_FUNC)&savvy_InternalDataFrame_to_view__impl, 4},
     {NULL, NULL, 0}};

r/sedonadb/man/sd_register_udf.Rd

@@ -29,6 +29,7 @@ arrow-array = { workspace = true }
 datafusion = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-expr = { workspace = true }
+datafusion-ffi = { workspace = true }
 savvy = "*"
 savvy-ffi = "*"
 sedona = { path = "../../../../rust/sedona" }

r/sedonadb/src/init.c

@@ -26,10 +26,15 @@ SEXP savvy_sedonadb_adbc_init_func__ffi(void);
 // methods and associated functions for InternalContext
 SEXP savvy_InternalContext_data_frame_from_array_stream__ffi(
     SEXP self__, SEXP c_arg__stream_xptr, SEXP c_arg__collect_now);
+SEXP savvy_InternalContext_data_frame_from_table_provider__ffi(
+    SEXP self__, SEXP c_arg__provider_xptr);
 SEXP savvy_InternalContext_deregister_table__ffi(SEXP self__,
                                                  SEXP c_arg__table_ref);
 SEXP savvy_InternalContext_new__ffi(void);
 SEXP savvy_InternalContext_read_parquet__ffi(SEXP self__, SEXP c_arg__paths);
+SEXP savvy_InternalContext_register_scalar_udf__ffi(
+    SEXP self__, SEXP c_arg__scalar_udf_xptr);
+SEXP savvy_InternalContext_scalar_udf_xptr__ffi(SEXP self__, SEXP c_arg__name);
 SEXP savvy_InternalContext_sql__ffi(SEXP self__, SEXP c_arg__query);
 SEXP savvy_InternalContext_view__ffi(SEXP self__, SEXP c_arg__table_ref);
 
@@ -43,11 +48,13 @@ SEXP savvy_InternalDataFrame_show__ffi(SEXP self__, SEXP c_arg__ctx,
                                        SEXP c_arg__width_chars,
                                        SEXP c_arg__ascii, SEXP c_arg__limit);
 SEXP savvy_InternalDataFrame_to_arrow_schema__ffi(SEXP self__, SEXP c_arg__out);
-SEXP savvy_InternalDataFrame_to_arrow_stream__ffi(SEXP self__, SEXP c_arg__out);
+SEXP savvy_InternalDataFrame_to_arrow_stream__ffi(
+    SEXP self__, SEXP c_arg__out, SEXP c_arg__requested_schema_xptr);
 SEXP savvy_InternalDataFrame_to_parquet__ffi(
     SEXP self__, SEXP c_arg__ctx, SEXP c_arg__path, SEXP c_arg__partition_by,
     SEXP c_arg__sort_by, SEXP c_arg__single_file_output,
     SEXP c_arg__overwrite_bbox_columns, SEXP c_arg__geoparquet_version);
+SEXP savvy_InternalDataFrame_to_provider__ffi(SEXP self__);
 SEXP savvy_InternalDataFrame_to_view__ffi(SEXP self__, SEXP c_arg__ctx,
                                           SEXP c_arg__table_ref,
                                           SEXP c_arg__overwrite);

r/sedonadb/src/rust/Cargo.toml

@@ -16,20 +16,19 @@
 // under the License.
 use std::sync::Arc;
 
-use arrow_array::{
-    ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream},
-    RecordBatchReader,
-};
+use arrow_array::RecordBatchReader;
 use arrow_schema::ArrowError;
 use datafusion::catalog::{MemTable, TableProvider};
-use savvy::{savvy, savvy_err, Result};
+use datafusion_ffi::udf::FFI_ScalarUDF;
+use savvy::{savvy, savvy_err, IntoExtPtrSexp, Result};
 
 use sedona::{context::SedonaContext, record_batch_reader_provider::RecordBatchReaderProvider};
 use sedona_geoparquet::provider::GeoParquetReadOptions;
 use tokio::runtime::Runtime;
 
 use crate::{
     dataframe::{new_data_frame, InternalDataFrame},
+    ffi::{import_array_stream, import_scalar_udf, import_table_provider, FFIScalarUdfR},
     runtime::wait_for_future_captured_r,
 };
 
@@ -94,14 +93,7 @@ impl InternalContext {
         stream_xptr: savvy::Sexp,
         collect_now: bool,
     ) -> savvy::Result<InternalDataFrame> {
-        let ffi_stream =
-            unsafe { savvy_ffi::R_ExternalPtrAddr(stream_xptr.0) as *mut FFI_ArrowArrayStream };
-        if ffi_stream.is_null() {
-            return Err(savvy_err!("external pointer to null in to_arrow_schema()"));
-        }
-
-        let stream = unsafe { FFI_ArrowArrayStream::from_raw(ffi_stream as _) };
-        let stream_reader = ArrowArrayStreamReader::try_new(stream)?;
+        let stream_reader = import_array_stream(stream_xptr)?;
 
         // Some readers are sensitive to being collected on the R thread or not, so
         // provide the option to collect everything immediately.
@@ -117,8 +109,37 @@ impl InternalContext {
         Ok(new_data_frame(inner, self.runtime.clone()))
     }
 
+    pub fn data_frame_from_table_provider(
+        &self,
+        provider_xptr: savvy::Sexp,
+    ) -> Result<InternalDataFrame> {
+        let provider = import_table_provider(provider_xptr)?;
+        let inner = self.inner.ctx.read_table(provider)?;
+        Ok(new_data_frame(inner, self.runtime.clone()))
+    }
+
     pub fn deregister_table(&self, table_ref: &str) -> savvy::Result<()> {
         self.inner.ctx.deregister_table(table_ref)?;
         Ok(())
     }
+
+    pub fn scalar_udf_xptr(&self, name: &str) -> savvy::Result<savvy::Sexp> {
+        if let Some(udf) = self.inner.ctx.state().scalar_functions().get(name) {
+            let ffi_scalar_udf: FFI_ScalarUDF = udf.clone().into();
+            let mut ffi_xptr = FFIScalarUdfR(ffi_scalar_udf).into_external_pointer();
+            unsafe { savvy_ffi::Rf_protect(ffi_xptr.0) };
+            ffi_xptr.set_class(vec!["datafusion_scalar_udf"])?;
+            unsafe { savvy_ffi::Rf_unprotect(1) };
+
+            Ok(ffi_xptr)
+        } else {
+            Err(savvy_err!("Scalar UDF '{name}' was not found"))
+        }
+    }
+
+    pub fn register_scalar_udf(&self, scalar_udf_xptr: savvy::Sexp) -> savvy::Result<()> {
+        let scalar_udf = import_scalar_udf(scalar_udf_xptr)?;
+        self.inner.ctx.register_udf(scalar_udf);
+        Ok(())
+    }
 }