From 08030f37fa34995af2ea43143922476ace82e39d Mon Sep 17 00:00:00 2001
From: Lordworms <48054792+Lordworms@users.noreply.github.com>
Date: Tue, 23 Apr 2024 12:38:37 -0500
Subject: [PATCH] Implement rewrite for EliminateOneUnion and EliminateJoin
 (#10184)

* implement rewrite for EliminateJoin

* implement rewrite for eliminate_on_union
---
 datafusion/expr/src/logical_plan/tree_node.rs |  2 +-
 datafusion/optimizer/src/eliminate_join.rs    | 51 ++++++++++++-------
 .../optimizer/src/eliminate_one_union.rs      | 30 +++++++----
 3 files changed, 54 insertions(+), 29 deletions(-)

diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs
index f5db5a2704bc..37a36c36ca53 100644
--- a/datafusion/expr/src/logical_plan/tree_node.rs
+++ b/datafusion/expr/src/logical_plan/tree_node.rs
@@ -363,7 +363,7 @@ impl TreeNode for LogicalPlan {
 
 /// Converts a `Arc<LogicalPlan>` without copying, if possible. Copies the plan
 /// if there is a shared reference
-fn unwrap_arc(plan: Arc<LogicalPlan>) -> LogicalPlan {
+pub fn unwrap_arc(plan: Arc<LogicalPlan>) -> LogicalPlan {
     Arc::try_unwrap(plan)
         // if None is returned, there is another reference to this
         // LogicalPlan, so we can not own it, and must clone instead
diff --git a/datafusion/optimizer/src/eliminate_join.rs b/datafusion/optimizer/src/eliminate_join.rs
index caf45dda9896..fea87e758790 100644
--- a/datafusion/optimizer/src/eliminate_join.rs
+++ b/datafusion/optimizer/src/eliminate_join.rs
@@ -18,7 +18,8 @@
 //! [`EliminateJoin`] rewrites `INNER JOIN` with `true`/`null`
 use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::{Result, ScalarValue};
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result, ScalarValue};
 use datafusion_expr::JoinType::Inner;
 use datafusion_expr::{
     logical_plan::{EmptyRelation, LogicalPlan},
@@ -39,38 +40,50 @@ impl EliminateJoin {
 impl OptimizerRule for EliminateJoin {
     fn try_optimize(
         &self,
-        plan: &LogicalPlan,
+        _plan: &LogicalPlan,
         _config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
+        internal_err!("Should have called EliminateJoin::rewrite")
+    }
+
+    fn name(&self) -> &str {
+        "eliminate_join"
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(ApplyOrder::TopDown)
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
         match plan {
             LogicalPlan::Join(join) if join.join_type == Inner && join.on.is_empty() => {
                 match join.filter {
                     Some(Expr::Literal(ScalarValue::Boolean(Some(true)))) => {
-                        Ok(Some(LogicalPlan::CrossJoin(CrossJoin {
-                            left: join.left.clone(),
-                            right: join.right.clone(),
-                            schema: join.schema.clone(),
+                        Ok(Transformed::yes(LogicalPlan::CrossJoin(CrossJoin {
+                            left: join.left,
+                            right: join.right,
+                            schema: join.schema,
                         })))
                     }
-                    Some(Expr::Literal(ScalarValue::Boolean(Some(false)))) => {
-                        Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
+                    Some(Expr::Literal(ScalarValue::Boolean(Some(false)))) => Ok(
+                        Transformed::yes(LogicalPlan::EmptyRelation(EmptyRelation {
                             produce_one_row: false,
-                            schema: join.schema.clone(),
-                        })))
-                    }
-                    _ => Ok(None),
+                            schema: join.schema,
+                        })),
+                    ),
+                    _ => Ok(Transformed::no(LogicalPlan::Join(join))),
                 }
             }
-            _ => Ok(None),
+            _ => Ok(Transformed::no(plan)),
         }
     }
 
-    fn name(&self) -> &str {
-        "eliminate_join"
-    }
-
-    fn apply_order(&self) -> Option<ApplyOrder> {
-        Some(ApplyOrder::TopDown)
+    fn supports_rewrite(&self) -> bool {
+        true
     }
 }
 
diff --git a/datafusion/optimizer/src/eliminate_one_union.rs b/datafusion/optimizer/src/eliminate_one_union.rs
index 95a3370ab1b5..11a9009cd96c 100644
--- a/datafusion/optimizer/src/eliminate_one_union.rs
+++ b/datafusion/optimizer/src/eliminate_one_union.rs
@@ -17,8 +17,8 @@
 
 //! [`EliminateOneUnion`]  eliminates single element `Union`
 use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
-use datafusion_expr::logical_plan::{LogicalPlan, Union};
+use datafusion_common::{internal_err, tree_node::Transformed, Result};
+use datafusion_expr::logical_plan::{tree_node::unwrap_arc, LogicalPlan, Union};
 
 use crate::optimizer::ApplyOrder;
 
@@ -36,21 +36,33 @@ impl EliminateOneUnion {
 impl OptimizerRule for EliminateOneUnion {
     fn try_optimize(
         &self,
-        plan: &LogicalPlan,
+        _plan: &LogicalPlan,
         _config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
-        match plan {
-            LogicalPlan::Union(Union { inputs, .. }) if inputs.len() == 1 => {
-                Ok(inputs.first().map(|input| input.as_ref().clone()))
-            }
-            _ => Ok(None),
-        }
+        internal_err!("Should have called EliminateOneUnion::rewrite")
     }
 
     fn name(&self) -> &str {
         "eliminate_one_union"
     }
 
+    fn supports_rewrite(&self) -> bool {
+        true
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
+        match plan {
+            LogicalPlan::Union(Union { mut inputs, .. }) if inputs.len() == 1 => {
+                Ok(Transformed::yes(unwrap_arc(inputs.pop().unwrap())))
+            }
+            _ => Ok(Transformed::no(plan)),
+        }
+    }
+
     fn apply_order(&self) -> Option<ApplyOrder> {
         Some(ApplyOrder::TopDown)
     }