Merge pull request #9571 from leiysky/join-order

chore(planner): Improve join reorder algorithm

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

@@ -16,10 +16,17 @@ use crate::optimizer::RuleID;
 use crate::optimizer::RuleSet;
 
 pub fn get_explore_rule_set() -> RuleSet {
+    join_rule_set_rs_b2()
+}
+
+/// Get rule set of join order RS-B2.
+/// Read paper "The Complexity of Transformation-Based Join Enumeration" for more details.
+fn join_rule_set_rs_b2() -> RuleSet {
     RuleSet::create_with_ids(vec![
         RuleID::CommuteJoin,
         RuleID::LeftAssociateJoin,
         RuleID::RightAssociateJoin,
+        RuleID::ExchangeJoin,
     ])
     .unwrap()
 }

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

@@ -21,12 +21,17 @@ use super::CascadesOptimizer;
 
 pub struct Scheduler {
     task_queue: VecDeque<Task>,
+
+    /// A counter to track the number of tasks
+    /// that have been scheduled.
+    scheduled_task_count: u64,
 }
 
 impl Scheduler {
     pub fn new() -> Self {
         Self {
             task_queue: Default::default(),
+            scheduled_task_count: 0,
         }
     }
 
@@ -38,8 +43,16 @@ impl Scheduler {
                 continue;
             }
             task.execute(optimizer, self)?;
+
+            // Update the counter
+            self.scheduled_task_count += 1;
         }
 
+        tracing::debug!(
+            "CascadesOptimizer: scheduled {} tasks",
+            self.scheduled_task_count
+        );
+
         Ok(())
     }
 

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

@@ -117,7 +117,7 @@ impl HeuristicOptimizer {
         for rule in rule_list.iter() {
             let mut state = TransformResult::new();
             if s_expr.match_pattern(rule.pattern()) && !s_expr.applied_rule(&rule.id()) {
-                s_expr.apply_rule(&rule.id());
+                s_expr.set_applied_rule(&rule.id());
                 rule.apply(&s_expr, &mut state)?;
                 if !state.results().is_empty() {
                     // Recursive optimize the result

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

@@ -18,6 +18,7 @@ use common_exception::Result;
 use super::group::Group;
 use crate::optimizer::memo::Memo;
 use crate::optimizer::pattern_extractor::PatternExtractor;
+use crate::optimizer::rule::AppliedRules;
 use crate::optimizer::rule::RulePtr;
 use crate::optimizer::rule::TransformResult;
 use crate::optimizer::SExpr;
@@ -36,6 +37,9 @@ pub struct MExpr {
 
     pub plan: RelOperator,
     pub children: Vec<IndexType>,
+
+    // Disable rules for current `MExpr`
+    pub applied_rules: AppliedRules,
 }
 
 impl MExpr {
@@ -44,12 +48,14 @@ impl MExpr {
         index: IndexType,
         plan: RelOperator,
         children: Vec<IndexType>,
+        applied_rules: AppliedRules,
     ) -> Self {
         MExpr {
             group_index,
             plan,
             children,
             index,
+            applied_rules,
         }
     }
 
@@ -87,6 +93,10 @@ impl MExpr {
         rule: &RulePtr,
         transform_state: &mut TransformResult,
     ) -> Result<()> {
+        if self.applied_rules.get(&rule.id()) {
+            return Ok(());
+        }
+
         let mut extractor = PatternExtractor::create();
         let exprs = extractor.extract(memo, self, rule.pattern())?;
 

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

@@ -109,6 +109,7 @@ impl Memo {
             self.group(group_index)?.num_exprs(),
             s_expr.plan,
             children_group,
+            s_expr.applied_rules,
         );
         self.insert_m_expr(group_index, m_expr)?;
 

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

@@ -35,6 +35,9 @@ use crate::optimizer::rule::rewrite::RulePushDownLimitSort;
 use crate::optimizer::rule::rewrite::RulePushDownLimitUnion;
 use crate::optimizer::rule::rewrite::RulePushDownSortScan;
 use crate::optimizer::rule::rewrite::RuleSplitAggregate;
+use crate::optimizer::rule::transform::RuleExchangeJoin;
+use crate::optimizer::rule::transform::RuleLeftExchangeJoin;
+use crate::optimizer::rule::transform::RuleRightExchangeJoin;
 use crate::optimizer::rule::RuleID;
 use crate::optimizer::rule::RulePtr;
 
@@ -70,6 +73,9 @@ impl RuleFactory {
             RuleID::CommuteJoin => Ok(Box::new(RuleCommuteJoin::new())),
             RuleID::LeftAssociateJoin => Ok(Box::new(RuleLeftAssociateJoin::new())),
             RuleID::RightAssociateJoin => Ok(Box::new(RuleRightAssociateJoin::new())),
+            RuleID::LeftExchangeJoin => Ok(Box::new(RuleLeftExchangeJoin::new())),
+            RuleID::RightExchangeJoin => Ok(Box::new(RuleRightExchangeJoin::new())),
+            RuleID::ExchangeJoin => Ok(Box::new(RuleExchangeJoin::new())),
         }
     }
 }

src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_join.rs

@@ -308,7 +308,7 @@ impl Rule for RulePushDownFilterJoin {
         if !need_push {
             return Ok(());
         }
-        result.apply_rule(&self.id);
+        result.set_applied_rule(&self.id);
         state.add_result(result);
 
         Ok(())

src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_limit_union.rs

@@ -100,7 +100,7 @@ impl Rule for RulePushDownLimitUnion {
         let mut result = SExpr::create_binary(union.into(), union_left_child, union_right_child);
         // Add original limit to top
         result = SExpr::create_unary(limit.into(), result);
-        result.apply_rule(&self.id);
+        result.set_applied_rule(&self.id);
         state.add_result(result);
 
         Ok(())

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

@@ -56,6 +56,9 @@ pub enum RuleID {
     CommuteJoin,
     LeftAssociateJoin,
     RightAssociateJoin,
+    LeftExchangeJoin,
+    RightExchangeJoin,
+    ExchangeJoin,
 }
 
 impl Display for RuleID {
@@ -83,6 +86,9 @@ impl Display for RuleID {
             RuleID::CommuteJoin => write!(f, "CommuteJoin"),
             RuleID::LeftAssociateJoin => write!(f, "LeftAssociateJoin"),
             RuleID::RightAssociateJoin => write!(f, "RightAssociateJoin"),
+            RuleID::LeftExchangeJoin => write!(f, "LeftExchangeJoin"),
+            RuleID::RightExchangeJoin => write!(f, "RightExchangeJoin"),
+            RuleID::ExchangeJoin => write!(f, "ExchangeJoin"),
         }
     }
 }

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

@@ -13,10 +13,16 @@
 // limitations under the License.
 
 mod rule_commute_join;
+mod rule_exchange_join;
 mod rule_left_associate_join;
+mod rule_left_exchange_join;
 mod rule_right_associate_join;
+mod rule_right_exchange_join;
 mod util;
 
 pub use rule_commute_join::RuleCommuteJoin;
+pub use rule_exchange_join::RuleExchangeJoin;
 pub use rule_left_associate_join::RuleLeftAssociateJoin;
+pub use rule_left_exchange_join::RuleLeftExchangeJoin;
 pub use rule_right_associate_join::RuleRightAssociateJoin;
+pub use rule_right_exchange_join::RuleRightExchangeJoin;

src/query/sql/src/planner/optimizer/rule/transform/rule_commute_join.rs

@@ -78,8 +78,17 @@ impl Rule for RuleCommuteJoin {
                 (join.left_conditions, join.right_conditions) =
                     (join.right_conditions, join.left_conditions);
                 join.join_type = join.join_type.opposite();
-                let result =
+                let mut result =
                     SExpr::create_binary(join.into(), right_child.clone(), left_child.clone());
+
+                // Disable the following rules for the generated expression
+                result.set_applied_rule(&RuleID::CommuteJoin);
+                result.set_applied_rule(&RuleID::LeftAssociateJoin);
+                result.set_applied_rule(&RuleID::LeftExchangeJoin);
+                result.set_applied_rule(&RuleID::RightAssociateJoin);
+                result.set_applied_rule(&RuleID::RightExchangeJoin);
+                result.set_applied_rule(&RuleID::ExchangeJoin);
+
                 state.add_result(result);
             }
             _ => {}