Set runtime on all descendents within the same library

Author: timsaucer

datafusion/ffi/src/execution_plan.rs

@@ -204,13 +204,67 @@ impl Clone for FFI_ExecutionPlan {
     }
 }
 
+/// Helper function to recursively identify any children that do not
+/// have a runtime set but should because they are local to this same
+/// library. This does imply a restriction that all execution plans
+/// in this chain that are within the same library use the same runtime.
+fn pass_runtime_to_children(
+    plan: &Arc<dyn ExecutionPlan>,
+    runtime: &Handle,
+) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+    println!("checking plan {:?}", plan.name());
+    let mut updated_children = false;
+    let plan_is_foreign = plan.as_any().is::<ForeignExecutionPlan>();
+
+    let children = plan
+        .children()
+        .into_iter()
+        .map(|child| {
+            let child = match pass_runtime_to_children(child, runtime)? {
+                Some(child) => {
+                    updated_children = true;
+                    child
+                }
+                None => Arc::clone(child),
+            };
+
+            // If the parent is foreign and the child is local to this library, then when
+            // 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 `with_new_children` below it will extract the
+            // FFI plan that does contain the runtime.
+            if plan_is_foreign && !child.as_any().is::<ForeignExecutionPlan>() {
+                updated_children = true;
+                let ffi_child = FFI_ExecutionPlan::new(child, Some(runtime.clone()));
+                let foreign_child = ForeignExecutionPlan::try_from(ffi_child);
+                foreign_child.map(|c| Arc::new(c) as Arc<dyn ExecutionPlan>)
+            } else {
+                Ok(child)
+            }
+        })
+        .collect::<Result<Vec<_>>>()?;
+    if updated_children {
+        Arc::clone(plan).with_new_children(children).map(Some)
+    } else {
+        Ok(None)
+    }
+}
+
 impl FFI_ExecutionPlan {
     /// This function is called on the provider's side.
-    pub fn new(plan: Arc<dyn ExecutionPlan>, runtime: Option<Handle>) -> Self {
+    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.as_any().downcast_ref::<ForeignExecutionPlan>() {
             return plan.plan.clone();
         }
 
+        if let Some(rt) = &runtime
+            && let Ok(Some(p)) = pass_runtime_to_children(&plan, rt)
+        {
+            plan = p;
+        }
+
         let private_data = Box::new(ExecutionPlanPrivateData { plan, runtime });
         Self {
             properties: properties_fn_wrapper,
@@ -278,28 +332,34 @@ impl TryFrom<&FFI_ExecutionPlan> for Arc<dyn ExecutionPlan> {
 
     fn try_from(plan: &FFI_ExecutionPlan) -> Result<Self, Self::Error> {
         if (plan.library_marker_id)() == crate::get_library_marker_id() {
-            return Ok(Arc::clone(plan.inner()));
+            Ok(Arc::clone(plan.inner()))
+        } else {
+            let plan = ForeignExecutionPlan::try_from(plan.clone())?;
+            Ok(Arc::new(plan))
         }
+    }
+}
 
+impl TryFrom<FFI_ExecutionPlan> for ForeignExecutionPlan {
+    type Error = DataFusionError;
+    fn try_from(plan: FFI_ExecutionPlan) -> Result<Self, Self::Error> {
         unsafe {
-            let name = (plan.name)(plan).into();
+            let name = (plan.name)(&plan).into();
 
-            let properties: PlanProperties = (plan.properties)(plan).try_into()?;
+            let properties: PlanProperties = (plan.properties)(&plan).try_into()?;
 
-            let children_rvec = (plan.children)(plan);
+            let children_rvec = (plan.children)(&plan);
             let children = children_rvec
                 .iter()
                 .map(<Arc<dyn ExecutionPlan>>::try_from)
                 .collect::<Result<Vec<_>>>()?;
 
-            let plan = ForeignExecutionPlan {
+            Ok(ForeignExecutionPlan {
                 name,
-                plan: plan.clone(),
+                plan,
                 properties,
                 children,
-            };
-
-            Ok(Arc::new(plan))
+            })
         }
     }
 }
@@ -366,11 +426,10 @@ impl ExecutionPlan for ForeignExecutionPlan {
     }
 }
 
-#[cfg(test)]
-pub(crate) mod tests {
-    use arrow::datatypes::{DataType, Field, Schema};
-    use datafusion::physical_plan::Partitioning;
-    use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+#[cfg(any(test, feature = "integration-tests"))]
+pub mod tests {
+    use datafusion_physical_plan::Partitioning;
+    use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
     use super::*;
 
@@ -384,7 +443,7 @@ pub(crate) mod tests {
         pub fn new(schema: arrow::datatypes::SchemaRef) -> Self {
             Self {
                 props: PlanProperties::new(
-                    datafusion::physical_expr::EquivalenceProperties::new(schema),
+                    datafusion_physical_expr::EquivalenceProperties::new(schema),
                     Partitioning::UnknownPartitioning(3),
                     EmissionType::Incremental,
                     Boundedness::Bounded,
@@ -442,8 +501,9 @@ pub(crate) mod tests {
 
     #[test]
     fn test_round_trip_ffi_execution_plan() -> Result<()> {
-        let schema =
-            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+            arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Float32, false),
+        ]));
 
         let original_plan = Arc::new(EmptyExec::new(schema));
         let original_name = original_plan.name().to_string();
@@ -455,7 +515,7 @@ pub(crate) mod tests {
 
         assert_eq!(original_name, foreign_plan.name());
 
-        let display = datafusion::physical_plan::display::DisplayableExecutionPlan::new(
+        let display = datafusion_physical_plan::display::DisplayableExecutionPlan::new(
             foreign_plan.as_ref(),
         );
 
@@ -470,8 +530,9 @@ pub(crate) mod tests {
 
     #[test]
     fn test_ffi_execution_plan_children() -> Result<()> {
-        let schema =
-            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+            arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Float32, false),
+        ]));
 
         // Version 1: Adding child to the foreign plan
         let child_plan = Arc::new(EmptyExec::new(Arc::clone(&schema)));
@@ -509,8 +570,9 @@ pub(crate) mod tests {
 
     #[test]
     fn test_ffi_execution_plan_local_bypass() {
-        let schema =
-            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
+            arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Float32, false),
+        ]));
 
         let plan = Arc::new(EmptyExec::new(schema));
 

datafusion/ffi/src/tests/mod.rs

@@ -39,6 +39,8 @@ use super::udf::FFI_ScalarUDF;
 use crate::catalog_provider::FFI_CatalogProvider;
 use crate::catalog_provider_list::FFI_CatalogProviderList;
 use crate::config::extension_options::FFI_ExtensionOptions;
+use crate::execution_plan::FFI_ExecutionPlan;
+use crate::execution_plan::tests::EmptyExec;
 use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use crate::tests::catalog::create_catalog_provider_list;
 use crate::udaf::FFI_AggregateUDF;
@@ -92,6 +94,8 @@ pub struct ForeignLibraryModule {
     /// Create extension options, for either ConfigOptions or TableOptions
     pub create_extension_options: extern "C" fn() -> FFI_ExtensionOptions,
 
+    pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan,
+
     pub version: extern "C" fn() -> u64,
 }
 
@@ -133,6 +137,13 @@ extern "C" fn construct_table_provider(
     }
 }
 
+pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan {
+    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+
+    let plan = Arc::new(EmptyExec::new(schema));
+    FFI_ExecutionPlan::new(plan, None)
+}
+
 #[export_root_module]
 /// This defines the entry point for using the module.
 pub fn get_foreign_library_module() -> ForeignLibraryModuleRef {
@@ -147,6 +158,7 @@ pub fn get_foreign_library_module() -> ForeignLibraryModuleRef {
         create_stddev_udaf: create_ffi_stddev_func,
         create_rank_udwf: create_ffi_rank_func,
         create_extension_options: config::create_extension_options,
+        create_empty_exec,
         version: super::version,
     }
     .leak_into_prefix()

datafusion/ffi/tests/ffi_execution_plan.rs

@@ -0,0 +1,91 @@
+#[cfg(feature = "integration-tests")]
+mod tests {
+    use arrow::datatypes::Field;
+    use arrow::datatypes::Schema;
+    use arrow_schema::DataType;
+    use datafusion_common::DataFusionError;
+    use datafusion_ffi::execution_plan::FFI_ExecutionPlan;
+    use datafusion_ffi::execution_plan::ForeignExecutionPlan;
+    use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec};
+    use datafusion_ffi::tests::utils::get_module;
+    use datafusion_physical_plan::ExecutionPlan;
+    use std::sync::Arc;
+
+    #[test]
+    fn test_ffi_execution_plan_new_sets_runtimes_on_children()
+    -> Result<(), DataFusionError> {
+        // We want to test the case where we have two libraries.
+        // Library A will have a foreign plan from Library B, called child_plan.
+        // Library A will add a plan called grandchild_plan under child_plan
+        // Library A will create a plan called parent_plan, that has child_plan
+        // under it. So we should have:
+        // parent_plan (local) -> child_plan (foreign) -> grandchild_plan (local)
+        // Then we want to turn parent_plan into a FFI plan.
+        // Verify that grandchild_plan also gets the same runtime as parent_plan.
+
+        let module = get_module()?;
+
+        fn generate_local_plan() -> Arc<dyn ExecutionPlan> {
+            let schema =
+                Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+
+            Arc::new(EmptyExec::new(schema))
+        }
+
+        let child_plan =
+            module
+                .create_empty_exec()
+                .ok_or(DataFusionError::NotImplemented(
+                    "External module failed to implement create_empty_exec".to_string(),
+                ))?();
+        let child_plan: Arc<dyn ExecutionPlan> = (&child_plan)
+            .try_into()
+            .expect("should be able create plan");
+        assert!(child_plan.as_any().is::<ForeignExecutionPlan>());
+
+        let grandchild_plan = generate_local_plan();
+
+        let child_plan = child_plan.with_new_children(vec![grandchild_plan])?;
+
+        unsafe {
+            // Originally the runtime is not set. We go through the unsafe casting
+            // of data here because the `inner()` function is private and this is
+            // only an integration test so we do not want to expose it.
+            let ffi_child = FFI_ExecutionPlan::new(Arc::clone(&child_plan), None);
+            let ffi_grandchild =
+                (ffi_child.children)(&ffi_child).into_iter().next().unwrap();
+
+            let grandchild_private_data =
+                ffi_grandchild.private_data as *const ExecutionPlanPrivateData;
+            assert!((*grandchild_private_data).runtime.is_none());
+        }
+
+        let parent_plan = generate_local_plan().with_new_children(vec![child_plan])?;
+
+        // Adding the grandchild beneath this FFI plan should get the runtime passed down.
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .build()
+            .unwrap();
+        let ffi_parent =
+            FFI_ExecutionPlan::new(parent_plan, Some(runtime.handle().clone()));
+
+        unsafe {
+            let ffi_child = (ffi_parent.children)(&ffi_parent)
+                .into_iter()
+                .next()
+                .unwrap();
+            let ffi_grandchild =
+                (ffi_child.children)(&ffi_child).into_iter().next().unwrap();
+            assert_eq!(
+                (ffi_grandchild.library_marker_id)(),
+                (ffi_parent.library_marker_id)()
+            );
+
+            let grandchild_private_data =
+                ffi_grandchild.private_data as *const ExecutionPlanPrivateData;
+            assert!((*grandchild_private_data).runtime.is_some());
+        }
+
+        Ok(())
+    }
+}