datafusion_ffi/

schema_provider.rs

// 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::ffi::c_void;
use std::sync::Arc;

use async_ffi::{FfiFuture, FutureExt};
use async_trait::async_trait;
use datafusion_catalog::{SchemaProvider, TableProvider};
use datafusion_common::error::{DataFusionError, Result};
use datafusion_proto::logical_plan::{
    DefaultLogicalExtensionCodec, LogicalExtensionCodec,
};
use stabby::string::String as SString;
use stabby::vec::Vec as SVec;
use tokio::runtime::Handle;

use crate::execution::FFI_TaskContextProvider;
use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
use crate::table_provider::{FFI_TableProvider, ForeignTableProvider};
use crate::util::{FFI_Option, FFI_Result};
use crate::{df_result, sresult_return};

/// A stable struct for sharing [`SchemaProvider`] across FFI boundaries.
#[repr(C)]
#[derive(Debug)]
pub struct FFI_SchemaProvider {
    pub owner_name: FFI_Option<SString>,

    pub table_names: unsafe extern "C" fn(provider: &Self) -> SVec<SString>,

    pub table: unsafe extern "C" fn(
        provider: &Self,
        name: SString,
    ) -> FfiFuture<
        FFI_Result<FFI_Option<FFI_TableProvider>>,
    >,

    pub register_table: unsafe extern "C" fn(
        provider: &Self,
        name: SString,
        table: FFI_TableProvider,
    )
        -> FFI_Result<FFI_Option<FFI_TableProvider>>,

    pub deregister_table:
        unsafe extern "C" fn(
            provider: &Self,
            name: SString,
        ) -> FFI_Result<FFI_Option<FFI_TableProvider>>,

    pub table_exist: unsafe extern "C" fn(provider: &Self, name: SString) -> bool,

    pub logical_codec: FFI_LogicalExtensionCodec,

    /// Used to create a clone on the provider of the execution plan. This should
    /// only need to be called by the receiver of the plan.
    pub clone: unsafe extern "C" fn(plan: &Self) -> Self,

    /// 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 [`ForeignSchemaProvider`] 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_SchemaProvider {}
unsafe impl Sync for FFI_SchemaProvider {}

struct ProviderPrivateData {
    provider: Arc<dyn SchemaProvider>,
    runtime: Option<Handle>,
}

impl FFI_SchemaProvider {
    unsafe fn inner(&self) -> &Arc<dyn SchemaProvider> {
        unsafe {
            let private_data = self.private_data as *const ProviderPrivateData;
            &(*private_data).provider
        }
    }

    unsafe fn runtime(&self) -> Option<Handle> {
        unsafe {
            let private_data = self.private_data as *const ProviderPrivateData;
            (*private_data).runtime.clone()
        }
    }
}

unsafe extern "C" fn table_names_fn_wrapper(
    provider: &FFI_SchemaProvider,
) -> SVec<SString> {
    unsafe {
        let provider = provider.inner();

        let table_names = provider.table_names();
        table_names.into_iter().map(|s| s.into()).collect()
    }
}

unsafe extern "C" fn table_fn_wrapper(
    provider: &FFI_SchemaProvider,
    name: SString,
) -> FfiFuture<FFI_Result<FFI_Option<FFI_TableProvider>>> {
    unsafe {
        let runtime = provider.runtime();
        let logical_codec = provider.logical_codec.clone();
        let provider = Arc::clone(provider.inner());

        async move {
            let table = sresult_return!(provider.table(name.as_str()).await)
                .map(|t| {
                    FFI_TableProvider::new_with_ffi_codec(t, true, runtime, logical_codec)
                })
                .into();

            FFI_Result::Ok(table)
        }
        .into_ffi()
    }
}

unsafe extern "C" fn register_table_fn_wrapper(
    provider: &FFI_SchemaProvider,
    name: SString,
    table: FFI_TableProvider,
) -> FFI_Result<FFI_Option<FFI_TableProvider>> {
    unsafe {
        let runtime = provider.runtime();
        let logical_codec = provider.logical_codec.clone();
        let provider = provider.inner();

        let table = Arc::new(ForeignTableProvider(table));

        let returned_table = sresult_return!(provider.register_table(name.into(), table))
            .map(|t| {
                FFI_TableProvider::new_with_ffi_codec(t, true, runtime, logical_codec)
            });

        FFI_Result::Ok(returned_table.into())
    }
}

unsafe extern "C" fn deregister_table_fn_wrapper(
    provider: &FFI_SchemaProvider,
    name: SString,
) -> FFI_Result<FFI_Option<FFI_TableProvider>> {
    unsafe {
        let runtime = provider.runtime();
        let logical_codec = provider.logical_codec.clone();
        let provider = provider.inner();

        let returned_table = sresult_return!(provider.deregister_table(name.as_str()))
            .map(|t| {
                FFI_TableProvider::new_with_ffi_codec(t, true, runtime, logical_codec)
            });

        FFI_Result::Ok(returned_table.into())
    }
}

unsafe extern "C" fn table_exist_fn_wrapper(
    provider: &FFI_SchemaProvider,
    name: SString,
) -> bool {
    unsafe { provider.inner().table_exist(name.as_str()) }
}

unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_SchemaProvider) {
    unsafe {
        debug_assert!(!provider.private_data.is_null());
        let private_data =
            Box::from_raw(provider.private_data as *mut ProviderPrivateData);
        drop(private_data);
        provider.private_data = std::ptr::null_mut();
    }
}

unsafe extern "C" fn clone_fn_wrapper(
    provider: &FFI_SchemaProvider,
) -> FFI_SchemaProvider {
    unsafe {
        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_SchemaProvider {
            owner_name: provider.owner_name.clone(),
            table_names: table_names_fn_wrapper,
            table: table_fn_wrapper,
            register_table: register_table_fn_wrapper,
            deregister_table: deregister_table_fn_wrapper,
            table_exist: table_exist_fn_wrapper,
            logical_codec: provider.logical_codec.clone(),
            clone: clone_fn_wrapper,
            release: release_fn_wrapper,
            version: super::version,
            private_data,
            library_marker_id: crate::get_library_marker_id,
        }
    }
}

impl Drop for FFI_SchemaProvider {
    fn drop(&mut self) {
        unsafe { (self.release)(self) }
    }
}

impl FFI_SchemaProvider {
    /// Creates a new [`FFI_SchemaProvider`].
    pub fn new(
        provider: Arc<dyn SchemaProvider>,
        runtime: Option<Handle>,
        task_ctx_provider: impl Into<FFI_TaskContextProvider>,
        logical_codec: Option<Arc<dyn LogicalExtensionCodec>>,
    ) -> Self {
        let task_ctx_provider = task_ctx_provider.into();
        let logical_codec =
            logical_codec.unwrap_or_else(|| Arc::new(DefaultLogicalExtensionCodec {}));
        let logical_codec = FFI_LogicalExtensionCodec::new(
            logical_codec,
            runtime.clone(),
            task_ctx_provider.clone(),
        );
        Self::new_with_ffi_codec(provider, runtime, logical_codec)
    }

    pub fn new_with_ffi_codec(
        provider: Arc<dyn SchemaProvider>,
        runtime: Option<Handle>,
        logical_codec: FFI_LogicalExtensionCodec,
    ) -> Self {
        if let Some(provider) = provider.downcast_ref::<ForeignSchemaProvider>() {
            return provider.0.clone();
        }

        let owner_name = provider.owner_name().map(|s| s.into()).into();
        let private_data = Box::new(ProviderPrivateData { provider, runtime });

        Self {
            owner_name,
            table_names: table_names_fn_wrapper,
            table: table_fn_wrapper,
            register_table: register_table_fn_wrapper,
            deregister_table: deregister_table_fn_wrapper,
            table_exist: table_exist_fn_wrapper,
            logical_codec,
            clone: clone_fn_wrapper,
            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,
        }
    }
}

/// This wrapper struct exists on the receiver side of the FFI interface, so it has
/// no guarantees about being able to access the data in `private_data`. Any functions
/// defined on this struct must only use the stable functions provided in
/// FFI_SchemaProvider to interact with the foreign table provider.
#[derive(Debug)]
pub struct ForeignSchemaProvider(pub FFI_SchemaProvider);

unsafe impl Send for ForeignSchemaProvider {}
unsafe impl Sync for ForeignSchemaProvider {}

impl From<&FFI_SchemaProvider> for Arc<dyn SchemaProvider> {
    fn from(provider: &FFI_SchemaProvider) -> Self {
        if (provider.library_marker_id)() == crate::get_library_marker_id() {
            return Arc::clone(unsafe { provider.inner() });
        }

        Arc::new(ForeignSchemaProvider(provider.clone())) as Arc<dyn SchemaProvider>
    }
}

impl Clone for FFI_SchemaProvider {
    fn clone(&self) -> Self {
        unsafe { (self.clone)(self) }
    }
}

#[async_trait]
impl SchemaProvider for ForeignSchemaProvider {
    fn owner_name(&self) -> Option<&str> {
        let name: Option<&SString> = self.0.owner_name.as_ref();
        name.map(|s| s.as_str())
    }

    fn table_names(&self) -> Vec<String> {
        unsafe {
            (self.0.table_names)(&self.0)
                .into_iter()
                .map(|s| s.into())
                .collect()
        }
    }

    async fn table(
        &self,
        name: &str,
    ) -> Result<Option<Arc<dyn TableProvider>>, DataFusionError> {
        unsafe {
            let table: Option<FFI_TableProvider> =
                df_result!((self.0.table)(&self.0, name.into()).await)?.into();

            let table = table.as_ref().map(<Arc<dyn TableProvider>>::from);

            Ok(table)
        }
    }

    fn register_table(
        &self,
        name: String,
        table: Arc<dyn TableProvider>,
    ) -> Result<Option<Arc<dyn TableProvider>>> {
        unsafe {
            let ffi_table = match table.downcast_ref::<ForeignTableProvider>() {
                Some(t) => t.0.clone(),
                None => FFI_TableProvider::new_with_ffi_codec(
                    table,
                    true,
                    None,
                    self.0.logical_codec.clone(),
                ),
            };

            let returned_provider: Option<FFI_TableProvider> =
                df_result!((self.0.register_table)(&self.0, name.into(), ffi_table))?
                    .into();

            Ok(returned_provider
                .map(|t| Arc::new(ForeignTableProvider(t)) as Arc<dyn TableProvider>))
        }
    }

    fn deregister_table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
        let returned_provider: Option<FFI_TableProvider> = unsafe {
            df_result!((self.0.deregister_table)(&self.0, name.into()))?.into()
        };

        Ok(returned_provider
            .map(|t| Arc::new(ForeignTableProvider(t)) as Arc<dyn TableProvider>))
    }

    /// Returns true if table exist in the schema provider, false otherwise.
    fn table_exist(&self, name: &str) -> bool {
        unsafe { (self.0.table_exist)(&self.0, name.into()) }
    }
}

#[cfg(test)]
mod tests {
    use arrow::datatypes::Schema;
    use datafusion::catalog::MemorySchemaProvider;
    use datafusion::datasource::empty::EmptyTable;

    use super::*;

    fn empty_table() -> Arc<dyn TableProvider> {
        Arc::new(EmptyTable::new(Arc::new(Schema::empty())))
    }

    #[tokio::test]
    async fn test_round_trip_ffi_schema_provider() {
        let schema_provider = Arc::new(MemorySchemaProvider::new());
        assert!(
            schema_provider
                .as_ref()
                .register_table("prior_table".to_string(), empty_table())
                .unwrap()
                .is_none()
        );

        let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx();

        let mut ffi_schema_provider =
            FFI_SchemaProvider::new(schema_provider, None, task_ctx_provider, None);
        ffi_schema_provider.library_marker_id = crate::mock_foreign_marker_id;

        let foreign_schema_provider: Arc<dyn SchemaProvider> =
            (&ffi_schema_provider).into();

        let prior_table_names = foreign_schema_provider.table_names();
        assert_eq!(prior_table_names.len(), 1);
        assert_eq!(prior_table_names[0], "prior_table");

        // Replace an existing table with one of the same name generates an error
        let returned_schema = foreign_schema_provider
            .register_table("prior_table".to_string(), empty_table());
        assert!(returned_schema.is_err());
        assert_eq!(foreign_schema_provider.table_names().len(), 1);

        // Add a new table
        let returned_schema = foreign_schema_provider
            .register_table("second_table".to_string(), empty_table())
            .expect("Unable to register table");
        assert!(returned_schema.is_none());
        assert_eq!(foreign_schema_provider.table_names().len(), 2);

        // Remove a table
        let returned_schema = foreign_schema_provider
            .deregister_table("prior_table")
            .expect("Unable to deregister table");
        assert!(returned_schema.is_some());
        assert_eq!(foreign_schema_provider.table_names().len(), 1);

        // Retrieve non-existent table
        let returned_schema = foreign_schema_provider
            .table("prior_table")
            .await
            .expect("Unable to query table");
        assert!(returned_schema.is_none());
        assert!(!foreign_schema_provider.table_exist("prior_table"));

        // Retrieve valid table
        let returned_schema = foreign_schema_provider
            .table("second_table")
            .await
            .expect("Unable to query table");
        assert!(returned_schema.is_some());
        assert!(foreign_schema_provider.table_exist("second_table"));
    }

    #[test]
    fn test_ffi_schema_provider_local_bypass() {
        let schema_provider = Arc::new(MemorySchemaProvider::new());

        let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx();
        let mut ffi_schema =
            FFI_SchemaProvider::new(schema_provider, None, task_ctx_provider, None);

        // Verify local libraries can be downcast to their original
        let foreign_schema: Arc<dyn SchemaProvider> = (&ffi_schema).into();
        assert!(
            foreign_schema
                .downcast_ref::<MemorySchemaProvider>()
                .is_some()
        );

        // Verify different library markers generate foreign providers
        ffi_schema.library_marker_id = crate::mock_foreign_marker_id;
        let foreign_schema: Arc<dyn SchemaProvider> = (&ffi_schema).into();
        assert!(
            foreign_schema
                .downcast_ref::<ForeignSchemaProvider>()
                .is_some()
        );
    }
}