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:
Ning Sun
2025-09-09 09:26:38 +08:00
committed by GitHub
parent 16febbd4c2
commit 948a6578fa
6 changed files with 86 additions and 8 deletions

View File

@@ -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)