(datafusion) 01/01: fix(ffi): don't discard downcast-delegating wrappers of foreign plans (#25749)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25749-991fd23dd0be0046af5945b8c3905612859b657a
in repository https://gitbox.apache.org/repos/asf/datafusion.git

commit a0a048f78a38f998ea2a88e32faa66db80048e4d
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Mon Sep 28 17:11:03 2026 +0000

    fix(ffi): don't discard downcast-delegating wrappers of foreign plans (#25749)
    
    ## Which issue does this PR close?
    
    - Closes #25717.
    
    ## Rationale for this change
    
    A transparent wrapper `ExecutionPlan` (one that implements
    `downcast_delegate`, such as `datafusion-tracing`'s `InstrumentedExec`)
    around a `ForeignExecutionPlan` is silently dropped when it is passed
    back across the FFI boundary. For example, installing the tracing
    instrumentation rule via FFI leaves the plan un-instrumented. Also, when
    such a wrapper is the child of a foreign plan, it does not receive the
    runtime handle.
    
    ## What changes are included in this PR?
    
    - `FFI_ExecutionPlan::new`: the "already foreign" shortcut now checks
    the concrete type via `Any` instead of the delegate-aware
    `downcast_ref`, so it only fires for a real `ForeignExecutionPlan`.
    - `pass_runtime_to_children`: the child check also uses `Any`, so a
    local wrapper child of a foreign parent is re-wrapped with the runtime.
    - `plan_is_foreign` intentionally stays delegate-aware, since a
    delegating wrapper typically forwards `replace_children` to the foreign
    plan and its children still need the runtime. A comment explains this.
    
    ## What is the testing strategy for this PR?
    
    - Added an optional `downcast_delegate` mode to the existing `EmptyExec`
    test plan (`with_downcast_delegate`).
    - Added `test_ffi_execution_plan_delegating_wrapper`, which checks that
    the wrapper survives `FFI_ExecutionPlan::new` and that a wrapper child
    of a foreign parent gets the runtime. Reverting either fix on its own
    makes the test fail.
    
    ## Are there any user-facing changes?
    
    No API changes. Wrapper plans that implement `downcast_delegate` are now
    kept when they cross the FFI boundary instead of being discarded.
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    Co-authored-by: Claude Opus 5.5 (1M context) <no...@anthropic.com>
    Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
 datafusion/ffi/src/execution_plan.rs | 68 ++++++++++++++++++++++++++++++++++--
 1 file changed, 65 insertions(+), 3 deletions(-)

diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs
index b242d1b275..6ef1f17cf6 100644
--- a/datafusion/ffi/src/execution_plan.rs
+++ b/datafusion/ffi/src/execution_plan.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::pin::Pin;
 use std::sync::Arc;
@@ -288,6 +289,9 @@ fn pass_runtime_to_children(
     runtime: &Handle,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
     let mut updated_children = false;
+    // This intentionally uses the delegate-aware `is`: a transparent wrapper
+    // around a `ForeignExecutionPlan` typically forwards `replace_children` to
+    // it, so its children will still cross the FFI boundary and need the runtime.
     let plan_is_foreign = plan.is::<ForeignExecutionPlan>();
 
     let children = plan
@@ -306,8 +310,12 @@ fn pass_runtime_to_children(
             // we called `children()` above we will get something other than a
             // `ForeignExecutionPlan`. In this case wrap the plan in a `ForeignExecutionPlan`
             // because when we call `replace_children` below it will extract the
-            // FFI plan that does contain the runtime.
-            if plan_is_foreign && !child.is::<ForeignExecutionPlan>() {
+            // FFI plan that does contain the runtime. Check the concrete type rather
+            // than the delegate-aware `is` so that a local wrapper which delegates to
+            // a `ForeignExecutionPlan` still receives the runtime.
+            if plan_is_foreign
+                && !(child.as_ref() as &dyn Any).is::<ForeignExecutionPlan>()
+            {
                 updated_children = true;
                 let ffi_child = FFI_ExecutionPlan::new(child, Some(runtime.clone()));
                 let foreign_child = ForeignExecutionPlan::try_from(ffi_child);
@@ -334,7 +342,13 @@ impl FFI_ExecutionPlan {
     pub fn new(mut plan: Arc<dyn ExecutionPlan>, runtime: Option<Handle>) -> Self {
         // Note to developers: `pass_runtime_to_children` relies on the logic here to
         // get the underlying FFI plan during calls to `new_with_children`.
-        if let Some(plan) = plan.downcast_ref::<ForeignExecutionPlan>() {
+        //
+        // Check the concrete type rather than using the delegate-aware
+        // `downcast_ref`, otherwise a wrapper whose `downcast_delegate` is a
+        // `ForeignExecutionPlan` would be discarded here.
+        if let Some(plan) =
+            (plan.as_ref() as &dyn Any).downcast_ref::<ForeignExecutionPlan>()
+        {
             return plan.plan.clone();
         }
 
@@ -571,6 +585,7 @@ pub mod tests {
         dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>,
         metrics: Option<MetricsSet>,
         statistics: Option<Statistics>,
+        downcast_delegate: bool,
     }
 
     impl EmptyExec {
@@ -587,6 +602,7 @@ pub mod tests {
                 dynamic_expressions: Vec::default(),
                 metrics: None,
                 statistics: None,
+                downcast_delegate: false,
             }
         }
 
@@ -600,6 +616,13 @@ pub mod tests {
             self
         }
 
+        /// Act as a transparent wrapper around `child` via `downcast_delegate`.
+        pub fn with_downcast_delegate(mut self, child: Arc<dyn ExecutionPlan>) -> Self {
+            self.children = vec![child];
+            self.downcast_delegate = true;
+            self
+        }
+
         pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self {
             self.props =
                 Arc::new(self.props.as_ref().clone().with_partitioning(partitioning));
@@ -658,6 +681,7 @@ pub mod tests {
                 dynamic_expressions: self.dynamic_expressions.clone(),
                 metrics: self.metrics.clone(),
                 statistics: self.statistics.clone(),
+                downcast_delegate: self.downcast_delegate,
             }))
         }
 
@@ -679,6 +703,12 @@ pub mod tests {
             unimplemented!()
         }
 
+        fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> {
+            self.downcast_delegate
+                .then(|| self.children.first().map(|c| c.as_ref()))
+                .flatten()
+        }
+
         fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> {
             self.dynamic_expressions.iter().map(Arc::clone).collect()
         }
@@ -1004,4 +1034,36 @@ pub mod tests {
         let foreign_plan: Arc<dyn ExecutionPlan> = (&ffi_plan).try_into().unwrap();
         assert!(foreign_plan.is::<ForeignExecutionPlan>());
     }
+
+    /// A wrapper whose `downcast_delegate` is a `ForeignExecutionPlan` must not be
+    /// mistaken for one: it must survive `FFI_ExecutionPlan::new` and still receive
+    /// the runtime when it is the local child of a foreign parent.
+    #[tokio::test]
+    async fn test_ffi_execution_plan_delegating_wrapper() -> Result<()> {
+        let schema = Arc::new(arrow::datatypes::Schema::empty());
+
+        let mut ffi_inner =
+            FFI_ExecutionPlan::new(Arc::new(EmptyExec::new(Arc::clone(&schema))), None);
+        ffi_inner.library_marker_id = crate::mock_foreign_marker_id;
+        let foreign_inner: Arc<dyn ExecutionPlan> =
+            Arc::new(ForeignExecutionPlan::try_from(ffi_inner)?);
+        let wrapper: Arc<dyn ExecutionPlan> = Arc::new(
+            EmptyExec::new(Arc::clone(&schema)).with_downcast_delegate(foreign_inner),
+        );
+        assert!(wrapper.is::<ForeignExecutionPlan>());
+
+        let ffi_wrapper = FFI_ExecutionPlan::new(Arc::clone(&wrapper), None);
+        assert!(ffi_wrapper.inner().downcast_delegate().is_some());
+
+        let parent = Arc::new(EmptyExec::new(schema)).replace_children(
+            vec![wrapper],
+            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+        )?;
+        let mut ffi_parent = FFI_ExecutionPlan::new(parent, None);
+        ffi_parent.library_marker_id = crate::mock_foreign_marker_id;
+        let foreign_parent: Arc<dyn ExecutionPlan> = (&ffi_parent).try_into()?;
+        assert!(pass_runtime_to_children(&foreign_parent, &Handle::current())?.is_some());
+
+        Ok(())
+    }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org