Move arrow wrappers into their own file

Author: timsaucer

datafusion/ffi/src/arrow_wrappers.rs

@@ -0,0 +1,53 @@
+use std::sync::Arc;
+
+use abi_stable::StableAbi;
+use arrow::{
+    datatypes::{Schema, SchemaRef},
+    ffi::{FFI_ArrowArray, FFI_ArrowSchema},
+};
+use log::error;
+
+/// This is a wrapper struct around FFI_ArrowSchema simply to indicate
+/// to the StableAbi macros that the underlying struct is FFI safe.
+#[repr(C)]
+#[derive(Debug, StableAbi)]
+pub struct WrappedSchema(#[sabi(unsafe_opaque_field)] pub FFI_ArrowSchema);
+
+impl From<SchemaRef> for WrappedSchema {
+    fn from(value: SchemaRef) -> Self {
+        let ffi_schema = match FFI_ArrowSchema::try_from(value.as_ref()) {
+            Ok(s) => s,
+            Err(e) => {
+                error!("Unable to convert DataFusion Schema to FFI_ArrowSchema in FFI_PlanProperties. {}", e);
+                FFI_ArrowSchema::empty()
+            }
+        };
+
+        WrappedSchema(ffi_schema)
+    }
+}
+
+impl From<WrappedSchema> for SchemaRef {
+    fn from(value: WrappedSchema) -> Self {
+        let schema = match Schema::try_from(&value.0) {
+            Ok(s) => s,
+            Err(e) => {
+                error!("Unable to convert from FFI_ArrowSchema to DataFusion Schema in FFI_PlanProperties. {}", e);
+                Schema::empty()
+            }
+        };
+        Arc::new(schema)
+    }
+}
+
+/// This is a wrapper struct for FFI_ArrowArray to indicate to StableAbi
+/// that the struct is FFI Safe. For convenience, we also include the
+/// schema needed to create a record batch from the array.
+#[repr(C)]
+#[derive(Debug, StableAbi)]
+pub struct WrappedArray {
+    #[sabi(unsafe_opaque_field)]
+    pub array: FFI_ArrowArray,
+
+    pub schema: WrappedSchema,
+}

datafusion/ffi/src/execution_plan.rs

@@ -21,7 +21,6 @@ use abi_stable::{
     std_types::{RResult, RString, RVec},
     StableAbi,
 };
-use arrow::ffi_stream::FFI_ArrowArrayStream;
 use datafusion::error::Result;
 use datafusion::{
     error::DataFusionError,
@@ -33,28 +32,36 @@ use crate::{
     plan_properties::FFI_PlanProperties, record_batch_stream::FFI_RecordBatchStream,
 };
 
+/// A stable struct for sharing a [`ExecutionPlan`] across FFI boundaries.
 #[repr(C)]
 #[derive(Debug, StableAbi)]
-pub struct WrappedArrayStream(#[sabi(unsafe_opaque_field)] pub FFI_ArrowArrayStream);
-
-#[repr(C)]
-#[derive(Debug, StableAbi)]
-#[allow(missing_docs)]
 #[allow(non_camel_case_types)]
 pub struct FFI_ExecutionPlan {
+    /// Return the plan properties
     pub properties: unsafe extern "C" fn(plan: &Self) -> FFI_PlanProperties,
 
+    /// Return a vector of children plans
     pub children: unsafe extern "C" fn(plan: &Self) -> RVec<FFI_ExecutionPlan>,
 
+    /// Return the plan name.
     pub name: unsafe extern "C" fn(plan: &Self) -> RString,
 
+    /// Execute the plan and return a record batch stream. Errors
+    /// will be returned as a string.
     pub execute: unsafe extern "C" fn(
         plan: &Self,
         partition: usize,
     ) -> RResult<FFI_RecordBatchStream, RString>,
 
+    /// 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,
+
+    /// Release the memory of the private data when it is no longer being used.
     pub release: unsafe extern "C" fn(arg: &mut Self),
+
+    /// Internal data. This is only to be accessed by the provider of the plan.
+    /// A [`ForeignExecutionPlan`] should never attempt to access this data.
     pub private_data: *mut c_void,
 }
 
@@ -154,6 +161,9 @@ impl Drop for FFI_ExecutionPlan {
     }
 }
 
+/// This struct is used to access an execution plan provided by a foreign
+/// library across a FFI boundary.
+///
 /// The ForeignExecutionPlan is to be used by the caller of the plan, so it has
 /// no knowledge or access to the private data. All interaction with the plan
 /// must occur through the functions defined in FFI_ExecutionPlan.

datafusion/ffi/src/lib.rs

@@ -17,6 +17,7 @@
 // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143
 #![deny(clippy::clone_on_ref_ptr)]
 
+pub mod arrow_wrappers;
 pub mod execution_plan;
 pub mod plan_properties;
 pub mod record_batch_stream;

datafusion/ffi/src/plan_properties.rs

@@ -24,10 +24,7 @@ use abi_stable::{
     },
     StableAbi,
 };
-use arrow::{
-    datatypes::{Schema, SchemaRef},
-    ffi::FFI_ArrowSchema,
-};
+use arrow::datatypes::SchemaRef;
 use datafusion::{
     error::{DataFusionError, Result},
     physical_expr::EquivalenceProperties,
@@ -42,9 +39,10 @@ use datafusion_proto::{
     },
     protobuf::{Partitioning, PhysicalSortExprNodeCollection},
 };
-use log::error;
 use prost::Message;
 
+use crate::arrow_wrappers::WrappedSchema;
+
 // TODO: should we just make ExecutionMode repr(C)?
 #[repr(C)]
 #[allow(non_camel_case_types)]
@@ -75,37 +73,6 @@ impl From<FFI_ExecutionMode> for ExecutionMode {
     }
 }
 
-#[repr(C)]
-#[derive(Debug, StableAbi)]
-pub struct WrappedSchema(#[sabi(unsafe_opaque_field)] pub FFI_ArrowSchema);
-
-impl From<SchemaRef> for WrappedSchema {
-    fn from(value: SchemaRef) -> Self {
-        let ffi_schema = match FFI_ArrowSchema::try_from(value.as_ref()) {
-            Ok(s) => s,
-            Err(e) => {
-                error!("Unable to convert DataFusion Schema to FFI_ArrowSchema in FFI_PlanProperties. {}", e);
-                FFI_ArrowSchema::empty()
-            }
-        };
-
-        WrappedSchema(ffi_schema)
-    }
-}
-
-impl From<WrappedSchema> for SchemaRef {
-    fn from(value: WrappedSchema) -> Self {
-        let schema = match Schema::try_from(&value.0) {
-            Ok(s) => s,
-            Err(e) => {
-                error!("Unable to convert from FFI_ArrowSchema to DataFusion Schema in FFI_PlanProperties. {}", e);
-                Schema::empty()
-            }
-        };
-        Arc::new(schema)
-    }
-}
-
 #[repr(C)]
 #[derive(Debug, StableAbi)]
 #[allow(missing_docs)]

datafusion/ffi/src/record_batch_stream.rs

@@ -24,7 +24,7 @@ use abi_stable::{
 use arrow::array::{Array, RecordBatch};
 use arrow::{
     array::{make_array, StructArray},
-    ffi::{from_ffi, to_ffi, FFI_ArrowArray},
+    ffi::{from_ffi, to_ffi},
 };
 use async_ffi::{ContextExt, FfiContext, FfiPoll};
 use datafusion::error::Result;
@@ -34,16 +34,7 @@ use datafusion::{
 };
 use futures::{Stream, TryStreamExt};
 
-use crate::plan_properties::WrappedSchema;
-
-#[repr(C)]
-#[derive(Debug, StableAbi)]
-pub struct WrappedArray {
-    #[sabi(unsafe_opaque_field)]
-    array: FFI_ArrowArray,
-
-    schema: WrappedSchema,
-}
+use crate::arrow_wrappers::{WrappedArray, WrappedSchema};
 
 #[repr(C)]
 #[derive(Debug, StableAbi)]

datafusion/ffi/src/table_provider.rs

@@ -42,7 +42,7 @@ use datafusion_proto::{
 use prost::Message;
 
 use crate::{
-    plan_properties::WrappedSchema,
+    arrow_wrappers::WrappedSchema,
     session_config::ForeignSessionConfig,
     table_source::{FFI_TableProviderFilterPushDown, FFI_TableType},
 };