fix: preserve provided arguments during FFI object construction

datafusion/ffi/src/proto/logical_extension_codec.rs

@@ -294,6 +294,15 @@ impl Drop for FFI_LogicalExtensionCodec {
 
 impl FFI_LogicalExtensionCodec {
     /// Creates a new [`FFI_LogicalExtensionCodec`].
+    ///
+    /// If `codec` is already foreign, this re-exports its original FFI handle
+    /// rather than adding another wrapper layer. The handle still adopts the
+    /// `task_ctx_provider` supplied here, so it is never silently discarded and
+    /// an imported codec can be rebound to a different session.
+    ///
+    /// `runtime` is only honored when a new wrapper is created. An
+    /// already-foreign handle keeps the runtime of the library that owns it,
+    /// because that value lives in private data this side cannot reach.
     pub fn new(
         codec: Arc<dyn LogicalExtensionCodec>,
         runtime: Option<Handle>,
@@ -302,7 +311,9 @@ impl FFI_LogicalExtensionCodec {
         if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
             .downcast_ref::<ForeignLogicalExtensionCodec>()
         {
-            return codec.0.clone();
+            let mut codec = codec.0.clone();
+            codec.task_ctx_provider = task_ctx_provider.into();
+            return codec;
         }
 
         let task_ctx_provider = task_ctx_provider.into();
@@ -507,7 +518,9 @@ mod tests {
     use datafusion_proto::logical_plan::LogicalExtensionCodec;
     use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 
-    use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+    use crate::proto::logical_extension_codec::{
+        FFI_LogicalExtensionCodec, ForeignLogicalExtensionCodec,
+    };
     use crate::proto::physical_extension_codec::tests::TestExtensionCodec;
 
     fn create_test_table() -> MemTable {
@@ -727,4 +740,61 @@ mod tests {
         let foreign_codec: Arc<dyn LogicalExtensionCodec> = (&ffi_codec).into();
         assert!(!arc_ptr_eq(&foreign_codec, &codec));
     }
+
+    /// Importing a codec and re-wrapping it with a different task context
+    /// provider must rebind the handle. See
+    /// <https://github.com/apache/datafusion/issues/24722>.
+    #[test]
+    fn ffi_logical_extension_codec_rebind_adopts_task_ctx_provider() {
+        let (_ctx_a, provider_a) = crate::util::tests::test_session_and_ctx();
+        let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx();
+
+        let mut ffi_codec = FFI_LogicalExtensionCodec::new(
+            Arc::new(TestExtensionCodec {}) as Arc<dyn LogicalExtensionCodec>,
+            None,
+            provider_a,
+        );
+        ffi_codec.library_marker_id = crate::mock_foreign_marker_id;
+
+        let imported: Arc<dyn LogicalExtensionCodec> = (&ffi_codec).into();
+        assert!(
+            (Arc::clone(&imported) as Arc<dyn std::any::Any>)
+                .downcast_ref::<ForeignLogicalExtensionCodec>()
+                .is_some()
+        );
+
+        let rebound = FFI_LogicalExtensionCodec::new(imported, None, provider_b);
+
+        let task_ctx: Arc<TaskContext> = (&rebound.task_ctx_provider)
+            .try_into()
+            .expect("rebound codec resolves");
+        assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id());
+    }
+
+    /// Because the provider is held as a `Weak`, a codec that cannot be rebound
+    /// forces callers to keep the original session alive. Once rebinding works,
+    /// dropping it must not invalidate the handle.
+    #[test]
+    fn ffi_logical_extension_codec_rebind_releases_original_session() {
+        let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx();
+
+        let rebound = {
+            let (ctx_a, provider_a) = crate::util::tests::test_session_and_ctx();
+            let mut ffi_codec = FFI_LogicalExtensionCodec::new(
+                Arc::new(TestExtensionCodec {}) as Arc<dyn LogicalExtensionCodec>,
+                None,
+                provider_a,
+            );
+            ffi_codec.library_marker_id = crate::mock_foreign_marker_id;
+            let imported: Arc<dyn LogicalExtensionCodec> = (&ffi_codec).into();
+            drop(ctx_a);
+
+            FFI_LogicalExtensionCodec::new(imported, None, provider_b)
+        };
+
+        let task_ctx: Arc<TaskContext> = (&rebound.task_ctx_provider)
+            .try_into()
+            .expect("rebound codec must not depend on the original session");
+        assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id());
+    }
 }

datafusion/ffi/src/proto/physical_extension_codec.rs

@@ -280,6 +280,15 @@ impl Drop for FFI_PhysicalExtensionCodec {
 
 impl FFI_PhysicalExtensionCodec {
     /// Creates a new [`FFI_PhysicalExtensionCodec`].
+    ///
+    /// If `codec` is already foreign, this re-exports its original FFI handle
+    /// rather than adding another wrapper layer. The handle still adopts the
+    /// `task_ctx_provider` supplied here, so it is never silently discarded and
+    /// an imported codec can be rebound to a different session.
+    ///
+    /// `runtime` is only honored when a new wrapper is created. An
+    /// already-foreign handle keeps the runtime of the library that owns it,
+    /// because that value lives in private data this side cannot reach.
     pub fn new(
         codec: Arc<dyn PhysicalExtensionCodec>,
         runtime: Option<Handle>,
@@ -288,7 +297,9 @@ impl FFI_PhysicalExtensionCodec {
         if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
             .downcast_ref::<ForeignPhysicalExtensionCodec>()
         {
-            return codec.0.clone();
+            let mut codec = codec.0.clone();
+            codec.task_ctx_provider = task_ctx_provider.into();
+            return codec;
         }
 
         let task_ctx_provider = task_ctx_provider.into();
@@ -449,7 +460,9 @@ pub(crate) mod tests {
     };
 
     use crate::execution_plan::tests::EmptyExec;
-    use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+    use crate::proto::physical_extension_codec::{
+        FFI_PhysicalExtensionCodec, ForeignPhysicalExtensionCodec,
+    };
 
     #[derive(Debug)]
     pub(crate) struct TestExtensionCodec;
@@ -710,4 +723,34 @@ pub(crate) mod tests {
         let foreign_codec: Arc<dyn PhysicalExtensionCodec> = (&ffi_codec).into();
         assert!(!arc_ptr_eq(&foreign_codec, &codec));
     }
+
+    /// Importing a codec and re-wrapping it with a different task context
+    /// provider must rebind the handle. See
+    /// <https://github.com/apache/datafusion/issues/24722>.
+    #[test]
+    fn ffi_physical_extension_codec_rebind_adopts_task_ctx_provider() {
+        let (_ctx_a, provider_a) = crate::util::tests::test_session_and_ctx();
+        let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx();
+
+        let mut ffi_codec = FFI_PhysicalExtensionCodec::new(
+            Arc::new(TestExtensionCodec {}) as Arc<dyn PhysicalExtensionCodec>,
+            None,
+            provider_a,
+        );
+        ffi_codec.library_marker_id = crate::mock_foreign_marker_id;
+
+        let imported: Arc<dyn PhysicalExtensionCodec> = (&ffi_codec).into();
+        assert!(
+            (Arc::clone(&imported) as Arc<dyn std::any::Any>)
+                .downcast_ref::<ForeignPhysicalExtensionCodec>()
+                .is_some()
+        );
+
+        let rebound = FFI_PhysicalExtensionCodec::new(imported, None, provider_b);
+
+        let task_ctx: Arc<TaskContext> = (&rebound.task_ctx_provider)
+            .try_into()
+            .expect("rebound codec resolves");
+        assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id());
+    }
 }

datafusion/ffi/src/query_planner.rs

@@ -442,4 +442,40 @@ mod tests {
 
         Ok(())
     }
+
+    // Control for https://github.com/apache/datafusion/issues/24722: this
+    // constructor adopts the supplied codecs on the already-foreign path.
+    #[test]
+    fn test_rebind_foreign_query_planner_adopts_codecs() {
+        use datafusion_execution::TaskContext;
+
+        let ctx_a = Arc::new(SessionContext::new());
+        let ctx_b = Arc::new(SessionContext::new());
+        let provider_b = Arc::clone(&ctx_b) as Arc<dyn TaskContextProvider>;
+
+        let mut ffi_a = create_ffi_query_planner(Arc::clone(&ctx_a));
+        ffi_a.library_marker_id = crate::mock_foreign_marker_id;
+        let imported: Arc<dyn QueryPlanner + Send + Sync> = (&ffi_a).into();
+        let any_ref: &dyn std::any::Any = imported.as_ref();
+        assert!(any_ref.downcast_ref::<ForeignQueryPlanner>().is_some());
+
+        let rebound = FFI_QueryPlanner::new_with_ffi_codecs(
+            imported,
+            FFI_LogicalExtensionCodec::new(
+                Arc::new(DefaultLogicalExtensionCodec {}),
+                None,
+                &provider_b,
+            ),
+            FFI_PhysicalExtensionCodec::new(
+                Arc::new(DefaultPhysicalExtensionCodec {}),
+                None,
+                &provider_b,
+            ),
+        );
+
+        let bound_to: Arc<TaskContext> = (&rebound.logical_codec.task_ctx_provider)
+            .try_into()
+            .unwrap();
+        assert_eq!(bound_to.session_id(), ctx_b.task_ctx().session_id());
+    }
 }

datafusion/ffi/src/table_provider.rs

@@ -564,14 +564,26 @@ impl FFI_TableProvider {
         )
     }
 
+    /// Creates an [`FFI_TableProvider`] using a prebuilt FFI logical codec.
+    ///
+    /// If `provider` is already foreign, this re-exports its original FFI
+    /// handle rather than adding another wrapper layer. The handle still adopts
+    /// the `logical_codec` supplied here, so it is never silently discarded and
+    /// an imported provider can be rebound to a different session.
+    ///
+    /// `runtime` is only honored when a new wrapper is created. An
+    /// already-foreign handle keeps the runtime of the library that owns it,
+    /// because that value lives in private data this side cannot reach.
     pub fn new_with_ffi_codec(
         provider: Arc<dyn TableProvider>,
         can_support_pushdown_filters: bool,
         runtime: Option<Handle>,
         logical_codec: FFI_LogicalExtensionCodec,
     ) -> Self {
         if let Some(provider) = provider.downcast_ref::<ForeignTableProvider>() {
-            return provider.0.clone();
+            let mut provider = provider.0.clone();
+            provider.logical_codec = logical_codec;
+            return provider;
         }
         let private_data = Box::new(ProviderPrivateData { provider, runtime });
 
@@ -1271,4 +1283,40 @@ mod tests {
 
         Ok(())
     }
+
+    /// Re-wrapping an imported provider with a rebuilt logical codec must adopt
+    /// that codec. See <https://github.com/apache/datafusion/issues/24722>.
+    #[test]
+    fn test_rebind_foreign_table_provider_adopts_logical_codec() -> Result<()> {
+        let (_ctx_a, provider_a) = crate::util::tests::test_session_and_ctx();
+        let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx();
+
+        let mut ffi_provider = FFI_TableProvider::new(
+            create_test_table_provider()?,
+            true,
+            None,
+            provider_a,
+            None,
+        );
+        ffi_provider.library_marker_id = crate::mock_foreign_marker_id;
+
+        let imported: Arc<dyn TableProvider> = (&ffi_provider).into();
+        assert!(imported.downcast_ref::<ForeignTableProvider>().is_some());
+
+        // Rebuild the codec against session B and re-wrap.
+        let codec_b = FFI_LogicalExtensionCodec::new(
+            Arc::new(DefaultLogicalExtensionCodec {}),
+            None,
+            provider_b,
+        );
+        let rebound =
+            FFI_TableProvider::new_with_ffi_codec(imported, true, None, codec_b);
+
+        let task_ctx: Arc<TaskContext> = (&rebound.logical_codec.task_ctx_provider)
+            .try_into()
+            .expect("rebound provider's codec resolves");
+        assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id());
+
+        Ok(())
+    }
 }