Implement new_with_children and repartitioned on execution plan in FFI

Author: timsaucer

datafusion-examples/examples/custom_data_source/adapter_serialization.rs

@@ -308,6 +308,10 @@ 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

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

datafusion-examples/examples/proto/expression_deduplication.rs

@@ -199,6 +199,10 @@ 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/ffi/src/execution_plan.rs

@@ -20,7 +20,8 @@ use std::pin::Pin;
 use std::sync::Arc;
 
 use abi_stable::StableAbi;
-use abi_stable::std_types::{RString, RVec};
+use abi_stable::std_types::{ROption, RResult, RString, RVec};
+use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::{DataFusionError, Result};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
@@ -29,11 +30,12 @@ use datafusion_physical_plan::{
 };
 use tokio::runtime::Handle;
 
+use crate::config::FFI_ConfigOptions;
 use crate::execution::FFI_TaskContext;
 use crate::plan_properties::FFI_PlanProperties;
 use crate::record_batch_stream::FFI_RecordBatchStream;
 use crate::util::FFIResult;
-use crate::{df_result, rresult};
+use crate::{df_result, rresult, rresult_return};
 
 /// A stable struct for sharing a [`ExecutionPlan`] across FFI boundaries.
 #[repr(C)]
@@ -45,6 +47,9 @@ pub struct FFI_ExecutionPlan {
     /// Return a vector of children plans
     pub children: unsafe extern "C" fn(plan: &Self) -> RVec<FFI_ExecutionPlan>,
 
+    pub with_new_children:
+        unsafe extern "C" fn(plan: &Self, children: RVec<Self>) -> FFIResult<Self>,
+
     /// Return the plan name.
     pub name: unsafe extern "C" fn(plan: &Self) -> RString,
 
@@ -56,6 +61,12 @@ pub struct FFI_ExecutionPlan {
         context: FFI_TaskContext,
     ) -> FFIResult<FFI_RecordBatchStream>,
 
+    pub repartitioned: unsafe extern "C" fn(
+        plan: &Self,
+        target_partitions: usize,
+        config: FFI_ConfigOptions,
+    ) -> FFIResult<ROption<FFI_ExecutionPlan>>,
+
     /// Used to create a clone on the provider of the execution plan. This should
     /// only need to be called by the receiver of the plan.
     pub clone: unsafe extern "C" fn(plan: &Self) -> Self,
@@ -86,6 +97,11 @@ impl FFI_ExecutionPlan {
         let private_data = self.private_data as *const ExecutionPlanPrivateData;
         unsafe { &(*private_data).plan }
     }
+
+    fn runtime(&self) -> Option<Handle> {
+        let private_data = self.private_data as *const ExecutionPlanPrivateData;
+        unsafe { (*private_data).runtime.clone() }
+    }
 }
 
 unsafe extern "C" fn properties_fn_wrapper(
@@ -97,37 +113,69 @@ unsafe extern "C" fn properties_fn_wrapper(
 unsafe extern "C" fn children_fn_wrapper(
     plan: &FFI_ExecutionPlan,
 ) -> RVec<FFI_ExecutionPlan> {
-    unsafe {
-        let private_data = plan.private_data as *const ExecutionPlanPrivateData;
-        let plan = &(*private_data).plan;
-        let runtime = &(*private_data).runtime;
+    let runtime = plan.runtime();
+    let plan = plan.inner();
 
-        let children: Vec<_> = plan
-            .children()
-            .into_iter()
-            .map(|child| FFI_ExecutionPlan::new(Arc::clone(child), runtime.clone()))
-            .collect();
+    let children: Vec<_> = plan
+        .children()
+        .into_iter()
+        .map(|child| FFI_ExecutionPlan::new(Arc::clone(child), runtime.clone()))
+        .collect();
 
-        children.into()
-    }
+    children.into()
+}
+
+unsafe extern "C" fn with_new_children_fn_wrapper(
+    plan: &FFI_ExecutionPlan,
+    children: RVec<FFI_ExecutionPlan>,
+) -> FFIResult<FFI_ExecutionPlan> {
+    let runtime = plan.runtime();
+    let plan = Arc::clone(plan.inner());
+    let children = rresult_return!(
+        children
+            .iter()
+            .map(<Arc<dyn ExecutionPlan>>::try_from)
+            .collect::<Result<Vec<_>>>()
+    );
+
+    let new_plan = rresult_return!(plan.with_new_children(children));
+
+    RResult::ROk(FFI_ExecutionPlan::new(new_plan, runtime))
 }
 
 unsafe extern "C" fn execute_fn_wrapper(
     plan: &FFI_ExecutionPlan,
     partition: usize,
     context: FFI_TaskContext,
 ) -> FFIResult<FFI_RecordBatchStream> {
-    unsafe {
-        let ctx = context.into();
-        let private_data = plan.private_data as *const ExecutionPlanPrivateData;
-        let plan = &(*private_data).plan;
-        let runtime = (*private_data).runtime.clone();
-
-        rresult!(
-            plan.execute(partition, ctx)
-                .map(|rbs| FFI_RecordBatchStream::new(rbs, runtime))
-        )
-    }
+    let ctx = context.into();
+    let runtime = plan.runtime();
+    let plan = plan.inner();
+
+    let _guard = runtime.as_ref().map(|rt| rt.enter());
+
+    rresult!(
+        plan.execute(partition, ctx)
+            .map(|rbs| FFI_RecordBatchStream::new(rbs, runtime))
+    )
+}
+
+unsafe extern "C" fn repartitioned_fn_wrapper(
+    plan: &FFI_ExecutionPlan,
+    target_partitions: usize,
+    config: FFI_ConfigOptions,
+) -> FFIResult<ROption<FFI_ExecutionPlan>> {
+    let maybe_config: Result<ConfigOptions, DataFusionError> = config.try_into();
+    let config = rresult_return!(maybe_config);
+    let runtime = plan.runtime();
+    let plan = plan.inner();
+
+    rresult!(
+        plan.repartitioned(target_partitions, &config)
+            .map(|maybe_plan| maybe_plan
+                .map(|plan| FFI_ExecutionPlan::new(plan, runtime))
+                .into())
+    )
 }
 
 unsafe extern "C" fn name_fn_wrapper(plan: &FFI_ExecutionPlan) -> RString {
@@ -145,12 +193,10 @@ unsafe extern "C" fn release_fn_wrapper(plan: &mut FFI_ExecutionPlan) {
 }
 
 unsafe extern "C" fn clone_fn_wrapper(plan: &FFI_ExecutionPlan) -> FFI_ExecutionPlan {
-    unsafe {
-        let private_data = plan.private_data as *const ExecutionPlanPrivateData;
-        let plan_data = &(*private_data);
+    let runtime = plan.runtime();
+    let plan = plan.inner();
 
-        FFI_ExecutionPlan::new(Arc::clone(&plan_data.plan), plan_data.runtime.clone())
-    }
+    FFI_ExecutionPlan::new(Arc::clone(plan), runtime)
 }
 
 impl Clone for FFI_ExecutionPlan {
@@ -162,13 +208,18 @@ impl Clone for FFI_ExecutionPlan {
 impl FFI_ExecutionPlan {
     /// This function is called on the provider's side.
     pub fn new(plan: Arc<dyn ExecutionPlan>, runtime: Option<Handle>) -> Self {
-        let private_data = Box::new(ExecutionPlanPrivateData { plan, runtime });
+        if let Some(plan) = plan.as_any().downcast_ref::<ForeignExecutionPlan>() {
+            return plan.plan.clone();
+        }
 
+        let private_data = Box::new(ExecutionPlanPrivateData { plan, runtime });
         Self {
             properties: properties_fn_wrapper,
             children: children_fn_wrapper,
+            with_new_children: with_new_children_fn_wrapper,
             name: name_fn_wrapper,
             execute: execute_fn_wrapper,
+            repartitioned: repartitioned_fn_wrapper,
             clone: clone_fn_wrapper,
             release: release_fn_wrapper,
             private_data: Box::into_raw(private_data) as *mut c_void,
@@ -275,12 +326,16 @@ impl ExecutionPlan for ForeignExecutionPlan {
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        Ok(Arc::new(ForeignExecutionPlan {
-            plan: self.plan.clone(),
-            name: self.name.clone(),
-            children,
-            properties: Arc::clone(&self.properties),
-        }))
+        unsafe {
+            let children = children
+                .into_iter()
+                .map(|child| FFI_ExecutionPlan::new(child, None))
+                .collect::<RVec<_>>();
+            let new_plan =
+                df_result!((self.plan.with_new_children)(&self.plan, children))?;
+
+            (&new_plan).try_into()
+        }
     }
 
     fn execute(
@@ -310,6 +365,22 @@ impl ExecutionPlan for ForeignExecutionPlan {
         }
         Ok(tnr)
     }
+
+    fn repartitioned(
+        &self,
+        target_partitions: usize,
+        config: &ConfigOptions,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        let config = config.into();
+        let maybe_plan: Option<FFI_ExecutionPlan> = df_result!(unsafe {
+            (self.plan.repartitioned)(&self.plan, target_partitions, config)
+        })?
+        .into();
+
+        maybe_plan
+            .map(|plan| <Arc<dyn ExecutionPlan>>::try_from(&plan))
+            .transpose()
+    }
 }
 
 #[cfg(test)]

datafusion/ffi/src/physical_expr/mod.rs

@@ -448,6 +448,10 @@ impl Drop for FFI_PhysicalExpr {
 impl From<Arc<dyn PhysicalExpr>> for FFI_PhysicalExpr {
     /// Creates a new [`FFI_PhysicalExpr`].
     fn from(expr: Arc<dyn PhysicalExpr>) -> Self {
+        if let Some(expr) = expr.as_any().downcast_ref::<ForeignPhysicalExpr>() {
+            return expr.expr.clone();
+        }
+
         let private_data = Box::new(PhysicalExprPrivateData { expr });
 
         Self {

datafusion/ffi/src/proto/physical_extension_codec.rs

@@ -111,14 +111,14 @@ unsafe impl Send for FFI_PhysicalExtensionCodec {}
 unsafe impl Sync for FFI_PhysicalExtensionCodec {}
 
 struct PhysicalExtensionCodecPrivateData {
-    provider: Arc<dyn PhysicalExtensionCodec>,
+    codec: Arc<dyn PhysicalExtensionCodec>,
     runtime: Option<Handle>,
 }
 
 impl FFI_PhysicalExtensionCodec {
     fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> {
         let private_data = self.private_data as *const PhysicalExtensionCodecPrivateData;
-        unsafe { &(*private_data).provider }
+        unsafe { &(*private_data).codec }
     }
 
     fn runtime(&self) -> &Option<Handle> {
@@ -132,6 +132,7 @@ unsafe extern "C" fn try_decode_fn_wrapper(
     buf: RSlice<u8>,
     inputs: RVec<FFI_ExecutionPlan>,
 ) -> FFIResult<FFI_ExecutionPlan> {
+    let runtime = codec.runtime().clone();
     let task_ctx: Arc<TaskContext> =
         rresult_return!((&codec.task_ctx_provider).try_into());
     let codec = codec.inner();
@@ -144,7 +145,7 @@ unsafe extern "C" fn try_decode_fn_wrapper(
     let plan =
         rresult_return!(codec.try_decode(buf.as_ref(), &inputs, task_ctx.as_ref()));
 
-    RResult::ROk(FFI_ExecutionPlan::new(plan, None))
+    RResult::ROk(FFI_ExecutionPlan::new(plan, runtime))
 }
 
 unsafe extern "C" fn try_encode_fn_wrapper(
@@ -240,11 +241,10 @@ unsafe extern "C" fn try_encode_udwf_fn_wrapper(
     RResult::ROk(bytes.into())
 }
 
-unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_PhysicalExtensionCodec) {
+unsafe extern "C" fn release_fn_wrapper(codec: &mut FFI_PhysicalExtensionCodec) {
     unsafe {
-        let private_data = Box::from_raw(
-            provider.private_data as *mut PhysicalExtensionCodecPrivateData,
-        );
+        let private_data =
+            Box::from_raw(codec.private_data as *mut PhysicalExtensionCodecPrivateData);
         drop(private_data);
     }
 }
@@ -267,13 +267,19 @@ impl Drop for FFI_PhysicalExtensionCodec {
 impl FFI_PhysicalExtensionCodec {
     /// Creates a new [`FFI_PhysicalExtensionCodec`].
     pub fn new(
-        provider: Arc<dyn PhysicalExtensionCodec + Send>,
+        codec: Arc<dyn PhysicalExtensionCodec + Send>,
         runtime: Option<Handle>,
         task_ctx_provider: impl Into<FFI_TaskContextProvider>,
     ) -> Self {
+        if let Some(codec) = codec
+            .as_any()
+            .downcast_ref::<ForeignPhysicalExtensionCodec>()
+        {
+            return codec.0.clone();
+        }
+
         let task_ctx_provider = task_ctx_provider.into();
-        let private_data =
-            Box::new(PhysicalExtensionCodecPrivateData { provider, runtime });
+        let private_data = Box::new(PhysicalExtensionCodecPrivateData { codec, runtime });
 
         Self {
             try_decode: try_decode_fn_wrapper,
@@ -306,11 +312,11 @@ unsafe impl Send for ForeignPhysicalExtensionCodec {}
 unsafe impl Sync for ForeignPhysicalExtensionCodec {}
 
 impl From<&FFI_PhysicalExtensionCodec> for Arc<dyn PhysicalExtensionCodec> {
-    fn from(provider: &FFI_PhysicalExtensionCodec) -> Self {
-        if (provider.library_marker_id)() == crate::get_library_marker_id() {
-            Arc::clone(provider.inner())
+    fn from(codec: &FFI_PhysicalExtensionCodec) -> Self {
+        if (codec.library_marker_id)() == crate::get_library_marker_id() {
+            Arc::clone(codec.inner())
         } else {
-            Arc::new(ForeignPhysicalExtensionCodec(provider.clone()))
+            Arc::new(ForeignPhysicalExtensionCodec(codec.clone()))
         }
     }
 }
@@ -403,6 +409,10 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec {
 
         Ok(())
     }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
 }
 
 #[cfg(test)]
@@ -558,6 +568,10 @@ pub(crate) mod tests {
 
             Ok(())
         }
+
+        fn as_any(&self) -> &dyn std::any::Any {
+            self
+        }
     }
 
     fn create_test_exec() -> Arc<dyn ExecutionPlan> {