FFI: return underlying trait type when converting from FFI structs (#18672)

## Which issue does this PR close?

- Addresses part of #18671
but does not close it.
- Minor: removes `return_type` in `FFI_ScalarUDF` since it will never be
called due to the `return_field_from_args` implementation.

## Rationale for this change

This PR adds the concept of a `library_marker_id` to the FFI crate. The
reason for this is described in the `README` text as part of the PR. In
our current use of the FFI library we get into issues where we round
trip FFI structs back to their original library that now have FFI
wrappers when they are no longer needed.

## What changes are included in this PR?

- Adds a method to find the memory address of a static constant in the
library.
- Adds a check in methods that are creating Foreign FFI structs to see
if we are actually in the local library or are actually foreign.
- Replaces the `From<> for Foreign` to `From<> for Arc<dyn ...>`. The
actual use case is essentially always to use these structs as their
implementation of the trait they back.
- Adds unit tests and a method to mock being in foreign code
- Removed an unused function call in `FFI_ScalarUDF`

## Are these changes tested?

- Code is tested against existing unit tests and with
`datafusion-python`.
- Added unit tests
- Coverage report compared to `main`:

<img width="1146" height="788" alt="coverage-report"
src="https://github.com/user-attachments/assets/f422ee20-35d6-4a10-ae96-c38b59e26acb"
/>


## Are there any user-facing changes?

This does change the API for the FFI functions.

---------

Co-authored-by: Renato Marroquin <[email protected]>

Author: timsaucer

Cargo.lock

@@ -24,6 +24,5 @@ publish = false
 [dependencies]
 abi_stable = "0.11.3"
 datafusion = { workspace = true }
-datafusion-ffi = { workspace = true }
 ffi_module_interface = { path = "../ffi_module_interface" }
 tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] }

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

@@ -23,7 +23,7 @@ use datafusion::{
 };
 
 use abi_stable::library::{development_utils::compute_library_path, RootModule};
-use datafusion_ffi::table_provider::ForeignTableProvider;
+use datafusion::datasource::TableProvider;
 use ffi_module_interface::TableProviderModuleRef;
 
 #[tokio::main]
@@ -49,13 +49,13 @@ async fn main() -> Result<()> {
             ))?();
 
     // 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();
+    // turn it into a `TableProvider`.
+    let foreign_table_provider: Arc<dyn TableProvider> = (&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))?;
+    ctx.register_table("external_table", foreign_table_provider)?;
     let df = ctx.table("external_table").await?;
     df.show().await?;
 

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

@@ -58,6 +58,7 @@ semver = "1.0.27"
 tokio = { workspace = true }
 
 [dev-dependencies]
+datafusion = { workspace = true, default-features = false, features = ["sql"] }
 doc-comment = { workspace = true }
 
 [features]

datafusion/ffi/Cargo.toml

@@ -101,6 +101,75 @@ In this crate we have a variety of structs which closely mimic the behavior of
 their internal counterparts. To see detailed notes about how to use them, see
 the example in `FFI_TableProvider`.
 
+## Memory Management
+
+One of the advantages of Rust is the ownership model, which means programmers
+_usually_ do not need to worry about memory management. When interacting with
+foreign code, this is not necessarily true. If you review the structures in
+this crate, you will find that many of them implement the `Drop` trait and
+perform a foreign call.
+
+Suppose we have a `FFI_CatalogProvider`, for example. This struct is safe to
+pass across the FFI boundary, so it may be owned by either the library that
+produces the underlying `CatalogProvider` or by another library that consumes
+it. If we look closer at the `FFI_CatalogProvider`, it has a pointer to
+some private data. That private data is only accessible on the producer's
+side. If you attempt to access it on the consumer's side, you may get
+segmentation faults or other bad behavior. Within that private data is the
+actual `Arc<dyn CatalogProvider`. That `Arc<>` must be freed, but if the
+`FFI_CatalogProvider` is only owned on the consumer's side, we have no way
+to access the private data and free it.
+
+To account for this, most structs in this crate have a `release` method that
+is used to clean up any privately held data. This calls into the producer's
+side, regardless of if it is called on either the local or foreign side.
+Most of the structs in this crate carry atomic reference counts to the
+underlying data, and this is straight forward. Some structs like the
+`FFI_Accumulator` contain an inner `Box<dyn Accumulator>`. The reason for
+this is that we need to be able to mutably access these based on the
+`Accumulator` trait definition. For these we have slightly more complicated
+release code based on whether it is being dropped on the local or foreign side.
+Traits that use a `Box<>` for their underlying data also cannot implement
+`Clone`.
+
+## Library Marker ID
+
+When reviewing the code, many of the structs in this crate contain a call to
+a `library_marker_id`. The purpose of this call is to determine if a library is
+accessing _local_ code through the FFI structs. Consider this example: you have
+a `primary` program that exposes functions to create a schema provider. You
+have a `secondary` library that exposes a function to create a catalog provider
+and the `secondary` library uses the schema provider of the `primary` program.
+From the point of view of the `secondary` library, the schema provider is
+foreign code.
+
+Now when we register the `secondary` library with the `primary` program as a
+catalog provider and we make calls to get a schema, the `secondary` library
+will return a FFI wrapped schema provider back to the `primary` program. In
+this case that schema provider is actually local code to the `primary` program
+except that it is wrapped in the FFI code!
+
+We work around this by the `library_marker_id` calls. What this does is it
+creates a global variable within each library and returns a `usize` address
+of that library. This is guaranteed to be unique for every library that contains
+FFI code. By comparing these `usize` addresses we can determine if a FFI struct
+is local or foreign.
+
+In our example of the schema provider, if you were to make a call in your
+primary program to get the schema provider, it would reach out to the foreign
+catalog provider and send back a `FFI_SchemaProvider` object. By then
+comparing the `library_marker_id` of this object to the `primary` program, we
+determine it is local code. This means it is safe to access the underlying
+private data.
+
+Users of the FFI code should not need to access these function. If you are
+implementing a new FFI struct, then it is recommended that you follow the
+established patterns for converting from FFI struct into the underlying
+traits. Specifically you should use `crate::get_library_marker_id` and in
+your unit tests you should override this with
+`crate::mock_foreign_marker_id` to force your test to create the foreign
+variant of your struct.
+
 [apache datafusion]: https://datafusion.apache.org/
 [api docs]: http://docs.rs/datafusion-ffi/latest
 [rust abi]: https://doc.rust-lang.org/reference/abi.html

datafusion/ffi/README.md

@@ -70,6 +70,11 @@ pub struct FFI_CatalogProvider {
     /// Internal data. This is only to be accessed by the provider of the plan.
     /// A [`ForeignCatalogProvider`] should never attempt to access this data.
     pub private_data: *mut c_void,
+
+    /// Utility to identify when FFI objects are accessed locally through
+    /// the foreign interface. See [`crate::get_library_marker_id`] and
+    /// the crate's `README.md` for more information.
+    pub library_marker_id: extern "C" fn() -> usize,
 }
 
 unsafe impl Send for FFI_CatalogProvider {}
@@ -116,7 +121,7 @@ unsafe extern "C" fn register_schema_fn_wrapper(
 ) -> RResult<ROption<FFI_SchemaProvider>, RString> {
     let runtime = provider.runtime();
     let provider = provider.inner();
-    let schema = Arc::new(ForeignSchemaProvider::from(schema));
+    let schema: Arc<dyn SchemaProvider + Send> = schema.into();
 
     let returned_schema =
         rresult_return!(provider.register_schema(name.as_str(), schema))
@@ -169,6 +174,7 @@ unsafe extern "C" fn clone_fn_wrapper(
         release: release_fn_wrapper,
         version: super::version,
         private_data,
+        library_marker_id: crate::get_library_marker_id,
     }
 }
 
@@ -195,6 +201,7 @@ impl FFI_CatalogProvider {
             release: release_fn_wrapper,
             version: super::version,
             private_data: Box::into_raw(private_data) as *mut c_void,
+            library_marker_id: crate::get_library_marker_id,
         }
     }
 }
@@ -209,9 +216,14 @@ pub struct ForeignCatalogProvider(pub(crate) FFI_CatalogProvider);
 unsafe impl Send for ForeignCatalogProvider {}
 unsafe impl Sync for ForeignCatalogProvider {}
 
-impl From<&FFI_CatalogProvider> for ForeignCatalogProvider {
+impl From<&FFI_CatalogProvider> for Arc<dyn CatalogProvider + Send> {
     fn from(provider: &FFI_CatalogProvider) -> Self {
-        Self(provider.clone())
+        if (provider.library_marker_id)() == crate::get_library_marker_id() {
+            return Arc::clone(unsafe { provider.inner() });
+        }
+
+        Arc::new(ForeignCatalogProvider(provider.clone()))
+            as Arc<dyn CatalogProvider + Send>
     }
 }
 
@@ -298,9 +310,10 @@ mod tests {
             .unwrap()
             .is_none());
 
-        let ffi_catalog = FFI_CatalogProvider::new(catalog, None);
+        let mut ffi_catalog = FFI_CatalogProvider::new(catalog, None);
+        ffi_catalog.library_marker_id = crate::mock_foreign_marker_id;
 
-        let foreign_catalog: ForeignCatalogProvider = (&ffi_catalog).into();
+        let foreign_catalog: Arc<dyn CatalogProvider + Send> = (&ffi_catalog).into();
 
         let prior_schema_names = foreign_catalog.schema_names();
         assert_eq!(prior_schema_names.len(), 1);
@@ -335,4 +348,26 @@ mod tests {
         let returned_schema = foreign_catalog.schema("second_schema");
         assert!(returned_schema.is_some());
     }
+
+    #[test]
+    fn test_ffi_catalog_provider_local_bypass() {
+        let catalog = Arc::new(MemoryCatalogProvider::new());
+
+        let mut ffi_catalog = FFI_CatalogProvider::new(catalog, None);
+
+        // Verify local libraries can be downcast to their original
+        let foreign_catalog: Arc<dyn CatalogProvider + Send> = (&ffi_catalog).into();
+        assert!(foreign_catalog
+            .as_any()
+            .downcast_ref::<MemoryCatalogProvider>()
+            .is_some());
+
+        // Verify different library markers generate foreign providers
+        ffi_catalog.library_marker_id = crate::mock_foreign_marker_id;
+        let foreign_catalog: Arc<dyn CatalogProvider + Send> = (&ffi_catalog).into();
+        assert!(foreign_catalog
+            .as_any()
+            .downcast_ref::<ForeignCatalogProvider>()
+            .is_some());
+    }
 }

datafusion/ffi/src/catalog_provider.rs

@@ -58,6 +58,11 @@ pub struct FFI_CatalogProviderList {
     /// Internal data. This is only to be accessed by the provider of the plan.
     /// A [`ForeignCatalogProviderList`] should never attempt to access this data.
     pub private_data: *mut c_void,
+
+    /// Utility to identify when FFI objects are accessed locally through
+    /// the foreign interface. See [`crate::get_library_marker_id`] and
+    /// the crate's `README.md` for more information.
+    pub library_marker_id: extern "C" fn() -> usize,
 }
 
 unsafe impl Send for FFI_CatalogProviderList {}
@@ -94,7 +99,7 @@ unsafe extern "C" fn register_catalog_fn_wrapper(
 ) -> ROption<FFI_CatalogProvider> {
     let runtime = provider.runtime();
     let provider = provider.inner();
-    let catalog = Arc::new(ForeignCatalogProvider::from(catalog));
+    let catalog: Arc<dyn CatalogProvider + Send> = catalog.into();
 
     provider
         .register_catalog(name.into(), catalog)
@@ -138,6 +143,7 @@ unsafe extern "C" fn clone_fn_wrapper(
         release: release_fn_wrapper,
         version: super::version,
         private_data,
+        library_marker_id: crate::get_library_marker_id,
     }
 }
 
@@ -163,6 +169,7 @@ impl FFI_CatalogProviderList {
             release: release_fn_wrapper,
             version: super::version,
             private_data: Box::into_raw(private_data) as *mut c_void,
+            library_marker_id: crate::get_library_marker_id,
         }
     }
 }
@@ -177,9 +184,14 @@ pub struct ForeignCatalogProviderList(FFI_CatalogProviderList);
 unsafe impl Send for ForeignCatalogProviderList {}
 unsafe impl Sync for ForeignCatalogProviderList {}
 
-impl From<&FFI_CatalogProviderList> for ForeignCatalogProviderList {
+impl From<&FFI_CatalogProviderList> for Arc<dyn CatalogProviderList + Send> {
     fn from(provider: &FFI_CatalogProviderList) -> Self {
-        Self(provider.clone())
+        if (provider.library_marker_id)() == crate::get_library_marker_id() {
+            return Arc::clone(unsafe { provider.inner() });
+        }
+
+        Arc::new(ForeignCatalogProviderList(provider.clone()))
+            as Arc<dyn CatalogProviderList + Send>
     }
 }
 
@@ -248,9 +260,11 @@ mod tests {
             .register_catalog("prior_catalog".to_owned(), prior_catalog)
             .is_none());
 
-        let ffi_catalog_list = FFI_CatalogProviderList::new(catalog_list, None);
+        let mut ffi_catalog_list = FFI_CatalogProviderList::new(catalog_list, None);
+        ffi_catalog_list.library_marker_id = crate::mock_foreign_marker_id;
 
-        let foreign_catalog_list: ForeignCatalogProviderList = (&ffi_catalog_list).into();
+        let foreign_catalog_list: Arc<dyn CatalogProviderList + Send> =
+            (&ffi_catalog_list).into();
 
         let prior_catalog_names = foreign_catalog_list.catalog_names();
         assert_eq!(prior_catalog_names.len(), 1);
@@ -280,4 +294,28 @@ mod tests {
         let returned_catalog = foreign_catalog_list.catalog("second_catalog");
         assert!(returned_catalog.is_some());
     }
+
+    #[test]
+    fn test_ffi_catalog_provider_list_local_bypass() {
+        let catalog_list = Arc::new(MemoryCatalogProviderList::new());
+
+        let mut ffi_catalog_list = FFI_CatalogProviderList::new(catalog_list, None);
+
+        // Verify local libraries can be downcast to their original
+        let foreign_catalog_list: Arc<dyn CatalogProviderList + Send> =
+            (&ffi_catalog_list).into();
+        assert!(foreign_catalog_list
+            .as_any()
+            .downcast_ref::<MemoryCatalogProviderList>()
+            .is_some());
+
+        // Verify different library markers generate foreign providers
+        ffi_catalog_list.library_marker_id = crate::mock_foreign_marker_id;
+        let foreign_catalog_list: Arc<dyn CatalogProviderList + Send> =
+            (&ffi_catalog_list).into();
+        assert!(foreign_catalog_list
+            .as_any()
+            .downcast_ref::<ForeignCatalogProviderList>()
+            .is_some());
+    }
 }