Add FunctionRegistry::register_udaf and FunctionRegistry::register_udwf

mo

Author: alamb

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

@@ -823,15 +823,7 @@ impl SessionContext {
     /// Any functions registered with the udf name or its aliases will be overwritten with this new function
     pub fn register_udf(&self, f: ScalarUDF) {
         let mut state = self.state.write();
-        let aliases = f.aliases();
-        for alias in aliases {
-            state
-                .scalar_functions
-                .insert(alias.to_string(), Arc::new(f.clone()));
-        }
-        state
-            .scalar_functions
-            .insert(f.name().to_string(), Arc::new(f));
+        state.register_udf(Arc::new(f)).ok();
     }
 
     /// Registers an aggregate UDF within this context.
@@ -842,10 +834,7 @@ impl SessionContext {
     /// - `SELECT MY_UDAF(x)...` will look for an aggregate named `"my_udaf"`
     /// - `SELECT "my_UDAF"(x)` will look for an aggregate named `"my_UDAF"`
     pub fn register_udaf(&self, f: AggregateUDF) {
-        self.state
-            .write()
-            .aggregate_functions
-            .insert(f.name().to_string(), Arc::new(f));
+        self.state.write().register_udaf(Arc::new(f)).ok();
     }
 
     /// Registers a window UDF within this context.
@@ -856,10 +845,7 @@ impl SessionContext {
     /// - `SELECT MY_UDWF(x)...` will look for a window function named `"my_udwf"`
     /// - `SELECT "my_UDWF"(x)` will look for a window function named `"my_UDWF"`
     pub fn register_udwf(&self, f: WindowUDF) {
-        self.state
-            .write()
-            .window_functions
-            .insert(f.name().to_string(), Arc::new(f));
+        self.state.write().register_udwf(Arc::new(f)).ok();
     }
 
     /// Creates a [`DataFrame`] for reading a data source.
@@ -1984,8 +1970,24 @@ impl FunctionRegistry for SessionState {
     }
 
     fn register_udf(&mut self, udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> {
+        let aliases = udf.aliases();
+        for alias in aliases {
+            self.scalar_functions.insert(alias.to_string(), udf.clone());
+        }
+
         Ok(self.scalar_functions.insert(udf.name().into(), udf))
     }
+
+    fn register_udaf(
+        &mut self,
+        udaf: Arc<AggregateUDF>,
+    ) -> Result<Option<Arc<AggregateUDF>>> {
+        Ok(self.aggregate_functions.insert(udaf.name().into(), udaf))
+    }
+
+    fn register_udwf(&mut self, udwf: Arc<WindowUDF>) -> Result<Option<Arc<WindowUDF>>> {
+        Ok(self.window_functions.insert(udwf.name().into(), udwf))
+    }
 }
 
 impl OptimizerConfig for SessionState {

datafusion/execution/src/registry.rs

@@ -27,13 +27,16 @@ pub trait FunctionRegistry {
     /// Set of all available udfs.
     fn udfs(&self) -> HashSet<String>;
 
-    /// Returns a reference to the udf named `name`.
+    /// Returns a reference to the user defined scalar function (udf) named
+    /// `name`.
     fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>;
 
-    /// Returns a reference to the udaf named `name`.
+    /// Returns a reference to the user defined table function (udaf) named
+    /// `name`.
     fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>>;
 
-    /// Returns a reference to the udwf named `name`.
+    /// Returns a reference to the user defined window function (udwf) named
+    /// `name`.
     fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>>;
 
     /// Registers a new [`ScalarUDF`], returning any previously registered
@@ -45,7 +48,26 @@ pub trait FunctionRegistry {
         not_impl_err!("Registering ScalarUDF")
     }
 
-    // TODO add register_udaf and register_udwf
+    /// Registers a new [`AggregateUDF`], returning any previously registered
+    /// implementation.
+    ///
+    /// Returns an error (the default) if the function can not be registered,
+    /// for example if the registry is read only.
+    fn register_udaf(
+        &mut self,
+        _udaf: Arc<AggregateUDF>,
+    ) -> Result<Option<Arc<AggregateUDF>>> {
+        not_impl_err!("Registering AggregateUDF")
+    }
+
+    /// Registers a new [`WindowUDF`], returning any previously registered
+    /// implementation.
+    ///
+    /// Returns an error (the default) if the function can not be registered,
+    /// for example if the registry is read only.
+    fn register_udwf(&mut self, _udaf: Arc<WindowUDF>) -> Result<Option<Arc<WindowUDF>>> {
+        not_impl_err!("Registering WindowUDF")
+    }
 }
 
 /// Serializer and deserializer registry for extensions like [UserDefinedLogicalNode].