diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index fe73c99ac25c..0e41a67c4d65 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -50,7 +50,7 @@ use crate::{ DropView, Explain, LogicalPlan, LogicalPlanBuilder, PlanType, SetVariable, TableSource, TableType, ToStringifiedPlan, UNNAMED_TABLE, }, - optimizer::analyzer::{Analyzer, AnalyzerRule}, + optimizer::analyzer::Analyzer, optimizer::optimizer::{Optimizer, OptimizerConfig, OptimizerRule}, physical_optimizer::optimizer::{PhysicalOptimizer, PhysicalOptimizerRule}, physical_plan::{udaf::AggregateUDF, udf::ScalarUDF, ExecutionPlan}, @@ -69,6 +69,7 @@ use datafusion_common::{ OwnedTableReference, SchemaReference, }; use datafusion_execution::registry::SerializerRegistry; +use datafusion_expr::analyzer::AnalyzerRuleRef; use datafusion_expr::{ logical_plan::{DdlStatement, Statement}, var_provider::is_system_variables, @@ -1702,12 +1703,6 @@ impl SessionState { self } - /// Override the `AnalyzerRule`s optimizer plan rules. - pub fn with_analyzer_rules(mut self, rules: Vec) -> Self { - self.analyzer = Analyzer::with_rules(rules); - self - } - /// Replace the entire list of [`OptimizerRule`]s used to optimize plans pub fn with_optimizer_rules( mut self, @@ -2242,23 +2237,14 @@ impl FunctionRegistry for SessionState { }) } - /// Add `analyzer_rule` to the `index` of the list of - /// `AnalyzerRule`s used to rewrite queries. - fn register_analyzer_rule( - &mut self, - analyzer_rule: AnalyzerRuleRef, - index: usize, - ) -> Result<()> { - if index > self.analyzer.rules.len() { - return exec_err!( - "index out of range, index: {index}, length: {:?}", - self.analyzer.rules.len() - ); - } - let mut split = self.analyzer.rules.split_off(index); - self.analyzer.rules.push_back(analyzer_rule); - self.analyzer.rules.append(&mut split); - Ok(()) + /// Override the `AnalyzerRule`s optimizer plan rules. + fn with_analyzer_rules(&mut self, rules: Vec) { + self.analyzer = Analyzer::with_rules(rules); + } + + /// return the existing analyzer rules + fn analyzer_rules(&self) -> Vec { + self.analyzer.rules.clone() } fn register_udf(&mut self, udf: Arc) -> Result>> { @@ -2314,21 +2300,6 @@ impl FunctionRegistry for SessionState { } Ok(udwf) } - - fn deregister_analyzer_rule(&mut self, index: usize) -> Result<()> { - if index > self.analyzer.rules.len() { - return exec_err!( - "index out of range, index: {index}, length: {:?}", - self.analyzer.rules.len() - ); - } - - let mut split = self.analyzer.rules.split_off(index); - // split off should includes the element at index - split.pop_front().unwrap(); - self.analyzer.rules.append(&mut split); - Ok(()) - } } impl OptimizerConfig for SessionState { diff --git a/datafusion/execution/src/registry.rs b/datafusion/execution/src/registry.rs index a05d69fc9118..ce2f6c5972d4 100644 --- a/datafusion/execution/src/registry.rs +++ b/datafusion/execution/src/registry.rs @@ -68,16 +68,13 @@ pub trait FunctionRegistry { not_impl_err!("Registering WindowUDF") } - /// Registers a new `AnalyzerRule` at the given index. - /// Returns an error if the index is out of bounds. - fn register_analyzer_rule( - &mut self, - _rule: AnalyzerRuleRef, - _index: usize, - ) -> Result<()> { - not_impl_err!("Registering AnalyzerRule") + /// Return the existing analyzer rules + fn analyzer_rules(&self) -> Vec { + vec![] } + fn with_analyzer_rules(&mut self, _rules: Vec) {} + /// Deregisters a [`ScalarUDF`], returning the implementation that was /// deregistered. /// @@ -104,12 +101,6 @@ pub trait FunctionRegistry { fn deregister_udwf(&mut self, _name: &str) -> Result>> { not_impl_err!("Deregistering WindowUDF") } - - /// Deregsiters an `AnalyzerRule` at the given index. - /// Returns an error if the index is out of bounds. - fn deregister_analyzer_rule(&mut self, _index: usize) -> Result<()> { - not_impl_err!("DeRegistering AnalyzerRule") - } } /// Serializer and deserializer registry for extensions like [UserDefinedLogicalNode]. diff --git a/datafusion/functions-array/src/lib.rs b/datafusion/functions-array/src/lib.rs index 2f1e5bffb86d..d4cd0eedabcf 100644 --- a/datafusion/functions-array/src/lib.rs +++ b/datafusion/functions-array/src/lib.rs @@ -39,7 +39,7 @@ mod utils; use analyzer::rewrite::OperatorToFunction; use datafusion_common::Result; use datafusion_execution::FunctionRegistry; -use datafusion_expr::{analyzer::AnalyzerRuleRef, ScalarUDF}; +use datafusion_expr::ScalarUDF; use log::debug; use std::sync::Arc; @@ -97,18 +97,12 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> { Ok(()) as Result<()> })?; - let rules: Vec<(AnalyzerRuleRef, usize)> = vec![ - // The order here should be adjusted based on the rules defined in the optimizer - // - // OperatorToFunction should be run before TypeCoercion, since it rewrite based on the argument types (List or Scalar), - // and TypeCoercion may cast the argument types from Scalar to List. - (Arc::new(OperatorToFunction::new()), 1), - ]; + let mut rules = registry.analyzer_rules(); + // OperatorToFunction should be run before TypeCoercion, since it rewrite based on the argument types (List or Scalar), + // and TypeCoercion may cast the argument types from Scalar to List. + rules.insert(1, Arc::new(OperatorToFunction::new())); - rules.into_iter().try_for_each(|(rule, index)| { - registry.register_analyzer_rule(rule, index)?; - Ok(()) as Result<()> - })?; + registry.with_analyzer_rules(rules); Ok(()) } diff --git a/datafusion/optimizer/src/analyzer/mod.rs b/datafusion/optimizer/src/analyzer/mod.rs index 6f0edec1a3c9..a3866969d3d5 100644 --- a/datafusion/optimizer/src/analyzer/mod.rs +++ b/datafusion/optimizer/src/analyzer/mod.rs @@ -36,14 +36,13 @@ use datafusion_expr::expr::InSubquery; use datafusion_expr::utils::inspect_expr_pre; use datafusion_expr::{Expr, LogicalPlan}; use log::debug; -use std::collections::LinkedList; use std::sync::Arc; /// A rule-based Analyzer. #[derive(Clone)] pub struct Analyzer { /// All rules to apply - pub rules: LinkedList, + pub rules: Vec, } impl Default for Analyzer { @@ -66,9 +65,7 @@ impl Analyzer { /// Create a new analyzer with the given rules pub fn with_rules(rules: Vec) -> Self { - Self { - rules: LinkedList::from_iter(rules), - } + Self { rules } } /// Analyze the logical plan by applying analyzer rules, and