From bd40a105f2dddb5d983ef24217db10753f0ac9d0 Mon Sep 17 00:00:00 2001
From: Yuchen Liang <70461588+yliang412@users.noreply.github.com>
Date: Wed, 27 Mar 2024 18:39:09 -0400
Subject: [PATCH] feat: display logical plan after heuristic (#141)

It might be helpful to see what the plan looks like after the heuristics
phase. It might help us to identify things that are not handled in the
system and help us make design decisions forward on implementing rules.

Signed-off-by: Yuchen Liang <yuchenl3@andrew.cmu.edu>
---
 optd-datafusion-bridge/src/lib.rs | 10 ++++++++++
 optd-sqlplannertest/src/lib.rs    | 10 ++++++++++
 2 files changed, 20 insertions(+)

diff --git a/optd-datafusion-bridge/src/lib.rs b/optd-datafusion-bridge/src/lib.rs
index 4bea295d..e6413653 100644
--- a/optd-datafusion-bridge/src/lib.rs
+++ b/optd-datafusion-bridge/src/lib.rs
@@ -233,6 +233,16 @@ impl OptdQueryPlanner {
 
         if optimizer.is_heuristic_enabled() {
             optd_rel = optimizer.heuristic_optimize(optd_rel);
+            if let Some(explains) = &mut explains {
+                explains.push(StringifiedPlan::new(
+                    PlanType::OptimizedLogicalPlan {
+                        optimizer_name: "optd-heuristic".to_string(),
+                    },
+                    PlanNode::from_rel_node(optd_rel.clone())
+                        .unwrap()
+                        .explain_to_string(None),
+                ))
+            }
         }
 
         let (group_id, optimized_rel, meta) = optimizer.cascades_optimize(optd_rel)?;
diff --git a/optd-sqlplannertest/src/lib.rs b/optd-sqlplannertest/src/lib.rs
index 2dbd6f54..d69ca31d 100644
--- a/optd-sqlplannertest/src/lib.rs
+++ b/optd-sqlplannertest/src/lib.rs
@@ -163,6 +163,16 @@ impl DatafusionDBMS {
                         .map(|x| &x[1])
                         .unwrap()
                 )?;
+            } else if subtask == "logical_optd_heuristic" {
+                writeln!(
+                    r,
+                    "{}",
+                    result
+                        .iter()
+                        .find(|x| x[0] == "logical_plan after optd-heuristic")
+                        .map(|x| &x[1])
+                        .unwrap()
+                )?;
             } else if subtask == "logical_optd" {
                 writeln!(
                     r,