diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index ed2c594f1bc02..2668421a93254 100644 --- a/datafusion/ffi/src/proto/logical_extension_codec.rs +++ b/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, runtime: Option, @@ -302,7 +311,9 @@ impl FFI_LogicalExtensionCodec { if let Some(codec) = (Arc::clone(&codec) as Arc) .downcast_ref::() { - 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 = (&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 + /// . + #[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, + None, + provider_a, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + + let imported: Arc = (&ffi_codec).into(); + assert!( + (Arc::clone(&imported) as Arc) + .downcast_ref::() + .is_some() + ); + + let rebound = FFI_LogicalExtensionCodec::new(imported, None, provider_b); + + let task_ctx: Arc = (&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, + None, + provider_a, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + let imported: Arc = (&ffi_codec).into(); + drop(ctx_a); + + FFI_LogicalExtensionCodec::new(imported, None, provider_b) + }; + + let task_ctx: Arc = (&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()); + } } diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index 95d2ed68a6ea3..ad8a208967c36 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/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, runtime: Option, @@ -288,7 +297,9 @@ impl FFI_PhysicalExtensionCodec { if let Some(codec) = (Arc::clone(&codec) as Arc) .downcast_ref::() { - 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 = (&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 + /// . + #[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, + None, + provider_a, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + + let imported: Arc = (&ffi_codec).into(); + assert!( + (Arc::clone(&imported) as Arc) + .downcast_ref::() + .is_some() + ); + + let rebound = FFI_PhysicalExtensionCodec::new(imported, None, provider_b); + + let task_ctx: Arc = (&rebound.task_ctx_provider) + .try_into() + .expect("rebound codec resolves"); + assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); + } } diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 6d895d65c5fc1..727e1e57819a7 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/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; + + 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 = (&ffi_a).into(); + let any_ref: &dyn std::any::Any = imported.as_ref(); + assert!(any_ref.downcast_ref::().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 = (&rebound.logical_codec.task_ctx_provider) + .try_into() + .unwrap(); + assert_eq!(bound_to.session_id(), ctx_b.task_ctx().session_id()); + } } diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index ee9377bff064e..f1eb6076b6783 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -402,6 +402,16 @@ 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, can_support_pushdown_filters: bool, @@ -409,7 +419,9 @@ impl FFI_TableProvider { logical_codec: FFI_LogicalExtensionCodec, ) -> Self { if let Some(provider) = provider.downcast_ref::() { - 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 }); @@ -903,4 +915,40 @@ mod tests { Ok(()) } + + /// Re-wrapping an imported provider with a rebuilt logical codec must adopt + /// that codec. See . + #[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 = (&ffi_provider).into(); + assert!(imported.downcast_ref::().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 = (&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(()) + } }