tune plan space budget

Signed-off-by: Alex Chi Z <[email protected]>

Author: skyzh

optd-core/src/cascades/optimizer.rs

@@ -27,7 +27,8 @@ pub type RuleId = usize;
 
 #[derive(Default, Clone, Debug)]
 pub struct OptimizerContext {
-    pub budget_used: bool,
+    pub budget_used_logical: bool,
+    pub budget_used_all: bool,
     pub rules_applied: usize,
 }
 
@@ -42,13 +43,20 @@ pub struct OptimizerProperties {
     pub disable_pruning: bool,
 }
 
+#[derive(Debug, Default)]
+pub struct CascadesStats {
+    pub rule_match_count: HashMap<usize, usize>,
+    pub rule_total_bindings: HashMap<usize, usize>,
+}
+
 pub struct CascadesOptimizer<T: NodeType, M: Memo<T> = NaiveMemo<T>> {
     memo: M,
     pub(super) tasks: VecDeque<Box<dyn Task<T, M>>>,
     explored_group: HashSet<GroupId>,
     explored_expr: HashSet<ExprId>,
     fired_rules: HashMap<ExprId, HashSet<RuleId>>,
     rules: Arc<[Arc<dyn Rule<T, Self>>]>,
+    pub stats: CascadesStats,
     disabled_rules: HashSet<usize>,
     cost: Arc<dyn CostModel<T, M>>,
     property_builders: Arc<[Box<dyn LogicalPropertyBuilderAny<T>>]>,
@@ -123,6 +131,7 @@ impl<T: NodeType> CascadesOptimizer<T, NaiveMemo<T>> {
             property_builders,
             prop,
             disabled_rules: HashSet::new(),
+            stats: CascadesStats::default(),
         }
     }
 
@@ -250,39 +259,63 @@ impl<T: NodeType, M: Memo<T>> CascadesOptimizer<T, M> {
         self.tasks
             .push_back(Box::new(OptimizeGroupTask::new(group_id, None)));
         // get the task from the stack
-        self.ctx.budget_used = false;
+        self.ctx.budget_used_logical = false;
+        self.ctx.budget_used_all = false;
         let plan_space_begin = self.memo.estimated_plan_space();
         let mut iter = 0;
         while let Some(task) = self.tasks.pop_back() {
             let new_tasks = task.execute(self)?;
             self.tasks.extend(new_tasks);
             iter += 1;
-            if !self.ctx.budget_used {
+            if !self.ctx.budget_used_logical {
                 let plan_space = self.memo.estimated_plan_space();
                 if let Some(partial_explore_space) = self.prop.partial_explore_space {
                     if plan_space - plan_space_begin > partial_explore_space {
                         println!(
                             "plan space size budget used, not applying logical rules any more. current plan space: {}",
                             plan_space
                         );
-                        self.ctx.budget_used = true;
+                        self.ctx.budget_used_logical = true;
                         if self.prop.panic_on_budget {
                             panic!("plan space size budget used");
                         }
                     }
-                } else if let Some(partial_explore_iter) = self.prop.partial_explore_iter {
+                }
+             }
+             if !self.ctx.budget_used_all {
+                if let Some(partial_explore_iter) = self.prop.partial_explore_iter {
                     if iter >= partial_explore_iter {
                         println!(
-                            "plan explore iter budget used, not applying logical rules any more. current plan space: {}",
-                            plan_space
+                            "plan explore iter budget used, not applying physical/logical rules any more if there's no winner. current iter: {}",
+                            iter
                         );
-                        self.ctx.budget_used = true;
+                        self.ctx.budget_used_all = true;
                         if self.prop.panic_on_budget {
                             panic!("plan space size budget used");
                         }
                     }
                 }
             }
+            if iter > 100000 && iter % 10000 == 0 {
+                println!("iter={}", iter);
+                println!("plan_space={}", self.memo.estimated_plan_space());
+                for (id, rule) in self.rules.iter().enumerate() {
+                    println!(
+                        "{}: matched={}, bindings={}",
+                        rule.name(),
+                        self.stats
+                            .rule_match_count
+                            .get(&id)
+                            .copied()
+                            .unwrap_or_default(),
+                        self.stats
+                            .rule_total_bindings
+                            .get(&id)
+                            .copied()
+                            .unwrap_or_default()
+                    );
+                }
+            }
         }
         Ok(())
     }

optd-core/src/cascades/tasks/apply_rule.rs

@@ -177,7 +177,11 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for ApplyRuleTask {
         let group_id = optimizer.get_group_id(self.expr_id);
         let mut tasks = vec![];
         let binding_exprs = match_and_pick_expr(rule.matcher(), self.expr_id, optimizer);
+        if !binding_exprs.is_empty() {
+            *optimizer.stats.rule_match_count.entry(self.rule_id).or_default() += 1;
+        }
         for binding in binding_exprs {
+            *optimizer.stats.rule_total_bindings.entry(self.rule_id).or_default() += 1;
             trace!(event = "before_apply_rule", task = "apply_rule", input_binding=%binding);
             let applied = rule.apply(optimizer, binding);
 

optd-core/src/cascades/tasks/optimize_expression.rs

@@ -42,6 +42,7 @@ fn top_matches<T: NodeType>(matcher: &RuleMatcher<T>, match_typ: T) -> bool {
 impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeExpressionTask {
     fn execute(&self, optimizer: &mut CascadesOptimizer<T, M>) -> Result<Vec<Box<dyn Task<T, M>>>> {
         let expr = optimizer.get_expr_memoed(self.expr_id);
+        let group_id = optimizer.get_group_id(self.expr_id);
         trace!(event = "task_begin", task = "optimize_expr", expr_id = %self.expr_id, expr = %expr);
         let mut tasks = vec![];
         for (rule_id, rule) in optimizer.rules().iter().enumerate() {
@@ -53,17 +54,27 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeExpressionTask {
                 continue;
             }
             // Skip transformation rules when budget is used
-            if optimizer.ctx.budget_used && !rule.is_impl_rule() {
+            if (optimizer.ctx.budget_used_logical || optimizer.ctx.budget_used_all)
+                && !rule.is_impl_rule()
+            {
                 continue;
             }
+            if optimizer.ctx.budget_used_all
+                && optimizer.get_group_info(group_id).winner.has_full_winner()
+            {
+                break;
+            }
             if top_matches(rule.matcher(), expr.typ.clone()) {
-                tasks.push(
-                    Box::new(ApplyRuleTask::new(rule_id, self.expr_id, self.exploring, self.upper_bound))
-                        as Box<dyn Task<T, M>>,
-                );
+                tasks.push(Box::new(ApplyRuleTask::new(
+                    rule_id,
+                    self.expr_id,
+                    self.exploring,
+                    self.upper_bound,
+                )) as Box<dyn Task<T, M>>);
                 for &input_group_id in &expr.children {
                     tasks.push(
-                        Box::new(ExploreGroupTask::new(input_group_id, self.upper_bound)) as Box<dyn Task<T, M>>
+                        Box::new(ExploreGroupTask::new(input_group_id, self.upper_bound))
+                            as Box<dyn Task<T, M>>,
                     );
                 }
             }

optd-datafusion-repr/src/lib.rs

@@ -102,7 +102,6 @@ impl DatafusionOptimizer {
             rule_wrappers.push(rule);
         }
         rule_wrappers.push(Arc::new(rules::FilterProjectTransposeRule::new()));
-        rule_wrappers.push(Arc::new(rules::FilterCrossJoinTransposeRule::new()));
         rule_wrappers.push(Arc::new(rules::FilterInnerJoinTransposeRule::new()));
         rule_wrappers.push(Arc::new(rules::FilterSortTransposeRule::new()));
         rule_wrappers.push(Arc::new(rules::FilterAggTransposeRule::new()));
@@ -150,7 +149,7 @@ impl DatafusionOptimizer {
                 OptimizerProperties {
                     panic_on_budget: false,
                     partial_explore_iter: Some(1 << 20),
-                    partial_explore_space: Some(1 << 10),
+                    partial_explore_space: None, // remove this in the future
                     disable_pruning: false,
                 },
             ),