DRAFT: Resolve function calls by name during planning

datafusion-examples/examples/rewrite_expr.rs

@@ -16,13 +16,14 @@
 // under the License.
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::execution::FunctionRegistry;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TreeNode};
-use datafusion_common::{plan_err, DataFusionError, Result, ScalarValue};
+use datafusion_common::{internal_err, plan_err, DataFusionError, Result, ScalarValue};
 use datafusion_expr::{
     AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource, WindowUDF,
 };
-use datafusion_optimizer::analyzer::{Analyzer, AnalyzerRule};
+use datafusion_optimizer::analyzer::{Analyzer, AnalyzerConfig, AnalyzerRule};
 use datafusion_optimizer::optimizer::Optimizer;
 use datafusion_optimizer::{utils, OptimizerConfig, OptimizerContext, OptimizerRule};
 use datafusion_sql::planner::{ContextProvider, SqlToRel};
@@ -32,6 +33,38 @@ use datafusion_sql::TableReference;
 use std::any::Any;
 use std::sync::Arc;
 
+struct ExamplesAnalyzerConfig<'a> {
+    config_options: &'a ConfigOptions,
+}
+
+impl<'a> FunctionRegistry for ExamplesAnalyzerConfig<'a> {
+    fn udfs(&self) -> std::collections::HashSet<String> {
+        std::collections::HashSet::new()
+    }
+
+    fn udf(&self, _name: &str) -> Result<Arc<ScalarUDF>> {
+        internal_err!("Mock Function Registry")
+    }
+
+    fn udaf(&self, _name: &str) -> Result<Arc<AggregateUDF>> {
+        internal_err!("Mock Function Registry")
+    }
+
+    fn udwf(&self, _name: &str) -> Result<Arc<WindowUDF>> {
+        internal_err!("Mock Function Registry")
+    }
+}
+
+impl<'a> AnalyzerConfig for ExamplesAnalyzerConfig<'a> {
+    fn function_registry(&self) -> &dyn FunctionRegistry {
+        self
+    }
+
+    fn options(&self) -> &ConfigOptions {
+        self.config_options
+    }
+}
+
 pub fn main() -> Result<()> {
     // produce a logical plan using the datafusion-sql crate
     let dialect = PostgreSqlDialect {};
@@ -50,8 +83,11 @@ pub fn main() -> Result<()> {
     // run the analyzer with our custom rule
     let config = OptimizerContext::default().with_skip_failing_rules(false);
     let analyzer = Analyzer::with_rules(vec![Arc::new(MyAnalyzerRule {})]);
+    let analyzer_config = ExamplesAnalyzerConfig {
+        config_options: config.options(),
+    };
     let analyzed_plan =
-        analyzer.execute_and_check(&logical_plan, config.options(), |_, _| {})?;
+        analyzer.execute_and_check(&logical_plan, &analyzer_config, |_, _| {})?;
     println!(
         "Analyzed Logical Plan:\n\n{}\n",
         analyzed_plan.display_indent()
@@ -80,7 +116,11 @@ fn observe(plan: &LogicalPlan, rule: &dyn OptimizerRule) {
 struct MyAnalyzerRule {}
 
 impl AnalyzerRule for MyAnalyzerRule {
-    fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result<LogicalPlan> {
+    fn analyze(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn AnalyzerConfig,
+    ) -> Result<LogicalPlan> {
         Self::analyze_plan(plan)
     }
 

datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs

@@ -349,12 +349,15 @@ mod tests {
     use arrow::datatypes::DataType::Decimal128;
     use arrow::datatypes::Schema;
     use arrow::datatypes::{DataType, Field};
+    use datafusion_common::internal_err;
     use datafusion_common::{config::ConfigOptions, TableReference, ToDFSchema};
     use datafusion_common::{DataFusionError, Result};
+    use datafusion_execution::FunctionRegistry;
     use datafusion_expr::{
         builder::LogicalTableSource, cast, col, lit, AggregateUDF, Expr, ScalarUDF,
         TableSource, WindowUDF,
     };
+    use datafusion_optimizer::analyzer::AnalyzerConfig;
     use datafusion_physical_expr::execution_props::ExecutionProps;
     use datafusion_physical_expr::{create_physical_expr, PhysicalExpr};
     use datafusion_sql::planner::ContextProvider;
@@ -370,6 +373,38 @@ mod tests {
     use std::ops::Rem;
     use std::sync::Arc;
 
+    struct TestAnalyzerConfig<'a> {
+        config_options: &'a ConfigOptions,
+    }
+
+    impl<'a> FunctionRegistry for TestAnalyzerConfig<'a> {
+        fn udfs(&self) -> HashSet<String> {
+            HashSet::new()
+        }
+
+        fn udf(&self, _name: &str) -> Result<Arc<ScalarUDF>> {
+            internal_err!("mock function registry")
+        }
+
+        fn udaf(&self, _name: &str) -> Result<Arc<AggregateUDF>> {
+            internal_err!("mock function registry")
+        }
+
+        fn udwf(&self, _name: &str) -> Result<Arc<WindowUDF>> {
+            internal_err!("mock function registry")
+        }
+    }
+
+    impl<'a> AnalyzerConfig for TestAnalyzerConfig<'a> {
+        fn function_registry(&self) -> &dyn datafusion_execution::FunctionRegistry {
+            self
+        }
+
+        fn options(&self) -> &ConfigOptions {
+            self.config_options
+        }
+    }
+
     struct PrimitiveTypeField {
         name: &'static str,
         physical_ty: PhysicalType,
@@ -1314,7 +1349,10 @@ mod tests {
         let analyzer = Analyzer::new();
         let optimizer = Optimizer::new();
         // analyze and optimize the logical plan
-        let plan = analyzer.execute_and_check(&plan, config.options(), |_, _| {})?;
+        let analyzer_config = TestAnalyzerConfig {
+            config_options: config.options(),
+        };
+        let plan = analyzer.execute_and_check(&plan, &analyzer_config, |_, _| {})?;
         let plan = optimizer.optimize(&plan, &config, |_, _| {})?;
         // convert the logical plan into a physical plan
         let exprs = plan.expressions();

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

@@ -104,7 +104,7 @@ use crate::catalog::information_schema::{InformationSchemaProvider, INFORMATION_
 use crate::catalog::listing_schema::ListingSchemaProvider;
 use crate::datasource::object_store::ObjectStoreUrl;
 use datafusion_optimizer::{
-    analyzer::{Analyzer, AnalyzerRule},
+    analyzer::{Analyzer, AnalyzerConfig, AnalyzerRule},
     OptimizerConfig,
 };
 use datafusion_sql::planner::object_name_to_table_reference;
@@ -1729,7 +1729,7 @@ impl SessionState {
             // analyze & capture output of each rule
             let analyzed_plan = match self.analyzer.execute_and_check(
                 e.plan.as_ref(),
-                self.options(),
+                self,
                 |analyzed_plan, analyzer| {
                     let analyzer_name = analyzer.name().to_string();
                     let plan_type = PlanType::AnalyzedLogicalPlan { analyzer_name };
@@ -1785,9 +1785,7 @@ impl SessionState {
                 logical_optimization_succeeded,
             }))
         } else {
-            let analyzed_plan =
-                self.analyzer
-                    .execute_and_check(plan, self.options(), |_, _| {})?;
+            let analyzed_plan = self.analyzer.execute_and_check(plan, self, |_, _| {})?;
             self.optimizer.optimize(&analyzed_plan, self, |_, _| {})
         }
     }
@@ -1875,6 +1873,16 @@ impl SessionState {
     }
 }
 
+impl AnalyzerConfig for SessionState {
+    fn function_registry(&self) -> &dyn FunctionRegistry {
+        self
+    }
+
+    fn options(&self) -> &ConfigOptions {
+        self.config_options()
+    }
+}
+
 struct SessionContextProvider<'a> {
     state: &'a SessionState,
     tables: HashMap<String, Arc<dyn TableSource>>,

datafusion/expr/src/expr_fn.rs

@@ -27,7 +27,8 @@ use crate::{
     aggregate_function, built_in_function, conditional_expressions::CaseBuilder,
     logical_plan::Subquery, AccumulatorFactoryFunction, AggregateUDF,
     BuiltinScalarFunction, Expr, LogicalPlan, Operator, ReturnTypeFunction,
-    ScalarFunctionImplementation, ScalarUDF, Signature, StateTypeFunction, Volatility,
+    ScalarFunctionDefinition, ScalarFunctionImplementation, ScalarUDF, Signature,
+    StateTypeFunction, Volatility,
 };
 use arrow::datatypes::DataType;
 use datafusion_common::{Column, Result};
@@ -1007,18 +1008,18 @@ pub fn create_udwf(
     )
 }
 
-/// Calls a named built in function
+/// Calls a named function
 /// ```
 /// use datafusion_expr::{col, lit, call_fn};
 ///
 /// // create the expression sin(x) < 0.2
 /// let expr = call_fn("sin", vec![col("x")]).unwrap().lt(lit(0.2));
 /// ```
 pub fn call_fn(name: impl AsRef<str>, args: Vec<Expr>) -> Result<Expr> {
-    match name.as_ref().parse::<BuiltinScalarFunction>() {
-        Ok(fun) => Ok(Expr::ScalarFunction(ScalarFunction::new(fun, args))),
-        Err(e) => Err(e),
-    }
+    Ok(Expr::ScalarFunction(ScalarFunction {
+        func_def: ScalarFunctionDefinition::Name(Arc::from(name.as_ref())),
+        args,
+    }))
 }
 
 #[cfg(test)]

datafusion/optimizer/Cargo.toml

@@ -43,6 +43,7 @@ arrow = { workspace = true }
 async-trait = { workspace = true }
 chrono = { workspace = true }
 datafusion-common = { workspace = true }
+datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-physical-expr = { path = "../physical-expr", version = "33.0.0", default-features = false }
 hashbrown = { version = "0.14", features = ["raw"] }

datafusion/optimizer/src/analyzer/count_wildcard_rule.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::analyzer::AnalyzerRule;
-use datafusion_common::config::ConfigOptions;
+
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
 use datafusion_common::Result;
 use datafusion_expr::expr::{AggregateFunction, AggregateFunctionDefinition, InSubquery};
@@ -29,6 +29,8 @@ use datafusion_expr::{
 };
 use std::sync::Arc;
 
+use super::AnalyzerConfig;
+
 /// Rewrite `Count(Expr:Wildcard)` to `Count(Expr:Literal)`.
 ///
 /// Resolves issue: <https://github.com/apache/arrow-datafusion/issues/5473>
@@ -42,7 +44,7 @@ impl CountWildcardRule {
 }
 
 impl AnalyzerRule for CountWildcardRule {
-    fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result<LogicalPlan> {
+    fn analyze(&self, plan: LogicalPlan, _: &dyn AnalyzerConfig) -> Result<LogicalPlan> {
         plan.transform_down(&analyze_internal)
     }