Add configurable UNION DISTINCT to FILTER rewrite optimization

Author: xiedeyantu

datafusion/common/src/config.rs

@@ -1168,6 +1168,37 @@ config_namespace! {
         /// closer to the leaf table scans, and push those projections down
         /// towards the leaf nodes.
         pub enable_leaf_expression_pushdown: bool, default = true
+
+        /// When set to true, the logical optimizer will rewrite `UNION DISTINCT`
+        /// branches that read from the same source and differ only by filter predicates
+        /// into a single branch with a combined filter.
+        ///
+        /// This optimization is conservative: it only applies when the `UNION DISTINCT`
+        /// branches share the same source and compatible wrapper nodes such as identical
+        /// projections or aliases.
+        ///
+        /// Example:
+        ///
+        /// Disabled (`false`, default):
+        /// ```text
+        /// Distinct:
+        ///   Union
+        ///     Projection: test.col_int32 AS a, test.col_uint32 AS b
+        ///       Filter: test.col_int32 = Int32(1)
+        ///         TableScan: test
+        ///     Projection: test.col_int32 AS a, test.col_uint32 AS b
+        ///       Filter: test.col_uint32 = UInt32(5)
+        ///         TableScan: test
+        /// ```
+        ///
+        /// Enabled (`true`):
+        /// ```text
+        /// Distinct:
+        ///   Projection: test.col_int32 AS a, test.col_uint32 AS b
+        ///     Filter: test.col_int32 = Int32(1) OR test.col_uint32 = UInt32(5)
+        ///       TableScan: test
+        /// ```
+        pub enable_unions_to_filter: bool, default = false
     }
 }
 

datafusion/optimizer/src/lib.rs

@@ -70,6 +70,7 @@ pub mod rewrite_set_comparison;
 pub mod scalar_subquery_to_join;
 pub mod simplify_expressions;
 pub mod single_distinct_to_groupby;
+pub mod unions_to_filter;
 pub mod utils;
 
 #[cfg(test)]

datafusion/optimizer/src/optimizer.rs

@@ -56,6 +56,7 @@ use crate::rewrite_set_comparison::RewriteSetComparison;
 use crate::scalar_subquery_to_join::ScalarSubqueryToJoin;
 use crate::simplify_expressions::SimplifyExpressions;
 use crate::single_distinct_to_groupby::SingleDistinctToGroupBy;
+use crate::unions_to_filter::UnionsToFilter;
 use crate::utils::log_plan;
 
 /// Transforms one [`LogicalPlan`] into another which computes the same results,
@@ -280,6 +281,7 @@ impl Optimizer {
         let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
             Arc::new(RewriteSetComparison::new()),
             Arc::new(OptimizeUnions::new()),
+            Arc::new(UnionsToFilter::new()),
             Arc::new(SimplifyExpressions::new()),
             Arc::new(ReplaceDistinctWithAggregate::new()),
             Arc::new(EliminateJoin::new()),