完善udf插件化代码

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -1080,11 +1080,7 @@ impl TryInto<Expr> for &protobuf::LogicalExprNode {
             }
             // argo engine add start
             ExprType::AggregateUdfExpr(expr) => {
-                let gpm = global_plugin_manager("").lock().unwrap();
-                let plugin_registrar = gpm.plugin_managers.get(&PluginEnum::UDF).unwrap();
-                if let Some(udf_plugin_manager) =
-                    plugin_registrar.as_any().downcast_ref::<UDFPluginManager>()
-                {
+                if let Some(udf_plugin_manager) = get_udf_plugin_manager("") {
                     let fun = udf_plugin_manager
                         .aggregate_udfs
                         .get(expr.fun_name.as_str())
@@ -1106,11 +1102,7 @@ impl TryInto<Expr> for &protobuf::LogicalExprNode {
                 }
             }
             ExprType::ScalarUdfProtoExpr(expr) => {
-                let gpm = global_plugin_manager("").lock().unwrap();
-                let plugin_registrar = gpm.plugin_managers.get(&PluginEnum::UDF).unwrap();
-                if let Some(udf_plugin_manager) =
-                    plugin_registrar.as_any().downcast_ref::<UDFPluginManager>()
-                {
+                if let Some(udf_plugin_manager) = get_udf_plugin_manager("") {
                     let fun = udf_plugin_manager
                         .scalar_udfs
                         .get(expr.fun_name.as_str())
@@ -1335,7 +1327,7 @@ impl TryInto<Field> for &protobuf::Field {
 use crate::serde::protobuf::ColumnStats;
 use datafusion::physical_plan::{aggregates, windows};
 use datafusion::plugin::plugin_manager::global_plugin_manager;
-use datafusion::plugin::udf::UDFPluginManager;
+use datafusion::plugin::udf::{get_udf_plugin_manager, UDFPluginManager};
 use datafusion::plugin::PluginEnum;
 use datafusion::prelude::{
     array, date_part, date_trunc, length, lower, ltrim, md5, rtrim, sha224, sha256,

ballista/rust/core/src/serde/physical_plan/from_proto.rs

@@ -83,7 +83,7 @@ use datafusion::physical_plan::{
     AggregateExpr, ColumnStatistics, ExecutionPlan, PhysicalExpr, Statistics, WindowExpr,
 };
 use datafusion::plugin::plugin_manager::global_plugin_manager;
-use datafusion::plugin::udf::UDFPluginManager;
+use datafusion::plugin::udf::{get_udf_plugin_manager, UDFPluginManager};
 use datafusion::plugin::PluginEnum;
 use datafusion::prelude::CsvReadOptions;
 use log::debug;
@@ -320,10 +320,7 @@ impl TryInto<Arc<dyn ExecutionPlan>> for &protobuf::PhysicalPlanNode {
                             ExprType::AggregateUdfExpr(agg_node) => {
                                 let name = agg_node.fun_name.as_str();
                                 let udaf_fun_name = &name[0..name.find('(').unwrap()];
-                                let gpm = global_plugin_manager("").lock().unwrap();
-                                let plugin_registrar = gpm.plugin_managers.get(&PluginEnum::UDF).unwrap();
-                                if let Some(udf_plugin_manager) = plugin_registrar.as_any().downcast_ref::<UDFPluginManager>()
-                                {
+                                if let Some(udf_plugin_manager) = get_udf_plugin_manager("") {
                                     let fun = udf_plugin_manager.aggregate_udfs.get(udaf_fun_name).ok_or_else(|| {
                                         proto_error(format!(
                                             "can not get udaf:{} from plugins!",
@@ -584,11 +581,7 @@ impl TryFrom<&protobuf::PhysicalExprNode> for Arc<dyn PhysicalExpr> {
             }
             // argo engine add.
             ExprType::ScalarUdfProtoExpr(e) => {
-                let gpm = global_plugin_manager("").lock().unwrap();
-                let plugin_registrar = gpm.plugin_managers.get(&PluginEnum::UDF).unwrap();
-                if let Some(udf_plugin_manager) =
-                    plugin_registrar.as_any().downcast_ref::<UDFPluginManager>()
-                {
+                if let Some(udf_plugin_manager) = get_udf_plugin_manager("") {
                     let fun = udf_plugin_manager
                         .scalar_udfs
                         .get(&e.fun_name)

datafusion/src/execution/context.rs

@@ -83,9 +83,7 @@ use crate::physical_plan::planner::DefaultPhysicalPlanner;
 use crate::physical_plan::udf::ScalarUDF;
 use crate::physical_plan::ExecutionPlan;
 use crate::physical_plan::PhysicalPlanner;
-use crate::plugin::plugin_manager::global_plugin_manager;
-use crate::plugin::udf::UDFPluginManager;
-use crate::plugin::PluginEnum;
+use crate::plugin::udf::get_udf_plugin_manager;
 use crate::sql::{
     parser::{DFParser, FileType},
     planner::{ContextProvider, SqlToRel},
@@ -198,13 +196,9 @@ impl ExecutionContext {
             })),
         };
 
-        let gpm = global_plugin_manager(config.plugin_dir.as_str());
-
         // register udf
-        let gpm_guard = gpm.lock().unwrap();
-        let plugin_registrar = gpm_guard.plugin_managers.get(&PluginEnum::UDF).unwrap();
         if let Some(udf_plugin_manager) =
-            plugin_registrar.as_any().downcast_ref::<UDFPluginManager>()
+            get_udf_plugin_manager(config.plugin_dir.as_str())
         {
             udf_plugin_manager
                 .scalar_udfs

datafusion/src/plugin/mod.rs

@@ -1,7 +1,26 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
 use crate::error::Result;
 use crate::plugin::udf::UDFPluginManager;
+use libloading::Library;
 use std::any::Any;
 use std::env;
+use std::sync::Arc;
 
 /// plugin manager
 pub mod plugin_manager;
@@ -47,17 +66,15 @@ pub struct PluginDeclaration {
 
     /// One of PluginEnum
     pub plugin_type: unsafe extern "C" fn() -> PluginEnum,
-
-    /// `register` is a function which impl PluginRegistrar. It will be call when plugin load.
-    pub register: unsafe extern "C" fn(&mut Box<dyn PluginRegistrar>),
 }
 
 /// Plugin Registrar , Every plugin need implement this trait
 pub trait PluginRegistrar: Send + Sync + 'static {
-    /// The implementer of the plug-in needs to call this interface to report his own information to the plug-in manager
-    fn register_plugin(&mut self, plugin: Box<dyn Plugin>) -> Result<()>;
+    /// # Safety
+    /// load plugin from library
+    unsafe fn load(&mut self, library: Arc<Library>) -> Result<()>;
 
-    /// Returns the plugin registrar as [`Any`](std::any::Any) so that it can be
+    /// Returns the plugin as [`Any`](std::any::Any) so that it can be
     /// downcast to a specific implementation.
     fn as_any(&self) -> &dyn Any;
 }
@@ -66,22 +83,12 @@ pub trait PluginRegistrar: Send + Sync + 'static {
 ///
 /// # Notes
 ///
-/// This works by automatically generating an `extern "C"` function with a
+/// This works by automatically generating an `extern "C"` function named `get_plugin_type` with a
 /// pre-defined signature and symbol name. And then generating a PluginDeclaration.
 /// Therefore you will only be able to declare one plugin per library.
 #[macro_export]
 macro_rules! declare_plugin {
-    ($plugin_type:expr, $curr_plugin_type:ty, $constructor:path) => {
-        #[no_mangle]
-        pub extern "C" fn register_plugin(
-            registrar: &mut Box<dyn $crate::plugin::PluginRegistrar>,
-        ) {
-            // make sure the constructor is the correct type.
-            let constructor: fn() -> $curr_plugin_type = $constructor;
-            let object = constructor();
-            registrar.register_plugin(Box::new(object)).unwrap();
-        }
-
+    ($plugin_type:expr) => {
         #[no_mangle]
         pub extern "C" fn get_plugin_type() -> $crate::plugin::PluginEnum {
             $plugin_type
@@ -93,7 +100,6 @@ macro_rules! declare_plugin {
                 rustc_version: $crate::plugin::RUSTC_VERSION,
                 core_version: $crate::plugin::CORE_VERSION,
                 plugin_type: get_plugin_type,
-                register: register_plugin,
             };
     };
 }

datafusion/src/plugin/plugin_manager.rs

@@ -1,3 +1,19 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
 use crate::error::{DataFusionError, Result};
 use crate::plugin::{PluginDeclaration, CORE_VERSION, RUSTC_VERSION};
 use crate::plugin::{PluginEnum, PluginRegistrar};
@@ -13,11 +29,14 @@ use once_cell::sync::OnceCell;
 /// To prevent the library from being loaded multiple times, we use once_cell defines a Arc<Mutex<GlobalPluginManager>>
 /// Because datafusion is a library, not a service, users may not need to load all plug-ins in the process.
 /// So fn global_plugin_manager return Arc<Mutex<GlobalPluginManager>>. In this way, users can load the required library through the load method of GlobalPluginManager when needed
+static INSTANCE: OnceCell<Arc<Mutex<GlobalPluginManager>>> = OnceCell::new();
+
+/// global_plugin_manager
 pub fn global_plugin_manager(
     plugin_path: &str,
 ) -> &'static Arc<Mutex<GlobalPluginManager>> {
-    static INSTANCE: OnceCell<Arc<Mutex<GlobalPluginManager>>> = OnceCell::new();
     INSTANCE.get_or_init(move || unsafe {
+        println!("====================init===================");
         let mut gpm = GlobalPluginManager::default();
         gpm.load(plugin_path).unwrap();
         Arc::new(Mutex::new(gpm))
@@ -38,6 +57,9 @@ impl GlobalPluginManager {
     /// # Safety
     /// find plugin file from `plugin_path` and load it .
     unsafe fn load(&mut self, plugin_path: &str) -> Result<()> {
+        if "".eq(plugin_path) {
+            return Ok(());
+        }
         // find library file from udaf_plugin_path
         info!("load plugin from dir:{}", plugin_path);
         println!("load plugin from dir:{}", plugin_path);
@@ -54,18 +76,19 @@ impl GlobalPluginManager {
 
             let library = Arc::new(library);
 
-            // get a pointer to the plugin_declaration symbol.
-            let dec = library
-                .get::<*mut PluginDeclaration>(b"plugin_declaration\0")
-                .map_err(|e| {
-                    DataFusionError::IoError(io::Error::new(
-                        io::ErrorKind::Other,
-                        format!("not found plugin_declaration in the library: {}", e),
-                    ))
-                })?
-                .read();
-
-            // version checks to prevent accidental ABI incompatibilities
+            let dec = library.get::<*mut PluginDeclaration>(b"plugin_declaration\0");
+            if dec.is_err() {
+                info!(
+                    "not found plugin_declaration in the library: {}",
+                    plugin_file.path().to_str().unwrap()
+                );
+                return Ok(());
+            }
+
+            let dec = dec.unwrap().read();
+
+            // ersion checks to prevent accidental ABI incompatibilities
+
             if dec.rustc_version != RUSTC_VERSION || dec.core_version != CORE_VERSION {
                 return Err(DataFusionError::IoError(io::Error::new(
                     io::ErrorKind::Other,
@@ -82,8 +105,7 @@ impl GlobalPluginManager {
                 }
                 Some(manager) => manager,
             };
-
-            (dec.register)(curr_plugin_manager);
+            curr_plugin_manager.load(library)?;
             self.plugin_files
                 .push(plugin_file.path().to_str().unwrap().to_string());
         }
@@ -112,17 +134,20 @@ impl GlobalPluginManager {
             if let Some(path) = item.path().extension() {
                 if let Some(suffix) = path.to_str() {
                     if suffix == "dylib" || suffix == "so" || suffix == "dll" {
-                        info!("load plugin from library file:{}", path.to_str().unwrap());
+                        info!(
+                            "load plugin from library file:{}",
+                            item.path().to_str().unwrap()
+                        );
                         println!(
                             "load plugin from library file:{}",
-                            path.to_str().unwrap()
+                            item.path().to_str().unwrap()
                         );
                         return Some(item);
                     }
                 }
             }
 
-            return None;
+            None
         }) {
             plugin_files.push(entry);
         }