Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions datafusion/core/src/execution/context/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1755,6 +1755,14 @@ impl FunctionRegistry for SessionContext {
) -> Result<()> {
self.state.write().register_expr_planner(expr_planner)
}

fn udafs(&self) -> HashSet<String> {
self.state.read().udafs()
}

fn udwfs(&self) -> HashSet<String> {
self.state.read().udwfs()
}
}

/// Create a new task context instance from SessionContext
Expand Down
8 changes: 8 additions & 0 deletions datafusion/core/src/execution/session_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1917,6 +1917,14 @@ impl FunctionRegistry for SessionState {
self.expr_planners.push(expr_planner);
Ok(())
}

fn udafs(&self) -> HashSet<String> {
self.aggregate_functions.keys().cloned().collect()
}

fn udwfs(&self) -> HashSet<String> {
self.window_functions.keys().cloned().collect()
}
}

impl OptimizerConfig for SessionState {
Expand Down
8 changes: 8 additions & 0 deletions datafusion/execution/src/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,14 @@ impl FunctionRegistry for TaskContext {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}

fn udafs(&self) -> HashSet<String> {
self.aggregate_functions.keys().cloned().collect()
}

fn udwfs(&self) -> HashSet<String> {
self.window_functions.keys().cloned().collect()
}
}

#[cfg(test)]
Expand Down
16 changes: 15 additions & 1 deletion datafusion/expr/src/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,15 @@ use std::sync::Arc;

/// A registry knows how to build logical expressions out of user-defined function' names
pub trait FunctionRegistry {
/// Set of all available udfs.
/// Returns names of all available scalar user defined functions.
fn udfs(&self) -> HashSet<String>;

/// Returns names of all available aggregate user defined functions.
fn udafs(&self) -> HashSet<String>;

/// Returns names of all available window user defined functions.
fn udwfs(&self) -> HashSet<String>;

/// Returns a reference to the user defined scalar function (udf) named
/// `name`.
fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>;
Expand Down Expand Up @@ -200,4 +206,12 @@ impl FunctionRegistry for MemoryFunctionRegistry {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}

fn udafs(&self) -> HashSet<String> {
self.udafs.keys().cloned().collect()
}

fn udwfs(&self) -> HashSet<String> {
self.udwfs.keys().cloned().collect()
}
}
8 changes: 8 additions & 0 deletions datafusion/proto/src/bytes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,14 @@ impl Serializeable for Expr {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}

fn udafs(&self) -> std::collections::HashSet<String> {
std::collections::HashSet::default()
}

fn udwfs(&self) -> std::collections::HashSet<String> {
std::collections::HashSet::default()
}
}
Expr::from_bytes_with_registry(&bytes, &PlaceHolderRegistry)?;

Expand Down
8 changes: 8 additions & 0 deletions datafusion/proto/src/bytes/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,4 +59,12 @@ impl FunctionRegistry for NoRegistry {
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}

fn udafs(&self) -> HashSet<String> {
HashSet::new()
}

fn udwfs(&self) -> HashSet<String> {
HashSet::new()
}
}
2 changes: 2 additions & 0 deletions datafusion/spark/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@
//! # impl FunctionRegistry for SessionContext {
//! # fn register_udf(&mut self, _udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> { Ok (None) }
//! # fn udfs(&self) -> HashSet<String> { unimplemented!() }
//! # fn udafs(&self) -> HashSet<String> { unimplemented!() }
//! # fn udwfs(&self) -> HashSet<String> { unimplemented!() }
//! # fn udf(&self, _name: &str) -> Result<Arc<ScalarUDF>> { unimplemented!() }
//! # fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {unimplemented!() }
//! # fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> { unimplemented!() }
Expand Down