fix: preserve computed projections in unions_to_filter by aoto-tech · Pull Request #25074 · apache/datafusion

aoto-tech · 2026-09-08T12:48:47Z

Which issue does this PR close?

Closes unions_to_filter drops computed projections and returns wrong results #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.

kumarUjjawal

Thank you @pushnanashi2 for the fix. I have left few comment for you. Please take a look. Some slt test would also be nice to have.

kumarUjjawal · 2026-09-08T14:12:13Z

-            LogicalPlan::Projection(Projection { input, .. }) => {
-                Arc::unwrap_or_clone(input)
-            }
+            LogicalPlan::Projection(_) => return None,


Removing SubqueryAlias has the same defect you just fixed for projections? SELECT x.a FROM t AS x WHERE x.a = 1 UNION SELECT x.a FROM t AS x WHERE x.a = 2 plans as a filter over the alias over TableScan: t, so the merged predicate still names x.a while the source only exposes t.a.

We should trip nothing and take the filter input as the source. GroupKey equality then keeps different projections apart and still merges identical ones, so views and derived tables keep the rewrite instead of losing it. That source needs the volatility test that the wrappers already get.

Thanks, keeping the filter input intact makes sense. I’ll update the implementation and add tests over the next day or two.

I’ll also check whether different underlying sources with otherwise identical scan metadata can end up in the same GroupKey, since TableScan::eq does not compare source. I haven’t reproduced this yet; I’ll investigate and report back.

kumarUjjawal · 2026-09-08T14:12:36Z

-            wrappers,
-        })),
+        other => {
+            let Some(source) = strip_passthrough_nodes(other) else {


peel_wrappers already consumed every Projection and SubqueryAlias before this match, so the plan that reaches this arm can never be either one.

strip_passthrough_nodes returns Some on the first iteration here, and the new debug message cannot fire. I would use other directly as the source.

Thanks for pointing this out. You're right: after peel_wrappers, other cannot be a Projection or SubqueryAlias, so strip_passthrough_nodes always returns Some(other) here. I’ll use other directly as the source and remove the redundant check and unreachable debug message

codecov-commenter · 2026-09-08T18:21:56Z

Codecov Report

❌ Patch coverage is 90.83558% with 34 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.89%. Comparing base (681705e) to head (40afd95).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/optimizer/src/unions_to_filter.rs	90.83%	1 Missing and 33 partials ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #25074    +/-   ##
========================================
  Coverage   81.88%   81.89%            
========================================
  Files        1133     1133            
  Lines      424522   424855   +333     
  Branches   424522   424855   +333     
========================================
+ Hits       347622   347936   +314     
+ Misses      56288    56285     -3     
- Partials    20612    20634    +22

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

aoto-tech · 2026-09-10T08:17:09Z

@kumarUjjawal san
Thanks for the review. I’ve pushed updates addressing both comments:

Kept the filter input intact, including Projection and SubqueryAlias. Different computed projections remain separate, while identical ones over the same source can still be merged.
Added recursive checks for volatile expressions and subqueries throughout the retained source.
Simplified the other arm to use other directly and removed strip_passthrough_nodes, including the redundant check and unreachable debug message.
Added optimizer unit tests and SQL regression tests covering computed projections and aliases, with the rewrite both enabled and disabled in the SQL tests.

Regarding the source-identity concern I mentioned earlier, GroupKey now also checks the corresponding TableSource instances using Arc::ptr_eq, since TableScan equality ignores source. This conservatively prevents distinct source instances with identical scan metadata from being merged. I added regression tests for both direct scans and scans nested below projections and aliases.

I also used unwrap() instead of ? in the newly added tests to address a CI coverage issue: when using ?, coverage was silently missed without an explicit error.

Please take another look when you have a chance.

kumarUjjawal · 2026-09-10T08:18:43Z

Thank you @aoto-tech I will take a look on the weekends.

kumarUjjawal

LGTM!

jayzhan211 · 2026-09-13T07:39:08Z

+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


It would be nice to add another explain plan to show the plan

Great suggestion! I’ll make that change.

I hadn’t noticed

jayzhan211

Thanks @aoto-tech , non-blocking suggestions

jayzhan211 · 2026-09-13T07:40:29Z

+    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").add(lit(100)).alias("amount")])


col("a") + lit(100) without std::ops::Add

The current test only checks that it returns 1 and 2, so it doesn’t prove that the rewrite still happens with the alias. I’ll add an EXPLAIN case for that.

aoto-tech · 2026-09-13T09:33:34Z

@jayzhan211 san
I’ve addressed your suggestions.

kumarUjjawal · 2026-09-13T16:38:09Z

Thank you @aoto-tech and @jayzhan211 for the review.

fix: preserve computed projections in unions_to_filter

56f24c4

github-actions Bot added the optimizer Optimizer rules label Sep 8, 2026

kumarUjjawal reviewed Sep 8, 2026

View reviewed changes

aoto-tech and others added 3 commits September 9, 2026 21:23

Merge branch 'main' into fix/unions-to-filter-computed-projection

cb8f98a

fix(optimizer): preserve union branch filter inputs

d101432

fix(optimizer): distinguish union branch table sources

359aa64

github-actions Bot added the sqllogictest SQL Logic Tests (.slt) label Sep 9, 2026

kumarUjjawal approved these changes Sep 11, 2026

View reviewed changes

jayzhan211 reviewed Sep 13, 2026

View reviewed changes

jayzhan211 approved these changes Sep 13, 2026

View reviewed changes

aoto-tech and others added 2 commits September 13, 2026 17:30

test(optimizer): address union rewrite review suggestions

b91ca60

Merge branch 'main' into fix/unions-to-filter-computed-projection

40afd95

kumarUjjawal added this pull request to the merge queue Sep 13, 2026

Merged via the queue into apache:main with commit c5257f0 Sep 13, 2026
41 checks passed

nuno-faria added the v56.0.0 label Oct 3, 2026

Conversation

aoto-tech commented Sep 8, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

kumarUjjawal left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

aoto-tech commented Sep 10, 2026

Uh oh!

kumarUjjawal commented Sep 10, 2026

Uh oh!

kumarUjjawal left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayzhan211 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aoto-tech commented Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kumarUjjawal commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

codecov-commenter commented Sep 8, 2026 •

edited

Loading

aoto-tech commented Sep 13, 2026 •

edited

Loading