fix: walk every input and through every aggregate in columns_read_above_aggregate

The old walk stopped at the first Aggregate found top-down and used
any() over the plan's inputs, which short-circuits once one input
returns true. So a read of the NULL-filled column above a stacked
(non-grouping-set) aggregate, in a later join input, or between
nested grouping sets was never collected, wrongly reporting
SafeToExtend and producing wrong rows instead of staying correlated.

Now a node counts as "below a grouping set" when any grouping-set
aggregate is below it specifically (not just any aggregate), every
input is visited unconditionally, and the walk continues through
every aggregate rather than stopping at the first one.

Adds the three repro queries as statement error cases in subquery.slt.

Author: namanjain24-sudo

datafusion/optimizer/src/decorrelate.rs

@@ -918,32 +918,39 @@ fn is_grouping_call(expr: &Expr) -> bool {
     matches!(expr, Expr::AggregateFunction(agg) if agg.func.name().eq_ignore_ascii_case("grouping"))
 }
 
-/// Columns read by a node strictly above the first `Aggregate` found in
+/// 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.
 ///
 /// [`PullUpCorrelatedExpr::f_up`] runs bottom-up, so by the time it visits an
 /// `Aggregate` it cannot yet tell whether a `HAVING` or `Projection` above it
 /// reads one of the columns a grouping set pull up would add. This walks the
-/// plan top-down instead, before the rewrite starts, and stops at the first
-/// `Aggregate` along each branch, collecting the columns every node above it
-/// reads in its own expressions. A nested `Subquery` is a different
-/// correlation scope and is skipped, the same way
+/// plan top-down instead, before the rewrite starts, collecting the columns
+/// every node above a grouping set `Aggregate` reads in its own expressions.
+/// It visits every input and walks through every `Aggregate` - a node that
+/// reads the pulled-up column may sit in a later input, above another
+/// (non-grouping-set) `Aggregate`, or between nested grouping sets, not only
+/// directly above the first `Aggregate` found. A nested `Subquery` is a
+/// different correlation scope and is skipped, the same way
 /// [`PullUpCorrelatedExpr::f_down`] skips it.
 fn columns_read_above_aggregate(plan: &LogicalPlan) -> BTreeSet<Column> {
     fn walk(plan: &LogicalPlan, above: &mut BTreeSet<Column>) -> bool {
         if matches!(plan, LogicalPlan::Subquery(_)) {
             return false;
         }
-        if matches!(plan, LogicalPlan::Aggregate(_)) {
-            return true;
+        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(_))))
     }
 
     let mut above = BTreeSet::new();

datafusion/sqllogictest/test_files/subquery.slt

@@ -2943,6 +2943,25 @@ SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k
 5 true
 NULL false
 
+# A filter above the grouping set `Aggregate`, through another `Aggregate`
+# stacked on top of it, reads the NULL filled column. `walk` must not stop at
+# the first `Aggregate` it finds (the outer `count(*)`) without recursing into
+# its input, or it never sees that read and wrongly decorrelates.
+statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression ScalarSubquery
+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;
+
+# The same read, but in a later join input rather than directly above the
+# grouping set `Aggregate`. `walk` must visit every input, not stop once one
+# of them turns up a grouping set `Aggregate` below it.
+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 (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;
+
+# The same read, above a second, nested grouping set `Aggregate` stacked on
+# top of the first. `walk` must keep walking through a grouping set
+# `Aggregate` too, not only a plain one.
+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 (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;
+
 # 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 holds k = 1 and is filtered out. On main this