WIP on execution plan

Author: timsaucer

src/ffi/execution_plan.rs

@@ -3,19 +3,105 @@ use std::{
     sync::Arc,
 };
 
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::{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 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);
+
+impl FFI_ExecutionPlan {
+
+    pub fn empty() -> Self {
+        Self {
+            private_data: std::ptr::null_mut(),
+        }
+    }
+}
+
+impl FFI_ExecutionPlan {
+    pub fn new(plan: Arc<dyn ExecutionPlan + Send>) -> Self {
+        let private_data = Box::new(ExecutionPlanPrivateData {
+            plan,
+            last_error: None,
+        });
+
+        Self {
+            private_data: Box::into_raw(private_data) as *mut c_void
+        }
+    }
+}
+
+impl ExecutionPlan for FFI_ExecutionPlan {
+    fn name(&self) -> &str {
+        todo!()
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        todo!()
+    }
+
+    fn properties(&self) -> &datafusion::physical_plan::PlanProperties {
+        self.properties
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        todo!()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+        todo!()
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<datafusion::execution::TaskContext>,
+    ) -> datafusion::error::Result<datafusion::execution::SendableRecordBatchStream> {
+        todo!()
+    }
+}
+
+impl DisplayAs for FFI_ExecutionPlan {
+    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)]
+#[allow(non_camel_case_types)]
+pub struct FFI_PlanProperties {
+    /// See [ExecutionPlanProperties::equivalence_properties]
+    pub eq_properties: EquivalenceProperties,
+    /// See [ExecutionPlanProperties::output_partitioning]
+    pub partitioning: Partitioning,
+    /// See [ExecutionPlanProperties::execution_mode]
+    pub execution_mode: ExecutionMode,
+    /// See [ExecutionPlanProperties::output_ordering]
+    output_ordering: Option<LexOrdering>,
+}

src/ffi/session_config.rs

@@ -1,6 +1,6 @@
-use std::ffi::{c_void, CString};
+use std::{ffi::{c_void, CString}, sync::Arc};
 
-use datafusion::prelude::SessionConfig;
+use datafusion::{catalog::Session, prelude::SessionConfig};
 
 #[repr(C)]
 #[derive(Debug)]
@@ -28,3 +28,19 @@ impl ExportedSessionConfig {
         unsafe { &mut *((*self.session).private_data as *mut SessionConfigPrivateData) }
     }
 }
+
+impl FFI_SessionConfig {
+    /// Creates a new [`FFI_TableProvider`].
+    pub fn new(session: &dyn Session) -> Self {
+        let config = session.config().clone();
+        let private_data = Box::new(SessionConfigPrivateData {
+            config,
+            last_error: None,
+        });
+
+        Self {
+            version: 2,
+            private_data: Box::into_raw(private_data) as *mut c_void,
+        }
+    }
+}

src/ffi/table_provider.rs

@@ -37,12 +37,12 @@ pub struct FFI_TableProvider {
     pub schema: Option<unsafe extern "C" fn(provider: *const FFI_TableProvider) -> FFI_ArrowSchema>,
     pub scan: Option<
         unsafe extern "C" fn(
-            provider: *mut FFI_TableProvider,
-            session_config: *mut FFI_SessionConfig,
+            provider: *const FFI_TableProvider,
+            session_config: *const FFI_SessionConfig,
             n_projections: c_int,
-            projections: *mut c_int,
+            projections: *const c_int,
             n_filters: c_int,
-            filters: *mut *const c_char,
+            filters: *const *const c_char,
             limit: c_int,
             out: *mut FFI_ExecutionPlan,
         ) -> c_int,
@@ -58,25 +58,20 @@ struct ProviderPrivateData {
     last_error: Option<CString>,
 }
 
-struct ExportedTableProvider {
-    provider: *mut FFI_TableProvider,
-}
-struct ConstExportedTableProvider {
-    provider: *const FFI_TableProvider,
-}
+struct ExportedTableProvider(*const FFI_TableProvider);
 
 // The callback used to get array schema
 unsafe extern "C" fn provider_schema(provider: *const FFI_TableProvider) -> FFI_ArrowSchema {
-    ConstExportedTableProvider { provider }.provider_schema()
+    ExportedTableProvider(provider).provider_schema()
 }
 
 unsafe extern "C" fn provider_scan(
-    provider: *mut FFI_TableProvider,
-    session_config: *mut FFI_SessionConfig,
+    provider: *const FFI_TableProvider,
+    session_config: *const FFI_SessionConfig,
     n_projections: c_int,
-    projections: *mut c_int,
+    projections: *const c_int,
     n_filters: c_int,
-    filters: *mut *const c_char,
+    filters: *const *const c_char,
     limit: c_int,
     mut out: *mut FFI_ExecutionPlan,
 ) -> c_int {
@@ -104,7 +99,7 @@ unsafe extern "C" fn provider_scan(
 
     let limit = limit.try_into().ok();
 
-    let plan = ExportedTableProvider { provider }.provider_scan(
+    let plan = ExportedTableProvider(provider).provider_scan(
         &session,
         maybe_projections,
         filters_vec,
@@ -120,9 +115,9 @@ unsafe extern "C" fn provider_scan(
     }
 }
 
-impl ConstExportedTableProvider {
+impl ExportedTableProvider {
     fn get_private_data(&self) -> &ProviderPrivateData {
-        unsafe { &*((*self.provider).private_data as *const ProviderPrivateData) }
+        unsafe { &*((*self.0).private_data as *const ProviderPrivateData) }
     }
 
     pub fn provider_schema(&self) -> FFI_ArrowSchema {
@@ -133,12 +128,6 @@ impl ConstExportedTableProvider {
         // so we expect it to always pass. Maybe some logging should be added.
         FFI_ArrowSchema::try_from(provider.schema().as_ref()).unwrap_or(FFI_ArrowSchema::empty())
     }
-}
-
-impl ExportedTableProvider {
-    fn get_private_data(&mut self) -> &mut ProviderPrivateData {
-        unsafe { &mut *((*self.provider).private_data as *mut ProviderPrivateData) }
-    }
 
     pub fn provider_scan(
         &mut self,
@@ -250,22 +239,46 @@ impl TableProvider for FFI_TableProvider {
     /// parallelized or distributed.
     async fn scan(
         &self,
-        _ctx: &dyn Session,
+        session: &dyn Session,
         projection: Option<&Vec<usize>>,
         filters: &[Expr],
         // limit can be used to reduce the amount scanned
         // from the datasource as a performance optimization.
         // If set, it contains the amount of rows needed by the `LogicalPlan`,
         // The datasource should return *at least* this number of rows if available.
-        _limit: Option<usize>,
+        limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let scan_fn = self.scan.ok_or(DataFusionError::NotImplemented(
             "Scan not defined on FFI_TableProvider".to_string(),
         ))?;
 
-        Err(datafusion::error::DataFusionError::NotImplemented(
-            "scan not implemented".to_string(),
-        ))
+        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_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())
+            .collect();
+
+        let limit = match limit {
+            Some(l) => l as c_int,
+            None => -1,
+        };
+
+        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)
+        };
+
+        match err_code {
+            0 => Ok(Arc::new(out)),
+            _ => Err(datafusion::error::DataFusionError::Internal("Unable to perform scan via FFI".to_string()))
+        }
     }
 
     /// Tests whether the table provider can make use of a filter expression