Add configurable UNION DISTINCT to FILTER rewrite optimization

datafusion/common/src/config.rs

@@ -1348,6 +1348,13 @@ 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 and only applies when the
+        /// branches share the same source and compatible wrapper nodes such as identical
+        /// projections or aliases.
+        pub enable_unions_to_filter: bool, default = false
     }
 }
 

datafusion/core/src/optimizer_rule_reference.md

@@ -39,28 +39,29 @@ Rule order matters. The default pipeline may change between releases.
 | ----- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------- |
 | 1     | `rewrite_set_comparison`                  | Rewrites `ANY` and `ALL` set-comparison subqueries into `EXISTS`-based boolean expressions with correct SQL NULL semantics. |
 | 2     | `optimize_unions`                         | Flattens nested unions and removes unions with a single input.                                                              |
-| 3     | `simplify_expressions`                    | Constant-folds and simplifies expressions while preserving output names.                                                    |
-| 4     | `replace_distinct_aggregate`              | Rewrites `DISTINCT` and `DISTINCT ON` operators into aggregate-based plans that later rules can optimize further.           |
-| 5     | `eliminate_join`                          | Replaces keyless inner joins with a literal `false` filter by an empty relation.                                            |
-| 6     | `decorrelate_predicate_subquery`          | Converts eligible `IN` and `EXISTS` predicate subqueries into semi or anti joins.                                           |
-| 7     | `scalar_subquery_to_join`                 | Rewrites eligible scalar subqueries into joins and adds schema-preserving projections.                                      |
-| 8     | `decorrelate_lateral_join`                | Rewrites eligible lateral joins into regular joins.                                                                         |
-| 9     | `extract_equijoin_predicate`              | Splits join filters into equijoin keys and residual predicates.                                                             |
-| 10    | `eliminate_duplicated_expr`               | Removes duplicate expressions from projections, aggregates, and similar operators.                                          |
-| 11    | `eliminate_filter`                        | Drops always-true filters and replaces always-false or NULL filters with empty relations.                                   |
-| 12    | `eliminate_cross_join`                    | Uses filter predicates to replace cross joins with inner joins when join keys can be found.                                 |
-| 13    | `eliminate_limit`                         | Removes no-op limits and simplifies trivial limit shapes.                                                                   |
-| 14    | `propagate_empty_relation`                | Pushes empty-relation knowledge upward so operators fed by no rows collapse early.                                          |
-| 15    | `filter_null_join_keys`                   | Adds `IS NOT NULL` filters to nullable equijoin keys that can never match.                                                  |
-| 16    | `eliminate_outer_join`                    | Rewrites outer joins to inner joins when later filters reject the NULL-extended rows.                                       |
-| 17    | `push_down_limit`                         | Moves literal limits closer to scans and unions and merges adjacent limits.                                                 |
-| 18    | `push_down_filter`                        | Moves filters as early as possible through filter-commutative operators.                                                    |
-| 19    | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans.                                             |
-| 20    | `eliminate_group_by_constant`             | Removes constant or functionally redundant expressions from `GROUP BY`.                                                     |
-| 21    | `common_sub_expression_eliminate`         | Computes repeated subexpressions once and reuses the result.                                                                |
-| 22    | `extract_leaf_expressions`                | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier.                      |
-| 23    | `push_down_leaf_projections`              | Pushes the helper projections created by leaf extraction toward leaf inputs.                                                |
-| 24    | `optimize_projections`                    | Prunes unused columns and removes unnecessary logical projections.                                                          |
+| 3     | `unions_to_filter`                        | Merges `UNION DISTINCT` branches that share the same source into a single filtered branch with a disjunctive predicate.     |
+| 4     | `simplify_expressions`                    | Constant-folds and simplifies expressions while preserving output names.                                                    |
+| 5     | `replace_distinct_aggregate`              | Rewrites `DISTINCT` and `DISTINCT ON` operators into aggregate-based plans that later rules can optimize further.           |
+| 6     | `eliminate_join`                          | Replaces keyless inner joins with a literal `false` filter by an empty relation.                                            |
+| 7     | `decorrelate_predicate_subquery`          | Converts eligible `IN` and `EXISTS` predicate subqueries into semi or anti joins.                                           |
+| 8     | `scalar_subquery_to_join`                 | Rewrites eligible scalar subqueries into joins and adds schema-preserving projections.                                      |
+| 9     | `decorrelate_lateral_join`                | Rewrites eligible lateral joins into regular joins.                                                                         |
+| 10    | `extract_equijoin_predicate`              | Splits join filters into equijoin keys and residual predicates.                                                             |
+| 11    | `eliminate_duplicated_expr`               | Removes duplicate expressions from projections, aggregates, and similar operators.                                          |
+| 12    | `eliminate_filter`                        | Drops always-true filters and replaces always-false or NULL filters with empty relations.                                   |
+| 13    | `eliminate_cross_join`                    | Uses filter predicates to replace cross joins with inner joins when join keys can be found.                                 |
+| 14    | `eliminate_limit`                         | Removes no-op limits and simplifies trivial limit shapes.                                                                   |
+| 15    | `propagate_empty_relation`                | Pushes empty-relation knowledge upward so operators fed by no rows collapse early.                                          |
+| 16    | `filter_null_join_keys`                   | Adds `IS NOT NULL` filters to nullable equijoin keys that can never match.                                                  |
+| 17    | `eliminate_outer_join`                    | Rewrites outer joins to inner joins when later filters reject the NULL-extended rows.                                       |
+| 18    | `push_down_limit`                         | Moves literal limits closer to scans and unions and merges adjacent limits.                                                 |
+| 19    | `push_down_filter`                        | Moves filters as early as possible through filter-commutative operators.                                                    |
+| 20    | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans.                                             |
+| 21    | `eliminate_group_by_constant`             | Removes constant or functionally redundant expressions from `GROUP BY`.                                                     |
+| 22    | `common_sub_expression_eliminate`         | Computes repeated subexpressions once and reuses the result.                                                                |
+| 23    | `extract_leaf_expressions`                | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier.                      |
+| 24    | `push_down_leaf_projections`              | Pushes the helper projections created by leaf extraction toward leaf inputs.                                                |
+| 25    | `optimize_projections`                    | Prunes unused columns and removes unnecessary logical projections.                                                          |
 
 ### Physical Optimizer Rules
 

datafusion/optimizer/Cargo.toml

@@ -83,3 +83,7 @@ harness = false
 [[bench]]
 name = "optimize_projections"
 harness = false
+
+[[bench]]
+name = "unions_to_filter"
+harness = false

datafusion/optimizer/benches/unions_to_filter.rs

@@ -0,0 +1,195 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Microbenchmarks for the [`UnionsToFilter`] optimizer rule.
+//!
+//! Three scenarios are covered:
+//!
+//! 1. **merge** – N branches over the *same* table, each with a simple
+//!    equality filter.  All branches should be merged into a single
+//!    `DISTINCT(Filter(OR …))` plan.
+//!
+//! 2. **no_merge** – N branches over *different* tables.  The rule must
+//!    recognise that no merge is possible and leave the plan unchanged.
+//!    This exercises the "cold path" without any rewrite work.
+//!
+//! 3. **merge_with_projection** – N branches over the same table but each
+//!    branch wraps the filter in a `Projection`.  This exercises the wrapper-
+//!    peeling and re-wrapping paths in addition to the core merge logic.
+//!
+//! To generate a flamegraph (requires `cargo-flamegraph`):
+//! ```text
+//! cargo flamegraph -p datafusion-optimizer --bench unions_to_filter \
+//!     --flamechart --root --profile profiling --freq 1000 -- --bench
+//! ```
+
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, logical_plan::table_scan};
+use datafusion_expr::{col, lit};
+use datafusion_optimizer::OptimizerContext;
+use datafusion_optimizer::unions_to_filter::UnionsToFilter;
+use datafusion_optimizer::{Optimizer, OptimizerRule};
+use std::hint::black_box;
+use std::sync::Arc;
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/// Build a three-column table scan for `name`.
+fn scan(name: &str) -> LogicalPlan {
+    let schema = Schema::new(vec![
+        Field::new("a", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+        Field::new("c", DataType::Int32, false),
+    ]);
+    table_scan(Some(name), &schema, None)
+        .unwrap()
+        .build()
+        .unwrap()
+}
+
+/// Build a `DISTINCT (UNION ALL …)` plan whose `n` branches all filter over
+/// the *same* table (`t`), so the rule can merge them.
+fn build_merge_plan(n: usize) -> LogicalPlan {
+    assert!(n >= 2);
+    let mut builder: Option<LogicalPlanBuilder> = None;
+    for i in 0..n {
+        let branch = LogicalPlanBuilder::from(scan("t"))
+            .filter(col("a").eq(lit(i as i32)))
+            .unwrap()
+            .build()
+            .unwrap();
+        builder = Some(match builder {
+            None => LogicalPlanBuilder::from(branch),
+            Some(b) => b.union(branch).unwrap(),
+        });
+    }
+    builder.unwrap().distinct().unwrap().build().unwrap()
+}
+
+/// Build a `DISTINCT (UNION ALL …)` plan whose `n` branches each filter over a
+/// *different* table, so no merge is possible.
+fn build_no_merge_plan(n: usize) -> LogicalPlan {
+    assert!(n >= 2);
+    let mut builder: Option<LogicalPlanBuilder> = None;
+    for i in 0..n {
+        let branch = LogicalPlanBuilder::from(scan(&format!("t{i}")))
+            .filter(col("a").eq(lit(i as i32)))
+            .unwrap()
+            .build()
+            .unwrap();
+        builder = Some(match builder {
+            None => LogicalPlanBuilder::from(branch),
+            Some(b) => b.union(branch).unwrap(),
+        });
+    }
+    builder.unwrap().distinct().unwrap().build().unwrap()
+}
+
+/// Build a `DISTINCT (UNION ALL …)` plan whose `n` branches each wrap the
+/// filter inside a `Projection` over the *same* table.
+fn build_merge_with_projection_plan(n: usize) -> LogicalPlan {
+    assert!(n >= 2);
+    let mut builder: Option<LogicalPlanBuilder> = None;
+    for i in 0..n {
+        let branch = LogicalPlanBuilder::from(scan("t"))
+            .filter(col("a").eq(lit(i as i32)))
+            .unwrap()
+            .project(vec![col("a"), col("b")])
+            .unwrap()
+            .build()
+            .unwrap();
+        builder = Some(match builder {
+            None => LogicalPlanBuilder::from(branch),
+            Some(b) => b.union(branch).unwrap(),
+        });
+    }
+    builder.unwrap().distinct().unwrap().build().unwrap()
+}
+
+/// Run the [`UnionsToFilter`] rule through the full [`Optimizer`] pipeline
+/// (single pass, feature flag enabled).
+fn run_optimizer(plan: &LogicalPlan, ctx: &OptimizerContext) -> LogicalPlan {
+    let rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> =
+        vec![Arc::new(UnionsToFilter::new())];
+    Optimizer::with_rules(rules)
+        .optimize(plan.clone(), ctx, |_, _| {})
+        .unwrap()
+}
+
+// ---------------------------------------------------------------------------
+// Benchmark functions
+// ---------------------------------------------------------------------------
+
+fn bench_merge(c: &mut Criterion) {
+    let mut options = ConfigOptions::default();
+    options.optimizer.enable_unions_to_filter = true;
+    let ctx =
+        OptimizerContext::new_with_config_options(Arc::new(options)).with_max_passes(1);
+
+    let mut group = c.benchmark_group("unions_to_filter/merge");
+    for n in [2, 8, 32, 128] {
+        let plan = build_merge_plan(n);
+        group.bench_with_input(BenchmarkId::from_parameter(n), &plan, |b, p| {
+            b.iter(|| black_box(run_optimizer(p, &ctx)));
+        });
+    }
+    group.finish();
+}
+
+fn bench_no_merge(c: &mut Criterion) {
+    let mut options = ConfigOptions::default();
+    options.optimizer.enable_unions_to_filter = true;
+    let ctx =
+        OptimizerContext::new_with_config_options(Arc::new(options)).with_max_passes(1);
+
+    let mut group = c.benchmark_group("unions_to_filter/no_merge");
+    for n in [2, 8, 32, 128] {
+        let plan = build_no_merge_plan(n);
+        group.bench_with_input(BenchmarkId::from_parameter(n), &plan, |b, p| {
+            b.iter(|| black_box(run_optimizer(p, &ctx)));
+        });
+    }
+    group.finish();
+}
+
+fn bench_merge_with_projection(c: &mut Criterion) {
+    let mut options = ConfigOptions::default();
+    options.optimizer.enable_unions_to_filter = true;
+    let ctx =
+        OptimizerContext::new_with_config_options(Arc::new(options)).with_max_passes(1);
+
+    let mut group = c.benchmark_group("unions_to_filter/merge_with_projection");
+    for n in [2, 8, 32, 128] {
+        let plan = build_merge_with_projection_plan(n);
+        group.bench_with_input(BenchmarkId::from_parameter(n), &plan, |b, p| {
+            b.iter(|| black_box(run_optimizer(p, &ctx)));
+        });
+    }
+    group.finish();
+}
+
+criterion_group!(
+    benches,
+    bench_merge,
+    bench_no_merge,
+    bench_merge_with_projection
+);
+criterion_main!(benches);

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()),