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

Author: milenkovicm

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

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

datafusion/core/src/execution/session_state.rs

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

datafusion/execution/src/task.rs

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

datafusion/expr/src/registry.rs

@@ -27,9 +27,25 @@ 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> {
+        // This default implementation is provided temporarily
+        // to maintain backward compatibility for the 50.1 release.
+        // It will be reverted to a required method in future versions.
+        HashSet::default()
+    }
+
+    /// Returns names of all available window user defined functions.
+    fn udwfs(&self) -> HashSet<String> {
+        // This default implementation is provided temporarily
+        // to maintain backward compatibility for the 50.1 release.
+        // It will be reverted to a required method in future versions.
+        HashSet::default()
+    }
+
     /// Returns a reference to the user defined scalar function (udf) named
     /// `name`.
     fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>;
@@ -200,4 +216,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()
+    }
 }

datafusion/proto/src/bytes/mod.rs

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

datafusion/proto/src/bytes/registry.rs

@@ -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()
+    }
 }

datafusion/spark/src/lib.rs

@@ -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!() }