fix(planner): Correct the mapping relationship between RuleID and DEFAULT_REWRITE_RULES

src/query/sql/Cargo.toml

@@ -48,8 +48,6 @@ dashmap = "5.4"
 educe = "0.4"
 globiter = "0.1"
 itertools = "0.10.5"
-num-derive = "0.3.3"
-num-traits = "0.2.15"
 once_cell = "1.15.0"
 opendal = { workspace = true }
 ordered-float = { workspace = true }

src/query/sql/src/planner/optimizer/cascades/cascade.rs

@@ -18,8 +18,9 @@ use std::sync::Arc;
 use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
 use common_exception::Result;
+use roaring::RoaringBitmap;
 
-use super::explore_rules::get_explore_rule_set;
+use super::explore_rules::calc_explore_rule_set;
 use crate::optimizer::cascades::scheduler::Scheduler;
 use crate::optimizer::cascades::tasks::OptimizeGroupTask;
 use crate::optimizer::cascades::tasks::Task;
@@ -29,36 +30,36 @@ use crate::optimizer::cost::DefaultCostModel;
 use crate::optimizer::format::display_memo;
 use crate::optimizer::memo::Memo;
 use crate::optimizer::rule::TransformResult;
-use crate::optimizer::RuleSet;
 use crate::optimizer::SExpr;
 use crate::IndexType;
-use crate::MetadataRef;
 
 /// A cascades-style search engine to enumerate possible alternations of a relational expression and
 /// find the optimal one.
 pub struct CascadesOptimizer {
-    pub(crate) memo: Memo,
-    pub(crate) cost_model: Box<dyn CostModel>,
+    pub memo: Memo,
+    pub explore_rule_set: roaring::RoaringBitmap,
+
+    pub cost_model: Box<dyn CostModel>,
+
     /// group index -> best cost context
-    pub(crate) best_cost_map: HashMap<IndexType, CostContext>,
-    pub(crate) explore_rule_set: RuleSet,
-    pub(crate) metadata: MetadataRef,
+    pub best_cost_map: HashMap<IndexType, CostContext>,
+    _ctx: Arc<dyn TableContext>,
 }
 
 impl CascadesOptimizer {
-    pub fn create(ctx: Arc<dyn TableContext>, metadata: MetadataRef) -> Result<Self> {
+    pub fn create(ctx: Arc<dyn TableContext>) -> Result<Self> {
         let enable_bushy_join = ctx.get_settings().get_enable_bushy_join()? != 0;
         let explore_rule_set = if ctx.get_settings().get_enable_cbo()? {
-            get_explore_rule_set(enable_bushy_join)
+            calc_explore_rule_set(enable_bushy_join)
         } else {
-            RuleSet::create()
+            RoaringBitmap::new()
         };
         Ok(CascadesOptimizer {
             memo: Memo::create(),
             cost_model: Box::new(DefaultCostModel),
             best_cost_map: HashMap::new(),
+            _ctx: ctx,
             explore_rule_set,
-            metadata,
         })
     }
 

src/query/sql/src/planner/optimizer/cascades/explore_rules.rs

@@ -12,30 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use roaring::RoaringBitmap;
+
 use crate::optimizer::RuleID;
-use crate::optimizer::RuleSet;
 
-pub fn get_explore_rule_set(enable_bushy_join: bool) -> RuleSet {
+pub fn calc_explore_rule_set(enable_bushy_join: bool) -> RoaringBitmap {
     if enable_bushy_join {
-        rule_set_rs_b2()
+        calc_join_rule_set_rs_b2()
     } else {
-        rule_set_rs_l1()
+        calc_join_rule_set_rs_l1()
     }
 }
 
 /// Get rule set of join order RS-B2, which may generate bushy trees.
 /// Read paper "The Complexity of Transformation-Based Join Enumeration" for more details.
-fn rule_set_rs_b2() -> RuleSet {
-    RuleSet::create_with_ids(vec![
-        RuleID::CommuteJoin,
-        RuleID::LeftAssociateJoin,
-        RuleID::RightAssociateJoin,
-        RuleID::ExchangeJoin,
-    ])
+fn calc_join_rule_set_rs_b2() -> RoaringBitmap {
+    (RuleID::CommuteJoin as u32..RuleID::CommuteJoinBaseTable as u32).collect::<RoaringBitmap>()
 }
 
 /// Get rule set of join order RS-L1, which will only generate left-deep trees.
 /// Read paper "The Complexity of Transformation-Based Join Enumeration" for more details.
-fn rule_set_rs_l1() -> RuleSet {
-    RuleSet::create_with_ids(vec![RuleID::CommuteJoinBaseTable, RuleID::LeftExchangeJoin])
+fn calc_join_rule_set_rs_l1() -> RoaringBitmap {
+    (RuleID::CommuteJoinBaseTable as u32..RuleID::RightExchangeJoin as u32)
+        .collect::<RoaringBitmap>()
 }

src/query/sql/src/planner/optimizer/cascades/tasks/apply_rule.rs

@@ -58,7 +58,7 @@ impl ApplyRuleTask {
         let group = optimizer.memo.group(self.target_group_index)?;
         let m_expr = group.m_expr(self.m_expr_index)?;
         let mut state = TransformResult::new();
-        let rule = RuleFactory::create_rule(self.rule_id, optimizer.metadata.clone())?;
+        let rule = RuleFactory::create().create_rule(self.rule_id, None)?;
         m_expr.apply_rule(&optimizer.memo, &rule, &mut state)?;
         optimizer.insert_from_transform_state(self.target_group_index, state)?;
 

src/query/sql/src/planner/optimizer/cascades/tasks/explore_expr.rs

@@ -22,6 +22,10 @@ use super::Task;
 use crate::optimizer::cascades::scheduler::Scheduler;
 use crate::optimizer::cascades::tasks::SharedCounter;
 use crate::optimizer::cascades::CascadesOptimizer;
+use crate::optimizer::RuleID;
+use crate::optimizer::RULE_FACTORY;
+use crate::plans::Operator;
+use crate::plans::RelOperator;
 use crate::IndexType;
 
 #[derive(Clone, Copy, Debug)]
@@ -140,6 +144,18 @@ impl ExploreExprTask {
         }
     }
 
+    fn calc_operator_rule_set(
+        &self,
+        optimizer: &CascadesOptimizer,
+        operator: &RelOperator,
+    ) -> roaring::RoaringBitmap {
+        unsafe {
+            operator.exploration_candidate_rules()
+                & (&RULE_FACTORY.exploration_rules)
+                & (&optimizer.explore_rule_set)
+        }
+    }
+
     fn explore_self(
         &mut self,
         optimizer: &mut CascadesOptimizer,
@@ -149,11 +165,11 @@ impl ExploreExprTask {
             .memo
             .group(self.group_index)?
             .m_expr(self.m_expr_index)?;
-        let rule_set = &optimizer.explore_rule_set;
+        let rule_set = self.calc_operator_rule_set(optimizer, &m_expr.plan);
 
         for rule_id in rule_set.iter() {
             let apply_rule_task = ApplyRuleTask::with_parent(
-                rule_id,
+                unsafe { std::mem::transmute::<u8, RuleID>(rule_id as u8) },
                 m_expr.group_index,
                 m_expr.index,
                 &self.ref_count,

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

@@ -20,37 +20,45 @@ use once_cell::sync::Lazy;
 
 use super::prune_unused_columns::UnusedColumnPruner;
 use crate::optimizer::heuristic::decorrelate::decorrelate_subquery;
+use crate::optimizer::rule::RulePtr;
 use crate::optimizer::rule::TransformResult;
 use crate::optimizer::ColumnSet;
-use crate::optimizer::RuleFactory;
 use crate::optimizer::RuleID;
 use crate::optimizer::SExpr;
+use crate::optimizer::RULE_FACTORY;
+use crate::plans::Operator;
+use crate::plans::RelOperator;
 use crate::BindContext;
 use crate::MetadataRef;
 
 pub static DEFAULT_REWRITE_RULES: Lazy<Vec<RuleID>> = Lazy::new(|| {
     vec![
+        // Filter
         RuleID::NormalizeDisjunctiveFilter,
         RuleID::NormalizeScalarFilter,
         RuleID::EliminateFilter,
-        RuleID::EliminateEvalScalar,
         RuleID::MergeFilter,
-        RuleID::MergeEvalScalar,
         RuleID::PushDownFilterUnion,
         RuleID::PushDownFilterAggregate,
+        RuleID::PushDownFilterSort,
+        RuleID::PushDownFilterEvalScalar,
+        RuleID::PushDownFilterJoin,
+        RuleID::PushDownFilterScan,
+        // EvalScalar
+        RuleID::EliminateEvalScalar,
+        RuleID::MergeEvalScalar,
+        // Limit
         RuleID::PushDownLimitUnion,
         RuleID::RulePushDownLimitExpression,
         RuleID::PushDownLimitSort,
         RuleID::PushDownLimitAggregate,
         RuleID::PushDownLimitOuterJoin,
         RuleID::PushDownLimitScan,
-        RuleID::PushDownFilterSort,
-        RuleID::PushDownFilterEvalScalar,
-        RuleID::PushDownFilterJoin,
+        // Agg
         RuleID::FoldCountAggregate,
         RuleID::SplitAggregate,
-        RuleID::PushDownFilterScan,
         RuleID::PushDownPrewhere, /* PushDownPrwhere should be after all rules except PushDownFilterScan */
+        // Sort
         RuleID::PushDownSortScan, // PushDownFilterScan should be after PushDownPrewhere
     ]
 });
@@ -113,13 +121,27 @@ impl HeuristicOptimizer {
         Ok(result)
     }
 
+    fn calc_operator_rule_set(&self, operator: &RelOperator) -> roaring::RoaringBitmap {
+        unsafe { operator.transformation_candidate_rules() & (&RULE_FACTORY.transformation_rules) }
+    }
+
+    fn get_rule(&self, rule_id: u32) -> Result<RulePtr> {
+        unsafe {
+            RULE_FACTORY.create_rule(
+                DEFAULT_REWRITE_RULES[rule_id as usize],
+                Some(self.metadata.clone()),
+            )
+        }
+    }
+
     /// Try to apply the rules to the expression.
     /// Return the final result that no rule can be applied.
     fn apply_transform_rules(&self, s_expr: &SExpr) -> Result<SExpr> {
         let mut s_expr = s_expr.clone();
+        let rule_set = self.calc_operator_rule_set(&s_expr.plan);
 
-        for rule_id in DEFAULT_REWRITE_RULES.iter() {
-            let rule = RuleFactory::create_rule(*rule_id, self.metadata.clone())?;
+        for rule_id in rule_set.iter() {
+            let rule = self.get_rule(rule_id)?;
             let mut state = TransformResult::new();
             if s_expr.match_pattern(rule.pattern()) && !s_expr.applied_rule(&rule.id()) {
                 s_expr.set_applied_rule(&rule.id());

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

@@ -43,4 +43,5 @@ pub use rule::try_push_down_filter_join;
 pub use rule::RuleFactory;
 pub use rule::RuleID;
 pub use rule::RuleSet;
+pub use rule::RULE_FACTORY;
 pub use s_expr::SExpr;

src/query/sql/src/planner/optimizer/optimizer.rs

@@ -147,9 +147,10 @@ pub fn optimize_query(
 ) -> Result<SExpr> {
     let contains_local_table_scan = contains_local_table_scan(&s_expr, &metadata);
 
-    let mut heuristic = HeuristicOptimizer::new(ctx.clone(), bind_context, metadata.clone());
+    let mut heuristic = HeuristicOptimizer::new(ctx.clone(), bind_context, metadata);
     let mut result = heuristic.optimize(s_expr)?;
-    let mut cascades = CascadesOptimizer::create(ctx.clone(), metadata)?;
+
+    let mut cascades = CascadesOptimizer::create(ctx.clone())?;
     result = cascades.optimize(result)?;
     // So far, we don't have ability to execute distributed query
     // with reading data from local tales(e.g. system tables).
@@ -176,10 +177,10 @@ fn get_optimized_memo(
     metadata: MetadataRef,
     bind_context: Box<BindContext>,
 ) -> Result<(Memo, HashMap<IndexType, CostContext>)> {
-    let mut heuristic = HeuristicOptimizer::new(ctx.clone(), bind_context, metadata.clone());
+    let mut heuristic = HeuristicOptimizer::new(ctx.clone(), bind_context, metadata);
     let result = heuristic.optimize(s_expr)?;
 
-    let mut cascades = CascadesOptimizer::create(ctx, metadata)?;
+    let mut cascades = CascadesOptimizer::create(ctx)?;
     cascades.optimize(result)?;
     Ok((cascades.memo, cascades.best_cost_map))
 }

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

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 use common_exception::Result;
+use once_cell::sync::Lazy;
+use roaring::RoaringBitmap;
 
 use super::rewrite::RuleEliminateEvalScalar;
 use super::rewrite::RuleFoldCountAggregate;
@@ -47,16 +49,36 @@ use crate::optimizer::rule::RuleID;
 use crate::optimizer::rule::RulePtr;
 use crate::MetadataRef;
 
-pub struct RuleFactory;
+// read only, so thread safe
+pub static mut RULE_FACTORY: Lazy<RuleFactory> = Lazy::new(RuleFactory::create);
+
+pub struct RuleFactory {
+    pub transformation_rules: roaring::RoaringBitmap,
+    pub exploration_rules: roaring::RoaringBitmap,
+}
 
 impl RuleFactory {
-    pub fn create_rule(id: RuleID, metadata: MetadataRef) -> Result<RulePtr> {
+    pub fn create() -> Self {
+        RuleFactory {
+            transformation_rules: (RuleID::NormalizeScalarFilter as u32
+                ..RuleID::CommuteJoin as u32)
+                .collect::<RoaringBitmap>(),
+            exploration_rules: (RuleID::CommuteJoin as u32..(RuleID::RightExchangeJoin as u32) + 1)
+                .collect::<RoaringBitmap>(),
+        }
+    }
+
+    pub fn create_rule(&self, id: RuleID, metadata: Option<MetadataRef>) -> Result<RulePtr> {
         match id {
             RuleID::EliminateEvalScalar => Ok(Box::new(RuleEliminateEvalScalar::new())),
             RuleID::PushDownFilterUnion => Ok(Box::new(RulePushDownFilterUnion::new())),
             RuleID::PushDownFilterEvalScalar => Ok(Box::new(RulePushDownFilterEvalScalar::new())),
-            RuleID::PushDownFilterJoin => Ok(Box::new(RulePushDownFilterJoin::new(metadata))),
-            RuleID::PushDownFilterScan => Ok(Box::new(RulePushDownFilterScan::new(metadata))),
+            RuleID::PushDownFilterJoin => {
+                Ok(Box::new(RulePushDownFilterJoin::new(metadata.unwrap())))
+            }
+            RuleID::PushDownFilterScan => {
+                Ok(Box::new(RulePushDownFilterScan::new(metadata.unwrap())))
+            }
             RuleID::PushDownFilterSort => Ok(Box::new(RulePushDownFilterSort::new())),
             RuleID::PushDownLimitUnion => Ok(Box::new(RulePushDownLimitUnion::new())),
             RuleID::PushDownLimitScan => Ok(Box::new(RulePushDownLimitScan::new())),
@@ -82,7 +104,7 @@ impl RuleFactory {
             RuleID::LeftExchangeJoin => Ok(Box::new(RuleLeftExchangeJoin::new())),
             RuleID::RightExchangeJoin => Ok(Box::new(RuleRightExchangeJoin::new())),
             RuleID::ExchangeJoin => Ok(Box::new(RuleExchangeJoin::new())),
-            RuleID::PushDownPrewhere => Ok(Box::new(RulePushDownPrewhere::new(metadata))),
+            RuleID::PushDownPrewhere => Ok(Box::new(RulePushDownPrewhere::new(metadata.unwrap()))),
         }
     }
 }

src/query/sql/src/planner/optimizer/rule/mod.rs

@@ -21,6 +21,7 @@ mod transform;
 mod transform_result;
 
 pub use factory::RuleFactory;
+pub use factory::RULE_FACTORY;
 pub use rewrite::try_push_down_filter_join;
 pub use rule::Rule;
 pub use rule::RuleID;