feat(plannertest): add disable pruning and dump memo flags (#223)

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

Author: skyzh

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

@@ -190,9 +190,10 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for ApplyRuleTask {
                                 as Box<dyn Task<T, M>>,
                         );
                     } else {
-                        tasks
-                            .push(Box::new(OptimizeInputsTask::new(expr_id, true))
-                                as Box<dyn Task<T, M>>);
+                        tasks.push(Box::new(OptimizeInputsTask::new(
+                            expr_id,
+                            !optimizer.prop.disable_pruning,
+                        )) as Box<dyn Task<T, M>>);
                     }
                     optimizer.unmark_expr_explored(expr_id);
                     trace!(event = "apply_rule", expr_id = %self.expr_id, rule_id = %self.rule_id, new_expr_id = %expr_id);

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

@@ -42,7 +42,10 @@ 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(OptimizeInputsTask::new(expr, true)) as Box<dyn Task<T, M>>);
+                tasks.push(Box::new(OptimizeInputsTask::new(
+                    expr,
+                    !optimizer.prop.disable_pruning,
+                )) as Box<dyn Task<T, M>>);
             }
         }
         trace!(event = "task_finish", task = "optimize_group", group_id = %self.group_id, exprs_cnt = exprs_cnt);

optd-sqlplannertest/src/lib.rs

@@ -121,6 +121,11 @@ impl DatafusionDBMS {
             } else {
                 optimizer.panic_on_explore_limit(false);
             }
+            if flags.disable_pruning {
+                optimizer.disable_pruning(true);
+            } else {
+                optimizer.disable_pruning(false);
+            }
             let rules = optimizer.rules();
             if flags.enable_logical_rules.is_empty() {
                 for r in 0..rules.len() {
@@ -190,6 +195,17 @@ impl DatafusionDBMS {
                 }
             }
         }
+        if flags.dump_memo_table {
+            let mut guard = self
+                .optd_optimizer
+                .as_ref()
+                .unwrap()
+                .optimizer
+                .lock()
+                .unwrap();
+            let optimizer = guard.as_mut().unwrap().optd_optimizer_mut();
+            optimizer.dump();
+        }
         Ok(result)
     }
 
@@ -326,6 +342,8 @@ struct TestFlags {
     enable_df_logical: bool,
     enable_logical_rules: Vec<String>,
     panic_on_budget: bool,
+    dump_memo_table: bool,
+    disable_pruning: bool,
 }
 
 /// Extract the flags from a task. The flags are specified in square brackets.
@@ -353,6 +371,10 @@ fn extract_flags(task: &str) -> Result<TestFlags> {
                 }
             } else if flag == "panic_on_budget" {
                 options.panic_on_budget = true;
+            } else if flag == "dump_memo_table" {
+                options.dump_memo_table = true;
+            } else if flag == "disable_pruning" {
+                options.disable_pruning = true;
             } else {
                 bail!("Unknown flag: {}", flag);
             }