unions_to_filter drops computed projections and returns wrong results

[aoto-tech]

Describe the bug

In DataFusion 55.0.0, built from commit a5c809f, enabling datafusion.optimizer.enable_unions_to_filter changes the result of a valid UNION DISTINCT query.

The query returns [1100, 1200] with the option disabled and [1000] with the option enabled. No error or warning is produced.

The option defaults to false, so the default configuration is not affected. The impact is incorrect query results when this opt-in optimizer rule is enabled.

To Reproduce

The following was reproduced with datafusion-cli built from commit a5c809f.

CREATE TABLE prices AS
SELECT 1000 AS amount;

SET datafusion.optimizer.enable_unions_to_filter = true;

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;

The derived tables deliberately reuse the base table name as their alias. Run the same SELECT with enable_unions_to_filter set to false as a control.

Observed results:

enable_unions_to_filter = false:
1100
1200

enable_unions_to_filter = true:
1000

With the option enabled, EXPLAIN shows that the UNION and both computed projections have disappeared:

Aggregate: groupBy=[[amount]], aggr=[[]]
Filter: prices.amount > Int64(0) OR prices.amount > Int64(1)
TableScan: prices projection=[amount]

Expected behavior

The query should return the same two rows regardless of the optimizer setting:

1100
1200

The optimizer should preserve amount + 100 and amount + 200 before merging the branch predicates.

Additional context

The cause appears to be the interaction between extract_branch and strip_passthrough_nodes in unions_to_filter.rs.

Projections above the branch filter are retained, while projections below the filter are removed without checking their expressions. In this query, both branches are reduced to the same TableScan: prices and the same outer projection, so they are merged into one scan with the combined predicate.

Relevant source:

datafusion/datafusion/optimizer/src/unions_to_filter.rs

Lines 161 to 210 in a5c809f

	fn extract_branch(plan: LogicalPlan) -> Result<Option<UnionBranch>> {
	let (wrappers, plan) = peel_wrappers(plan);
	// Volatile or subquery expressions in the projection must not be merged:
	// they are evaluated once per branch in the original plan but would be
	// evaluated once per combined row after the rewrite, which can change the
	// output row set.
	if !wrapper_projections_are_safe(&wrappers) {
	debug!(
	"unions_to_filter skipped: projection wrapper contains volatile expression or subquery"
	);
	return Ok(None);
	}
	match plan {
	LogicalPlan::Filter(Filter {
	predicate, input, ..
	}) => {
	if !is_mergeable_predicate(&predicate) {
	debug!(
	"unions_to_filter skipped: branch predicate contains volatility or a subquery"
	);
	return Ok(None);
	}
	Ok(Some(UnionBranch {
	source: strip_passthrough_nodes(Arc::unwrap_or_clone(input)),
	predicate,
	wrappers,
	}))
	}
	// 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
	// row restriction (LIMIT) or rely on an order guarantee that UNION does
	// not preserve (ORDER BY). Bail out to leave the UNION unchanged.
	LogicalPlan::Limit(_) => {
	debug!("unions_to_filter skipped: branch contains LIMIT");
	Ok(None)
	}
	LogicalPlan::Sort(_) => {
	debug!("unions_to_filter skipped: branch contains ORDER BY / SORT");
	Ok(None)
	}
	other => Ok(Some(UnionBranch {
	source: strip_passthrough_nodes(other),
	predicate: Expr::Literal(
	datafusion_common::ScalarValue::Boolean(Some(true)),
	None,
	),
	wrappers,
	})),

datafusion/datafusion/optimizer/src/unions_to_filter.rs

Lines 281 to 293 in a5c809f

	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,
	};
	}
	}

I searched existing DataFusion issues and pull requests for unions_to_filter and the computed-projection case but did not find a matching report.

I investigated the source code and ran the reproducer myself. I used AI to cross-check my findings and proofread the English because I am not a native English speaker.

A regression test should cover two branches over the same base table with different computed projections below their filters, with the optimization both enabled and disabled.