From 42adbba9a2dd70a1cad056597c85170fa6ad48c0 Mon Sep 17 00:00:00 2001 From: naman Date: Sun, 27 Sep 2026 02:32:15 +0530 Subject: [PATCH 1/5] fix: decorrelate grouping sets that leave out the correlated column when safe correlated column, even when nothing would ever observe the difference. Only reject now when a set is empty, the aggregate computes GROUPING/GROUPING_ID, or something above the aggregate (checked ahead of time, since the rewrite runs bottom-up and hasn't visited it yet) reads the column that would change. Fixes #25708 --- datafusion/optimizer/src/decorrelate.rs | 248 +++++++++++++----- .../optimizer/src/decorrelate_lateral_join.rs | 4 +- .../src/decorrelate_predicate_subquery.rs | 3 +- .../optimizer/src/scalar_subquery_to_join.rs | 4 +- .../sqllogictest/test_files/subquery.slt | 30 ++- 5 files changed, 219 insertions(+), 70 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index e4577d6eb0c4c..a96f5a5961ff8 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -36,8 +36,8 @@ use datafusion_expr::utils::{ collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, }; use datafusion_expr::{ - BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, FetchType, - LogicalPlan, LogicalPlanBuilder, Operator, SkipType, expr, lit, + Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, + FetchType, LogicalPlan, LogicalPlanBuilder, Operator, SkipType, expr, lit, }; /// This struct rewrite the sub query plan by pull up the correlated @@ -91,6 +91,13 @@ pub struct PullUpCorrelatedExpr { /// The list is cleared when the pull up passes a node that can put a NULL /// back into such a column: an outer join, a union or a grouping set. pub correlated_filters: Vec, + /// Columns read, above the subquery's `Aggregate`, by a node the rewrite + /// has not reached yet when it visits that `Aggregate` (the rewrite runs + /// bottom-up, so a `Projection` or `HAVING` above it is only visited + /// afterwards). Populated once, from the original subquery plan, by + /// [`Self::with_column_refs_above_aggregate`], since [`Self::f_up`] has no + /// way to compute it on its own. + column_refs_above_aggregate: BTreeSet, } impl Default for PullUpCorrelatedExpr { @@ -113,6 +120,7 @@ impl PullUpCorrelatedExpr { pull_up_having_expr: None, pulled_up_scalar_agg: false, correlated_filters: Vec::new(), + column_refs_above_aggregate: BTreeSet::new(), } } @@ -135,6 +143,43 @@ impl PullUpCorrelatedExpr { self.exists_sub_query = exists_sub_query; self } + + /// Mark the plan as one whose correlated expressions cannot be pulled up + /// and stop descending into it + fn unsupported(&mut self, plan: LogicalPlan) -> Result> { + self.can_pull_up = false; + Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump)) + } + + /// Record the columns read above the subquery's `Aggregate`, computed + /// from `subquery_plan` (the plan as it is before this rewrite runs). + /// + /// A grouping set that omits a correlated column only needs to reject the + /// pull up when something reads the column the pull up would fill in, so + /// [`Self::f_up`] needs this computed ahead of time. See + /// [`columns_read_above_aggregate`]. + pub fn with_column_refs_above_aggregate( + mut self, + subquery_plan: &LogicalPlan, + ) -> Self { + self.column_refs_above_aggregate = columns_read_above_aggregate(subquery_plan); + self + } +} + +/// The outcome of trying to fold the pull up's columns into a grouping set +/// `Aggregate`. See [`PullUpCorrelatedExpr::grouping_sets_pull_up_outcome`]. +#[derive(Debug, PartialEq, Eq)] +enum GroupingSetsPullUp { + /// Every set already groups by the columns; the aggregate's `group_expr` + /// stays as it is. + NoColumnsToAdd, + /// Some set leaves a required column out, but extending every set with it + /// changes nothing the rest of the query can observe. + SafeToExtend, + /// Some set leaves a required column out, and extending every set with it + /// would change what the query returns. + Unsafe, } /// Used to indicate the unmatched rows from the inner(subquery) table after the left out Join @@ -437,42 +482,54 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { &mut local_correlated_cols, ); - // A grouping set cannot take the columns the pull up adds. - // `LogicalPlanBuilder::aggregate` cross joins a plain group - // expression with the sets that are already there, 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 correlated filter matches nothing. - // The join that replaces the filter cannot bring those rows back, - // so the subquery stays correlated unless every set already groups - // by each column the pull up would add. + // A grouping set that omits a column the pull up adds either + // keeps its sets as they are (every set already groups by the + // column), or gains the column in every set, `ROLLUP(i.k)` + // (`GROUPING SETS ((i.k), ())`) becoming + // `GROUPING SETS ((i.k), (i.k, i.k))`. The latter loses the + // empty set, and with it the grand total row the subquery + // returns for every outer row, including the rows whose + // correlated filter matches nothing, so it is only safe when + // no set is empty and nothing reads the column the fill + // changes from `NULL` to the correlated value. let mut missing_exprs = if aggregate .group_expr .iter() .any(|expr| matches!(expr, Expr::GroupingSet(_))) { - if self.grouping_sets_cover_pull_up_cols( - &aggregate.group_expr, - &local_correlated_cols, - ) { - // Every set already groups by them, so the sets stay as - // they are. Adding the columns again would repeat them - // inside every set. - aggregate.group_expr.to_vec() - } else { - self.can_pull_up = false; - // The rewrite still runs, the same way the - // `can_pull_over_aggregation` case above does. The callers - // read `can_pull_up` only after the whole rewrite has - // finished, and the nodes above this one still expect the - // pulled up columns in its output, so leaving them out here - // would fail the rewrite with a schema error instead. They - // drop this plan and keep the correlated subquery. - self.collect_missing_exprs( - &aggregate.group_expr, - &local_correlated_cols, - )? + match self + .grouping_sets_pull_up_outcome(aggregate, &local_correlated_cols) + { + GroupingSetsPullUp::NoColumnsToAdd => { + // Every set already groups by them, so the sets stay as + // they are. Adding the columns again would repeat them + // inside every set. + aggregate.group_expr.to_vec() + } + GroupingSetsPullUp::SafeToExtend => { + // No set is empty and nothing reads the columns the + // fill changes from `NULL` to the correlated value, + // so extending every set is invisible to the rest of + // the query. + self.collect_missing_exprs( + &aggregate.group_expr, + &local_correlated_cols, + )? + } + GroupingSetsPullUp::Unsafe => { + self.can_pull_up = false; + // The rewrite still runs, the same way the + // `can_pull_over_aggregation` case above does. The callers + // read `can_pull_up` only after the whole rewrite has + // finished, and the nodes above this one still expect the + // pulled up columns in its output, so leaving them out here + // would fail the rewrite with a schema error instead. They + // drop this plan and keep the correlated subquery. + self.collect_missing_exprs( + &aggregate.group_expr, + &local_correlated_cols, + )? + } } } else { // add missing columns to Aggregation's group expressions @@ -618,26 +675,41 @@ impl PullUpCorrelatedExpr { Ok(Transformed::yes(plan.recompute_schema()?)) } - /// Whether the pull up can add its columns to `group_expr` without changing - /// what the aggregate returns. + /// Whether the pull up can add its columns to a grouping set `Aggregate`, + /// and whether it needs to. /// - /// `true` when `group_expr` holds no grouping set, and when every set of every - /// grouping set it holds already groups by each column - /// [`Self::collect_missing_exprs`] would add. In the second case the pull up - /// adds nothing and the aggregate keeps the sets it has. + /// [`GroupingSetsPullUp::NoColumnsToAdd`] when every set of every grouping + /// set already groups by each column [`Self::collect_missing_exprs`] + /// would add: the pull up adds nothing and the aggregate keeps the sets + /// it has. /// - /// `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. + /// Otherwise some set leaves a required column out, which fills that + /// column with `NULL` in the set's rows. Adding the column to every set + /// (`LogicalPlanBuilder::aggregate` cross joins a plain group expression + /// with the sets that are already there) gives it a value instead. This + /// is only invisible to the rest of the query when: + /// - no set is empty. `ROLLUP(i.k)`, `GROUPING SETS ((i.k), ())`, always + /// holds one, and turns into `GROUPING SETS ((i.k), (i.k, i.k))`, + /// losing the empty set and the grand total row it yields for every + /// outer row, including the rows the correlated filter matches + /// nothing for. The join that replaces the filter cannot bring that + /// row back. + /// - the aggregate does not compute `GROUPING`/`GROUPING_ID`, which + /// reads whether a row's set groups by a column. + /// - nothing above the aggregate reads the column, which + /// `column_refs_above_aggregate` was populated with ahead of time (this + /// rewrite runs bottom-up, so a `HAVING` or `Projection` above the + /// aggregate is only visited after it). /// - /// A non-empty set that leaves a column out fills it with NULL. Adding the - /// column would give it a value instead, which a `HAVING` or a projection - /// above the aggregate can read, so such a set is rejected as well. - fn grouping_sets_cover_pull_up_cols( + /// When one of these fails, [`GroupingSetsPullUp::Unsafe`] is returned: + /// the columns are still added, so the schema the nodes above expect + /// stays correct, but the caller must give up on decorrelating. + fn grouping_sets_pull_up_outcome( &self, - group_expr: &[Expr], + aggregate: &Aggregate, correlated_subquery_cols: &BTreeSet, - ) -> bool { + ) -> GroupingSetsPullUp { + let group_expr = &aggregate.group_expr; let grouping_sets = group_expr .iter() .filter_map(|expr| match expr { @@ -646,7 +718,7 @@ impl PullUpCorrelatedExpr { }) .collect::>(); if grouping_sets.is_empty() { - return true; + return GroupingSetsPullUp::NoColumnsToAdd; } // The same columns `collect_missing_exprs` appends: the correlated columns @@ -662,18 +734,37 @@ impl PullUpCorrelatedExpr { .any(|expr| matches!(expr, Expr::Column(c) if c == *col)) }); if required_cols.is_empty() { - return true; + return GroupingSetsPullUp::NoColumnsToAdd; } - grouping_sets.iter().all(|grouping_set| match grouping_set { - GroupingSet::Rollup(_) | GroupingSet::Cube(_) => false, - GroupingSet::GroupingSets(sets) => sets.iter().all(|set| { - required_cols.iter().all(|col| { - set.iter() - .any(|expr| matches!(expr, Expr::Column(c) if c == *col)) - }) - }), - }) + let already_covered = + grouping_sets.iter().all(|grouping_set| match grouping_set { + GroupingSet::Rollup(_) | GroupingSet::Cube(_) => false, + GroupingSet::GroupingSets(sets) => sets.iter().all(|set| { + required_cols.iter().all(|col| { + set.iter() + .any(|expr| matches!(expr, Expr::Column(c) if c == *col)) + }) + }), + }); + if already_covered { + return GroupingSetsPullUp::NoColumnsToAdd; + } + + let has_empty_set = grouping_sets.iter().any(|grouping_set| match grouping_set { + GroupingSet::Rollup(_) | GroupingSet::Cube(_) => true, + GroupingSet::GroupingSets(sets) => sets.iter().any(|set| set.is_empty()), + }); + let reads_grouping_fn = aggregate.aggr_expr.iter().any(is_grouping_call); + let read_above = required_cols + .iter() + .any(|col| self.column_refs_above_aggregate.contains(*col)); + + if has_empty_set || reads_grouping_fn || read_above { + GroupingSetsPullUp::Unsafe + } else { + GroupingSetsPullUp::SafeToExtend + } } fn collect_missing_exprs( @@ -828,6 +919,45 @@ fn can_pullup_over_aggregation(expr: &Expr) -> bool { } } +/// Whether `expr` computes `GROUPING`/`GROUPING_ID`, which reads whether a +/// row's grouping set groups by a particular column. +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 +/// `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 +/// [`PullUpCorrelatedExpr::f_down`] skips it. +fn columns_read_above_aggregate(plan: &LogicalPlan) -> BTreeSet { + fn walk(plan: &LogicalPlan, above: &mut BTreeSet) -> bool { + if matches!(plan, LogicalPlan::Subquery(_)) { + return false; + } + if matches!(plan, LogicalPlan::Aggregate(_)) { + return true; + } + 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 + } + + let mut above = BTreeSet::new(); + walk(plan, &mut above); + above +} + fn collect_local_correlated_cols( plan: &LogicalPlan, all_cols_map: &HashMap>, diff --git a/datafusion/optimizer/src/decorrelate_lateral_join.rs b/datafusion/optimizer/src/decorrelate_lateral_join.rs index a8df5e69e3f33..65a3b63e55998 100644 --- a/datafusion/optimizer/src/decorrelate_lateral_join.rs +++ b/datafusion/optimizer/src/decorrelate_lateral_join.rs @@ -107,7 +107,9 @@ fn rewrite_internal(join: Join) -> Result> { // Walk the subquery plan bottom-up, extracting correlated filter // predicates into join conditions and converting ungrouped aggregates // into group-by aggregates keyed on the correlation columns. - let mut pull_up = PullUpCorrelatedExpr::new().with_need_handle_count_bug(true); + let mut pull_up = PullUpCorrelatedExpr::new() + .with_need_handle_count_bug(true) + .with_column_refs_above_aggregate(subquery_plan); let rewritten_subquery = subquery_plan.clone().rewrite(&mut pull_up).data()?; if !pull_up.can_pull_up { return Ok(Transformed::no(LogicalPlan::Join(join))); diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 88a929654bea8..532333ed09bac 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -735,7 +735,8 @@ fn build_join( ) -> Result> { let mut pull_up = PullUpCorrelatedExpr::new() .with_in_predicate_opt(in_predicate_opt.cloned()) - .with_exists_sub_query(in_predicate_opt.is_none()); + .with_exists_sub_query(in_predicate_opt.is_none()) + .with_column_refs_above_aggregate(subquery); let new_plan = subquery.clone().rewrite(&mut pull_up).data()?; if !pull_up.can_pull_up { diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index c6374660ef1f9..4648267dceb28 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -351,7 +351,9 @@ fn build_join( // join with `Boolean(true)`) when the // `enable_physical_uncorrelated_scalar_subquery` option is disabled. let subquery_plan = subquery.subquery.as_ref(); - let mut pull_up = PullUpCorrelatedExpr::new().with_need_handle_count_bug(true); + let mut pull_up = PullUpCorrelatedExpr::new() + .with_need_handle_count_bug(true) + .with_column_refs_above_aggregate(subquery_plan); let decorrelated_subquery = subquery_plan.clone().rewrite(&mut pull_up).data()?; if !pull_up.can_pull_up { return Ok(None); diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index c1fd8564976ac..e781fedd46564 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2907,6 +2907,10 @@ DROP TABLE metrics; # `GROUPING SETS ((k), ())`, turns into `GROUPING SETS ((k), (k, 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. +# +# Follow-up (#25708): a non-empty set that leaves the correlated column out is +# safe to extend when nothing above the aggregate reads that column, so those +# queries decorrelate again below. statement ok CREATE TABLE gs_outer(k INT) AS VALUES (1), (2), (NULL), (4), (5); @@ -2927,14 +2931,17 @@ SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k # 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 set fills it with 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; +# where the set fills it with NULL. Nothing above this `EXISTS` reads `k` +# though (it only checks for a row), so the difference is invisible and the +# subquery still decorrelates. +query IB +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 ORDER BY gs_outer.k; +---- +1 true +2 true +4 false +5 true +NULL false # 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 @@ -2943,6 +2950,13 @@ SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k 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 same sets with a `GROUPING(k)` read out of the aggregate itself: for the +# `(j)` set it is 1 (k is not one of its columns) before the pull up, and 0 +# (k is now one of its columns in every set) after, so the subquery stays +# correlated even though nothing above the aggregate reads `k` directly. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression ScalarSubquery +SELECT gs_outer.k, (SELECT GROUPING(gs_inner.k) FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j)) LIMIT 1) FROM gs_outer; + # The correlated column does not have to appear in the grouping set at all. 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 ROLLUP(gs_inner.j)) FROM gs_outer; From 0e5367489d4c545d13250b7b386bc57d8e12533b Mon Sep 17 00:00:00 2001 From: naman Date: Sun, 27 Sep 2026 03:37:36 +0530 Subject: [PATCH 2/5] fix: drop private intra-doc link that failed the rustdoc CI check --- datafusion/optimizer/src/decorrelate.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index a96f5a5961ff8..0b651f7e2c4f9 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -157,7 +157,7 @@ impl PullUpCorrelatedExpr { /// A grouping set that omits a correlated column only needs to reject the /// pull up when something reads the column the pull up would fill in, so /// [`Self::f_up`] needs this computed ahead of time. See - /// [`columns_read_above_aggregate`]. + /// `columns_read_above_aggregate`. pub fn with_column_refs_above_aggregate( mut self, subquery_plan: &LogicalPlan, From f02e137196f0ee14a8b9d822440f49639c774019 Mon Sep 17 00:00:00 2001 From: naman Date: Mon, 5 Oct 2026 13:48:24 +0530 Subject: [PATCH 3/5] fix: remove unused unsupported() left over from the main rebase Resolving the rebase conflict in decorrelate.rs kept this method alongside main's independent rename of the same method to stop_pull_up() (with a plain Transformed return type). All call sites already use stop_pull_up(), so this was dead code. --- datafusion/optimizer/src/decorrelate.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0b651f7e2c4f9..e40993a74c132 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -36,8 +36,8 @@ use datafusion_expr::utils::{ collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, }; use datafusion_expr::{ - Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, - FetchType, LogicalPlan, LogicalPlanBuilder, Operator, SkipType, expr, lit, + Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, FetchType, + LogicalPlan, LogicalPlanBuilder, Operator, SkipType, expr, lit, }; /// This struct rewrite the sub query plan by pull up the correlated @@ -144,13 +144,6 @@ impl PullUpCorrelatedExpr { self } - /// Mark the plan as one whose correlated expressions cannot be pulled up - /// and stop descending into it - fn unsupported(&mut self, plan: LogicalPlan) -> Result> { - self.can_pull_up = false; - Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump)) - } - /// Record the columns read above the subquery's `Aggregate`, computed /// from `subquery_plan` (the plan as it is before this rewrite runs). /// From f85611207e9f7f66a3545fc92d9a2cb1fe6138de Mon Sep 17 00:00:00 2001 From: namanjain24-sudo <180642416+namanjain24-sudo@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:21:47 +0530 Subject: [PATCH 4/5] 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. --- datafusion/optimizer/src/decorrelate.rs | 23 ++++++++++++------- .../sqllogictest/test_files/subquery.slt | 19 +++++++++++++++ 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index e40993a74c132..b039d3ab46b2e 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/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 { fn walk(plan: &LogicalPlan, above: &mut BTreeSet) -> 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(); diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index e781fedd46564..65b7f20a9b2db 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/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 From 39d9b5805c24c28bdb4f5201a59b8a6d1362512c Mon Sep 17 00:00:00 2001 From: naman Date: Fri, 16 Oct 2026 00:54:59 +0530 Subject: [PATCH 5/5] fix: catch a GROUPING/GROUPING_ID read via __grouping_id, not the rewritten-away call is_grouping_call checked aggr_expr for a GROUPING/GROUPING_ID call, but ResolveGroupingFunction has already rewritten those into bit tests on __grouping_id in a node above the aggregate by the time this rule runs, so aggr_expr never holds one and the check never fires. Appending the correlated column to every set then shifts those bits, silently returning wrong rows instead of blocking decorrelation. Fix: look for __grouping_id directly in column_refs_above_aggregate, the same set already used to catch an ordinary column read above the aggregate. Removes the now-dead is_grouping_call. Adds the two repro queries as statement error cases in subquery.slt. --- datafusion/optimizer/src/decorrelate.rs | 17 +++++++++-------- .../sqllogictest/test_files/subquery.slt | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index b039d3ab46b2e..33033360daf33 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -748,12 +748,19 @@ impl PullUpCorrelatedExpr { GroupingSet::Rollup(_) | GroupingSet::Cube(_) => true, GroupingSet::GroupingSets(sets) => sets.iter().any(|set| set.is_empty()), }); - let reads_grouping_fn = aggregate.aggr_expr.iter().any(is_grouping_call); + // `GROUPING`/`GROUPING_ID` are resolved to bit tests on `__grouping_id` + // above the aggregate by `ResolveGroupingFunction` before this rule + // runs, so `aggregate.aggr_expr` never holds a `GROUPING` call by the + // time this is reached; adding a column shifts those bits instead. + let reads_grouping_id = self + .column_refs_above_aggregate + .iter() + .any(|col| col.name == Aggregate::INTERNAL_GROUPING_ID); let read_above = required_cols .iter() .any(|col| self.column_refs_above_aggregate.contains(*col)); - if has_empty_set || reads_grouping_fn || read_above { + if has_empty_set || reads_grouping_id || read_above { GroupingSetsPullUp::Unsafe } else { GroupingSetsPullUp::SafeToExtend @@ -912,12 +919,6 @@ fn can_pullup_over_aggregation(expr: &Expr) -> bool { } } -/// Whether `expr` computes `GROUPING`/`GROUPING_ID`, which reads whether a -/// row's grouping set groups by a particular column. -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 that has a grouping set `Aggregate` below it in /// `plan`, which is the subquery's original, not yet rewritten, plan. /// diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 65b7f20a9b2db..b5fc66ecf62d8 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2976,6 +2976,25 @@ SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression ScalarSubquery SELECT gs_outer.k, (SELECT GROUPING(gs_inner.k) FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j)) LIMIT 1) FROM gs_outer; +# `GROUPING`/`GROUPING_ID` are rewritten by `ResolveGroupingFunction` into bit +# tests on `__grouping_id` in a node above the aggregate before this rule +# runs, so the aggregate's own `aggr_expr` never holds a raw `GROUPING` call. +# Reading `GROUPING(j)` out of the `(j)`/`(x)` sets, as a HAVING or through an +# outer projection, reads `__grouping_id` above the aggregate without reading +# `j` or `x` themselves, so this is only caught by looking for +# `__grouping_id` in `column_refs_above_aggregate` directly. +statement ok +CREATE TABLE gs_inner3(k INT, j INT, x INT) AS VALUES (1, 10, 100), (NULL, 20, 200), (5, 30, 300), (2, 40, 400); + +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_inner3 WHERE gs_inner3.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner3.j), (gs_inner3.x)) HAVING GROUPING(gs_inner3.j) = 1) FROM gs_outer ORDER BY gs_outer.k; + +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 GROUPING(gs_inner3.j) AS g FROM gs_inner3 WHERE gs_inner3.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner3.j), (gs_inner3.x))) t WHERE t.g = 1) FROM gs_outer ORDER BY gs_outer.k; + +statement ok +DROP TABLE gs_inner3; + # The correlated column does not have to appear in the grouping set at all. 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 ROLLUP(gs_inner.j)) FROM gs_outer;