diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 7ac0dfa78215..c0107bc00ecc 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -60,7 +60,7 @@ jobs: with: rust-version: stable - name: Prepare cargo build - run: cargo check --profile ci --all-targets + run: cargo check --profile ci --all-targets --features integration-tests # cargo check common, functions and substrait with no default features linux-cargo-check-no-default-features: @@ -92,8 +92,8 @@ jobs: - name: Check workspace in debug mode run: cargo check --profile ci --all-targets --workspace - - name: Check workspace with avro,json features - run: cargo check --profile ci --workspace --benches --features avro,json + - name: Check workspace with additional features + run: cargo check --profile ci --workspace --benches --features avro,json,integration-tests - name: Check Cargo.lock for datafusion-cli run: | @@ -185,7 +185,7 @@ jobs: with: rust-version: stable - name: Run tests (excluding doctests) - run: cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --workspace --lib --tests --bins --features avro,json,backtrace + run: cargo test --profile ci --exclude datafusion-examples --exclude ffi_example_table_provider --exclude datafusion-benchmarks --workspace --lib --tests --bins --features avro,json,backtrace,integration-tests - name: Verify Working Directory Clean run: git diff --exit-code @@ -417,7 +417,7 @@ jobs: - name: Run tests (excluding doctests) shell: bash run: | - cargo test --profile ci --lib --tests --bins --features avro,json,backtrace + cargo test --profile ci --lib --tests --bins --features avro,json,backtrace,integration-tests cd datafusion-cli cargo test --profile ci --lib --tests --bins --all-features diff --git a/ci/scripts/rust_clippy.sh b/ci/scripts/rust_clippy.sh index f5c8b61e1c06..01eb6e710a2d 100755 --- a/ci/scripts/rust_clippy.sh +++ b/ci/scripts/rust_clippy.sh @@ -18,6 +18,6 @@ # under the License. set -ex -cargo clippy --all-targets --workspace --features avro,pyarrow -- -D warnings +cargo clippy --all-targets --workspace --features avro,pyarrow,integration-tests -- -D warnings cd datafusion-cli cargo clippy --all-targets --all-features -- -D warnings diff --git a/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs b/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs index c7eea8a8070b..c37d8f835ce8 100644 --- a/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs +++ b/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] diff --git a/datafusion/ffi/Cargo.toml b/datafusion/ffi/Cargo.toml index a0179ec44d7f..1a6248322d1c 100644 --- a/datafusion/ffi/Cargo.toml +++ b/datafusion/ffi/Cargo.toml @@ -33,10 +33,13 @@ workspace = true [lib] name = "datafusion_ffi" path = "src/lib.rs" +crate-type = ["cdylib", "rlib"] [dependencies] abi_stable = "0.11.3" arrow = { workspace = true, features = ["ffi"] } +arrow-array = { workspace = true } +arrow-schema = { workspace = true } async-ffi = { version = "0.5.0", features = ["abi_stable"] } async-trait = { workspace = true } datafusion = { workspace = true, default-features = false } @@ -44,7 +47,11 @@ 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 } + +[features] +integration-tests = [] diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 5ab321cc0114..a8c2f42fe251 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/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, pub context: Arc, + pub runtime: Option>, } 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, context: Arc) -> Self { - let private_data = Box::new(ExecutionPlanPrivateData { plan, context }); + pub fn new( + plan: Arc, + context: Arc, + runtime: Option>, + ) -> 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()?; diff --git a/datafusion/ffi/src/lib.rs b/datafusion/ffi/src/lib.rs index 8e09780edf03..bef36b1ddd48 100644 --- a/datafusion/ffi/src/lib.rs +++ b/datafusion/ffi/src/lib.rs @@ -26,5 +26,18 @@ pub mod session_config; pub mod table_provider; pub mod table_source; +#[cfg(feature = "integration-tests")] +pub mod tests; + +/// Returns the major version of the FFI implementation. If the API evolves, +/// we use the major version to identify compatibility over the unsafe +/// boundary. This call is intended to be used by implementers to validate +/// they have compatible libraries. +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); diff --git a/datafusion/ffi/src/record_batch_stream.rs b/datafusion/ffi/src/record_batch_stream.rs index c944e56c5cde..878ac24f6765 100644 --- a/datafusion/ffi/src/record_batch_stream.rs +++ b/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,12 +59,27 @@ pub struct FFI_RecordBatchStream { pub private_data: *mut c_void, } +pub struct RecordBatchStreamPrivateData { + pub rbs: SendableRecordBatchStream, + pub runtime: Option>, +} + impl From for FFI_RecordBatchStream { fn from(stream: SendableRecordBatchStream) -> Self { + Self::new(stream, None) + } +} + +impl FFI_RecordBatchStream { + pub fn new(stream: SendableRecordBatchStream, runtime: Option>) -> 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, } } } @@ -71,7 +87,8 @@ impl From for FFI_RecordBatchStream { 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>> { - 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) diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index b229d908d10d..183dfc8755d1 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/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, + runtime: Option>, } 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, can_support_pushdown_filters: bool, + runtime: Option>, ) -> 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(); diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs new file mode 100644 index 000000000000..38ddd13952b0 --- /dev/null +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -0,0 +1,282 @@ +// 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. + +//! This is an example of an async table provider that will call functions on +//! the tokio runtime of the library providing the function. Since we cannot +//! share a tokio runtime across the ffi boundary and the producer and consumer +//! may have different runtimes, we need to store a reference to the runtime +//! and enter it during streaming calls. The entering of the runtime will +//! occur by the datafusion_ffi crate during the streaming calls. This code +//! serves as an integration test of this feature. If we do not correctly +//! access the runtime, then you will get a panic when trying to do operations +//! such as spawning a tokio task. + +use std::{any::Any, fmt::Debug, sync::Arc}; + +use crate::table_provider::FFI_TableProvider; +use arrow_array::RecordBatch; +use arrow_schema::Schema; +use async_trait::async_trait; +use datafusion::{ + catalog::{Session, TableProvider}, + error::{DataFusionError, Result}, + execution::RecordBatchStream, + physical_expr::EquivalenceProperties, + physical_plan::{ExecutionPlan, Partitioning}, + prelude::Expr, +}; +use futures::Stream; +use tokio::{ + runtime::Runtime, + sync::{broadcast, mpsc}, +}; + +use super::create_record_batch; + +#[derive(Debug)] +pub struct AsyncTableProvider { + batch_request: mpsc::Sender, + shutdown: mpsc::Sender, + batch_receiver: broadcast::Receiver>, + _join_handle: Option>, +} + +fn async_table_provider_thread( + mut shutdown: mpsc::Receiver, + mut batch_request: mpsc::Receiver, + batch_sender: broadcast::Sender>, + tokio_rt: mpsc::Sender>, +) { + let runtime = Arc::new( + tokio::runtime::Builder::new_current_thread() + .build() + .expect("Unable to create tokio runtime"), + ); + let _runtime_guard = runtime.enter(); + tokio_rt + .blocking_send(Arc::clone(&runtime)) + .expect("Unable to send tokio runtime back to main thread"); + + runtime.block_on(async move { + let mut num_received = 0; + while let Some(true) = batch_request.recv().await { + let record_batch = match num_received { + 0 => Some(create_record_batch(1, 5)), + 1 => Some(create_record_batch(6, 1)), + 2 => Some(create_record_batch(7, 5)), + _ => None, + }; + num_received += 1; + + if batch_sender.send(record_batch).is_err() { + break; + } + } + }); + + let _ = shutdown.blocking_recv(); +} + +pub fn start_async_provider() -> (AsyncTableProvider, Arc) { + let (batch_request_tx, batch_request_rx) = mpsc::channel(10); + let (record_batch_tx, record_batch_rx) = broadcast::channel(10); + let (tokio_rt_tx, mut tokio_rt_rx) = mpsc::channel(10); + let (shutdown_tx, shutdown_rx) = mpsc::channel(1); + + // It is important that we are not using tokio to spawn here. We want this + // other thread to create it's own runtime, which is similar to a model used + // in datafusion-python and probably other places. This will let us test that + // we do correctly enter the runtime of the foreign provider. + let join_handle = Some(std::thread::spawn(move || { + async_table_provider_thread( + shutdown_rx, + batch_request_rx, + record_batch_tx, + tokio_rt_tx, + ) + })); + + let tokio_rt = tokio_rt_rx + .blocking_recv() + .expect("Unable to receive tokio runtime from spawned thread"); + + let table_provider = AsyncTableProvider { + shutdown: shutdown_tx, + batch_request: batch_request_tx, + batch_receiver: record_batch_rx, + _join_handle: join_handle, + }; + + (table_provider, tokio_rt) +} + +#[async_trait] +impl TableProvider for AsyncTableProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> Arc { + super::create_test_schema() + } + + fn table_type(&self) -> datafusion::logical_expr::TableType { + datafusion::logical_expr::TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(Arc::new(AsyncTestExecutionPlan::new( + self.batch_request.clone(), + self.batch_receiver.resubscribe(), + ))) + } +} + +impl Drop for AsyncTableProvider { + fn drop(&mut self) { + self.shutdown + .blocking_send(false) + .expect("Unable to call shutdown on spawned thread.") + } +} + +#[derive(Debug)] +struct AsyncTestExecutionPlan { + properties: datafusion::physical_plan::PlanProperties, + batch_request: mpsc::Sender, + batch_receiver: broadcast::Receiver>, +} + +impl AsyncTestExecutionPlan { + pub fn new( + batch_request: mpsc::Sender, + batch_receiver: broadcast::Receiver>, + ) -> Self { + Self { + properties: datafusion::physical_plan::PlanProperties::new( + EquivalenceProperties::new(super::create_test_schema()), + Partitioning::UnknownPartitioning(3), + datafusion::physical_plan::execution_plan::EmissionType::Incremental, + datafusion::physical_plan::execution_plan::Boundedness::Bounded, + ), + batch_request, + batch_receiver, + } + } +} + +impl ExecutionPlan for AsyncTestExecutionPlan { + fn name(&self) -> &str { + "async test execution plan" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::default() + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(AsyncTestRecordBatchStream { + batch_request: self.batch_request.clone(), + batch_receiver: self.batch_receiver.resubscribe(), + })) + } +} + +impl datafusion::physical_plan::DisplayAs for AsyncTestExecutionPlan { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + // Do nothing, just a test + Ok(()) + } +} + +struct AsyncTestRecordBatchStream { + batch_request: mpsc::Sender, + batch_receiver: broadcast::Receiver>, +} + +impl RecordBatchStream for AsyncTestRecordBatchStream { + fn schema(&self) -> arrow_schema::SchemaRef { + super::create_test_schema() + } +} + +impl Stream for AsyncTestRecordBatchStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let mut this = self.as_mut(); + + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { + // Nothing to do. We just need to simulate an async + // task running + }); + + if let Err(e) = this.batch_request.try_send(true) { + return std::task::Poll::Ready(Some(Err(DataFusionError::Execution( + format!("Unable to send batch request, {}", e), + )))); + } + + match this.batch_receiver.blocking_recv() { + Ok(batch) => match batch { + Some(batch) => std::task::Poll::Ready(Some(Ok(batch))), + None => std::task::Poll::Ready(None), + }, + Err(e) => std::task::Poll::Ready(Some(Err(DataFusionError::Execution( + format!("Unable receive record batch: {}", e), + )))), + } + } +} + +pub(crate) fn create_async_table_provider() -> FFI_TableProvider { + let (table_provider, tokio_rt) = start_async_provider(); + FFI_TableProvider::new(Arc::new(table_provider), true, Some(tokio_rt)) +} diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs new file mode 100644 index 000000000000..d2e865d6c2a1 --- /dev/null +++ b/datafusion/ffi/src/tests/mod.rs @@ -0,0 +1,99 @@ +// 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. + +use std::sync::Arc; + +use abi_stable::{ + declare_root_module_statics, export_root_module, + library::{LibraryError, RootModule}, + package_version_strings, + prefix_type::PrefixTypeTrait, + sabi_types::VersionStrings, + StableAbi, +}; + +use super::table_provider::FFI_TableProvider; +use arrow_array::RecordBatch; +use async_provider::create_async_table_provider; +use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::record_batch, +}; +use sync_provider::create_sync_table_provider; + +mod async_provider; +mod sync_provider; + +#[repr(C)] +#[derive(StableAbi)] +#[sabi(kind(Prefix(prefix_ref = TableProviderModuleRef)))] +/// This struct defines the module interfaces. It is to be shared by +/// both the module loading program and library that implements the +/// module. It is possible to move this definition into the loading +/// program and reference it in the modules, but this example shows +/// how a user may wish to separate these concerns. +pub struct TableProviderModule { + /// Constructs the table provider + pub create_table: extern "C" fn(synchronous: bool) -> FFI_TableProvider, + + pub version: extern "C" fn() -> u64, +} + +impl RootModule for TableProviderModuleRef { + declare_root_module_statics! {TableProviderModuleRef} + const BASE_NAME: &'static str = "datafusion_ffi"; + const NAME: &'static str = "datafusion_ffi"; + const VERSION_STRINGS: VersionStrings = package_version_strings!(); + + fn initialization(self) -> Result { + Ok(self) + } +} + +fn create_test_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Float64, true), + ])) +} + +pub fn create_record_batch(start_value: i32, num_values: usize) -> RecordBatch { + let end_value = start_value + num_values as i32; + let a_vals: Vec = (start_value..end_value).collect(); + let b_vals: Vec = a_vals.iter().map(|v| *v as f64).collect(); + + record_batch!(("a", Int32, a_vals), ("b", Float64, b_vals)).unwrap() +} + +/// Here we only wish to create a simple table provider as an example. +/// We create an in-memory table and convert it to it's FFI counterpart. +extern "C" fn construct_table_provider(synchronous: bool) -> FFI_TableProvider { + match synchronous { + true => create_sync_table_provider(), + false => create_async_table_provider(), + } +} + +#[export_root_module] +/// This defines the entry point for using the module. +pub fn get_simple_memory_table() -> TableProviderModuleRef { + TableProviderModule { + create_table: construct_table_provider, + version: super::version, + } + .leak_into_prefix() +} diff --git a/datafusion/ffi/src/tests/sync_provider.rs b/datafusion/ffi/src/tests/sync_provider.rs new file mode 100644 index 000000000000..ff85e0b15b39 --- /dev/null +++ b/datafusion/ffi/src/tests/sync_provider.rs @@ -0,0 +1,39 @@ +// 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. + +use std::sync::Arc; + +use crate::table_provider::FFI_TableProvider; +use datafusion::datasource::MemTable; + +use super::{create_record_batch, create_test_schema}; + +pub(crate) fn create_sync_table_provider() -> FFI_TableProvider { + let schema = create_test_schema(); + + // It is useful to create these as multiple record batches + // so that we can demonstrate the FFI stream. + let batches = vec![ + create_record_batch(1, 5), + create_record_batch(6, 1), + create_record_batch(7, 5), + ]; + + let table_provider = MemTable::try_new(schema, vec![batches]).unwrap(); + + FFI_TableProvider::new(Arc::new(table_provider), true, None) +} diff --git a/datafusion/ffi/tests/table_provider.rs b/datafusion/ffi/tests/table_provider.rs new file mode 100644 index 000000000000..9169c9f4221c --- /dev/null +++ b/datafusion/ffi/tests/table_provider.rs @@ -0,0 +1,135 @@ +// 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. + +/// Add an additional module here for convenience to scope this to only +/// when the feature integtation-tests is built +#[cfg(feature = "integration-tests")] +mod tests { + + use abi_stable::library::RootModule; + use datafusion::error::{DataFusionError, Result}; + use datafusion::prelude::SessionContext; + use datafusion_ffi::table_provider::ForeignTableProvider; + use datafusion_ffi::tests::TableProviderModuleRef; + use std::path::Path; + use std::sync::Arc; + + /// Compute the path to the library. It would be preferable to simply use + /// abi_stable::library::development_utils::compute_library_path however + /// our current CI pipeline has a `ci` profile that we need to use to + /// find the library. + pub fn compute_library_path( + target_path: &Path, + ) -> std::io::Result { + let debug_dir = target_path.join("debug"); + let release_dir = target_path.join("release"); + let ci_dir = target_path.join("ci"); + + let debug_path = M::get_library_path(&debug_dir.join("deps")); + let release_path = M::get_library_path(&release_dir.join("deps")); + let ci_path = M::get_library_path(&ci_dir.join("deps")); + + let all_paths = vec![ + (debug_dir.clone(), debug_path), + (release_dir, release_path), + (ci_dir, ci_path), + ]; + + let best_path = all_paths + .into_iter() + .filter(|(_, path)| path.exists()) + .filter_map(|(dir, path)| path.metadata().map(|m| (dir, m)).ok()) + .filter_map(|(dir, meta)| meta.modified().map(|m| (dir, m)).ok()) + .max_by_key(|(_, date)| *date) + .map(|(dir, _)| dir) + .unwrap_or(debug_dir); + + Ok(best_path) + } + + /// It is important that this test is in the `tests` directory and not in the + /// library directory so we can verify we are building a dynamic library and + /// testing it via a different executable. + #[cfg(feature = "integration-tests")] + async fn test_table_provider(synchronous: bool) -> Result<()> { + let expected_version = datafusion_ffi::version(); + + let crate_root = Path::new(env!("CARGO_MANIFEST_DIR")); + let target_dir = crate_root + .parent() + .expect("Failed to find crate parent") + .parent() + .expect("Failed to find workspace root") + .join("target"); + + // Find the location of the library. This is specific to the build environment, + // so you will need to change the approach here based on your use case. + // let target: &std::path::Path = "../../../../target/".as_ref(); + let library_path = + compute_library_path::(target_dir.as_path()) + .map_err(|e| DataFusionError::External(Box::new(e)))? + .join("deps"); + + // Load the module + let table_provider_module = + TableProviderModuleRef::load_from_directory(&library_path) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + assert_eq!( + table_provider_module + .version() + .expect("Unable to call version on FFI module")(), + expected_version + ); + + // By calling the code below, the table provided will be created within + // the module's code. + let ffi_table_provider = table_provider_module.create_table().ok_or( + DataFusionError::NotImplemented( + "External table provider failed to implement create_table".to_string(), + ), + )?(synchronous); + + // In order to access the table provider within this executable, we need to + // turn it into a `ForeignTableProvider`. + let foreign_table_provider: ForeignTableProvider = (&ffi_table_provider).into(); + + let ctx = SessionContext::new(); + + // Display the data to show the full cycle works. + ctx.register_table("external_table", Arc::new(foreign_table_provider))?; + let df = ctx.table("external_table").await?; + let results = df.collect().await?; + + assert_eq!(results.len(), 3); + assert_eq!(results[0], datafusion_ffi::tests::create_record_batch(1, 5)); + assert_eq!(results[1], datafusion_ffi::tests::create_record_batch(6, 1)); + assert_eq!(results[2], datafusion_ffi::tests::create_record_batch(7, 5)); + + Ok(()) + } + + #[tokio::test] + async fn async_test_table_provider() -> Result<()> { + test_table_provider(false).await + } + + #[tokio::test] + async fn sync_test_table_provider() -> Result<()> { + test_table_provider(true).await + } +}