FFI support for versions and alternate tokio runtimes

Cargo.toml

@@ -27,6 +27,7 @@ members = [
     "datafusion/expr-common",
     "datafusion/execution",
     "datafusion/ffi",
+    "datafusion/ffitest",
     "datafusion/functions",
     "datafusion/functions-aggregate",
     "datafusion/functions-aggregate-common",

datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs

@@ -53,7 +53,7 @@ extern "C" fn construct_simple_table_provider() -> FFI_TableProvider {
 
     let table_provider = MemTable::try_new(schema, vec![batches]).unwrap();
 
-    FFI_TableProvider::new(Arc::new(table_provider), true)
+    FFI_TableProvider::new(Arc::new(table_provider), true, None)
 }
 
 #[export_root_module]

datafusion/ffi/Cargo.toml

@@ -44,7 +44,8 @@ datafusion-proto = { workspace = true }
 futures = { workspace = true }
 log = { workspace = true }
 prost = { workspace = true }
+semver = "1.0.24"
+tokio = { workspace = true }
 
 [dev-dependencies]
 doc-comment = { workspace = true }
-tokio = { workspace = true }

datafusion/ffi/src/execution_plan.rs

@@ -27,6 +27,7 @@ use datafusion::{
     execution::{SendableRecordBatchStream, TaskContext},
     physical_plan::{DisplayAs, ExecutionPlan, PlanProperties},
 };
+use tokio::runtime::Runtime;
 
 use crate::{
     plan_properties::FFI_PlanProperties, record_batch_stream::FFI_RecordBatchStream,
@@ -71,6 +72,7 @@ unsafe impl Sync for FFI_ExecutionPlan {}
 pub struct ExecutionPlanPrivateData {
     pub plan: Arc<dyn ExecutionPlan>,
     pub context: Arc<TaskContext>,
+    pub runtime: Option<Arc<Runtime>>,
 }
 
 unsafe extern "C" fn properties_fn_wrapper(
@@ -88,11 +90,14 @@ unsafe extern "C" fn children_fn_wrapper(
     let private_data = plan.private_data as *const ExecutionPlanPrivateData;
     let plan = &(*private_data).plan;
     let ctx = &(*private_data).context;
+    let runtime = &(*private_data).runtime;
 
     let children: Vec<_> = plan
         .children()
         .into_iter()
-        .map(|child| FFI_ExecutionPlan::new(Arc::clone(child), Arc::clone(ctx)))
+        .map(|child| {
+            FFI_ExecutionPlan::new(Arc::clone(child), Arc::clone(ctx), runtime.clone())
+        })
         .collect();
 
     children.into()
@@ -105,9 +110,10 @@ unsafe extern "C" fn execute_fn_wrapper(
     let private_data = plan.private_data as *const ExecutionPlanPrivateData;
     let plan = &(*private_data).plan;
     let ctx = &(*private_data).context;
+    let runtime = (*private_data).runtime.as_ref().map(Arc::clone);
 
     match plan.execute(partition, Arc::clone(ctx)) {
-        Ok(rbs) => RResult::ROk(rbs.into()),
+        Ok(rbs) => RResult::ROk(FFI_RecordBatchStream::new(rbs, runtime)),
         Err(e) => RResult::RErr(
             format!("Error occurred during FFI_ExecutionPlan execute: {}", e).into(),
         ),
@@ -129,7 +135,11 @@ unsafe extern "C" fn clone_fn_wrapper(plan: &FFI_ExecutionPlan) -> FFI_Execution
     let private_data = plan.private_data as *const ExecutionPlanPrivateData;
     let plan_data = &(*private_data);
 
-    FFI_ExecutionPlan::new(Arc::clone(&plan_data.plan), Arc::clone(&plan_data.context))
+    FFI_ExecutionPlan::new(
+        Arc::clone(&plan_data.plan),
+        Arc::clone(&plan_data.context),
+        plan_data.runtime.clone(),
+    )
 }
 
 impl Clone for FFI_ExecutionPlan {
@@ -140,8 +150,16 @@ impl Clone for FFI_ExecutionPlan {
 
 impl FFI_ExecutionPlan {
     /// This function is called on the provider's side.
-    pub fn new(plan: Arc<dyn ExecutionPlan>, context: Arc<TaskContext>) -> Self {
-        let private_data = Box::new(ExecutionPlanPrivateData { plan, context });
+    pub fn new(
+        plan: Arc<dyn ExecutionPlan>,
+        context: Arc<TaskContext>,
+        runtime: Option<Arc<Runtime>>,
+    ) -> Self {
+        let private_data = Box::new(ExecutionPlanPrivateData {
+            plan,
+            context,
+            runtime,
+        });
 
         Self {
             properties: properties_fn_wrapper,
@@ -357,7 +375,7 @@ mod tests {
         let original_plan = Arc::new(EmptyExec::new(schema));
         let original_name = original_plan.name().to_string();
 
-        let local_plan = FFI_ExecutionPlan::new(original_plan, ctx.task_ctx());
+        let local_plan = FFI_ExecutionPlan::new(original_plan, ctx.task_ctx(), None);
 
         let foreign_plan: ForeignExecutionPlan = (&local_plan).try_into()?;
 

datafusion/ffi/src/lib.rs

@@ -26,5 +26,14 @@ pub mod session_config;
 pub mod table_provider;
 pub mod table_source;
 
+/// Returns the major version of the FFI implementation. If the API evolves,
+/// we use the major version to identify compatibility over the unsafe
+/// boundary.
+pub extern "C" fn version() -> u64 {
+    let version_str = env!("CARGO_PKG_VERSION");
+    let version = semver::Version::parse(version_str).expect("Invalid version string");
+    version.major
+}
+
 #[cfg(doctest)]
 doc_comment::doctest!("../README.md", readme_example_test);

datafusion/ffi/src/record_batch_stream.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{ffi::c_void, task::Poll};
+use std::{ffi::c_void, sync::Arc, task::Poll};
 
 use abi_stable::{
     std_types::{ROption, RResult, RString},
@@ -33,6 +33,7 @@ use datafusion::{
     execution::{RecordBatchStream, SendableRecordBatchStream},
 };
 use futures::{Stream, TryStreamExt};
+use tokio::runtime::Runtime;
 
 use crate::arrow_wrappers::{WrappedArray, WrappedSchema};
 
@@ -58,20 +59,36 @@ pub struct FFI_RecordBatchStream {
     pub private_data: *mut c_void,
 }
 
+pub struct RecordBatchStreamPrivateData {
+    pub rbs: SendableRecordBatchStream,
+    pub runtime: Option<Arc<Runtime>>,
+}
+
 impl From<SendableRecordBatchStream> for FFI_RecordBatchStream {
     fn from(stream: SendableRecordBatchStream) -> Self {
+        Self::new(stream, None)
+    }
+}
+
+impl FFI_RecordBatchStream {
+    pub fn new(stream: SendableRecordBatchStream, runtime: Option<Arc<Runtime>>) -> Self {
+        let private_data = Box::into_raw(Box::new(RecordBatchStreamPrivateData {
+            rbs: stream,
+            runtime,
+        })) as *mut c_void;
         FFI_RecordBatchStream {
             poll_next: poll_next_fn_wrapper,
             schema: schema_fn_wrapper,
-            private_data: Box::into_raw(Box::new(stream)) as *mut c_void,
+            private_data,
         }
     }
 }
 
 unsafe impl Send for FFI_RecordBatchStream {}
 
 unsafe extern "C" fn schema_fn_wrapper(stream: &FFI_RecordBatchStream) -> WrappedSchema {
-    let stream = stream.private_data as *const SendableRecordBatchStream;
+    let private_data = stream.private_data as *const RecordBatchStreamPrivateData;
+    let stream = &(*private_data).rbs;
 
     (*stream).schema().into()
 }
@@ -106,7 +123,10 @@ unsafe extern "C" fn poll_next_fn_wrapper(
     stream: &FFI_RecordBatchStream,
     cx: &mut FfiContext,
 ) -> FfiPoll<ROption<RResult<WrappedArray, RString>>> {
-    let stream = stream.private_data as *mut SendableRecordBatchStream;
+    let private_data = stream.private_data as *mut RecordBatchStreamPrivateData;
+    let stream = &mut (*private_data).rbs;
+
+    let _guard = (*private_data).runtime.as_ref().map(|rt| rt.enter());
 
     let poll_result = cx.with_context(|std_cx| {
         (*stream)

datafusion/ffi/src/table_provider.rs

@@ -40,6 +40,7 @@ use datafusion_proto::{
     protobuf::LogicalExprList,
 };
 use prost::Message;
+use tokio::runtime::Runtime;
 
 use crate::{
     arrow_wrappers::WrappedSchema,
@@ -139,6 +140,9 @@ pub struct FFI_TableProvider {
     /// Release the memory of the private data when it is no longer being used.
     pub release: unsafe extern "C" fn(arg: &mut Self),
 
+    /// Return the major DataFusion version number of this provider.
+    pub version: unsafe extern "C" fn() -> u64,
+
     /// Internal data. This is only to be accessed by the provider of the plan.
     /// A [`ForeignExecutionPlan`] should never attempt to access this data.
     pub private_data: *mut c_void,
@@ -149,6 +153,7 @@ unsafe impl Sync for FFI_TableProvider {}
 
 struct ProviderPrivateData {
     provider: Arc<dyn TableProvider + Send>,
+    runtime: Option<Arc<Runtime>>,
 }
 
 unsafe extern "C" fn schema_fn_wrapper(provider: &FFI_TableProvider) -> WrappedSchema {
@@ -216,6 +221,7 @@ unsafe extern "C" fn scan_fn_wrapper(
     let private_data = provider.private_data as *mut ProviderPrivateData;
     let internal_provider = &(*private_data).provider;
     let session_config = session_config.clone();
+    let runtime = &(*private_data).runtime;
 
     async move {
         let config = match ForeignSessionConfig::try_from(&session_config) {
@@ -261,7 +267,11 @@ unsafe extern "C" fn scan_fn_wrapper(
             Err(e) => return RResult::RErr(e.to_string().into()),
         };
 
-        RResult::ROk(FFI_ExecutionPlan::new(plan, ctx.task_ctx()))
+        RResult::ROk(FFI_ExecutionPlan::new(
+            plan,
+            ctx.task_ctx(),
+            runtime.clone(),
+        ))
     }
     .into_ffi()
 }
@@ -273,9 +283,11 @@ unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_TableProvider) {
 
 unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_TableProvider {
     let old_private_data = provider.private_data as *const ProviderPrivateData;
+    let runtime = (*old_private_data).runtime.clone();
 
     let private_data = Box::into_raw(Box::new(ProviderPrivateData {
         provider: Arc::clone(&(*old_private_data).provider),
+        runtime,
     })) as *mut c_void;
 
     FFI_TableProvider {
@@ -285,6 +297,7 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_Table
         supports_filters_pushdown: provider.supports_filters_pushdown,
         clone: clone_fn_wrapper,
         release: release_fn_wrapper,
+        version: super::version,
         private_data,
     }
 }
@@ -300,8 +313,9 @@ impl FFI_TableProvider {
     pub fn new(
         provider: Arc<dyn TableProvider + Send>,
         can_support_pushdown_filters: bool,
+        runtime: Option<Arc<Runtime>>,
     ) -> Self {
-        let private_data = Box::new(ProviderPrivateData { provider });
+        let private_data = Box::new(ProviderPrivateData { provider, runtime });
 
         Self {
             schema: schema_fn_wrapper,
@@ -313,6 +327,7 @@ impl FFI_TableProvider {
             },
             clone: clone_fn_wrapper,
             release: release_fn_wrapper,
+            version: super::version,
             private_data: Box::into_raw(private_data) as *mut c_void,
         }
     }
@@ -463,7 +478,7 @@ mod tests {
         let provider =
             Arc::new(MemTable::try_new(schema, vec![vec![batch1], vec![batch2]])?);
 
-        let ffi_provider = FFI_TableProvider::new(provider, true);
+        let ffi_provider = FFI_TableProvider::new(provider, true, None);
 
         let foreign_table_provider: ForeignTableProvider = (&ffi_provider).into();
 

datafusion/ffitest/Cargo.toml

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "datafusion-ffi-test"
+description = "Integration tests for DataFusion FFI"
+version = { workspace = true }
+edition = { workspace = true }
+homepage = { workspace = true }
+repository = { workspace = true }
+license = { workspace = true }
+authors = { workspace = true }
+publish = false
+# Specify MSRV here as `cargo msrv` doesn't support workspace version
+rust-version = "1.80.1"
+
+[lints]
+workspace = true
+
+[lib]
+name = "datafusion_ffi_test"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+abi_stable = "0.11.3"
+arrow = { workspace = true }
+arrow-array = { workspace = true }
+arrow-schema = { workspace = true }
+async-trait = { workspace = true }
+datafusion = { workspace = true, default-features = false }
+datafusion-ffi = { workspace = true }
+futures = { workspace = true }
+tokio = { workspace = true }