Skip to content

Commit 95357ed

Browse files
authored
feat: expose udafs and udwfs methods on FunctionRegistry (#17650)
* expose udafs and udwfs method on `FunctionRegistry` * fix doc test * add default implementations not to trigger backward incompatible change for others
1 parent 52690c6 commit 95357ed

File tree

7 files changed

+67
-1
lines changed

7 files changed

+67
-1
lines changed

datafusion/core/src/execution/context/mod.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1755,6 +1755,14 @@ impl FunctionRegistry for SessionContext {
17551755
) -> Result<()> {
17561756
self.state.write().register_expr_planner(expr_planner)
17571757
}
1758+
1759+
fn udafs(&self) -> HashSet<String> {
1760+
self.state.read().udafs()
1761+
}
1762+
1763+
fn udwfs(&self) -> HashSet<String> {
1764+
self.state.read().udwfs()
1765+
}
17581766
}
17591767

17601768
/// Create a new task context instance from SessionContext

datafusion/core/src/execution/session_state.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1917,6 +1917,14 @@ impl FunctionRegistry for SessionState {
19171917
self.expr_planners.push(expr_planner);
19181918
Ok(())
19191919
}
1920+
1921+
fn udafs(&self) -> HashSet<String> {
1922+
self.aggregate_functions.keys().cloned().collect()
1923+
}
1924+
1925+
fn udwfs(&self) -> HashSet<String> {
1926+
self.window_functions.keys().cloned().collect()
1927+
}
19201928
}
19211929

19221930
impl OptimizerConfig for SessionState {

datafusion/execution/src/task.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,14 @@ impl FunctionRegistry for TaskContext {
201201
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
202202
vec![]
203203
}
204+
205+
fn udafs(&self) -> HashSet<String> {
206+
self.aggregate_functions.keys().cloned().collect()
207+
}
208+
209+
fn udwfs(&self) -> HashSet<String> {
210+
self.window_functions.keys().cloned().collect()
211+
}
204212
}
205213

206214
#[cfg(test)]

datafusion/expr/src/registry.rs

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,25 @@ use std::sync::Arc;
2727

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

33+
/// Returns names of all available aggregate user defined functions.
34+
fn udafs(&self) -> HashSet<String> {
35+
// This default implementation is provided temporarily
36+
// to maintain backward compatibility for the 50.1 release.
37+
// It will be reverted to a required method in future versions.
38+
HashSet::default()
39+
}
40+
41+
/// Returns names of all available window user defined functions.
42+
fn udwfs(&self) -> HashSet<String> {
43+
// This default implementation is provided temporarily
44+
// to maintain backward compatibility for the 50.1 release.
45+
// It will be reverted to a required method in future versions.
46+
HashSet::default()
47+
}
48+
3349
/// Returns a reference to the user defined scalar function (udf) named
3450
/// `name`.
3551
fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>;
@@ -200,4 +216,12 @@ impl FunctionRegistry for MemoryFunctionRegistry {
200216
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
201217
vec![]
202218
}
219+
220+
fn udafs(&self) -> HashSet<String> {
221+
self.udafs.keys().cloned().collect()
222+
}
223+
224+
fn udwfs(&self) -> HashSet<String> {
225+
self.udwfs.keys().cloned().collect()
226+
}
203227
}

datafusion/proto/src/bytes/mod.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,14 @@ impl Serializeable for Expr {
170170
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
171171
vec![]
172172
}
173+
174+
fn udafs(&self) -> std::collections::HashSet<String> {
175+
std::collections::HashSet::default()
176+
}
177+
178+
fn udwfs(&self) -> std::collections::HashSet<String> {
179+
std::collections::HashSet::default()
180+
}
173181
}
174182
Expr::from_bytes_with_registry(&bytes, &PlaceHolderRegistry)?;
175183

datafusion/proto/src/bytes/registry.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,4 +59,12 @@ impl FunctionRegistry for NoRegistry {
5959
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
6060
vec![]
6161
}
62+
63+
fn udafs(&self) -> HashSet<String> {
64+
HashSet::new()
65+
}
66+
67+
fn udwfs(&self) -> HashSet<String> {
68+
HashSet::new()
69+
}
6270
}

datafusion/spark/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@
5353
//! # impl FunctionRegistry for SessionContext {
5454
//! # fn register_udf(&mut self, _udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> { Ok (None) }
5555
//! # fn udfs(&self) -> HashSet<String> { unimplemented!() }
56+
//! # fn udafs(&self) -> HashSet<String> { unimplemented!() }
57+
//! # fn udwfs(&self) -> HashSet<String> { unimplemented!() }
5658
//! # fn udf(&self, _name: &str) -> Result<Arc<ScalarUDF>> { unimplemented!() }
5759
//! # fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {unimplemented!() }
5860
//! # fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> { unimplemented!() }

0 commit comments

Comments
 (0)