refactor(core): prune based on upper bound (#265)

skyzh · skyzh · commit ef3beaee7255 · 2024-12-17T21:34:06.000-05:00
Signed-off-by: Alex Chi Z &lt;iskyzh@gmail.com&gt;

tune plan space budget

Signed-off-by: Alex Chi Z &lt;iskyzh@gmail.com&gt;
diff --git a/optd-core/src/cascades/memo.rs b/optd-core/src/cascades/memo.rs
@@ -352,7 +352,7 @@ impl<T: NodeType> NaiveMemo<T> {
     }
 
     fn verify_integrity(&self) {
-        if cfg!(debug_assertions) {
+        if false {
             let num_of_exprs = self.expr_id_to_expr_node.len();
             assert_eq!(num_of_exprs, self.expr_node_to_expr_id.len());
             assert_eq!(num_of_exprs, self.expr_id_to_group_id.len());
diff --git a/optd-core/src/cascades/optimizer.rs b/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(),
         }
     }
 
@@ -248,41 +257,65 @@ impl<T: NodeType, M: Memo<T>> CascadesOptimizer<T, M> {
     fn fire_optimize_tasks(&mut self, group_id: GroupId) -> Result<()> {
         trace!(event = "fire_optimize_tasks", root_group_id = %group_id);
         self.tasks
-            .push_back(Box::new(OptimizeGroupTask::new(group_id)));
+            .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(())
     }
diff --git a/optd-core/src/cascades/tasks/apply_rule.rs b/optd-core/src/cascades/tasks/apply_rule.rs
@@ -21,14 +21,21 @@ pub struct ApplyRuleTask {
     rule_id: RuleId,
     expr_id: ExprId,
     exploring: bool,
+    upper_bound: Option<f64>,
 }
 
 impl ApplyRuleTask {
-    pub fn new(rule_id: RuleId, expr_id: ExprId, exploring: bool) -> Self {
+    pub fn new(
+        rule_id: RuleId,
+        expr_id: ExprId,
+        exploring: bool,
+        upper_bound: Option<f64>,
+    ) -> Self {
         Self {
             rule_id,
             expr_id,
             exploring,
+            upper_bound,
         }
     }
 }
@@ -170,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);
 
@@ -181,13 +192,14 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for ApplyRuleTask {
                     let typ = expr.unwrap_typ();
                     if typ.is_logical() {
                         tasks.push(
-                            Box::new(OptimizeExpressionTask::new(expr_id, self.exploring))
+                            Box::new(OptimizeExpressionTask::new(expr_id, self.exploring, self.upper_bound))
                                 as Box<dyn Task<T, M>>,
                         );
                     } else {
                         tasks.push(Box::new(OptimizeInputsTask::new(
                             expr_id,
                             !optimizer.prop.disable_pruning,
+                            self.upper_bound
                         )) as Box<dyn Task<T, M>>);
                     }
                     optimizer.unmark_expr_explored(expr_id);
diff --git a/optd-core/src/cascades/tasks/explore_group.rs b/optd-core/src/cascades/tasks/explore_group.rs
@@ -14,11 +14,15 @@ use crate::nodes::NodeType;
 
 pub struct ExploreGroupTask {
     group_id: GroupId,
+    upper_bound: Option<f64>,
 }
 
 impl ExploreGroupTask {
-    pub fn new(group_id: GroupId) -> Self {
-        Self { group_id }
+    pub fn new(group_id: GroupId, upper_bound: Option<f64>) -> Self {
+        Self {
+            group_id,
+            upper_bound,
+        }
     }
 }
 
@@ -36,7 +40,7 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for ExploreGroupTask {
             let typ = optimizer.get_expr_memoed(expr).typ.clone();
             if typ.is_logical() {
                 tasks
-                    .push(Box::new(OptimizeExpressionTask::new(expr, true)) as Box<dyn Task<T, M>>);
+                    .push(Box::new(OptimizeExpressionTask::new(expr, true, self.upper_bound)) as Box<dyn Task<T, M>>);
             }
         }
         optimizer.mark_group_explored(self.group_id);
diff --git a/optd-core/src/cascades/tasks/optimize_expression.rs b/optd-core/src/cascades/tasks/optimize_expression.rs
@@ -16,11 +16,16 @@ use crate::rules::RuleMatcher;
 pub struct OptimizeExpressionTask {
     expr_id: ExprId,
     exploring: bool,
+    upper_bound: Option<f64>,
 }
 
 impl OptimizeExpressionTask {
-    pub fn new(expr_id: ExprId, exploring: bool) -> Self {
-        Self { expr_id, exploring }
+    pub fn new(expr_id: ExprId, exploring: bool, upper_bound: Option<f64>) -> Self {
+        Self {
+            expr_id,
+            exploring,
+            upper_bound,
+        }
     }
 }
 
@@ -37,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() {
@@ -48,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))
-                        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)) as Box<dyn Task<T, M>>
+                        Box::new(ExploreGroupTask::new(input_group_id, self.upper_bound))
+                            as Box<dyn Task<T, M>>,
                     );
                 }
             }
diff --git a/optd-core/src/cascades/tasks/optimize_group.rs b/optd-core/src/cascades/tasks/optimize_group.rs
@@ -15,11 +15,15 @@ use crate::nodes::NodeType;
 
 pub struct OptimizeGroupTask {
     group_id: GroupId,
+    upper_bound: Option<f64>,
 }
 
 impl OptimizeGroupTask {
-    pub fn new(group_id: GroupId) -> Self {
-        Self { group_id }
+    pub fn new(group_id: GroupId, upper_bound: Option<f64>) -> Self {
+        Self {
+            group_id,
+            upper_bound,
+        }
     }
 }
 
@@ -37,7 +41,7 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeGroupTask {
         for &expr in &exprs {
             let typ = optimizer.get_expr_memoed(expr).typ.clone();
             if typ.is_logical() {
-                tasks.push(Box::new(OptimizeExpressionTask::new(expr, false)) as Box<dyn Task<T, M>>);
+                tasks.push(Box::new(OptimizeExpressionTask::new(expr, false, self.upper_bound)) as Box<dyn Task<T, M>>);
             }
         }
         for &expr in &exprs {
@@ -46,6 +50,7 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeGroupTask {
                 tasks.push(Box::new(OptimizeInputsTask::new(
                     expr,
                     !optimizer.prop.disable_pruning,
+                    self.upper_bound
                 )) as Box<dyn Task<T, M>>);
             }
         }
diff --git a/optd-core/src/cascades/tasks/optimize_inputs.rs b/optd-core/src/cascades/tasks/optimize_inputs.rs
diff --git a/optd-datafusion-repr/src/lib.rs b/optd-datafusion-repr/src/lib.rs

Original file line number	Diff line number	Diff line change
`@@ -352,7 +352,7 @@ impl<T: NodeType> NaiveMemo<T> {`
`352`	`352`	`}`
`353`	`353`
`354`	`354`	`fn verify_integrity(&self) {`
`355`		`- if cfg!(debug_assertions) {`
	`355`	`+ if false {`
`356`	`356`	`let num_of_exprs = self.expr_id_to_expr_node.len();`
`357`	`357`	`assert_eq!(num_of_exprs, self.expr_node_to_expr_id.len());`
`358`	`358`	`assert_eq!(num_of_exprs, self.expr_id_to_group_id.len());`
Original file line number	Diff line number	Diff line change
`@@ -14,11 +14,15 @@ use crate::nodes::NodeType;`
`14`	`14`
`15`	`15`	`pub struct ExploreGroupTask {`
`16`	`16`	`group_id: GroupId,`
	`17`	`+ upper_bound: Option<f64>,`
`17`	`18`	`}`
`18`	`19`
`19`	`20`	`impl ExploreGroupTask {`
`20`		`- pub fn new(group_id: GroupId) -> Self {`
`21`		`- Self { group_id }`
	`21`	`+ pub fn new(group_id: GroupId, upper_bound: Option<f64>) -> Self {`
	`22`	`+ Self {`
	`23`	`+ group_id,`
	`24`	`+ upper_bound,`
	`25`	`+ }`
`22`	`26`	`}`
`23`	`27`	`}`
`24`	`28`
`@@ -36,7 +40,7 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for ExploreGroupTask {`
`36`	`40`	`let typ = optimizer.get_expr_memoed(expr).typ.clone();`
`37`	`41`	`if typ.is_logical() {`
`38`	`42`	`tasks`
`39`		`- .push(Box::new(OptimizeExpressionTask::new(expr, true)) as Box<dyn Task<T, M>>);`
	`43`	`+ .push(Box::new(OptimizeExpressionTask::new(expr, true, self.upper_bound)) as Box<dyn Task<T, M>>);`
`40`	`44`	`}`
`41`	`45`	`}`
`42`	`46`	`optimizer.mark_group_explored(self.group_id);`
Original file line number	Diff line number	Diff line change
`@@ -15,11 +15,15 @@ use crate::nodes::NodeType;`
`15`	`15`
`16`	`16`	`pub struct OptimizeGroupTask {`
`17`	`17`	`group_id: GroupId,`
	`18`	`+ upper_bound: Option<f64>,`
`18`	`19`	`}`
`19`	`20`
`20`	`21`	`impl OptimizeGroupTask {`
`21`		`- pub fn new(group_id: GroupId) -> Self {`
`22`		`- Self { group_id }`
	`22`	`+ pub fn new(group_id: GroupId, upper_bound: Option<f64>) -> Self {`
	`23`	`+ Self {`
	`24`	`+ group_id,`
	`25`	`+ upper_bound,`
	`26`	`+ }`
`23`	`27`	`}`
`24`	`28`	`}`
`25`	`29`
`@@ -37,7 +41,7 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeGroupTask {`
`37`	`41`	`for &expr in &exprs {`
`38`	`42`	`let typ = optimizer.get_expr_memoed(expr).typ.clone();`
`39`	`43`	`if typ.is_logical() {`
`40`		`- tasks.push(Box::new(OptimizeExpressionTask::new(expr, false)) as Box<dyn Task<T, M>>);`
	`44`	`+ tasks.push(Box::new(OptimizeExpressionTask::new(expr, false, self.upper_bound)) as Box<dyn Task<T, M>>);`
`41`	`45`	`}`
`42`	`46`	`}`
`43`	`47`	`for &expr in &exprs {`
`@@ -46,6 +50,7 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeGroupTask {`
`46`	`50`	`tasks.push(Box::new(OptimizeInputsTask::new(`
`47`	`51`	`expr,`
`48`	`52`	`!optimizer.prop.disable_pruning,`
	`53`	`+ self.upper_bound`
`49`	`54`	`)) as Box<dyn Task<T, M>>);`
`50`	`55`	`}`
`51`	`56`	`}`