diff --git a/datafusion/optimizer/src/unions_to_filter.rs b/datafusion/optimizer/src/unions_to_filter.rs index 158fd358287fe..2eda0daba1308 100644 --- a/datafusion/optimizer/src/unions_to_filter.rs +++ b/datafusion/optimizer/src/unions_to_filter.rs @@ -20,12 +20,14 @@ use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::Result; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion_common::tree_node::{ + Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, +}; use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema; use datafusion_expr::logical_plan::builder::LogicalPlanBuilder; use datafusion_expr::utils::disjunction; use datafusion_expr::{ - Distinct, Expr, Filter, LogicalPlan, Projection, SubqueryAlias, Union, + Distinct, Expr, Filter, LogicalPlan, Projection, SubqueryAlias, TableSource, Union, }; use log::debug; use std::sync::Arc; @@ -113,9 +115,11 @@ fn try_rewrite_distinct_union(plan: LogicalPlan) -> Result> return Ok(None); }; + let table_sources = collect_table_sources(&branch.source)?; let key = GroupKey { source: branch.source, wrappers: branch.wrappers, + table_sources, }; if let Some((_, conds)) = grouped.iter_mut().find(|(k, _)| k == &key) { conds.push(branch.predicate); @@ -172,7 +176,7 @@ fn extract_branch(plan: LogicalPlan) -> Result> { return Ok(None); } - match plan { + let (source, predicate) = match plan { LogicalPlan::Filter(Filter { predicate, input, .. }) => { @@ -182,11 +186,10 @@ fn extract_branch(plan: LogicalPlan) -> Result> { ); return Ok(None); } - Ok(Some(UnionBranch { - source: strip_passthrough_nodes(Arc::unwrap_or_clone(input)), - predicate, - wrappers, - })) + // Keep the input intact: projections compute values and aliases + // define the qualifiers used by the predicate. GroupKey compares + // the retained plans before merging their predicates. + (Arc::unwrap_or_clone(input), predicate) } // A Limit or Sort node changes the row-set semantics of the branch. // Merging two such branches into one would silently drop the per-branch @@ -194,29 +197,56 @@ fn extract_branch(plan: LogicalPlan) -> Result> { // not preserve (ORDER BY). Bail out to leave the UNION unchanged. LogicalPlan::Limit(_) => { debug!("unions_to_filter skipped: branch contains LIMIT"); - Ok(None) + return Ok(None); } LogicalPlan::Sort(_) => { debug!("unions_to_filter skipped: branch contains ORDER BY / SORT"); - Ok(None) + return Ok(None); } - other => Ok(Some(UnionBranch { - source: strip_passthrough_nodes(other), - predicate: Expr::Literal( - datafusion_common::ScalarValue::Boolean(Some(true)), - None, - ), - wrappers, - })), + other => ( + other, + Expr::Literal(datafusion_common::ScalarValue::Boolean(Some(true)), None), + ), + }; + + if !source_is_safe(&source)? { + debug!( + "unions_to_filter skipped: source contains volatile expression or subquery" + ); + return Ok(None); } + + Ok(Some(UnionBranch { + source, + predicate, + wrappers, + })) } -#[derive(Debug, Clone, PartialEq, Eq)] struct GroupKey { source: LogicalPlan, wrappers: Vec, + // LogicalPlan equality deliberately ignores TableScan::source. Retain the + // corresponding Arcs so structurally identical scans over different data + // cannot be grouped together. + table_sources: Vec>, } +impl PartialEq for GroupKey { + fn eq(&self, other: &Self) -> bool { + self.source == other.source + && self.wrappers == other.wrappers + && self.table_sources.len() == other.table_sources.len() + && self + .table_sources + .iter() + .zip(&other.table_sources) + .all(|(left, right)| Arc::ptr_eq(left, right)) + } +} + +impl Eq for GroupKey {} + #[derive(Debug, Clone, PartialEq, Eq)] enum Wrapper { Projection { @@ -278,20 +308,6 @@ fn wrap_branch(mut plan: LogicalPlan, wrappers: &[Wrapper]) -> Result LogicalPlan { - loop { - plan = match plan { - LogicalPlan::Projection(Projection { input, .. }) => { - Arc::unwrap_or_clone(input) - } - LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => { - Arc::unwrap_or_clone(input) - } - other => return other, - }; - } -} - fn align_plan_to_schema( plan: LogicalPlan, schema: datafusion_common::DFSchemaRef, @@ -322,6 +338,35 @@ fn is_mergeable_predicate(expr: &Expr) -> bool { !expr.is_volatile() && !expr_contains_subquery(expr) } +/// Check every expression in the retained source, including its descendants. +/// Merging branches also merges their source evaluations, so the same +/// restrictions as for projection wrappers apply throughout the source. +fn source_is_safe(source: &LogicalPlan) -> Result { + let mut safe = true; + source.apply(|node| { + node.apply_expressions(|expr| { + if is_mergeable_predicate(expr) { + Ok(TreeNodeRecursion::Continue) + } else { + safe = false; + Ok(TreeNodeRecursion::Stop) + } + }) + })?; + Ok(safe) +} + +fn collect_table_sources(source: &LogicalPlan) -> Result>> { + let mut table_sources = vec![]; + source.apply(|node| { + if let LogicalPlan::TableScan(scan) = node { + table_sources.push(Arc::clone(&scan.source)); + } + Ok(TreeNodeRecursion::Continue) + })?; + Ok(table_sources) +} + /// Returns `true` when every projection expression in `wrappers` is both /// non-volatile and subquery-free. /// @@ -340,7 +385,10 @@ fn wrapper_projections_are_safe(wrappers: &[Wrapper]) -> bool { fn expr_contains_subquery(expr: &Expr) -> bool { expr.exists(|e| match e { - Expr::ScalarSubquery(_) | Expr::Exists(_) | Expr::InSubquery(_) => Ok(true), + Expr::ScalarSubquery(_) + | Expr::Exists(_) + | Expr::InSubquery(_) + | Expr::SetComparison(_) => Ok(true), _ => Ok(false), }) .expect("boolean expression walk is infallible") @@ -353,10 +401,11 @@ mod tests { use crate::assert_optimized_plan_eq_snapshot; use crate::test::test_table_scan_with_name; use arrow::datatypes::DataType; - use datafusion_common::Result; + use datafusion_common::{Result, Spans}; + use datafusion_expr::expr::{SetComparison, SetQuantifier}; use datafusion_expr::{ - ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, - Volatility, col, lit, + ColumnarValue, Expr, Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, + Signature, Subquery, Volatility, col, lit, }; macro_rules! assert_optimized_plan_equal { @@ -406,12 +455,238 @@ mod tests { ScalarUDF::new_from_impl(VolatileTestUdf).call(vec![]) } + #[test] + fn set_comparison_is_detected_as_subquery() { + let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("t2").unwrap()) + .project(vec![col("b")]) + .unwrap() + .build() + .unwrap(); + let expr = Expr::SetComparison(SetComparison::new( + Box::new(col("a")), + Subquery { + subquery: Arc::new(subquery), + outer_ref_columns: vec![], + spans: Spans::new(), + }, + Operator::Gt, + SetQuantifier::Any, + )); + + assert!(expr_contains_subquery(&expr)); + } + + fn assert_not_rewritten(plan: LogicalPlan) { + let mut options = datafusion_common::config::ConfigOptions::default(); + options.optimizer.enable_unions_to_filter = true; + let optimizer_ctx = OptimizerContext::new_with_config_options(Arc::new(options)); + let result = UnionsToFilter::new() + .rewrite(plan.clone(), &optimizer_ctx) + .unwrap(); + assert!(!result.transformed); + assert_eq!(result.data, plan); + } + + #[test] + fn rewrite_union_distinct_preserves_alias_below_filter() -> Result<()> { + let scan = test_table_scan_with_name("t").unwrap(); + let left_source = LogicalPlanBuilder::from(scan.clone()) + .alias("x") + .unwrap() + .build() + .unwrap(); + let right_source = LogicalPlanBuilder::from(scan) + .alias("x") + .unwrap() + .build() + .unwrap(); + let left = LogicalPlanBuilder::from(left_source) + .filter(col("x.a").eq(lit(1))) + .unwrap() + .project(vec![col("x.a")]) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(right_source) + .filter(col("x.a").eq(lit(2))) + .unwrap() + .project(vec![col("x.a")]) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right) + .unwrap() + .build() + .unwrap(); + + assert_optimized_plan_equal!(plan, @r" + Distinct: + Projection: x.a + Projection: x.a + Filter: x.a = Int32(1) OR x.a = Int32(2) + SubqueryAlias: x + TableScan: t + ") + .unwrap(); + Ok(()) + } + + #[test] + fn rewrite_union_distinct_matching_computed_projection_below_filter() -> Result<()> { + let scan = test_table_scan_with_name("prices").unwrap(); + let left_source = LogicalPlanBuilder::from(scan.clone()) + .project(vec![(col("a") + lit(100)).alias("amount")]) + .unwrap() + .alias("prices") + .unwrap() + .build() + .unwrap(); + let right_source = LogicalPlanBuilder::from(scan) + .project(vec![(col("a") + lit(100)).alias("amount")]) + .unwrap() + .alias("prices") + .unwrap() + .build() + .unwrap(); + let left = LogicalPlanBuilder::from(left_source) + .filter(col("amount").gt(lit(0))) + .unwrap() + .project(vec![col("amount")]) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(right_source) + .filter(col("amount").gt(lit(1))) + .unwrap() + .project(vec![col("amount")]) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right) + .unwrap() + .build() + .unwrap(); + + assert_optimized_plan_equal!(plan, @r" + Distinct: + Projection: prices.amount + Projection: prices.amount + Filter: prices.amount > Int32(0) OR prices.amount > Int32(1) + SubqueryAlias: prices + Projection: prices.a + Int32(100) AS amount + TableScan: prices + ") + .unwrap(); + Ok(()) + } + + #[test] + fn keep_union_distinct_with_volatile_source_projection() { + for nested in [false, true] { + let mut source = + LogicalPlanBuilder::from(test_table_scan_with_name("t").unwrap()) + .project(vec![(volatile_expr() + lit(1.0_f64)).alias("v"), col("a")]) + .unwrap() + .alias("x") + .unwrap(); + if nested { + // The volatile projection is below another Filter, so checking + // just the source node's expressions would miss it. + source = source.filter(col("x.v").gt(lit(0.5_f64))).unwrap(); + } + let source = source.build().unwrap(); + let left = LogicalPlanBuilder::from(source.clone()) + .filter(col("x.a").eq(lit(1))) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(source) + .filter(col("x.a").eq(lit(2))) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right) + .unwrap() + .build() + .unwrap(); + assert_not_rewritten(plan); + } + } + + #[test] + fn keep_union_distinct_with_volatile_source_without_filter() { + let source = LogicalPlanBuilder::from(test_table_scan_with_name("t").unwrap()) + .project(vec![volatile_expr().alias("v"), col("a")]) + .unwrap() + // Keep the volatile projection inside the retained source rather + // than letting peel_wrappers classify it as an outer wrapper. + .distinct() + .unwrap() + .build() + .unwrap(); + let left = LogicalPlanBuilder::from(source.clone()) + .project(vec![col("a")]) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(source) + .filter(col("a").eq(lit(2))) + .unwrap() + .project(vec![col("a")]) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right) + .unwrap() + .build() + .unwrap(); + assert_not_rewritten(plan); + } + + #[test] + fn keep_union_distinct_with_subquery_in_source() { + let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("t2").unwrap()) + .project(vec![col("b")]) + .unwrap() + .build() + .unwrap(); + let source = LogicalPlanBuilder::from(test_table_scan_with_name("t").unwrap()) + .project(vec![ + datafusion_expr::scalar_subquery(Arc::new(subquery)).alias("v"), + col("a"), + ]) + .unwrap() + .build() + .unwrap(); + let left = LogicalPlanBuilder::from(source.clone()) + .filter(col("a").eq(lit(1))) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(source) + .filter(col("a").eq(lit(2))) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right) + .unwrap() + .build() + .unwrap(); + assert_not_rewritten(plan); + } + #[test] fn rewrite_union_distinct_same_source_filters() -> Result<()> { - let left = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let scan = test_table_scan_with_name("t")?; + let left = LogicalPlanBuilder::from(scan.clone()) .filter(col("a").eq(lit(1)))? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let right = LogicalPlanBuilder::from(scan) .filter(col("a").eq(lit(2)))? .build()?; @@ -428,6 +703,82 @@ mod tests { Ok(()) } + #[test] + fn keep_union_distinct_same_metadata_different_root_sources() -> Result<()> { + let left_scan_plan = test_table_scan_with_name("t")?; + let right_scan_plan = test_table_scan_with_name("t")?; + + // TableScan equality compares metadata but deliberately ignores its provider. + assert_eq!(left_scan_plan, right_scan_plan); + let (LogicalPlan::TableScan(left_scan), LogicalPlan::TableScan(right_scan)) = + (&left_scan_plan, &right_scan_plan) + else { + panic!("test helper must return TableScan plans"); + }; + assert!(!Arc::ptr_eq(&left_scan.source, &right_scan.source)); + + let left = LogicalPlanBuilder::from(left_scan_plan) + .filter(col("a").eq(lit(1)))? + .build()?; + let right = LogicalPlanBuilder::from(right_scan_plan) + .filter(col("a").eq(lit(2)))? + .build()?; + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right)? + .build()?; + + assert_not_rewritten(plan); + Ok(()) + } + + #[test] + fn keep_union_distinct_same_metadata_different_nested_sources() { + let left_scan_plan = test_table_scan_with_name("t").unwrap(); + let right_scan_plan = test_table_scan_with_name("t").unwrap(); + assert_eq!(left_scan_plan, right_scan_plan); + let left_sources = collect_table_sources(&left_scan_plan).unwrap(); + let right_sources = collect_table_sources(&right_scan_plan).unwrap(); + assert_eq!(left_sources.len(), 1); + assert_eq!(right_sources.len(), 1); + assert!(!Arc::ptr_eq(&left_sources[0], &right_sources[0])); + + // Verify collection through the Projection/SubqueryAlias source shape + // retained below a branch Filter. + let left_source = LogicalPlanBuilder::from(left_scan_plan) + .project(vec![col("a")]) + .unwrap() + .alias("x") + .unwrap() + .build() + .unwrap(); + let right_source = LogicalPlanBuilder::from(right_scan_plan) + .project(vec![col("a")]) + .unwrap() + .alias("x") + .unwrap() + .build() + .unwrap(); + assert_eq!(left_source, right_source); + + let left = LogicalPlanBuilder::from(left_source) + .filter(col("x.a").eq(lit(1))) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(right_source) + .filter(col("x.a").eq(lit(2))) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right) + .unwrap() + .build() + .unwrap(); + + assert_not_rewritten(plan); + } + #[test] fn keep_union_distinct_different_sources() -> Result<()> { let left = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) @@ -454,10 +805,11 @@ mod tests { #[test] fn keep_union_distinct_with_volatile_predicate() -> Result<()> { - let left = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let scan = test_table_scan_with_name("t")?; + let left = LogicalPlanBuilder::from(scan.clone()) .filter(volatile_expr().gt(lit(0.5_f64)))? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let right = LogicalPlanBuilder::from(scan) .filter(col("a").eq(lit(2)))? .build()?; @@ -478,10 +830,11 @@ mod tests { #[test] fn rewrite_union_distinct_with_matching_projection_prefix() -> Result<()> { - let left = LogicalPlanBuilder::from(test_table_scan_with_name("emp")?) + let scan = test_table_scan_with_name("emp")?; + let left = LogicalPlanBuilder::from(scan.clone()) .project(vec![col("a").alias("mgr"), col("b").alias("comm")])? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("emp")?) + let right = LogicalPlanBuilder::from(scan) .filter(col("b").eq(lit(5)))? .project(vec![col("a").alias("mgr"), col("b").alias("comm")])? .build()?; @@ -499,17 +852,55 @@ mod tests { Ok(()) } + #[test] + fn keep_union_distinct_with_computed_projection_below_filter() -> Result<()> { + let scan = test_table_scan_with_name("prices")?; + let left = LogicalPlanBuilder::from(scan.clone()) + .project(vec![(col("a") + lit(100)).alias("amount")])? + .alias("prices")? + .filter(col("amount").gt(lit(0)))? + .project(vec![col("amount")])? + .build()?; + let right = LogicalPlanBuilder::from(scan) + .project(vec![(col("a") + lit(200)).alias("amount")])? + .alias("prices")? + .filter(col("amount").gt(lit(1)))? + .project(vec![col("amount")])? + .build()?; + + let plan = LogicalPlanBuilder::from(left) + .union_distinct(right)? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Distinct: + Union + Projection: prices.amount + Filter: prices.amount > Int32(0) + SubqueryAlias: prices + Projection: prices.a + Int32(100) AS amount + TableScan: prices + Projection: prices.amount + Filter: prices.amount > Int32(1) + SubqueryAlias: prices + Projection: prices.a + Int32(200) AS amount + TableScan: prices + ")?; + Ok(()) + } + /// A volatile expression in the **projection** (SELECT list) must block the /// rewrite. Each original branch evaluates it independently; merging them /// would evaluate it once per combined row, changing the row set. #[test] fn keep_union_distinct_with_volatile_projection() -> Result<()> { // Both branches project volatile_test() AS v over the same source. - let left = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let scan = test_table_scan_with_name("t")?; + let left = LogicalPlanBuilder::from(scan.clone()) .filter(col("a").eq(lit(1)))? .project(vec![volatile_expr().alias("v"), col("a")])? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let right = LogicalPlanBuilder::from(scan) .filter(col("a").eq(lit(2)))? .project(vec![volatile_expr().alias("v"), col("a")])? .build()?; @@ -546,11 +937,12 @@ mod tests { ); let sq = scalar_subquery(subquery_plan); - let left = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let scan = test_table_scan_with_name("t")?; + let left = LogicalPlanBuilder::from(scan.clone()) .filter(col("a").eq(lit(1)))? .project(vec![sq.clone().alias("sub"), col("a")])? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("t")?) + let right = LogicalPlanBuilder::from(scan) .filter(col("a").eq(lit(2)))? .project(vec![sq.alias("sub"), col("a")])? .build()?; @@ -591,11 +983,12 @@ mod tests { /// single branch would lose the per-branch LIMIT semantics. #[test] fn keep_union_distinct_with_limit_branches() -> Result<()> { - let left = LogicalPlanBuilder::from(test_table_scan_with_name("emp")?) + let scan = test_table_scan_with_name("emp")?; + let left = LogicalPlanBuilder::from(scan.clone()) .project(vec![col("a").alias("mgr"), col("b").alias("comm")])? .limit(0, Some(2))? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("emp")?) + let right = LogicalPlanBuilder::from(scan) .project(vec![col("a").alias("mgr"), col("b").alias("comm")])? .limit(0, Some(2))? .build()?; @@ -622,11 +1015,12 @@ mod tests { /// in the result; merging the branches would silently discard the Sort. #[test] fn keep_union_distinct_with_sort_branches() -> Result<()> { - let left = LogicalPlanBuilder::from(test_table_scan_with_name("emp")?) + let scan = test_table_scan_with_name("emp")?; + let left = LogicalPlanBuilder::from(scan.clone()) .project(vec![col("a").alias("mgr"), col("b").alias("comm")])? .sort(vec![col("a").sort(true, true)])? .build()?; - let right = LogicalPlanBuilder::from(test_table_scan_with_name("emp")?) + let right = LogicalPlanBuilder::from(scan) .project(vec![col("a").alias("mgr"), col("b").alias("comm")])? .sort(vec![col("a").sort(true, true)])? .build()?; diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 7ec48ed21713c..9d953f5075425 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -302,6 +302,40 @@ physical_plan statement ok set datafusion.optimizer.enable_unions_to_filter = false; +statement ok +CREATE TABLE prices AS SELECT 1000 AS amount; + +# Establish the results with the rewrite disabled before repeating these +# queries with it enabled. The derived tables deliberately reuse the table +# and column names, but compute different values below their filters. +query I rowsort +SELECT prices.amount +FROM (SELECT amount + 100 AS amount FROM prices) AS prices +WHERE amount > 0 +UNION +SELECT prices.amount +FROM (SELECT amount + 200 AS amount FROM prices) AS prices +WHERE amount > 1 +---- +1100 +1200 + +query I rowsort +SELECT x.id FROM t1 AS x WHERE x.id = 1 +UNION +SELECT x.id FROM t1 AS x WHERE x.id = 2 +---- +1 +2 + +query I rowsort +SELECT p.amount FROM (SELECT id + 100 AS amount FROM t1) AS p WHERE p.amount = 101 +UNION +SELECT p.amount FROM (SELECT id + 100 AS amount FROM t1) AS p WHERE p.amount = 102 +---- +101 +102 + query TT EXPLAIN SELECT id, name FROM t1 WHERE id = 1 UNION SELECT id, name FROM t1 WHERE id = 2 ---- @@ -349,6 +383,61 @@ SELECT id, name FROM t1 WHERE id = 1 UNION SELECT id, name FROM t1 WHERE id = 2 1 Alex 2 Bob +# Regression: different computed projections below the branch filters must +# preserve both values when unions_to_filter is enabled. +query I rowsort +SELECT prices.amount +FROM (SELECT amount + 100 AS amount FROM prices) AS prices +WHERE amount > 0 +UNION +SELECT prices.amount +FROM (SELECT amount + 200 AS amount FROM prices) AS prices +WHERE amount > 1 +---- +1100 +1200 + +# Regression: predicates qualified with a table alias must still resolve +# against the source below the merged filter. +query I rowsort +SELECT x.id FROM t1 AS x WHERE x.id = 1 +UNION +SELECT x.id FROM t1 AS x WHERE x.id = 2 +---- +1 +2 + +query TT +EXPLAIN SELECT x.id FROM t1 AS x WHERE x.id = 1 +UNION +SELECT x.id FROM t1 AS x WHERE x.id = 2 +---- +logical_plan +01)Aggregate: groupBy=[[id]], aggr=[[]] +02)--SubqueryAlias: x +03)----Filter: t1.id = Int32(1) OR t1.id = Int32(2) +04)------TableScan: t1 projection=[id] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[] +02)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------FilterExec: id@0 = 1 OR id@0 = 2 +06)----------DataSourceExec: partitions=1, partition_sizes=[1] + +# Matching computed projections remain eligible for the rewrite and must +# return the same rows as they do with the rewrite disabled. +query I rowsort +SELECT p.amount FROM (SELECT id + 100 AS amount FROM t1) AS p WHERE p.amount = 101 +UNION +SELECT p.amount FROM (SELECT id + 100 AS amount FROM t1) AS p WHERE p.amount = 102 +---- +101 +102 + +statement ok +DROP TABLE prices; + statement ok set datafusion.optimizer.enable_unions_to_filter = false;