diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index e6f235da4dd4..b81f5cb44708 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -83,6 +83,9 @@ pub struct PullUpCorrelatedExpr { pub pull_up_having_expr: Option, /// whether we have converted a scalar aggregation into a group aggregation. When unnesting /// lateral joins, we need to produce a left outer join in such cases. + /// + /// An `Aggregate` above stops the pull up when this is set. `f_up` visits + /// the aggregate that sets it before any node above. pub pulled_up_scalar_agg: bool, /// Every correlated conjunct that a `Filter` of the subquery applies, /// before `remove_duplicated_filter` drops the ones that the `IN` @@ -556,6 +559,14 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { LogicalPlan::Aggregate(aggregate) if self.in_predicate_opt.is_some() || !self.join_filters.is_empty() => { + // A pulled up scalar aggregate below gives no row for unmatched + // outer rows, so this aggregate would see an empty input. Stop + // the rewrite here, since rebuilding the nodes above can fail + // with a schema error. See https://github.com/apache/datafusion/issues/26508 + if self.pulled_up_scalar_agg { + self.can_pull_up = false; + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop)); + } // If the aggregation is from a distinct it will not change the result for // exists/in subqueries so we can still pull up all predicates. let is_distinct = aggregate.aggr_expr.is_empty(); @@ -647,8 +658,6 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { } } if aggregate.group_expr.is_empty() { - // TODO: how do we handle the case where we have pulled multiple aggregations? For example, - // a group agg with a scalar agg as child. self.pulled_up_scalar_agg = true; } let new_plan = LogicalPlanBuilder::from((*aggregate.input).clone()) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index a4474c4b0388..45ec0345f75a 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -4003,3 +4003,76 @@ DROP TABLE unnest_outer; statement ok DROP TABLE unnest_inner; + +# An aggregate over a correlated scalar aggregate stays correlated. +# https://github.com/apache/datafusion/issues/26508 +statement ok +CREATE TABLE nested_agg_outer(k INT) AS VALUES (1), (2); + +statement ok +CREATE TABLE nested_agg_inner(id INT, v INT) AS VALUES (1, 10); + +# Errors rather than returning wrong rows. Expected `(1, 1), (2, 1)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT count(*) FROM (SELECT max(i.v) FROM nested_agg_inner AS i WHERE i.id = o.k)) FROM nested_agg_outer AS o ORDER BY o.k; + +# Errors rather than returning wrong rows. Expected `(1, 1), (2, 0)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT max(n) FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k)) FROM nested_agg_outer AS o ORDER BY o.k; + +# IN. Errors rather than returning wrong rows. Expected `2`. +query error Physical plan does not support logical expression InSubquery +SELECT o.k FROM nested_agg_outer AS o WHERE 0 IN (SELECT max(n) FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k)) ORDER BY o.k; + +# LATERAL. Errors rather than returning wrong rows. Expected `(1, 1), (2, 0)`. +query error Physical plan does not support logical expression OuterReferenceColumn +SELECT o.k, s.m FROM nested_agg_outer AS o LEFT JOIN LATERAL (SELECT max(n) AS m FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k)) AS s ON true ORDER BY o.k; + +# EXISTS. Errors rather than returning wrong rows. Expected `1, 2`. +query error Physical plan does not support logical expression Exists +SELECT o.k FROM nested_agg_outer AS o WHERE EXISTS (SELECT max(n) FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k)) ORDER BY o.k; + +# A projection between the aggregates. Errors rather than returning wrong rows. Expected `(1, 2), (2, 1)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT max(n2) FROM (SELECT n + 1 AS n2 FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k))) FROM nested_agg_outer AS o ORDER BY o.k; + +# A filter between the aggregates. Errors rather than returning wrong rows. +# Expected `(1, 1), (2, 1)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT count(*) FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k) WHERE n >= 0) FROM nested_agg_outer AS o ORDER BY o.k; + +# A GROUP BY-only aggregate over the scalar aggregate. +# Errors rather than returning wrong rows. Expected `(1, 1), (2, 1)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT count(*) FROM (SELECT n FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k) GROUP BY n)) FROM nested_agg_outer AS o ORDER BY o.k; + +# A grouped aggregate over the scalar aggregate. +# Errors rather than returning wrong rows. Expected `(1, 1), (2, 1)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT max(c) FROM (SELECT n, count(*) AS c FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k) GROUP BY n)) FROM nested_agg_outer AS o ORDER BY o.k; + +# Returned the right rows before, but is blocked too, since telling it apart +# needs the outer aggregate evaluated on the inner one's empty input row. +# Expected `(1, 10), (2, NULL)`. +query error Physical plan does not support logical expression ScalarSubquery +SELECT o.k, (SELECT max(m) FROM (SELECT max(i.v) AS m FROM nested_agg_inner AS i WHERE i.id = o.k)) FROM nested_agg_outer AS o ORDER BY o.k; + +# A projection over the scalar aggregate is still pulled up. +query II +SELECT o.k, (SELECT n + 1 FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k)) FROM nested_agg_outer AS o ORDER BY o.k; +---- +1 2 +2 1 + +# An aggregate over a grouped aggregate is still pulled up. +query II +SELECT o.k, (SELECT max(c) FROM (SELECT i.id, count(*) AS c FROM nested_agg_inner AS i WHERE i.id = o.k GROUP BY i.id)) FROM nested_agg_outer AS o ORDER BY o.k; +---- +1 1 +2 NULL + +statement ok +DROP TABLE nested_agg_outer; + +statement ok +DROP TABLE nested_agg_inner;