If we have foreign objects, get the underlying ffi object instead of wrapping

Author: timsaucer

datafusion-examples/examples/custom_data_source/adapter_serialization.rs

@@ -308,10 +308,6 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec {
             "try_encode not used - adapter wrapping happens in serialize_physical_plan"
         )
     }
-
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
 }
 
 impl PhysicalProtoConverterExtension for AdapterPreservingCodec {

datafusion-examples/examples/proto/composed_extension_codec.rs

@@ -158,10 +158,6 @@ impl PhysicalExtensionCodec for ParentPhysicalExtensionCodec {
             internal_err!("Not supported")
         }
     }
-
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
 }
 
 #[derive(Debug)]
@@ -236,8 +232,4 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec {
             internal_err!("Not supported")
         }
     }
-
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
 }

datafusion-examples/examples/proto/expression_deduplication.rs

@@ -199,10 +199,6 @@ impl PhysicalExtensionCodec for CachingCodec {
     ) -> Result<()> {
         datafusion::common::not_impl_err!("No custom extension nodes")
     }
-
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
 }
 
 impl PhysicalProtoConverterExtension for CachingCodec {

datafusion/catalog/src/table.rs

@@ -486,7 +486,7 @@ pub trait TableProviderFactory: Debug + Sync + Send {
 }
 
 /// A trait for table function implementations
-pub trait TableFunctionImpl: Debug + Sync + Send {
+pub trait TableFunctionImpl: Debug + Sync + Send + Any {
     /// Create a table provider
     fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>>;
 }

datafusion/ffi/src/catalog_provider.rs

@@ -250,6 +250,11 @@ impl FFI_CatalogProvider {
         runtime: Option<Handle>,
         logical_codec: FFI_LogicalExtensionCodec,
     ) -> Self {
+        if let Some(provider) = provider.as_any().downcast_ref::<ForeignCatalogProvider>()
+        {
+            return provider.0.clone();
+        }
+
         let private_data = Box::new(ProviderPrivateData { provider, runtime });
 
         Self {

datafusion/ffi/src/catalog_provider_list.rs

@@ -212,6 +212,13 @@ impl FFI_CatalogProviderList {
         runtime: Option<Handle>,
         logical_codec: FFI_LogicalExtensionCodec,
     ) -> Self {
+        if let Some(provider) = provider
+            .as_any()
+            .downcast_ref::<ForeignCatalogProviderList>()
+        {
+            return provider.0.clone();
+        }
+
         let private_data = Box::new(ProviderPrivateData { provider, runtime });
 
         Self {

datafusion/ffi/src/proto/logical_extension_codec.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::any::Any;
 use std::ffi::c_void;
 use std::sync::Arc;
 
@@ -296,6 +297,12 @@ impl FFI_LogicalExtensionCodec {
         runtime: Option<Handle>,
         task_ctx_provider: impl Into<FFI_TaskContextProvider>,
     ) -> Self {
+        if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
+            .downcast_ref::<ForeignLogicalExtensionCodec>()
+        {
+            return codec.0.clone();
+        }
+
         let task_ctx_provider = task_ctx_provider.into();
         let private_data = Box::new(LogicalExtensionCodecPrivateData { codec, runtime });
 

datafusion/ffi/src/proto/physical_extension_codec.rs

@@ -15,9 +15,6 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::ffi::c_void;
-use std::sync::Arc;
-
 use abi_stable::StableAbi;
 use abi_stable::std_types::{RResult, RSlice, RStr, RVec};
 use datafusion_common::error::Result;
@@ -27,6 +24,9 @@ use datafusion_expr::{
 };
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use std::any::Any;
+use std::ffi::c_void;
+use std::sync::Arc;
 use tokio::runtime::Handle;
 
 use crate::execution::FFI_TaskContextProvider;
@@ -271,8 +271,7 @@ impl FFI_PhysicalExtensionCodec {
         runtime: Option<Handle>,
         task_ctx_provider: impl Into<FFI_TaskContextProvider>,
     ) -> Self {
-        if let Some(codec) = codec
-            .as_any()
+        if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
             .downcast_ref::<ForeignPhysicalExtensionCodec>()
         {
             return codec.0.clone();
@@ -409,10 +408,6 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec {
 
         Ok(())
     }
-
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
 }
 
 #[cfg(test)]
@@ -568,10 +563,6 @@ pub(crate) mod tests {
 
             Ok(())
         }
-
-        fn as_any(&self) -> &dyn std::any::Any {
-            self
-        }
     }
 
     fn create_test_exec() -> Arc<dyn ExecutionPlan> {

datafusion/ffi/src/schema_provider.rs

@@ -259,6 +259,11 @@ impl FFI_SchemaProvider {
         runtime: Option<Handle>,
         logical_codec: FFI_LogicalExtensionCodec,
     ) -> Self {
+        if let Some(provider) = provider.as_any().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 });
 

datafusion/ffi/src/session/mod.rs

@@ -321,6 +321,10 @@ impl FFI_SessionRef {
         runtime: Option<Handle>,
         logical_codec: FFI_LogicalExtensionCodec,
     ) -> Self {
+        if let Some(session) = session.as_any().downcast_ref::<ForeignSession>() {
+            return session.session.clone();
+        }
+
         let private_data = Box::new(SessionPrivateData { session, runtime });
 
         Self {