First pass of apache#10296

emgeee · Apr 30, 2024 · 22609aa · 22609aa
1 parent dd56837
commit 22609aa
Showing 1 changed file with 31 additions and 17 deletions.
diff --git a/datafusion/optimizer/src/eliminate_nested_union.rs b/datafusion/optimizer/src/eliminate_nested_union.rs
@@ -18,8 +18,10 @@
 //! [`EliminateNestedUnion`]: flattens nested `Union` to a single `Union`
 use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result};
 use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
+use datafusion_expr::logical_plan::tree_node::unwrap_arc;
 use datafusion_expr::{Distinct, LogicalPlan, Union};
 use std::sync::Arc;
 
@@ -37,17 +39,37 @@ impl EliminateNestedUnion {
 impl OptimizerRule for EliminateNestedUnion {
     fn try_optimize(
         &self,
-        plan: &LogicalPlan,
+        _plan: &LogicalPlan,
         _config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
+        internal_err!("Should have called EliminateFilter::rewrite")
+    }
+
+    fn name(&self) -> &str {
+        "eliminate_nested_union"
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(ApplyOrder::BottomUp)
+    }
+
+    fn supports_rewrite(&self) -> bool {
+        true
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
         match plan {
             LogicalPlan::Union(Union { inputs, schema }) => {
                 let inputs = inputs
                     .iter()
                     .flat_map(extract_plans_from_union)
                     .collect::<Vec<_>>();
 
-                Ok(Some(LogicalPlan::Union(Union {
+                Ok(Transformed::yes(LogicalPlan::Union(Union {
                     inputs,
                     schema: schema.clone(),
                 })))
@@ -60,26 +82,18 @@ impl OptimizerRule for EliminateNestedUnion {
                         .flat_map(extract_plans_from_union)
                         .collect::<Vec<_>>();
 
-                    Ok(Some(LogicalPlan::Distinct(Distinct::All(Arc::new(
-                        LogicalPlan::Union(Union {
+                    Ok(Transformed::yes(LogicalPlan::Distinct(Distinct::All(
+                        Arc::new(LogicalPlan::Union(Union {
                             inputs,
                             schema: schema.clone(),
-                        }),
-                    )))))
+                        })),
+                    ))))
                 }
-                _ => Ok(None),
+                _ => Ok(Transformed::no(unwrap_arc(plan))),
             },
-            _ => Ok(None),
+            _ => Ok(Transformed::no(plan)),
         }
     }
-
-    fn name(&self) -> &str {
-        "eliminate_nested_union"
-    }
-
-    fn apply_order(&self) -> Option<ApplyOrder> {
-        Some(ApplyOrder::BottomUp)
-    }
 }
 
 fn extract_plans_from_union(plan: &Arc<LogicalPlan>) -> Vec<Arc<LogicalPlan>> {