diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index e4577d6eb0c4c..b039d3ab46b2e 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -36,7 +36,7 @@ use datafusion_expr::utils::{ collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, }; use datafusion_expr::{ - BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, FetchType, + Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, FetchType, LogicalPlan, LogicalPlanBuilder, Operator, SkipType, expr, lit, }; @@ -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,36 @@ impl PullUpCorrelatedExpr { self.exists_sub_query = exists_sub_query; self } + + /// 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 +475,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 +668,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 +711,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 +727,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 +912,52 @@ 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. +/// +/// [`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, 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; + } + let mut found_below = false; + for child in plan.inputs() { + found_below |= 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(); + 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..65b7f20a9b2db 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,36 @@ 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. +# 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 + +# 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 gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j))) FROM gs_outer; +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 @@ -2943,6 +2969,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;