adriangb commented on code in PR #26544: URL: https://github.com/apache/datafusion/pull/26544#discussion_r4336661693
########## datafusion/sqllogictest/test_files/subquery.slt: ########## @@ -3627,3 +3627,45 @@ 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); + +# 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 Review Comment: The same file has a precedent for this kind of test (around lines 692-716): it uses `query error`, matches only part of the message, and has a comment saying it "errors rather than returning wrong rows". Can these tests (and the extra cases suggested in the other review) do the same? For example: ``` # Errors rather than returning wrong rows. Expected `(1, 1), (2, 1)`. query error Physical plan does not support logical expression ScalarSubquery ``` If the `DataFusion error: This feature is not implemented:` prefix changes, the full-message match breaks. The comment also tells readers that the error is a stop-gap. ########## datafusion/sqllogictest/test_files/subquery.slt: ########## @@ -3627,3 +3627,45 @@ 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); + +# 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; + +# 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; + +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 + +# 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 Review Comment: The guard also stops two more shapes that have no test: a `GROUP BY`-only aggregate (the `is_distinct` path) over the scalar aggregate, and an aggregate with `GROUP BY` over the scalar aggregate. Both lose the row for unmatched outer rows in the same way, so blocking them is correct. I ran both on this branch and they give the error below. Tests would keep that behaviour in place: ```suggestion 2 NULL # 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; ``` ########## datafusion/sqllogictest/test_files/subquery.slt: ########## @@ -3627,3 +3627,45 @@ 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); + +# Expected `(1, 1), (2, 1)` and `(1, 1), (2, 0)`. Review Comment: nit: This one comment gives the expected results for the two queries below. Please put one comment above each query, so readers can see which result goes with which query. ########## datafusion/optimizer/src/decorrelate.rs: ########## @@ -505,6 +505,12 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { 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; + } Review Comment: nit: Two statements now update `can_pull_up` for the same "can we pull past this aggregate" rule. You could fold this into the line above: ```rust self.can_pull_up = self.can_pull_up && self.can_pull_over_aggregation && !self.pulled_up_scalar_agg; ``` (this changes nothing for the distinct case, because the new check is outside the `if !is_distinct` block and must stay there.) Also, `pulled_up_scalar_agg` now controls this decision, and the check works only because the child's `f_up` sets the flag first. A short doc comment on the field would make both points clear. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
