refine optimizer

Author: leiysky

.licenserc.yaml

@@ -34,5 +34,6 @@ header:
     - "**/*.toml"
     - "**/*.lock"
     - "**/*.yapf"
+    - "**/*.test"
 
   comment: on-failure

Cargo.lock

@@ -127,6 +127,7 @@ walkdir = "2.3.2"
 [dev-dependencies]
 clickhouse-driver = { git = "https://github.com/datafuse-extras/clickhouse_driver", rev = "cf978da" }
 criterion = "0.3.5"
+goldenfile = "1.3.0"
 jwt-simple = "0.10.9"
 maplit = "1.0.2"
 mysql_async = "0.29.0"

query/Cargo.toml

@@ -194,7 +194,7 @@ impl ExecuteState {
                 executor: executor.clone(),
                 block_buffer,
             });
-            interpreter.execute(None).await.unwrap();
+            interpreter.execute(None).await?;
 
             Ok(executor)
         } else {

query/src/servers/http/v1/query/execute_state.rs

@@ -82,7 +82,7 @@ use crate::sql::plans::AggregatePlan;
 use crate::sql::plans::AndExpr;
 use crate::sql::plans::CrossApply;
 use crate::sql::plans::EvalScalar;
-use crate::sql::plans::FilterPlan;
+use crate::sql::plans::Filter;
 use crate::sql::plans::LimitPlan;
 use crate::sql::plans::PhysicalHashJoin;
 use crate::sql::plans::PhysicalScan;
@@ -315,7 +315,7 @@ impl PipelineBuilder {
     fn build_filter(
         &mut self,
         ctx: Arc<QueryContext>,
-        filter: &FilterPlan,
+        filter: &Filter,
         input_schema: DataSchemaRef,
         pipeline: &mut NewPipeline,
     ) -> Result<DataSchemaRef> {

query/src/sql/exec/mod.rs

@@ -28,7 +28,7 @@ pub fn check_physical(expression: &SExpr) -> bool {
     }
 
     for child in expression.children() {
-        if !child.plan().is_physical() {
+        if !check_physical(child) {
             return false;
         }
     }

query/src/sql/exec/util.rs

@@ -16,12 +16,27 @@ mod implement;
 mod rule_list;
 
 use common_exception::Result;
+use lazy_static::lazy_static;
 
+use super::rule::RuleID;
 use crate::sql::optimizer::heuristic::implement::HeuristicImplementor;
-use crate::sql::optimizer::heuristic::rule_list::RuleList;
+pub use crate::sql::optimizer::heuristic::rule_list::RuleList;
 use crate::sql::optimizer::rule::TransformState;
 use crate::sql::optimizer::SExpr;
 
+lazy_static! {
+    pub static ref DEFAULT_REWRITE_RULES: Vec<RuleID> = vec![
+        RuleID::EliminateFilter,
+        RuleID::EliminateEvalScalar,
+        RuleID::EliminateProject,
+        RuleID::MergeFilter,
+        RuleID::MergeEvalScalar,
+        RuleID::MergeProject,
+        RuleID::PushDownFilterEvalScalar,
+        RuleID::PushDownFilterProject,
+    ];
+}
+
 /// A heuristic query optimizer. It will apply specific transformation rules in order and
 /// implement the logical plans with default implementation rules.
 pub struct HeuristicOptimizer {
@@ -30,15 +45,16 @@ pub struct HeuristicOptimizer {
 }
 
 impl HeuristicOptimizer {
-    pub fn create() -> Result<Self> {
-        Ok(HeuristicOptimizer {
-            rules: RuleList::create(vec![])?,
+    pub fn new(rules: RuleList) -> Self {
+        HeuristicOptimizer {
+            rules,
             implementor: HeuristicImplementor::new(),
-        })
+        }
     }
 
     pub fn optimize(&mut self, expression: SExpr) -> Result<SExpr> {
-        let result = self.optimize_expression(&expression)?;
+        let optimized = self.optimize_expression(&expression)?;
+        let result = self.implement_expression(&optimized)?;
         Ok(result)
     }
 
@@ -53,24 +69,41 @@ impl HeuristicOptimizer {
         Ok(result)
     }
 
-    fn apply_transform_rules(&self, s_expr: &SExpr, rule_list: &RuleList) -> Result<SExpr> {
-        let mut result = s_expr.clone();
+    fn implement_expression(&self, s_expr: &SExpr) -> Result<SExpr> {
+        let mut implemented_children = Vec::with_capacity(s_expr.arity());
+        for expr in s_expr.children() {
+            implemented_children.push(self.implement_expression(expr)?);
+        }
+        let implemented_expr = SExpr::create(s_expr.plan().clone(), implemented_children, None);
+        // Implement expression with Implementor
+        let mut state = TransformState::new();
+        self.implementor.implement(&implemented_expr, &mut state)?;
+        let result = if !state.results().is_empty() {
+            state.results()[0].clone()
+        } else {
+            implemented_expr
+        };
+        Ok(result)
+    }
 
+    // Return `None` if no rules matched
+    fn apply_transform_rules(&self, s_expr: &SExpr, rule_list: &RuleList) -> Result<SExpr> {
+        let mut s_expr = s_expr.clone();
         for rule in rule_list.iter() {
             let mut state = TransformState::new();
-            rule.apply(&result, &mut state)?;
-            if !state.results().is_empty() {
-                result = state.results()[0].clone();
-            }
-        }
+            if s_expr.match_pattern(rule.pattern()) && !s_expr.applied_rule(&rule.id()) {
+                rule.apply(&s_expr, &mut state)?;
+                s_expr.apply_rule(&rule.id());
+                if !state.results().is_empty() {
+                    // Recursive optimize the result
+                    let result = &state.results()[0];
+                    let optimized_result = self.optimize_expression(result)?;
 
-        // Implement expression with Implementor
-        let mut state = TransformState::new();
-        self.implementor.implement(s_expr, &mut state)?;
-        if !state.results().is_empty() {
-            result = state.results()[0].clone();
+                    return Ok(optimized_result);
+                }
+            }
         }
 
-        Ok(result)
+        Ok(s_expr.clone())
     }
 }

query/src/sql/optimizer/heuristic/mod.rs

@@ -25,6 +25,7 @@ mod s_expr;
 
 use common_exception::Result;
 pub use heuristic::HeuristicOptimizer;
+pub use heuristic::DEFAULT_REWRITE_RULES;
 pub use m_expr::MExpr;
 pub use memo::Memo;
 pub use optimize_context::OptimizeContext;
@@ -34,11 +35,14 @@ pub use property::PhysicalProperty;
 pub use property::RelExpr;
 pub use property::RelationalProperty;
 pub use property::RequiredProperty;
+pub use rule::RuleFactory;
 pub use s_expr::SExpr;
 
 use super::plans::Plan;
-use crate::sql::optimizer::rule::RuleID;
+pub use crate::sql::optimizer::heuristic::RuleList;
+pub use crate::sql::optimizer::rule::RuleID;
 use crate::sql::optimizer::rule::RuleSet;
+use crate::sql::MetadataRef;
 
 pub fn optimize(plan: Plan) -> Result<Plan> {
     match plan {
@@ -47,7 +51,7 @@ pub fn optimize(plan: Plan) -> Result<Plan> {
             bind_context,
             metadata,
         } => Ok(Plan::Query {
-            s_expr: optimize_query(s_expr)?,
+            s_expr: optimize_query(s_expr, metadata.clone())?,
             bind_context,
             metadata,
         }),
@@ -74,8 +78,9 @@ pub fn optimize(plan: Plan) -> Result<Plan> {
     }
 }
 
-pub fn optimize_query(expression: SExpr) -> Result<SExpr> {
-    let mut heuristic = HeuristicOptimizer::create()?;
+pub fn optimize_query(expression: SExpr, _metadata: MetadataRef) -> Result<SExpr> {
+    let rules = RuleList::create(DEFAULT_REWRITE_RULES.clone())?;
+    let mut heuristic = HeuristicOptimizer::new(rules);
     let s_expr = heuristic.optimize(expression)?;
     // TODO: enable cascades optimizer
     // let mut cascades = CascadesOptimizer::create(ctx);

query/src/sql/optimizer/mod.rs

@@ -14,6 +14,14 @@
 
 use common_exception::Result;
 
+use super::rewrite::RuleEliminateEvalScalar;
+use super::rewrite::RulePushDownFilterEvalScalar;
+use super::rewrite::RulePushDownFilterProject;
+use crate::sql::optimizer::rule::rewrite::RuleEliminateFilter;
+use crate::sql::optimizer::rule::rewrite::RuleEliminateProject;
+use crate::sql::optimizer::rule::rewrite::RuleMergeEvalScalar;
+use crate::sql::optimizer::rule::rewrite::RuleMergeFilter;
+use crate::sql::optimizer::rule::rewrite::RuleMergeProject;
 use crate::sql::optimizer::rule::rule_implement_get::RuleImplementGet;
 use crate::sql::optimizer::rule::rule_implement_hash_join::RuleImplementHashJoin;
 use crate::sql::optimizer::rule::RuleID;
@@ -28,8 +36,16 @@ impl RuleFactory {
 
     pub fn create_rule(&self, id: RuleID) -> Result<RulePtr> {
         match id {
-            RuleID::ImplementGet => Ok(Box::new(RuleImplementGet::create())),
-            RuleID::ImplementHashJoin => Ok(Box::new(RuleImplementHashJoin::create())),
+            RuleID::ImplementGet => Ok(Box::new(RuleImplementGet::new())),
+            RuleID::ImplementHashJoin => Ok(Box::new(RuleImplementHashJoin::new())),
+            RuleID::EliminateEvalScalar => Ok(Box::new(RuleEliminateEvalScalar::new())),
+            RuleID::PushDownFilterProject => Ok(Box::new(RulePushDownFilterProject::new())),
+            RuleID::PushDownFilterEvalScalar => Ok(Box::new(RulePushDownFilterEvalScalar::new())),
+            RuleID::EliminateFilter => Ok(Box::new(RuleEliminateFilter::new())),
+            RuleID::EliminateProject => Ok(Box::new(RuleEliminateProject::new())),
+            RuleID::MergeProject => Ok(Box::new(RuleMergeProject::new())),
+            RuleID::MergeEvalScalar => Ok(Box::new(RuleMergeEvalScalar::new())),
+            RuleID::MergeFilter => Ok(Box::new(RuleMergeFilter::new())),
         }
     }
 }

query/src/sql/optimizer/rule/factory.rs

@@ -12,17 +12,22 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::fmt::Display;
+use std::fmt::Formatter;
+
 use common_exception::Result;
 
 use crate::sql::optimizer::SExpr;
 
 mod factory;
+mod rewrite;
 mod rule_implement_get;
 mod rule_implement_hash_join;
 mod rule_set;
 mod transform_state;
 
 pub use factory::RuleFactory;
+pub use rule_set::AppliedRules;
 pub use rule_set::RuleSet;
 pub use transform_state::TransformState;
 
@@ -31,32 +36,39 @@ pub type RulePtr = Box<dyn Rule>;
 pub trait Rule {
     fn id(&self) -> RuleID;
 
-    fn apply(&self, expression: &SExpr, state: &mut TransformState) -> Result<()>;
+    fn apply(&self, s_expr: &SExpr, state: &mut TransformState) -> Result<()>;
 
     fn pattern(&self) -> &SExpr;
 }
 
-#[derive(Copy, Clone, Eq, PartialEq, Hash)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 pub enum RuleID {
+    PushDownFilterProject,
+    PushDownFilterEvalScalar,
+    EliminateEvalScalar,
+    EliminateFilter,
+    EliminateProject,
+    MergeProject,
+    MergeEvalScalar,
+    MergeFilter,
+
     ImplementGet,
     ImplementHashJoin,
 }
 
-impl RuleID {
-    pub fn name(&self) -> &'static str {
-        match self {
-            RuleID::ImplementGet => "ImplementGet",
-            RuleID::ImplementHashJoin => "ImplementHashJoin",
-        }
-    }
-
-    /// Unique integral id
-    /// TODO: maybe use a macro like https://docs.rs/iota/0.2.2/iota/ to implement this?
-    #[allow(dead_code)]
-    pub fn uid(&self) -> u32 {
+impl Display for RuleID {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            RuleID::ImplementGet => 0,
-            RuleID::ImplementHashJoin => 1,
+            RuleID::ImplementGet => write!(f, "ImplementGet"),
+            RuleID::ImplementHashJoin => write!(f, "ImplementHashJoin"),
+            RuleID::PushDownFilterProject => write!(f, "PushDownFilterProject"),
+            RuleID::PushDownFilterEvalScalar => write!(f, "PushDownFilterEvalScalar"),
+            RuleID::EliminateEvalScalar => write!(f, "EliminateEvalScalar"),
+            RuleID::EliminateFilter => write!(f, "EliminateFilter"),
+            RuleID::EliminateProject => write!(f, "EliminateProject"),
+            RuleID::MergeProject => write!(f, "MergeProject"),
+            RuleID::MergeEvalScalar => write!(f, "MergeEvalScalar"),
+            RuleID::MergeFilter => write!(f, "MergeFilter"),
         }
     }
 }