diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7bbeed2f1..4f503f8d5 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -53,14 +53,13 @@ use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; -use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ create_logical_extension_capsule, create_physical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx, get_tokio_runtime, physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future, - wait_for_future, + table_provider_factory_from_pycapsule, wait_for_future, }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; @@ -713,30 +712,22 @@ impl PySessionContext { pub fn register_table_factory( &self, format: &str, - mut factory: Bound<'_, PyAny>, + factory: Bound<'_, PyAny>, ) -> PyDataFusionResult<()> { - if factory.hasattr("__datafusion_table_provider_factory__")? { + let factory: Arc = if factory + .hasattr("__datafusion_table_provider_factory__")? + || factory.cast::().is_ok() + { let py = factory.py(); let ffi = self.ffi_logical_codec(); let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; - factory = factory - .getattr("__datafusion_table_provider_factory__")? - .call1((codec_capsule,))?; - } - - let factory: Arc = - if let Ok(capsule) = factory.cast::().map_err(py_datafusion_err) { - let data: NonNull = capsule - .pointer_checked(Some(c"datafusion_table_provider_factory"))? - .cast(); - let factory = unsafe { data.as_ref() }; - factory.into() - } else { - Arc::new(RustWrappedPyTableProviderFactory::new( - factory.into(), - self.ffi_logical_codec(), - )) - }; + table_provider_factory_from_pycapsule(&factory, (codec_capsule,))? + } else { + Arc::new(RustWrappedPyTableProviderFactory::new( + factory.into(), + self.ffi_logical_codec(), + )) + }; let st = self.ctx.state_ref(); let mut lock = st.write(); diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 9327d7f2f..f483a18ff 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -20,6 +20,7 @@ use std::ptr::NonNull; use std::sync::{Arc, OnceLock}; use std::time::Duration; +use datafusion::catalog::TableProviderFactory; use datafusion::datasource::TableProvider; use datafusion::execution::TaskContext; use datafusion::execution::context::SessionContext; @@ -30,6 +31,7 @@ use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_ffi::table_provider::FFI_TableProvider; +use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; use pyo3::prelude::*; @@ -249,6 +251,44 @@ pub fn create_physical_extension_capsule<'py>( /// instead. #[macro_export] macro_rules! from_pycapsule { + ($fn_name:ident, $capsule_name:literal, $ffi_type:ty, $output_type:ty, call_args) => { + pub fn $fn_name<'py, A>( + obj: &$crate::pyo3::Bound<'py, $crate::pyo3::PyAny>, + args: A, + ) -> $crate::pyo3::PyResult> + where + A: $crate::pyo3::call::PyCallArgs<'py>, + { + use $crate::pyo3::prelude::*; + use $crate::pyo3::types::PyCapsule; + + let mut obj = obj.clone(); + if obj.hasattr(concat!("__", $capsule_name, "__"))? { + obj = obj + .getattr(concat!("__", $capsule_name, "__"))? + .call1(args)?; + } + let capsule = obj.cast::().map_err(|_| { + $crate::errors::py_datafusion_err(concat!( + "Invalid ", + $capsule_name, + ". Does not contain PyCapsule object." + )) + })?; + $crate::validate_pycapsule(&capsule, $capsule_name)?; + + let expected_name = std::ffi::CString::new($capsule_name) + .expect("capsule name must not contain interior NUL bytes"); + let data: std::ptr::NonNull<$ffi_type> = capsule + .pointer_checked(Some(expected_name.as_c_str()))? + .cast(); + let output_obj = unsafe { data.as_ref() }; + let output_obj: std::sync::Arc<$output_type> = output_obj.into(); + + Ok(output_obj) + } + }; + ($fn_name:ident, $capsule_name:literal, $ffi_type:ty, $output_type:ty) => { pub fn $fn_name( obj: &$crate::pyo3::Bound<$crate::pyo3::PyAny>, @@ -340,6 +380,14 @@ from_pycapsule!( dyn PhysicalOptimizerRule + Send + Sync ); +from_pycapsule!( + table_provider_factory_from_pycapsule, + "datafusion_table_provider_factory", + FFI_TableProviderFactory, + dyn TableProviderFactory, + call_args +); + try_from_pycapsule!( task_context_from_pycapsule, "datafusion_task_context_provider",