tmp3

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

s3select/query/src/function/mod.rs

@@ -0,0 +1 @@
+pub mod simple_func_manager;

s3select/query/src/function/simple_func_manager.rs

@@ -0,0 +1,63 @@
+use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
+
+use api::query::function::FunctionMetadataManager;
+use api::{QueryError, QueryResult};
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+
+pub type SimpleFunctionMetadataManagerRef = Arc<SimpleFunctionMetadataManager>;
+
+#[derive(Debug, Default)]
+pub struct SimpleFunctionMetadataManager {
+    /// Scalar functions that are registered with the context
+    pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+    /// Aggregate functions registered in the context
+    pub aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    /// Window functions registered in the context
+    pub window_functions: HashMap<String, Arc<WindowUDF>>,
+}
+
+impl FunctionMetadataManager for SimpleFunctionMetadataManager {
+    fn register_udf(&mut self, f: ScalarUDF) -> QueryResult<()> {
+        self.scalar_functions.insert(f.inner().name().to_uppercase(), Arc::new(f));
+        Ok(())
+    }
+
+    fn register_udaf(&mut self, f: AggregateUDF) -> QueryResult<()> {
+        self.aggregate_functions.insert(f.inner().name().to_uppercase(), Arc::new(f));
+        Ok(())
+    }
+
+    fn register_udwf(&mut self, f: WindowUDF) -> QueryResult<()> {
+        self.window_functions.insert(f.inner().name().to_uppercase(), Arc::new(f));
+        Ok(())
+    }
+
+    fn udf(&self, name: &str) -> QueryResult<Arc<ScalarUDF>> {
+        let result = self.scalar_functions.get(&name.to_uppercase());
+
+        result
+            .cloned()
+            .ok_or_else(|| QueryError::FunctionExists { name: name.to_string() })
+    }
+
+    fn udaf(&self, name: &str) -> QueryResult<Arc<AggregateUDF>> {
+        let result = self.aggregate_functions.get(&name.to_uppercase());
+
+        result
+            .cloned()
+            .ok_or_else(|| QueryError::FunctionNotExists { name: name.to_string() })
+    }
+
+    fn udwf(&self, name: &str) -> QueryResult<Arc<WindowUDF>> {
+        let result = self.window_functions.get(&name.to_uppercase());
+
+        result
+            .cloned()
+            .ok_or_else(|| QueryError::FunctionNotExists { name: name.to_string() })
+    }
+
+    fn udfs(&self) -> HashSet<String> {
+        self.scalar_functions.keys().cloned().collect()
+    }
+}