ffi_physical_optimizer_stabby_upgrade

Author: coderfender

datafusion/ffi/src/physical_optimizer.rs

@@ -18,31 +18,30 @@
 use std::ffi::c_void;
 use std::sync::Arc;
 
-use abi_stable::StableAbi;
-use abi_stable::std_types::{RResult, RStr};
 use async_trait::async_trait;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::ExecutionPlan;
+use stabby::string::String as SString;
 use tokio::runtime::Handle;
 
 use crate::config::FFI_ConfigOptions;
 use crate::execution_plan::FFI_ExecutionPlan;
-use crate::util::FFIResult;
-use crate::{df_result, rresult_return};
+use crate::util::{FFI_Result, FFIResult};
+use crate::{df_result, sresult_return};
 
 /// A stable struct for sharing [`PhysicalOptimizerRule`] across FFI boundaries.
 #[repr(C)]
-#[derive(Debug, StableAbi)]
+#[derive(Debug)]
 pub struct FFI_PhysicalOptimizerRule {
     pub optimize: unsafe extern "C" fn(
         &Self,
         plan: &FFI_ExecutionPlan,
         config: FFI_ConfigOptions,
     ) -> FFIResult<FFI_ExecutionPlan>,
 
-    pub name: unsafe extern "C" fn(&Self) -> RStr,
+    pub name: unsafe extern "C" fn(&Self) -> SString,
 
     pub schema_check: unsafe extern "C" fn(&Self) -> bool,
 
@@ -92,14 +91,14 @@ unsafe extern "C" fn optimize_fn_wrapper(
 ) -> FFIResult<FFI_ExecutionPlan> {
     let runtime = rule.runtime();
     let rule = rule.inner();
-    let plan: Arc<dyn ExecutionPlan> = rresult_return!(plan.try_into());
-    let config = rresult_return!(ConfigOptions::try_from(config));
-    let optimized_plan = rresult_return!(rule.optimize(plan, &config));
+    let plan: Arc<dyn ExecutionPlan> = sresult_return!(plan.try_into());
+    let config = sresult_return!(ConfigOptions::try_from(config));
+    let optimized_plan = sresult_return!(rule.optimize(plan, &config));
 
-    RResult::ROk(FFI_ExecutionPlan::new(optimized_plan, runtime))
+    FFI_Result::Ok(FFI_ExecutionPlan::new(optimized_plan, runtime))
 }
 
-unsafe extern "C" fn name_fn_wrapper(rule: &FFI_PhysicalOptimizerRule) -> RStr<'_> {
+unsafe extern "C" fn name_fn_wrapper(rule: &FFI_PhysicalOptimizerRule) -> SString {
     let rule = rule.inner();
     rule.name().into()
 }
@@ -153,7 +152,7 @@ impl FFI_PhysicalOptimizerRule {
         if let Some(rule) = (Arc::clone(&rule) as Arc<dyn std::any::Any>)
             .downcast_ref::<ForeignPhysicalOptimizerRule>()
         {
-            return rule.0.clone();
+            return rule.rule.clone();
         }
 
         let private_data = Box::new(RulePrivateData { rule, runtime });
@@ -177,7 +176,10 @@ impl FFI_PhysicalOptimizerRule {
 /// defined on this struct must only use the stable functions provided in
 /// FFI_PhysicalOptimizerRule to interact with the foreign rule.
 #[derive(Debug)]
-pub struct ForeignPhysicalOptimizerRule(pub FFI_PhysicalOptimizerRule);
+pub struct ForeignPhysicalOptimizerRule {
+    name: String,
+    rule: FFI_PhysicalOptimizerRule,
+}
 
 unsafe impl Send for ForeignPhysicalOptimizerRule {}
 unsafe impl Sync for ForeignPhysicalOptimizerRule {}
@@ -188,8 +190,11 @@ impl From<&FFI_PhysicalOptimizerRule> for Arc<dyn PhysicalOptimizerRule + Send +
             return Arc::clone(rule.inner());
         }
 
-        Arc::new(ForeignPhysicalOptimizerRule(rule.clone()))
-            as Arc<dyn PhysicalOptimizerRule + Send + Sync>
+        let name: String = unsafe { (rule.name)(rule).into() };
+        Arc::new(ForeignPhysicalOptimizerRule {
+            name,
+            rule: rule.clone(),
+        }) as Arc<dyn PhysicalOptimizerRule + Send + Sync>
     }
 }
 
@@ -210,16 +215,16 @@ impl PhysicalOptimizerRule for ForeignPhysicalOptimizerRule {
         let plan = FFI_ExecutionPlan::new(plan, None);
 
         let optimized_plan =
-            unsafe { df_result!((self.0.optimize)(&self.0, &plan, config_options))? };
+            unsafe { df_result!((self.rule.optimize)(&self.rule, &plan, config_options))? };
         (&optimized_plan).try_into()
     }
 
     fn name(&self) -> &str {
-        unsafe { (self.0.name)(&self.0).as_str() }
+        &self.name
     }
 
     fn schema_check(&self) -> bool {
-        unsafe { (self.0.schema_check)(&self.0) }
+        unsafe { (self.rule.schema_check)(&self.rule) }
     }
 }
 
@@ -341,9 +346,9 @@ mod tests {
         let ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None);
         let cloned = ffi_rule.clone();
 
-        assert_eq!(unsafe { (ffi_rule.name)(&ffi_rule).as_str() }, unsafe {
-            (cloned.name)(&cloned).as_str()
-        });
+        let name1: String = unsafe { (ffi_rule.name)(&ffi_rule).into() };
+        let name2: String = unsafe { (cloned.name)(&cloned).into() };
+        assert_eq!(name1, name2);
 
         Ok(())
     }
@@ -363,10 +368,8 @@ mod tests {
 
         // Now wrap the foreign rule back into FFI - should not double-wrap
         let re_wrapped = FFI_PhysicalOptimizerRule::new(foreign_rule, None);
-        assert_eq!(
-            unsafe { (re_wrapped.name)(&re_wrapped).as_str() },
-            "no_op_rule"
-        );
+        let name: String = unsafe { (re_wrapped.name)(&re_wrapped).into() };
+        assert_eq!(name, "no_op_rule");
 
         Ok(())
     }