support func

Signed-off-by: junxiang Mu <1948535941@qq.com>

s3select/query/src/function/simple_func_manager.rs

@@ -1,13 +1,15 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use api::query::function::FunctionMetadataManager;
 use api::{QueryError, QueryResult};
+use datafusion::execution::SessionStateDefaults;
 use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+use tracing::debug;
 
 pub type SimpleFunctionMetadataManagerRef = Arc<SimpleFunctionMetadataManager>;
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct SimpleFunctionMetadataManager {
     /// Scalar functions that are registered with the context
     pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
@@ -17,19 +19,53 @@ pub struct SimpleFunctionMetadataManager {
     pub window_functions: HashMap<String, Arc<WindowUDF>>,
 }
 
+impl Default for SimpleFunctionMetadataManager {
+    fn default() -> Self {
+        let mut func_meta_manager = Self {
+            scalar_functions: Default::default(),
+            aggregate_functions: Default::default(),
+            window_functions: Default::default(),
+        };
+        SessionStateDefaults::default_scalar_functions().into_iter().for_each(|udf| {
+            let existing_udf = func_meta_manager.register_udf(udf.clone());
+            if let Ok(()) = existing_udf {
+                debug!("Overwrote an existing UDF: {}", udf.name());
+            }
+        });
+
+        SessionStateDefaults::default_aggregate_functions()
+            .into_iter()
+            .for_each(|udaf| {
+                let existing_udaf = func_meta_manager.register_udaf(udaf.clone());
+                if let Ok(()) = existing_udaf {
+                    debug!("Overwrote an existing UDAF: {}", udaf.name());
+                }
+            });
+
+        SessionStateDefaults::default_window_functions().into_iter().for_each(|udwf| {
+            let existing_udwf = func_meta_manager.register_udwf(udwf.clone());
+            if let Ok(()) = existing_udwf {
+                debug!("Overwrote an existing UDWF: {}", udwf.name());
+            }
+        });
+
+        func_meta_manager
+    }
+}
+
 impl FunctionMetadataManager for SimpleFunctionMetadataManager {
-    fn register_udf(&mut self, f: ScalarUDF) -> QueryResult<()> {
-        self.scalar_functions.insert(f.inner().name().to_uppercase(), Arc::new(f));
+    fn register_udf(&mut self, f: Arc<ScalarUDF>) -> QueryResult<()> {
+        self.scalar_functions.insert(f.inner().name().to_uppercase(), f);
         Ok(())
     }
 
-    fn register_udaf(&mut self, f: AggregateUDF) -> QueryResult<()> {
-        self.aggregate_functions.insert(f.inner().name().to_uppercase(), Arc::new(f));
+    fn register_udaf(&mut self, f: Arc<AggregateUDF>) -> QueryResult<()> {
+        self.aggregate_functions.insert(f.inner().name().to_uppercase(), f);
         Ok(())
     }
 
-    fn register_udwf(&mut self, f: WindowUDF) -> QueryResult<()> {
-        self.window_functions.insert(f.inner().name().to_uppercase(), Arc::new(f));
+    fn register_udwf(&mut self, f: Arc<WindowUDF>) -> QueryResult<()> {
+        self.window_functions.insert(f.inner().name().to_uppercase(), f);
         Ok(())
     }
 
@@ -57,7 +93,13 @@ impl FunctionMetadataManager for SimpleFunctionMetadataManager {
             .ok_or_else(|| QueryError::FunctionNotExists { name: name.to_string() })
     }
 
-    fn udfs(&self) -> HashSet<String> {
+    fn udfs(&self) -> Vec<String> {
         self.scalar_functions.keys().cloned().collect()
     }
+    fn udafs(&self) -> Vec<String> {
+        self.aggregate_functions.keys().cloned().collect()
+    }
+    fn udwfs(&self) -> Vec<String> {
+        self.window_functions.keys().cloned().collect()
+    }
 }