fix: decorrelate grouping sets that leave out the correlated column when safe

[namanjain24-sudo]

Which issue does this PR close?

• Closes Decorrelate subqueries whose grouping sets leave out the correlated column #25708.

Rationale for this change

#25529 fixed a wrong-results bug where a correlated filter under a grouping-set Aggregate was pulled up into the aggregate's grouping sets, but the guard it added is more conservative than it needs to be: it rejects the pull up whenever any grouping set omits the correlated column, even in cases like

SELECT o.k FROM o
WHERE EXISTS (
  SELECT 1 FROM i WHERE i.k = o.k
  GROUP BY GROUPING SETS ((i.k), (i.j))
)

where nothing above the aggregate ever reads i.k, so filling it in changes nothing the query can observe. These queries now fail to plan (This feature is not implemented: ... does not support logical expression Exists) instead of decorrelating.

What changes are included in this PR?

• datafusion/optimizer/src/decorrelate.rs: grouping_sets_cover_pull_up_cols becomes grouping_sets_pull_up_outcome, returning a three-way GroupingSetsPullUp (NoColumnsToAdd / SafeToExtend / Unsafe) instead of a bool. A set that omits a required column is now only Unsafe when:
  • some grouping set is empty (ROLLUP/CUBE always hold one), or
  • the aggregate itself computes GROUPING/GROUPING_ID, or
  • a node above the aggregate reads the column, per a new column_refs_above_aggregate set.
• Since PullUpCorrelatedExpr::f_up runs bottom-up, it hasn't visited whatever sits above the Aggregate (a HAVING, a Projection) by the time it needs to decide. column_refs_above_aggregate is instead computed once, top-down, over the original (not yet rewritten) subquery plan by a new columns_read_above_aggregate helper, and threaded in via a new with_column_refs_above_aggregate builder call at all three call sites (scalar_subquery_to_join.rs, decorrelate_predicate_subquery.rs, decorrelate_lateral_join.rs).
• datafusion/sqllogictest/test_files/subquery.slt: the EXISTS case above now decorrelates and asserts correct results; added a case showing a GROUPING(k) read still keeps the subquery correlated (correctly), alongside the existing HAVING/empty-set cases that stay correlated.

What is the testing strategy for this PR?

Extended datafusion/sqllogictest/test_files/subquery.slt's existing #25519 regression block:

• The previously-rejected EXISTS ... GROUPING SETS ((i.k), (i.j)) case is now a query IB asserting the correct results.
• Added a case with GROUPING(k) read out of the aggregate, asserting it still correctly stays correlated.
• The HAVING i.k IS NULL and every empty-set (ROLLUP/CUBE/explicit ()) case are unchanged and still assert they stay correlated.

Ran the full datafusion-optimizer unit + integration suite (905 + 26 tests) and the full sqllogictest suite (523 files), all green.

Are there any user-facing changes?

More EXISTS/IN/lateral-join subqueries with a correlated filter under a grouping-set aggregate now decorrelate into a join instead of failing to plan. No public API changes.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.23810% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.77%. Comparing base (c3ef346) to head (f856112).
⚠️ Report is 72 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/optimizer/src/decorrelate.rs	94.80%	3 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25781      +/-   ##
==========================================
+ Coverage   82.66%   82.77%   +0.10%     
==========================================
  Files        1147     1147              
  Lines      446357   450639    +4282     
  Branches   446357   450639    +4282     
==========================================
+ Hits       368971   372994    +4023     
+ Misses      54997    54932      -65     
- Partials    22389    22713     +324     

☔ 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.

[namanjain24-sudo]

Still open for review whenever someone has bandwidth.

[jayzhan211]

Thanks @namanjain24-sudo , here is a suggestion

[jayzhan211]

walk returns at the first Aggregate top-down, and any skips later inputs once one holds an aggregate. So reads of the NULL-filled column above a stacked aggregate, in a later join input, or between nested grouping sets are never collected. The result is SafeToExtend and wrong rows; on main these queries fail to plan instead.

Repro (tables from subquery.slt):

-- expected (1,1) (2,1) (4,0) (5,1) (NULL,0); returns 0 for every row
SELECT gs_outer.k, (SELECT count(*) FROM (SELECT gs_inner.k AS kk FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j))) t WHERE t.kk IS NULL) FROM gs_outer ORDER BY gs_outer.k;

-- expected true for 1, 2, 5; returns false for every row
SELECT gs_outer.k, EXISTS (SELECT 1 FROM (SELECT count(*) AS c FROM gs_inner) a JOIN (SELECT gs_inner.k FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j))) b ON a.c > 0 WHERE b.k IS NULL) FROM gs_outer ORDER BY gs_outer.k;

-- expected true for 1, 2, 5; returns false for every row
SELECT gs_outer.k, EXISTS (SELECT 1 FROM (SELECT gs_inner.k AS kk, gs_inner.j AS jj FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j))) t GROUP BY GROUPING SETS ((t.kk), (t.kk, t.jj)) HAVING t.kk IS NULL) FROM gs_outer ORDER BY gs_outer.k;

Fix: walk every input and through every aggregate, and count a node as "above" when any grouping-set aggregate is below it. With this, the three queries stay correlated, the ((k), (j)) EXISTS case still decorrelates, and the full slt suite passes; please add the three as statement error cases.

-/// Columns read by a node strictly above the first `Aggregate` found in
-/// `plan`, which is the subquery's original, not yet rewritten, plan.
+/// Columns read by a node that has a grouping set `Aggregate` below it in
+/// `plan`, which is the subquery's original, not yet rewritten, plan.
@@
     fn walk(plan: &LogicalPlan, above: &mut BTreeSet<Column>) -> bool {
         if matches!(plan, LogicalPlan::Subquery(_)) {
             return false;
         }
-        if matches!(plan, LogicalPlan::Aggregate(_)) {
-            return true;
+        // Visit every input, and walk through every aggregate: a node above
+        // a grouping set may sit in a later input, or above another aggregate.
+        let mut found_below = false;
+        for child in plan.inputs() {
+            found_below |= walk(child, above);
         }
-        let found_below = plan.inputs().into_iter().any(|child| walk(child, above));
         if found_below {
             for expr in plan.expressions() {
                 above.extend(expr.column_refs().into_iter().cloned());
             }
         }
         found_below
+            || matches!(plan, LogicalPlan::Aggregate(aggregate) if aggregate
+                .group_expr
+                .iter()
+                .any(|expr| matches!(expr, Expr::GroupingSet(_))))
     }

[namanjain24-sudo]

Applied exactly as suggested. Verified by reverting locally and running the three repros: query 1 returns count=0 for every row, queries 2 and 3 return false for every row — exactly the wrong output described. With the fix, all three correctly error instead, and the existing ((k), (j)) EXISTS case still decorrelates. Added the three as statement error cases in subquery.slt; full suite passes.