mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-05-25 01:10:37 +00:00
feat: add udtf (table function) registration (#6922)
* feat: work-in-progress udtf support Signed-off-by: Ning Sun <sunning@greptime.com> * feat: add table function support Signed-off-by: Ning Sun <sunning@greptime.com> * test: resolve test error Signed-off-by: Ning Sun <sunning@greptime.com> --------- Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
@@ -16,6 +16,7 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, LazyLock, RwLock};
|
||||
|
||||
use datafusion::catalog::TableFunction;
|
||||
use datafusion_expr::AggregateUDF;
|
||||
|
||||
use crate::admin::AdminFunction;
|
||||
@@ -42,6 +43,7 @@ use crate::system::SystemFunction;
|
||||
pub struct FunctionRegistry {
|
||||
functions: RwLock<HashMap<String, ScalarFunctionFactory>>,
|
||||
aggregate_functions: RwLock<HashMap<String, AggregateUDF>>,
|
||||
table_functions: RwLock<HashMap<String, Arc<TableFunction>>>,
|
||||
}
|
||||
|
||||
impl FunctionRegistry {
|
||||
@@ -87,6 +89,15 @@ impl FunctionRegistry {
|
||||
.insert(func.name().to_string(), func);
|
||||
}
|
||||
|
||||
/// Register a table function
|
||||
pub fn register_table_function(&self, func: TableFunction) {
|
||||
let _ = self
|
||||
.table_functions
|
||||
.write()
|
||||
.unwrap()
|
||||
.insert(func.name().to_string(), Arc::new(func));
|
||||
}
|
||||
|
||||
pub fn get_function(&self, name: &str) -> Option<ScalarFunctionFactory> {
|
||||
self.functions.read().unwrap().get(name).cloned()
|
||||
}
|
||||
@@ -106,6 +117,15 @@ impl FunctionRegistry {
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn table_functions(&self) -> Vec<Arc<TableFunction>> {
|
||||
self.table_functions
|
||||
.read()
|
||||
.unwrap()
|
||||
.values()
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Returns true if an aggregate function with the given name exists in the registry.
|
||||
pub fn is_aggr_func_exist(&self, name: &str) -> bool {
|
||||
self.aggregate_functions.read().unwrap().contains_key(name)
|
||||
|
||||
Reference in New Issue
Block a user