Working through execution plan FFI

Author: timsaucer

src/context.rs

@@ -59,7 +59,9 @@ use datafusion::datasource::listing::{
 };
 use datafusion::datasource::TableProvider;
 use datafusion::datasource::{provider, MemTable};
-use datafusion::execution::context::{DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext};
+use datafusion::execution::context::{
+    DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext,
+};
 use datafusion::execution::disk_manager::DiskManagerConfig;
 use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, UnboundedMemoryPool};
 use datafusion::execution::options::ReadOptions;

src/ffi/execution_plan.rs

@@ -1,36 +1,84 @@
 use std::{
     ffi::{c_void, CString},
+    ptr::null,
     sync::Arc,
 };
 
-use datafusion::{physical_expr::{EquivalenceProperties, LexOrdering}, physical_plan::{DisplayAs, ExecutionMode, ExecutionPlan, Partitioning, PlanProperties}};
+use datafusion::error::Result;
+use datafusion::{
+    error::DataFusionError,
+    parquet::file::properties,
+    physical_expr::{EquivalenceProperties, LexOrdering},
+    physical_plan::{DisplayAs, ExecutionMode, ExecutionPlan, Partitioning, PlanProperties},
+};
 
 #[repr(C)]
 #[derive(Debug)]
 #[allow(missing_docs)]
 #[allow(non_camel_case_types)]
 pub struct FFI_ExecutionPlan {
-    pub properties: Option<unsafe extern "C" fn(provider: *const FFI_ExecutionPlan) -> FFI_ArrowSchema>,
+    pub properties:
+        Option<unsafe extern "C" fn(plan: *const FFI_ExecutionPlan) -> FFI_PlanProperties>,
+    pub children: Option<
+        unsafe extern "C" fn(
+            plan: *const FFI_ExecutionPlan,
+            num_children: &mut usize,
+            out: &mut *const FFI_ExecutionPlan,
+        ) -> i32,
+    >,
 
     pub private_data: *mut c_void,
 }
 
-unsafe impl Send for FFI_ExecutionPlan {}
-unsafe impl Sync for FFI_ExecutionPlan {}
-
 pub struct ExecutionPlanPrivateData {
     pub plan: Arc<dyn ExecutionPlan + Send>,
     pub last_error: Option<CString>,
 }
 
-struct ExportedExecutionPlan(*const FFI_ExecutionPlan);
+unsafe extern "C" fn properties_fn_wrapper(plan: *const FFI_ExecutionPlan) -> FFI_PlanProperties {
+    let private_data = (*plan).private_data as *const ExecutionPlanPrivateData;
+    let properties = (*private_data).plan.properties();
+    properties.into()
+}
 
-impl FFI_ExecutionPlan {
+unsafe extern "C" fn children_fn_wrapper(
+    plan: *const FFI_ExecutionPlan,
+    num_children: &mut usize,
+    out: &mut *const FFI_ExecutionPlan,
+) -> i32 {
+    let private_data = (*plan).private_data as *const ExecutionPlanPrivateData;
+
+    let children = (*private_data).plan.children();
+    *num_children = children.len();
+    let children: Vec<FFI_ExecutionPlan> = children
+        .into_iter()
+        .map(|child| FFI_ExecutionPlan::new(child.clone()))
+        .collect();
+    *out = children.as_ptr();
+
+    0
+}
 
-    pub fn empty() -> Self {
-        Self {
-            private_data: std::ptr::null_mut(),
-        }
+// Since the trait ExecutionPlan requires borrowed values, we wrap our FFI.
+// This struct exists on the consumer side (datafusion-python, for example) and not
+// in the provider's side.
+#[derive(Debug)]
+pub struct ExportedExecutionPlan {
+    plan: *const FFI_ExecutionPlan,
+    properties: PlanProperties,
+    children: Vec<Arc<dyn ExecutionPlan>>,
+}
+
+unsafe impl Send for ExportedExecutionPlan {}
+unsafe impl Sync for ExportedExecutionPlan {}
+
+impl DisplayAs for ExportedExecutionPlan {
+    fn fmt_as(
+        &self,
+        t: datafusion::physical_plan::DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
+        todo!()
     }
 }
 
@@ -42,12 +90,63 @@ impl FFI_ExecutionPlan {
         });
 
         Self {
-            private_data: Box::into_raw(private_data) as *mut c_void
+            properties: Some(properties_fn_wrapper),
+            children: Some(children_fn_wrapper),
+            private_data: Box::into_raw(private_data) as *mut c_void,
+        }
+    }
+
+    pub fn empty() -> Self {
+        Self {
+            properties: None,
+            children: None,
+            private_data: std::ptr::null_mut(),
         }
     }
 }
 
-impl ExecutionPlan for FFI_ExecutionPlan {
+impl ExportedExecutionPlan {
+    pub fn new(plan: *const FFI_ExecutionPlan) -> Result<Self> {
+        let properties = unsafe {
+            let properties_fn = (*plan).properties.ok_or(DataFusionError::NotImplemented(
+                "properties not implemented on FFI_ExecutionPlan".to_string(),
+            ))?;
+            properties_fn(plan).into()
+        };
+
+        let children = unsafe {
+            let children_fn = (*plan).children.ok_or(DataFusionError::NotImplemented(
+                "children not implemented on FFI_ExecutionPlan".to_string(),
+            ))?;
+            let mut num_children = 0;
+            let mut children_ptr: *const FFI_ExecutionPlan = null();
+
+            if children_fn(plan, &mut num_children, &mut children_ptr) != 0 {
+                return Err(DataFusionError::Plan(
+                    "Error getting children for FFI_ExecutionPlan".to_string(),
+                ));
+            }
+
+            let ffi_vec = Vec::from_raw_parts(&mut children_ptr, num_children, num_children);
+            let maybe_children: Result<Vec<_>> = ffi_vec
+                .into_iter()
+                .map(|child| {
+                    ExportedExecutionPlan::new(child).map(|c| Arc::new(c) as Arc<dyn ExecutionPlan>)
+                })
+                .collect();
+
+            maybe_children?
+        };
+
+        Ok(Self {
+            plan,
+            properties,
+            children,
+        })
+    }
+}
+
+impl ExecutionPlan for ExportedExecutionPlan {
     fn name(&self) -> &str {
         todo!()
     }
@@ -57,11 +156,11 @@ impl ExecutionPlan for FFI_ExecutionPlan {
     }
 
     fn properties(&self) -> &datafusion::physical_plan::PlanProperties {
-        self.properties
+        &self.properties
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        todo!()
+        self.children.iter().collect()
     }
 
     fn with_new_children(
@@ -81,16 +180,15 @@ impl ExecutionPlan for FFI_ExecutionPlan {
 }
 
 impl DisplayAs for FFI_ExecutionPlan {
-    fn fmt_as(&self, t: datafusion::physical_plan::DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+    fn fmt_as(
+        &self,
+        t: datafusion::physical_plan::DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
         todo!()
     }
 }
 
-
-
-
-
-
 #[repr(C)]
 #[derive(Debug)]
 #[allow(missing_docs)]
@@ -104,4 +202,16 @@ pub struct FFI_PlanProperties {
     pub execution_mode: ExecutionMode,
     /// See [ExecutionPlanProperties::output_ordering]
     output_ordering: Option<LexOrdering>,
-}
+}
+
+impl From<&PlanProperties> for FFI_PlanProperties {
+    fn from(value: &PlanProperties) -> Self {
+        todo!()
+    }
+}
+
+impl From<FFI_PlanProperties> for PlanProperties {
+    fn from(value: FFI_PlanProperties) -> Self {
+        todo!()
+    }
+}

src/ffi/session_config.rs

@@ -1,4 +1,7 @@
-use std::{ffi::{c_void, CString}, sync::Arc};
+use std::{
+    ffi::{c_void, CString},
+    sync::Arc,
+};
 
 use datafusion::{catalog::Session, prelude::SessionConfig};
 
@@ -43,4 +46,4 @@ impl FFI_SessionConfig {
             private_data: Box::into_raw(private_data) as *mut c_void,
         }
     }
-}
+}

src/ffi/table_provider.rs

@@ -23,7 +23,7 @@ use datafusion::{
 use tokio::runtime::Runtime;
 
 use super::{
-    execution_plan::{ExecutionPlanPrivateData, FFI_ExecutionPlan},
+    execution_plan::{ExecutionPlanPrivateData, ExportedExecutionPlan, FFI_ExecutionPlan},
     session_config::{FFI_SessionConfig, SessionConfigPrivateData},
 };
 use datafusion::error::Result;
@@ -150,14 +150,15 @@ impl ExportedTableProvider {
         let runtime = Runtime::new().unwrap();
         let plan = runtime.block_on(provider.scan(session, projections, &filter_exprs, limit))?;
 
-        let plan_ptr = Box::new(ExecutionPlanPrivateData {
-            plan,
-            last_error: None,
-        });
+        // let plan_ptr = Box::new(ExecutionPlanPrivateData {
+        //     plan,
+        //     last_error: None,
+        // });
 
-        Ok(FFI_ExecutionPlan {
-            private_data: Box::into_raw(plan_ptr) as *mut c_void,
-        })
+        // Ok(FFI_ExecutionPlan {
+        //     private_data: Box::into_raw(plan_ptr) as *mut c_void,
+        // })
+        Ok(FFI_ExecutionPlan::new(plan))
     }
 }
 
@@ -255,14 +256,17 @@ impl TableProvider for FFI_TableProvider {
         let session_config = FFI_SessionConfig::new(session);
 
         let n_projections = projection.map(|p| p.len()).unwrap_or(0) as c_int;
-        let projections: Vec<c_int> = projection.map(|p| p.iter().map(|v| *v as c_int).collect()).unwrap_or_default();
+        let projections: Vec<c_int> = projection
+            .map(|p| p.iter().map(|v| *v as c_int).collect())
+            .unwrap_or_default();
         let projections_ptr = projections.as_ptr();
 
         let n_filters = filters.len() as c_int;
-        let filters: Vec<CString> = filters.iter().filter_map(|f| CString::new(f.to_string()).ok()).collect();
-        let filters_ptr: Vec<*const i8> = filters.iter()
-            .map(|s| s.as_ptr())
+        let filters: Vec<CString> = filters
+            .iter()
+            .filter_map(|f| CString::new(f.to_string()).ok())
             .collect();
+        let filters_ptr: Vec<*const i8> = filters.iter().map(|s| s.as_ptr()).collect();
 
         let limit = match limit {
             Some(l) => l as c_int,
@@ -272,13 +276,26 @@ impl TableProvider for FFI_TableProvider {
         let mut out = FFI_ExecutionPlan::empty();
 
         let err_code = unsafe {
-            scan_fn(self, &session_config, n_projections, projections_ptr, n_filters, filters_ptr.as_ptr(), limit, &mut out)
+            scan_fn(
+                self,
+                &session_config,
+                n_projections,
+                projections_ptr,
+                n_filters,
+                filters_ptr.as_ptr(),
+                limit,
+                &mut out,
+            )
         };
 
-        match err_code {
-            0 => Ok(Arc::new(out)),
-            _ => Err(datafusion::error::DataFusionError::Internal("Unable to perform scan via FFI".to_string()))
+        if 0 != err_code {
+            return Err(datafusion::error::DataFusionError::Internal(
+                "Unable to perform scan via FFI".to_string(),
+            ));
         }
+
+        let plan = ExportedExecutionPlan::new(&out)?;
+        Ok(Arc::new(plan))
     }
 
     /// Tests whether the table provider can make use of a filter expression