You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/13 16:47:44 UTC
(datafusion) branch main updated: fix: preserve computed projections in unions_to_filter (#25074)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new c5257f0540 fix: preserve computed projections in unions_to_filter (#25074)
c5257f0540 is described below
commit c5257f0540b9152f4aa2fcb7edad3760100f3fd9
Author: YukiAoto <84...@users.noreply.github.com>
AuthorDate: Sun Sep 13 16:38:13 2026 +0000
fix: preserve computed projections in unions_to_filter (#25074)
## Which issue does this PR close?
- Closes #25071
## Rationale for this change
With `enable_unions_to_filter` enabled, `UNION DISTINCT` branches
containing different computed projections below their filters can be
incorrectly treated as equivalent. This produces wrong result values and
row counts without an error.
## What changes are included in this PR?
`unions_to_filter` now skips the rewrite when a projection appears below
a branch filter. Alias nodes remain safe to strip. A regression test
covers branches computing `a + 100` and `a + 200` below their filters.
## What is the testing strategy for this PR?
- `cargo +1.97.0-x86_64-pc-windows-gnullvm test --locked -p
datafusion-optimizer unions_to_filter --lib` (9 passed)
- The standalone reproducer returns `[1100, 1200]` with the optimizer
setting both disabled and enabled after the fix.
- `git diff --check`
## Are there any user-facing changes?
This prevents incorrect query results when the opt-in `unions_to_filter`
optimizer rule is enabled.
---
datafusion/optimizer/src/unions_to_filter.rs | 496 ++++++++++++++++++++++++---
datafusion/sqllogictest/test_files/union.slt | 89 +++++
2 files changed, 534 insertions(+), 51 deletions(-)
diff --git a/datafusion/optimizer/src/unions_to_filter.rs b/datafusion/optimizer/src/unions_to_filter.rs
index 158fd35828..2eda0daba1 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<Option<LogicalPlan>>
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<Option<UnionBranch>> {
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<Option<UnionBranch>> {
);
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<Option<UnionBranch>> {
// 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<Wrapper>,
+ // LogicalPlan equality deliberately ignores TableScan::source. Retain the
+ // corresponding Arcs so structurally identical scans over different data
+ // cannot be grouped together.
+ table_sources: Vec<Arc<dyn TableSource>>,
}
+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<LogicalPla
Ok(plan)
}
-fn strip_passthrough_nodes(mut plan: LogicalPlan) -> 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<bool> {
+ 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<Vec<Arc<dyn TableSource>>> {
+ 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 7ec48ed217..9d953f5075 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org