fix: keep a correlated filter below an aggregate with a grouping set

[namanjain24-sudo]

Which issue does this PR close?

• Closes Wrong results: a correlated filter under an aggregate with a grouping set is pulled above the aggregate #25519.

Rationale for this change

A correlated subquery whose filter sits below an aggregate with a grouping set
returns wrong results, with no error and no warning.

PullUpCorrelatedExpr moves the correlated filter above the aggregate and adds
the correlated column to the aggregate's group expressions.
LogicalPlanBuilder::aggregate cross joins a plain group expression with the
sets a grouping set already holds, so ROLLUP(i.k), which is
GROUPING SETS ((i.k), ()), becomes GROUPING SETS ((i.k), (i.k, i.k)). The
empty set is gone, and with it the grand total row the subquery returns for
every outer row, including the rows whose filter matches nothing. The join that
replaces the filter cannot bring those rows back.

On main at b4a8c824b4, with datafusion-cli:

CREATE TABLE o(k INT) AS VALUES (1), (2), (NULL), (4), (5);
CREATE TABLE i(k INT) AS VALUES (1), (NULL), (5), (2);

SELECT o.k, EXISTS (SELECT 1 FROM i WHERE i.k = o.k GROUP BY ROLLUP(i.k)) AS e FROM o ORDER BY o.k;
SELECT o.k, o.k IN (SELECT i.k FROM i WHERE i.k = o.k GROUP BY ROLLUP(i.k)) AS m FROM o ORDER BY o.k;

o.k	EXISTS on main	IN on main	correct, per the issue
1	true	true	true / true
2	true	true	true / true
4	false	false	true / NULL
5	true	true	true / true
NULL	false	false	true / NULL

The physical plan for the EXISTS query shows the loss directly. ROLLUP(k)
has two sets, and both of them now group by k:

AggregateExec: group_by: (k), (k), mode: Partial

What changes are included in this PR?

datafusion/optimizer/src/decorrelate.rs, in the Aggregate arm of
PullUpCorrelatedExpr::f_up:

• When the group expressions hold a grouping set, the pull up no longer adds its
  columns to them. If any set would be missing one, can_pull_up is set to
  false and the subquery stays correlated, which is what the issue asks for
  instead of a wrong result.
• ROLLUP and CUBE always contain the empty set, so they are only safe when
  the pull up has nothing to add. An explicit GROUPING SETS is checked set by
  set.
• When every set already groups by each column the pull up needs, the aggregate
  keeps its sets untouched and the subquery still decorrelates. Before this PR
  that case appended the column anyway and produced
  GROUPING SETS ((c, c), (c, b, c)); it now stays GROUPING SETS ((c), (c, b)).
• The columns checked are the ones collect_missing_exprs would append: the
  correlated columns and the columns of a pulled up HAVING, minus the ones the
  group expressions already list on their own.

The rewrite still runs to completion in the unsupported case, the same way the
existing can_pull_over_aggregation case does. The three callers
(decorrelate_predicate_subquery, scalar_subquery_to_join,
decorrelate_lateral_join) read can_pull_up only after the whole rewrite has
finished, and the nodes above the aggregate still expect the pulled up columns
in its output, so returning early there fails the rewrite with a schema error
rather than declining the transform.

What is the testing strategy for this PR?

Five unit tests, in the two rules that reach this code path:

• decorrelate_predicate_subquery.rs: ROLLUP under EXISTS, CUBE under
  IN, a GROUPING SETS whose second set groups by another column, and the
  covering GROUPING SETS that must still decorrelate.
• scalar_subquery_to_join.rs: ROLLUP under a correlated scalar subquery.

Twelve queries in datafusion/sqllogictest/test_files/subquery.slt cover the
end-to-end behaviour: ROLLUP, CUBE, an explicit set list holding (), a set
list whose sets group by different columns, a grouping set that does not mention
the correlated column at all, and the IN, NOT EXISTS and correlated scalar
forms; plus the four that must keep working, a covering GROUPING SETS with one
and with two sets, a plain GROUP BY, and an uncorrelated ROLLUP subquery.

Ablation, with only decorrelate.rs reverted to main and every test kept (the
tests live in other files, so the revert does not delete them): all five unit
tests fail, and seven of the eight statement error cases in subquery.slt
fail with "query is expected to fail, but actually succeed". The eighth, the
correlated scalar subquery, already errors on main.

On the branch: cargo test -p datafusion-optimizer passes, 901 + 26 + 5 tests;
the full sqllogictest suite passes, 521 files; cargo clippy -p datafusion-optimizer --all-targets -- -D warnings and cargo fmt --all -- --check are clean.

Are there any user-facing changes?

Yes, and it is worth a look before merging. The queries above stop returning
wrong rows, but they do not start returning right ones: leaving the subquery
correlated means the physical planner rejects it with

This feature is not implemented: Physical plan does not support logical expression Exists(...)

That is the behaviour #25519 asks for, and it matches how the rule already
declines a Union, Sort or Extension that holds an outer reference, and a
Limit that holds one outside an EXISTS. Supporting these
plans properly needs the aggregate to keep its sets through decorrelation, which
is a larger change than this one.

Queries whose grouping set already groups by the correlated column in every set
are unaffected, other than no longer repeating that column inside each set.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 75.95628% with 44 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (b4a8c82) to head (f763afd).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
...on/optimizer/src/decorrelate_predicate_subquery.rs	70.70%	3 Missing and 26 partials ⚠️
datafusion/optimizer/src/decorrelate.rs	85.18%	5 Missing and 3 partials ⚠️
...atafusion/optimizer/src/scalar_subquery_to_join.rs	76.66%	2 Missing and 5 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25529      +/-   ##
==========================================
- Coverage   82.38%   82.38%   -0.01%     
==========================================
  Files        1138     1138              
  Lines      434328   434671     +343     
  Branches   434328   434671     +343     
==========================================
+ Hits       357824   358103     +279     
- Misses      54872    54886      +14     
- Partials    21632    21682      +50     

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

[jayzhan211]

Thanks @namanjain24-sudo! One non-blocking suggestion

[jayzhan211]

Comment gives the empty-set rationale, which doesn't apply: ((k), (j)) has no empty set. The reason is the NULL fill — the pull up turns (j) into (j, k), so gs_inner.k is non-NULL where the original set fills it with NULL. Suggested wording plus a case that shows it (main returns false for every row; should be true for 1, 2, 5):

-# A set that groups by another column does not carry the correlated column either.
+# A set that leaves out the correlated column fills it with NULL. The pull up
+# would turn `(j)` into `(j, k)`, and `k` would then carry a value in the rows
+# where the subquery returns NULL. Anything above the aggregate that reads `k`
+# sees the difference, so the subquery stays correlated.
+#
+# Known limitation: when nothing reads `k`, as here, the pull up was correct
+# before this guard and the query now fails to plan. Telling the two cases apart
+# needs the correlated column added to each set under an alias.
 statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists
 SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j))) FROM gs_outer;
+
+# The same sets with a HAVING that reads the NULL filled column. For k = 1 the
+# `(j)` set yields the row `(NULL, 10)`, which passes the HAVING, so EXISTS is
+# true. With `(j, k)` that row has `k = 1` and is filtered out.
+statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists
+SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j)) HAVING gs_inner.k IS NULL) FROM gs_outer;

The doc on grouping_sets_cover_pull_up_cols has the same gap, it only explains the empty set:

     /// `ROLLUP` and `CUBE` always contain the empty set, which yields a row for
     /// outer rows the correlated filter matches nothing for, so they are only safe
     /// when there is nothing to add.
+    ///
+    /// A non-empty set that leaves a column out fills it with NULL. Adding the
+    /// column would give it a value that a HAVING or a projection above the
+    /// aggregate can read, so such a set is rejected as well.

The alias-based fix is fine as a follow-up.

[namanjain24-sudo]

Thanks, you're right that the comment gave the wrong reason, and the NULL fill is the real one. Applied both your wording and the doc paragraph in f763afd, with the HAVING case added.

I checked the rows behind it first, for k = 1:

GROUPING SETS ((k), (j))       -> (1, NULL), (NULL, 10)
  + HAVING k IS NULL           -> (NULL, 10)          EXISTS true
GROUPING SETS ((k), (j, k))    -> (1, 10),   (1, NULL)
  + HAVING k IS NULL           -> no rows             EXISTS false

One detail on your note: false for every row is the HAVING query. The one already in the file answers true, true, false, true, false for 1, 2, 4, 5, NULL on main, which is correct, so it is only the "fails to plan now" limitation, as your comment says. The HAVING variant is the one main gets wrong, and it is in the file now with that stated.

Agreed on the alias-based fix as a follow-up.