Initial commit for example on using FFI Table provider in rust as a module loading system

Author: timsaucer

Cargo.toml

@@ -47,6 +47,9 @@ members = [
     "datafusion/substrait",
     "datafusion/wasmtest",
     "datafusion-examples",
+    "datafusion-examples/examples/ffi/ffi_example_table_provider",
+    "datafusion-examples/examples/ffi/ffi_module_interface",
+    "datafusion-examples/examples/ffi/ffi_module_loader",
     "test-utils",
     "benchmarks",
 ]

datafusion-examples/examples/ffi/ffi_example_table_provider/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "ffi_example_table_provider"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+datafusion = { workspace = true }
+datafusion-ffi = { workspace = true }
+abi_stable = "0.11.3"
+arrow = { workspace = true }
+arrow-array = { workspace = true }
+arrow-schema = { workspace = true }
+ffi_module_interface = { path = "../ffi_module_interface" }
+
+[lib]
+name = "ffi_example_table_provider"
+crate-type = ["cdylib",'rlib']

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

@@ -0,0 +1,44 @@
+use std::sync::Arc;
+
+use abi_stable::{export_root_module, prefix_type::PrefixTypeTrait};
+use arrow_array::RecordBatch;
+use datafusion::{
+    arrow::datatypes::{DataType, Field, Schema},
+    common::record_batch,
+    datasource::MemTable,
+};
+use datafusion_ffi::table_provider::FFI_TableProvider;
+use ffi_module_interface::{TableProviderModule, TableProviderModuleRef};
+
+fn create_record_batch(start_value: i32, num_values: usize) -> RecordBatch {
+    let end_value = start_value + num_values as i32;
+    let a_vals: Vec<i32> = (start_value..end_value).collect();
+    let b_vals: Vec<f64> = a_vals.iter().map(|v| *v as f64).collect();
+
+    record_batch!(("a", Int32, a_vals), ("b", Float64, b_vals)).unwrap()
+}
+
+extern "C" fn construct_simple_table_provider() -> FFI_TableProvider {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Int32, true),
+        Field::new("b", DataType::Float64, true),
+    ]));
+
+    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)
+}
+
+#[export_root_module]
+pub fn get_simple_memory_table() -> TableProviderModuleRef {
+    TableProviderModule {
+        create_table: construct_simple_table_provider,
+    }
+    .leak_into_prefix()
+}

datafusion-examples/examples/ffi/ffi_module_interface/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "ffi_module_interface"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+abi_stable = "0.11.3"
+datafusion-ffi = { workspace = true }

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

@@ -0,0 +1,27 @@
+use abi_stable::{
+    declare_root_module_statics,
+    library::{LibraryError, RootModule},
+    package_version_strings,
+    sabi_types::VersionStrings,
+    StableAbi,
+};
+use datafusion_ffi::table_provider::FFI_TableProvider;
+
+#[repr(C)]
+#[derive(StableAbi)]
+#[sabi(kind(Prefix(prefix_ref = TableProviderModuleRef)))]
+pub struct TableProviderModule {
+    /// Constructs the table provider
+    pub create_table: extern "C" fn() -> FFI_TableProvider,
+}
+
+impl RootModule for TableProviderModuleRef {
+    declare_root_module_statics! {TableProviderModuleRef}
+    const BASE_NAME: &'static str = "ffi_example_table_provider";
+    const NAME: &'static str = "ffi_example_table_provider";
+    const VERSION_STRINGS: VersionStrings = package_version_strings!();
+
+    fn initialization(self) -> Result<Self, LibraryError> {
+        Ok(self)
+    }
+}

datafusion-examples/examples/ffi/ffi_module_loader/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "ffi_module_loader"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] }
+datafusion = { workspace = true }
+datafusion-ffi = { workspace = true }
+ffi_module_interface = { path = "../ffi_module_interface" }
+abi_stable = "0.11.3"

datafusion-examples/examples/ffi/ffi_module_loader/src/main.rs

@@ -0,0 +1,39 @@
+use std::sync::Arc;
+
+use datafusion::{
+    error::{DataFusionError, Result},
+    prelude::SessionContext,
+};
+
+use abi_stable::library::{development_utils::compute_library_path, RootModule};
+use datafusion_ffi::table_provider::ForeignTableProvider;
+use ffi_module_interface::TableProviderModuleRef;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let target: &std::path::Path = "../../../../target/".as_ref();
+    let library_path = compute_library_path::<TableProviderModuleRef>(target).unwrap();
+
+    let table_provider_module =
+        TableProviderModuleRef::load_from_directory(&library_path)
+            .unwrap_or_else(|e| panic!("{}", e));
+
+    let ffi_table_provider =
+        table_provider_module
+            .create_table()
+            .ok_or(DataFusionError::NotImplemented(
+                "External table provider failed to implement create_table".to_string(),
+            ))?();
+
+    let foreign_table_provider: ForeignTableProvider = (&ffi_table_provider).into();
+
+    let ctx = SessionContext::new();
+
+    ctx.register_table("external_table", Arc::new(foreign_table_provider))?;
+
+    let df = ctx.table("external_table").await?;
+
+    df.show().await?;
+
+    Ok(())
+}