add pre-check

Author: xiedeyantu

datafusion/optimizer/src/unions_to_filter.rs

@@ -18,10 +18,9 @@
 //! Rewrites `UNION DISTINCT` branches that differ only by filter predicates
 //! into a single filtered branch plus `DISTINCT`.
 
-use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_common::Result;
-use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
 use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
 use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
 use datafusion_expr::utils::disjunction;
@@ -47,10 +46,6 @@ impl OptimizerRule for UnionsToFilter {
         "unions_to_filter"
     }
 
-    fn apply_order(&self) -> Option<ApplyOrder> {
-        Some(ApplyOrder::BottomUp)
-    }
-
     fn supports_rewrite(&self) -> bool {
         true
     }
@@ -64,6 +59,24 @@ impl OptimizerRule for UnionsToFilter {
             return Ok(Transformed::no(plan));
         }
 
+        // Fast pre-check: if the plan tree has no Distinct::All node at all we can
+        // skip the expensive bottom-up rewrite_with_subqueries traversal entirely.
+        // This matters for large UNION ALL plans (e.g. TPC-DS Q4) where the rule
+        // can never fire and the traversal overhead is otherwise measurable.
+        if !plan.exists(|p| Ok(matches!(p, LogicalPlan::Distinct(Distinct::All(_)))))? {
+            return Ok(Transformed::no(plan));
+        }
+
+        plan.rewrite_with_subqueries(&mut UnionsToFilterRewriter)
+    }
+}
+
+struct UnionsToFilterRewriter;
+
+impl TreeNodeRewriter for UnionsToFilterRewriter {
+    type Node = LogicalPlan;
+
+    fn f_up(&mut self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
         match plan {
             LogicalPlan::Distinct(Distinct::All(input)) => {
                 let inner = Arc::unwrap_or_clone(input);
@@ -370,6 +383,10 @@ mod tests {
     struct VolatileTestUdf;
 
     impl ScalarUDFImpl for VolatileTestUdf {
+        fn as_any(&self) -> &dyn std::any::Any {
+            self
+        }
+
         fn name(&self) -> &str {
             "volatile_test"
         }