From 6af238b724bc7bf1f8815b128bc0ff542e0aa096 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 20 Oct 2026 13:49:21 -0700 Subject: [PATCH 1/3] fix: keep a subquery correlated when an aggregate sits above a pulled up scalar aggregate --- datafusion/optimizer/src/decorrelate.rs | 11 ++++- .../sqllogictest/test_files/subquery.slt | 45 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0a1db9b2128c6..ab384b3742239 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -505,6 +505,15 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { if !is_distinct { self.can_pull_up = self.can_pull_up && self.can_pull_over_aggregation; } + // A scalar aggregate below returns one row for every outer row, + // but once pulled up it returns no row for the outer rows whose + // correlated filter matches nothing. This aggregate would then + // run over an empty input instead of that one row, and the count + // bug handling only covers the aggregate the join reads. + // See https://github.com/apache/datafusion/issues/26508. + if self.pulled_up_scalar_agg { + self.can_pull_up = false; + } let mut local_correlated_cols = BTreeSet::new(); collect_local_correlated_cols( &plan, @@ -588,8 +597,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 39b3961c82451..6e97b84f66dac 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -3627,3 +3627,48 @@ DROP TABLE unnest_outer; statement ok DROP TABLE unnest_inner; + +# An aggregate over a correlated scalar aggregate. The scalar aggregate returns +# one row for every outer row, so the aggregate above it never sees an empty +# input. 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); + +# The pull up would give no row for `k = 2`, so these stay correlated. The +# right answers are `(1, 1), (2, 1)` and `(1, 1), (2, 0)`. +statement error DataFusion error: This feature is not implemented: 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; + +statement error DataFusion error: This feature is not implemented: 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; + +# The same for IN and LATERAL. The right answers are `2` and `(1, 1), (2, 0)`. +statement error DataFusion error: This feature is not implemented: 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; + +statement error DataFusion error: This feature is not implemented: 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; + +# 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 + +# A grouped aggregate below returns no row for `k = 2` before the pull up too, +# so an aggregate over it 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; From d438ca560357b0b06d82d465859f35c84a4ef8f7 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 20 Oct 2026 16:49:44 -0700 Subject: [PATCH 2/3] shorten comments --- datafusion/optimizer/src/decorrelate.rs | 9 +++------ datafusion/sqllogictest/test_files/subquery.slt | 13 +++++-------- 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index ab384b3742239..11bf593600bfc 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -505,12 +505,9 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { if !is_distinct { self.can_pull_up = self.can_pull_up && self.can_pull_over_aggregation; } - // A scalar aggregate below returns one row for every outer row, - // but once pulled up it returns no row for the outer rows whose - // correlated filter matches nothing. This aggregate would then - // run over an empty input instead of that one row, and the count - // bug handling only covers the aggregate the join reads. - // See https://github.com/apache/datafusion/issues/26508. + // A pulled up scalar aggregate below gives no row for unmatched + // outer rows, so this aggregate would see an empty input. + // See https://github.com/apache/datafusion/issues/26508 if self.pulled_up_scalar_agg { self.can_pull_up = false; } diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 6e97b84f66dac..6e485a011fa48 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -3628,24 +3628,22 @@ DROP TABLE unnest_outer; statement ok DROP TABLE unnest_inner; -# An aggregate over a correlated scalar aggregate. The scalar aggregate returns -# one row for every outer row, so the aggregate above it never sees an empty -# input. https://github.com/apache/datafusion/issues/26508 +# 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); -# The pull up would give no row for `k = 2`, so these stay correlated. The -# right answers are `(1, 1), (2, 1)` and `(1, 1), (2, 0)`. +# Expected `(1, 1), (2, 1)` and `(1, 1), (2, 0)`. statement error DataFusion error: This feature is not implemented: 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; statement error DataFusion error: This feature is not implemented: 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; -# The same for IN and LATERAL. The right answers are `2` and `(1, 1), (2, 0)`. +# IN and LATERAL. Expected `2` and `(1, 1), (2, 0)`. statement error DataFusion error: This feature is not implemented: 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; @@ -3659,8 +3657,7 @@ SELECT o.k, (SELECT n + 1 FROM (SELECT count(*) AS n FROM nested_agg_inner AS i 1 2 2 1 -# A grouped aggregate below returns no row for `k = 2` before the pull up too, -# so an aggregate over it is still pulled up. +# 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; ---- From 8aa946e4b32d647727a883929405dd8eafdd6dbe Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 21 Oct 2026 08:40:20 -0700 Subject: [PATCH 3/3] Stop the rewrite at the aggregate, add review test cases --- datafusion/optimizer/src/decorrelate.rs | 17 +++++--- .../sqllogictest/test_files/subquery.slt | 43 ++++++++++++++++--- 2 files changed, 48 insertions(+), 12 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 11bf593600bfc..de2f60befa098 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -76,6 +76,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` @@ -499,18 +502,20 @@ 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(); if !is_distinct { self.can_pull_up = self.can_pull_up && self.can_pull_over_aggregation; } - // A pulled up scalar aggregate below gives no row for unmatched - // outer rows, so this aggregate would see an empty input. - // See https://github.com/apache/datafusion/issues/26508 - if self.pulled_up_scalar_agg { - self.can_pull_up = false; - } let mut local_correlated_cols = BTreeSet::new(); collect_local_correlated_cols( &plan, diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 6e485a011fa48..22a7227f8eb1a 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -3636,20 +3636,51 @@ 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); -# Expected `(1, 1), (2, 1)` and `(1, 1), (2, 0)`. -statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression ScalarSubquery +# 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; -statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression ScalarSubquery +# 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 and LATERAL. Expected `2` and `(1, 1), (2, 0)`. -statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression InSubquery +# 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; -statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression OuterReferenceColumn +# 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;