feature: dump and visualize the memo as a graph

optd-core/src/cascades/optimizer.rs

@@ -5,9 +5,13 @@
 
 use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
 use std::fmt::Display;
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::io::Write;
+use std::path::Path;
 use std::sync::Arc;
 
 use anyhow::Result;
+use itertools::Itertools;
 use tracing::trace;
 
 use super::memo::{ArcMemoPlanNode, GroupInfo, Memo};
@@ -40,6 +44,8 @@ pub struct OptimizerProperties {
     pub partial_explore_space: Option<usize>,
     /// Disable pruning during optimization.
     pub disable_pruning: bool,
+    /// Dump the memo as 0000.dot, 0001.dot, ... files to this path
+    pub dot_file_path: Option<String>,
 }
 
 pub struct CascadesOptimizer<T: NodeType, M: Memo<T> = NaiveMemo<T>> {
@@ -54,6 +60,7 @@ pub struct CascadesOptimizer<T: NodeType, M: Memo<T> = NaiveMemo<T>> {
     property_builders: Arc<[Box<dyn LogicalPropertyBuilderAny<T>>]>,
     pub ctx: OptimizerContext,
     pub prop: OptimizerProperties,
+    pub next_dot_file: usize,
 }
 
 /// `RelNode` only contains the representation of the plan nodes. Sometimes, we need more context,
@@ -123,6 +130,7 @@ impl<T: NodeType> CascadesOptimizer<T, NaiveMemo<T>> {
             property_builders,
             prop,
             disabled_rules: HashSet::new(),
+            next_dot_file: 0,
         }
     }
 
@@ -170,6 +178,93 @@ impl<T: NodeType, M: Memo<T>> CascadesOptimizer<T, M> {
         self.disabled_rules.contains(&rule_id)
     }
 
+    pub fn dump_dot(&self, writer: &mut dyn Write) -> Result<(), std::io::Error> {
+        let memo = self.memo();
+
+        // Collect all groups in a predictable iteration order
+        let groups: Vec<GroupId> = memo.get_all_group_ids().iter().sorted().cloned().collect();
+
+        writeln!(writer, "digraph Memo {{")?;
+        writeln!(
+            writer,
+            "compound=true; ranksep=1.0; node [colorscheme=set312];"
+        )?;
+        for group_id in groups.iter() {
+            let group = memo.get_group(*group_id);
+            writeln!(writer, "subgraph cluster_{} {{", group_id.0)?;
+            writeln!(writer, "rank=source;")?;
+            writeln!(writer, "edge [style=invis];")?;
+            writeln!(
+                writer,
+                "g{} [shape=plaintext,label=\"group_id=!{}\"];",
+                group_id.0, group_id.0
+            )?;
+            for expr_id in memo.get_all_exprs_in_group(*group_id).iter() {
+                let expr = memo.get_expr_memoed(*expr_id);
+                let mut s = DefaultHasher::new();
+                expr.typ.hash(&mut s);
+                let color = (s.finish() % 11) + 1; // %11 looks better than %12! :-)
+                let shape = if expr.typ.is_logical() { "oval" } else { "box" };
+                let rules = match self.fired_rules.get(expr_id) {
+                    None => 0,
+                    Some(v) => v.len(),
+                };
+                writeln!(
+                    writer,
+                    "e{} [shape={},label=\"{}: {:?} ({})\",style=filled,color={}]",
+                    expr_id.0, shape, expr_id.0, expr.typ, rules, color
+                )?;
+                writeln!(writer, "g{} -> e{};", group_id.0, expr_id.0)?;
+            }
+            writeln!(writer, "}}");
+        }
+        for group_id in groups.iter() {
+            for expr_id in memo.get_all_exprs_in_group(*group_id).iter() {
+                let expr = memo.get_expr_memoed(*expr_id);
+                for child in expr.children.iter() {
+                    writeln!(
+                        writer,
+                        "e{} -> g{} [lhead=\"cluster_{}\"];",
+                        expr_id.0, child.0, child.0
+                    )?;
+                }
+                let mut next_pred: usize = 0;
+                for pred_id in expr.predicates.iter() {
+                    let pred = memo.get_pred(*pred_id);
+                    let id = next_pred;
+                    self.dump_dot_pred(writer, &pred, expr_id.0, &mut next_pred);
+                    writeln!(writer, "e{} -> p{}_{};", expr_id.0, expr_id.0, id)?;
+                }
+            }
+        }
+        writeln!(writer, "}}")
+    }
+
+    fn dump_dot_pred(
+        &self,
+        writer: &mut dyn Write,
+        pred: &ArcPredNode<T>,
+        base: usize,
+        next_pred: &mut usize,
+    ) -> Result<(), std::io::Error> {
+        let mut s = DefaultHasher::new();
+        pred.typ.hash(&mut s);
+        let color = (s.finish() % 11) + 1;
+        let id = *next_pred;
+        *next_pred += 1;
+        writeln!(
+            writer,
+            "p{}_{} [shape=diamond,label=\"{:?}\",penwidth=3,color={}]",
+            base, id, pred.typ, color
+        )?;
+        for child in pred.children.iter() {
+            let child_id = *next_pred;
+            self.dump_dot_pred(writer, child, base, next_pred)?;
+            writeln!(writer, "p{}_{} -> p{}_{};", base, id, base, child_id)?;
+        }
+        Ok(())
+    }
+
     pub fn dump(&self) {
         for group_id in self.memo.get_all_group_ids() {
             let winner_str = match &self.memo.get_group_info(group_id).winner {

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

@@ -3,11 +3,12 @@
 // Use of this source code is governed by an MIT-style license that can be found in the LICENSE file or at
 // https://opensource.org/licenses/MIT.
 
+use std::path::Path;
 use std::sync::Arc;
 
 use anyhow::Result;
 use itertools::Itertools;
-use tracing::trace;
+use tracing::{debug, trace};
 
 use super::Task;
 use crate::cascades::memo::ArcMemoPlanNode;
@@ -164,6 +165,14 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for ApplyRuleTask {
             return Ok(vec![]);
         }
 
+        if let Some(pathname) = &optimizer.prop.dot_file_path {
+            let path = Path::new(pathname).join(format!("{:#08}.dot", optimizer.next_dot_file));
+            debug!("dumping memo to {:?}", path);
+            optimizer.next_dot_file += 1;
+            let mut writer = std::fs::File::create(path).unwrap();
+            optimizer.dump_dot(&mut writer);
+        }
+
         let rule = optimizer.rules()[self.rule_id].clone();
 
         trace!(event = "task_begin", task = "apply_rule", expr_id = %self.expr_id, rule_id = %self.rule_id, rule = %rule.name());

optd-datafusion-repr/src/lib.rs

@@ -152,6 +152,7 @@ impl DatafusionOptimizer {
                     partial_explore_iter: Some(1 << 20),
                     partial_explore_space: Some(1 << 10),
                     disable_pruning: false,
+                    dot_file_path: None,
                 },
             ),
             heuristic_optimizer: HeuristicsOptimizer::new_with_rules(