[Ballista] Add ballista plugin manager and UDF plugin

ballista/rust/core/Cargo.toml

@@ -42,7 +42,10 @@ datafusion = { path = "../../../datafusion/core", version = "7.0.0" }
 datafusion-proto = { path = "../../../datafusion/proto", version = "7.0.0" }
 futures = "0.3"
 hashbrown = "0.12"
+
+libloading = "0.7.3"
 log = "0.4"
+once_cell = "1.9.0"
 
 parking_lot = "0.12"
 parse_arg = "0.1.3"
@@ -53,9 +56,11 @@ sqlparser = "0.15"
 tokio = "1.0"
 tonic = "0.6"
 uuid = { version = "0.8", features = ["v4"] }
+walkdir = "2.3.2"
 
 [dev-dependencies]
 tempfile = "3"
 
 [build-dependencies]
+rustc_version = "0.4.0"
 tonic-build = { version = "0.6" }

ballista/rust/core/build.rs

@@ -20,6 +20,8 @@ fn main() -> Result<(), String> {
     println!("cargo:rerun-if-env-changed=FORCE_REBUILD");
 
     println!("cargo:rerun-if-changed=proto/ballista.proto");
+    let version = rustc_version::version().unwrap();
+    println!("cargo:rustc-env=RUSTC_VERSION={}", version);
     println!("cargo:rerun-if-changed=proto/datafusion.proto");
     tonic_build::configure()
         .extern_path(".datafusion", "::datafusion_proto::protobuf")

ballista/rust/core/src/config.rs

@@ -34,6 +34,8 @@ pub const BALLISTA_REPARTITION_AGGREGATIONS: &str = "ballista.repartition.aggreg
 pub const BALLISTA_REPARTITION_WINDOWS: &str = "ballista.repartition.windows";
 pub const BALLISTA_PARQUET_PRUNING: &str = "ballista.parquet.pruning";
 pub const BALLISTA_WITH_INFORMATION_SCHEMA: &str = "ballista.with_information_schema";
+/// give a plugin files dir, and then the dynamic library files in this dir will be load when scheduler state init.
+pub const BALLISTA_PLUGIN_DIR: &str = "ballista.plugin_dir";
 
 pub type ParseResult<T> = result::Result<T, String>;
 
@@ -139,6 +141,9 @@ impl BallistaConfig {
                     .parse::<bool>()
                     .map_err(|e| format!("{:?}", e))?;
             }
+            DataType::Utf8 => {
+                val.to_string();
+            }
             _ => {
                 return Err(format!("not support data type: {}", data_type));
             }
@@ -171,6 +176,9 @@ impl BallistaConfig {
             ConfigEntry::new(BALLISTA_WITH_INFORMATION_SCHEMA.to_string(),
                 "Sets whether enable information_schema".to_string(),
                 DataType::Boolean,Some("false".to_string())),
+            ConfigEntry::new(BALLISTA_PLUGIN_DIR.to_string(),
+                             "Sets the plugin dir".to_string(),
+                             DataType::Utf8,Some("".to_string())),
         ];
         entries
             .iter()
@@ -186,6 +194,10 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_DEFAULT_SHUFFLE_PARTITIONS)
     }
 
+    pub fn default_plugin_dir(&self) -> String {
+        self.get_string_setting(BALLISTA_PLUGIN_DIR)
+    }
+
     pub fn default_batch_size(&self) -> usize {
         self.get_usize_setting(BALLISTA_DEFAULT_BATCH_SIZE)
     }
@@ -233,6 +245,17 @@ impl BallistaConfig {
             v.parse::<bool>().unwrap()
         }
     }
+    fn get_string_setting(&self, key: &str) -> String {
+        if let Some(v) = self.settings.get(key) {
+            // infallible because we validate all configs in the constructor
+            v.to_string()
+        } else {
+            let entries = Self::valid_entries();
+            // infallible because we validate all configs in the constructor
+            let v = entries.get(key).unwrap().default_value.as_ref().unwrap();
+            v.to_string()
+        }
+    }
 }
 
 // an enum used to configure the scheduler policy
@@ -266,6 +289,7 @@ mod tests {
         let config = BallistaConfig::new()?;
         assert_eq!(2, config.default_shuffle_partitions());
         assert!(!config.default_with_information_schema());
+        assert_eq!("", config.default_plugin_dir().as_str());
         Ok(())
     }
 
@@ -284,6 +308,7 @@ mod tests {
     fn custom_config_invalid() -> Result<()> {
         let config = BallistaConfig::builder()
             .set(BALLISTA_DEFAULT_SHUFFLE_PARTITIONS, "true")
+            .set(BALLISTA_PLUGIN_DIR, "test_dir")
             .build();
         assert!(config.is_err());
         assert_eq!("General(\"Failed to parse user-supplied value 'ballista.shuffle.partitions' for configuration setting 'true': ParseIntError { kind: InvalidDigit }\")", format!("{:?}", config.unwrap_err()));
@@ -293,7 +318,6 @@ mod tests {
             .build();
         assert!(config.is_err());
         assert_eq!("General(\"Failed to parse user-supplied value 'ballista.with_information_schema' for configuration setting '123': ParseBoolError\")", format!("{:?}", config.unwrap_err()));
-
         Ok(())
     }
 }

ballista/rust/core/src/lib.rs

@@ -27,6 +27,8 @@ pub mod config;
 pub mod error;
 pub mod event_loop;
 pub mod execution_plans;
+/// some plugins
+pub mod plugin;
 pub mod utils;
 
 #[macro_use]

ballista/rust/core/src/plugin/mod.rs

@@ -0,0 +1,127 @@
+// 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;
+/// udf plugin
+pub mod udf;
+
+/// CARGO_PKG_VERSION
+pub static CORE_VERSION: &str = env!("CARGO_PKG_VERSION");
+/// RUSTC_VERSION
+pub static RUSTC_VERSION: &str = env!("RUSTC_VERSION");
+
+/// Top plugin trait
+pub trait Plugin {
+    /// Returns the plugin as [`Any`](std::any::Any) so that it can be
+    /// downcast to a specific implementation.
+    fn as_any(&self) -> &dyn Any;
+}
+
+/// The enum of Plugin
+#[derive(PartialEq, std::cmp::Eq, std::hash::Hash, Copy, Clone)]
+pub enum PluginEnum {
+    /// UDF/UDAF plugin
+    UDF,
+}
+
+impl PluginEnum {
+    /// new a struct which impl the PluginRegistrar trait
+    pub fn init_plugin_manager(&self) -> Box<dyn PluginRegistrar> {
+        match self {
+            PluginEnum::UDF => Box::new(UDFPluginManager::default()),
+        }
+    }
+}
+
+/// Every plugin need a PluginDeclaration
+#[derive(Copy, Clone)]
+pub struct PluginDeclaration {
+    /// Rust doesn’t have a stable ABI, meaning different compiler versions can generate incompatible code.
+    /// For these reasons, the UDF plug-in must be compiled using the same version of rustc as datafusion.
+    pub rustc_version: &'static str,
+
+    /// core version of the plugin. The plugin's core_version need same as plugin manager.
+    pub core_version: &'static str,
+
+    /// One of PluginEnum
+    pub plugin_type: unsafe extern "C" fn() -> PluginEnum,
+}
+
+/// Plugin Registrar , Every plugin need implement this trait
+pub trait PluginRegistrar: Send + Sync + 'static {
+    /// # Safety
+    /// load plugin from library
+    unsafe fn load(&mut self, library: Arc<Library>) -> Result<()>;
+
+    /// Returns the plugin as [`Any`](std::any::Any) so that it can be
+    /// downcast to a specific implementation.
+    fn as_any(&self) -> &dyn Any;
+}
+
+/// Declare a plugin's PluginDeclaration.
+///
+/// # Notes
+///
+/// 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) => {
+        #[no_mangle]
+        pub extern "C" fn get_plugin_type() -> $crate::plugin::PluginEnum {
+            $plugin_type
+        }
+
+        #[no_mangle]
+        pub static plugin_declaration: $crate::plugin::PluginDeclaration =
+            $crate::plugin::PluginDeclaration {
+                rustc_version: $crate::plugin::RUSTC_VERSION,
+                core_version: $crate::plugin::CORE_VERSION,
+                plugin_type: get_plugin_type,
+            };
+    };
+}
+
+/// get the plugin dir
+pub fn plugin_dir() -> String {
+    let current_exe_dir = match env::current_exe() {
+        Ok(exe_path) => exe_path.display().to_string(),
+        Err(_e) => "".to_string(),
+    };
+
+    // If current_exe_dir contain `deps` the root dir is the parent dir
+    // eg: /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/deps/plugins_app-067452b3ff2af70e
+    // the plugin dir is /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/deps
+    // else eg: /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/plugins_app
+    // the plugin dir is /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/
+    if current_exe_dir.contains("/deps/") {
+        let i = current_exe_dir.find("/deps/").unwrap();
+        String::from(&current_exe_dir.as_str()[..i + 6])
+    } else {
+        let i = current_exe_dir.rfind('/').unwrap();
+        String::from(&current_exe_dir.as_str()[..i])
+    }
+}