feat: function name hints for UDFs

Author: SteveLauC

datafusion-cli/Cargo.lock

@@ -30,6 +30,7 @@ use datafusion_sql::sqlparser::dialect::PostgreSqlDialect;
 use datafusion_sql::sqlparser::parser::Parser;
 use datafusion_sql::TableReference;
 use std::any::Any;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 pub fn main() -> Result<()> {
@@ -223,6 +224,18 @@ impl ContextProvider for MyContextProvider {
     fn options(&self) -> &ConfigOptions {
         &self.options
     }
+
+    fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>> {
+        HashMap::new()
+    }
+
+    fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>> {
+        HashMap::new()
+    }
+
+    fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>> {
+        HashMap::new()
+    }
 }
 
 struct MyTableSource {

datafusion-examples/examples/rewrite_expr.rs

@@ -1993,6 +1993,18 @@ impl<'a> ContextProvider for SessionContextProvider<'a> {
     fn options(&self) -> &ConfigOptions {
         self.state.config_options()
     }
+
+    fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>> {
+        self.state.scalar_functions.clone()
+    }
+
+    fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>> {
+        self.state.aggregate_functions().clone()
+    }
+
+    fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>> {
+        self.state.window_functions().clone()
+    }
 }
 
 impl FunctionRegistry for SessionState {

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

@@ -17,13 +17,12 @@
 
 //! Function module contains typing and signature for built-in and user defined functions.
 
-use crate::{Accumulator, BuiltinScalarFunction, PartitionEvaluator, Signature};
-use crate::{AggregateFunction, BuiltInWindowFunction, ColumnarValue};
+use crate::{
+    Accumulator, BuiltinScalarFunction, ColumnarValue, PartitionEvaluator, Signature,
+};
 use arrow::datatypes::DataType;
-use datafusion_common::utils::datafusion_strsim;
 use datafusion_common::Result;
 use std::sync::Arc;
-use strum::IntoEnumIterator;
 
 /// Scalar function
 ///
@@ -75,33 +74,3 @@ pub fn return_type(
 pub fn signature(fun: &BuiltinScalarFunction) -> Signature {
     fun.signature()
 }
-
-/// Suggest a valid function based on an invalid input function name
-pub fn suggest_valid_function(input_function_name: &str, is_window_func: bool) -> String {
-    let valid_funcs = if is_window_func {
-        // All aggregate functions and builtin window functions
-        AggregateFunction::iter()
-            .map(|func| func.to_string())
-            .chain(BuiltInWindowFunction::iter().map(|func| func.to_string()))
-            .collect()
-    } else {
-        // All scalar functions and aggregate functions
-        BuiltinScalarFunction::iter()
-            .map(|func| func.to_string())
-            .chain(AggregateFunction::iter().map(|func| func.to_string()))
-            .collect()
-    };
-    find_closest_match(valid_funcs, input_function_name)
-}
-
-/// Find the closest matching string to the target string in the candidates list, using edit distance(case insensitve)
-/// Input `candidates` must not be empty otherwise it will panic
-fn find_closest_match(candidates: Vec<String>, target: &str) -> String {
-    let target = target.to_lowercase();
-    candidates
-        .into_iter()
-        .min_by_key(|candidate| {
-            datafusion_strsim::levenshtein(&candidate.to_lowercase(), &target)
-        })
-        .expect("No candidates provided.") // Panic if `candidates` argument is empty
-}

datafusion/expr/src/function.rs

@@ -418,6 +418,18 @@ impl ContextProvider for MyContextProvider {
     fn options(&self) -> &ConfigOptions {
         &self.options
     }
+
+    fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>> {
+        HashMap::new()
+    }
+
+    fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>> {
+        HashMap::new()
+    }
+
+    fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>> {
+        HashMap::new()
+    }
 }
 
 struct MyTableSource {

datafusion/optimizer/tests/optimizer_integration.rs

@@ -43,6 +43,7 @@ datafusion-common = { workspace = true, default-features = true }
 datafusion-expr = { workspace = true }
 log = { workspace = true }
 sqlparser = { workspace = true }
+strum = { version = "0.26.1", features = ["derive"] }
 
 [dev-dependencies]
 ctor = { workspace = true }

datafusion/sql/Cargo.toml

@@ -131,4 +131,16 @@ impl ContextProvider for MyContextProvider {
     fn options(&self) -> &ConfigOptions {
         &self.options
     }
+
+    fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>> {
+        HashMap::new()
+    }
+
+    fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>> {
+        HashMap::new()
+    }
+
+    fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>> {
+        HashMap::new()
+    }
 }

datafusion/sql/examples/sql.rs

@@ -20,20 +20,63 @@ use arrow_schema::DataType;
 use datafusion_common::{
     not_impl_err, plan_datafusion_err, plan_err, DFSchema, Dependency, Result,
 };
-use datafusion_expr::expr::{ScalarFunction, Unnest};
-use datafusion_expr::function::suggest_valid_function;
 use datafusion_expr::window_frame::{check_window_frame, regularize_window_order_by};
 use datafusion_expr::{
-    expr, AggregateFunction, BuiltinScalarFunction, Expr, ExprSchemable, WindowFrame,
-    WindowFunctionDefinition,
+    expr, AggregateFunction, Expr, ExprSchemable, WindowFrame, WindowFunctionDefinition,
+};
+use datafusion_expr::{
+    expr::{ScalarFunction, Unnest},
+    BuiltInWindowFunction, BuiltinScalarFunction,
 };
 use sqlparser::ast::{
     Expr as SQLExpr, Function as SQLFunction, FunctionArg, FunctionArgExpr, WindowType,
 };
 use std::str::FromStr;
+use strum::IntoEnumIterator;
 
 use super::arrow_cast::ARROW_CAST_NAME;
 
+/// Suggest a valid function based on an invalid input function name
+pub fn suggest_valid_function(
+    input_function_name: &str,
+    is_window_func: bool,
+    ctx: &dyn ContextProvider,
+) -> String {
+    let valid_funcs = if is_window_func {
+        // All aggregate functions and builtin window functions
+        AggregateFunction::iter()
+            .map(|func| func.to_string())
+            .chain(BuiltInWindowFunction::iter().map(|func| func.to_string()))
+            .collect()
+    } else {
+        // All scalar functions and aggregate functions
+        let mut funcs = Vec::new();
+
+        funcs.extend(BuiltinScalarFunction::iter().map(|func| func.to_string()));
+        funcs.extend(ctx.udfs().into_keys());
+        funcs.extend(AggregateFunction::iter().map(|func| func.to_string()));
+        funcs.extend(ctx.udafs().into_keys());
+
+        funcs
+    };
+    find_closest_match(valid_funcs, input_function_name)
+}
+
+/// Find the closest matching string to the target string in the candidates list, using edit distance(case insensitve)
+/// Input `candidates` must not be empty otherwise it will panic
+fn find_closest_match(candidates: Vec<String>, target: &str) -> String {
+    let target = target.to_lowercase();
+    candidates
+        .into_iter()
+        .min_by_key(|candidate| {
+            datafusion_common::utils::datafusion_strsim::levenshtein(
+                &candidate.to_lowercase(),
+                &target,
+            )
+        })
+        .expect("No candidates provided.") // Panic if `candidates` argument is empty
+}
+
 impl<'a, S: ContextProvider> SqlToRel<'a, S> {
     pub(super) fn sql_function_to_expr(
         &self,
@@ -211,7 +254,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         }
 
         // Could not find the relevant function, so return an error
-        let suggested_func_name = suggest_valid_function(&name, is_function_window);
+        let suggested_func_name =
+            suggest_valid_function(&name, is_function_window, self.context_provider);
         plan_err!("Invalid function '{name}'.\nDid you mean '{suggested_func_name}'?")
     }
 

datafusion/sql/src/expr/function.rs

@@ -951,6 +951,18 @@ mod tests {
         fn get_window_meta(&self, _name: &str) -> Option<Arc<WindowUDF>> {
             None
         }
+
+        fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>> {
+            HashMap::new()
+        }
+
+        fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>> {
+            HashMap::new()
+        }
+
+        fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>> {
+            HashMap::new()
+        }
     }
 
     fn create_table_source(fields: Vec<Field>) -> Arc<dyn TableSource> {

datafusion/sql/src/expr/mod.rs

@@ -85,6 +85,10 @@ pub trait ContextProvider {
 
     /// Get configuration options
     fn options(&self) -> &ConfigOptions;
+
+    fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>>;
+    fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>>;
+    fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>>;
 }
 
 /// SQL parser options

datafusion/sql/src/planner.rs

@@ -2907,6 +2907,18 @@ impl ContextProvider for MockContextProvider {
     ) -> Result<Arc<dyn TableSource>> {
         Ok(Arc::new(EmptyTable::new(schema)))
     }
+
+    fn udfs(&self) -> HashMap<String, Arc<ScalarUDF>> {
+        self.udfs.clone()
+    }
+
+    fn udafs(&self) -> HashMap<String, Arc<AggregateUDF>> {
+        self.udafs.clone()
+    }
+
+    fn udwfs(&self) -> HashMap<String, Arc<WindowUDF>> {
+        HashMap::new()
+    }
 }
 
 #[test]