[comphead]

Which issue does this PR close?

• Part of Wrong results: correlated scalar subquery with nested aggregates #26508.

Rationale for this change

An aggregate without GROUP BY returns one row even when its input is empty. When such an aggregate is in a correlated subquery and another aggregate sits above it, DataFusion returns wrong results for outer rows that match no inner rows:

SELECT x.a,
       (SELECT count(*) FROM (SELECT max(y.b) FROM (VALUES (1, 10)) AS y(a, b) WHERE y.a = x.a)) AS v1,
       (SELECT max(n) FROM (SELECT count(*) AS n FROM (VALUES (1, 10)) AS y(a, b) WHERE y.a = x.a)) AS v2
FROM (VALUES (1), (2)) AS x(a)

For a = 2 this returns v1 = 0, v2 = NULL instead of v1 = 1, v2 = 0, which DuckDB, SQLite and PostgreSQL return. The same shape gives wrong results under IN and LEFT JOIN LATERAL too.

When the correlated filter is pulled up, the inner scalar aggregate becomes a grouped aggregate, which returns no row for the outer rows it does not match. The aggregate above then runs over an empty input. The count bug handling only fixes the aggregate the join reads, which here is the outer one, so the inner row is never restored. The existing TODO next to pulled_up_scalar_agg described this case.

What changes are included in this PR?

PullUpCorrelatedExpr stops the pull up when it reaches an Aggregate after a scalar aggregate has already been pulled up. The subquery stays correlated, the same way a window over an aggregate with count bug handling does today. The rewrite stops at that Aggregate, so the nodes above it are not rebuilt over the pulled up plan. Before, the rebuild could fail with a schema error. Such a query now fails with Physical plan does not support logical expression ScalarSubquery instead of returning wrong rows.

A projection over the scalar aggregate (SELECT n + 1 FROM (SELECT count(*) AS n ...)) and an aggregate over a grouped aggregate are still pulled up.

What is the testing strategy for this PR?

New sqllogictest cases at the end of subquery.slt cover the two queries from the issue, the IN, LATERAL and EXISTS forms, a projection, a filter, a GROUP BY-only aggregate and a grouped aggregate between the aggregates, a max over max query that is now blocked, and the two shapes that must still be pulled up. All sqllogictest files, the datafusion-optimizer unit tests and the core_integration sql:: tests (which include the TPC-DS planning tests) pass locally.

Are there any user-facing changes?

Queries with this shape now return a "not implemented" error instead of wrong results. Some queries of this shape returned correct rows before and now fail too, for example max over max. Telling them apart needs the outer aggregate evaluated on the inner one's empty-input row. No API changes.

[jayzhan211]

Thanks @comphead

nit: In decorrelate.rs around line 511, putting a filter between the two aggregates gives a confusing error instead of the not-implemented one:

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;
-- Schema error: Schema contains duplicate unqualified field name "count(Int64(1))"

The first optimizer pass now leaves the subquery correlated, and the second pass rewrites a reshaped plan that fails with a schema error. With datafusion.optimizer.max_passes = 1 it gives the expected Physical plan does not support logical expression ScalarSubquery. It's still an error (main returned (2, 0) here, which was wrong), so it would be nice to add it as a test case even if the message stays as is.

[jayzhan211]

This also blocks some shapes that were correct before and now fail to plan. For example, max over max, sum over sum and count(m) over max all returned the right rows on main, and now hit 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;
-- main: (1, 10), (2, NULL)   this PR: error

Blocking makes sense to me, since a precise rule would have to evaluate the outer aggregate on the inner aggregate's empty-input row. Could you mention this in the PR description and add a test that pins the new error, so it's clear this is an intended trade-off?

[comphead]

Added the max over max case as a pinned error, with a comment saying it returned the right rows before and why it is blocked too. The PR description now lists this trade-off under user-facing changes.

[jayzhan211]

nit: Two more shapes that returned wrong rows on main are now blocked too: EXISTS, and a projection between the two aggregates. It would be good to cover them here:

+# EXISTS. Expected `1, 2`.
+statement error DataFusion error: This feature is not implemented: 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. Expected `(1, 2), (2, 1)`.
+statement error DataFusion error: This feature is not implemented: 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;

[comphead]

Added both.

[adriangb]

Thanks for the fix!

We now return a "not implemented" error, not the result that #26508 expects (v1 = 1, v2 = 0). Can we change "Closes #26508" to "Part of #26508", or open a follow-up issue for the real decorrelation? If not, the issue closes while the queries still fail.

[adriangb]

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:

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.

[comphead]

The check moved to the top of the Aggregate branch. It now sets can_pull_up = false and returns early with TreeNodeRecursion::Stop, which fixes the schema error jayzhan211 found (details in the top-level comment). That leaves no second update of can_pull_up to fold. The check is still outside the is_distinct block. I also added the doc comment on pulled_up_scalar_agg.

[adriangb]

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.

[comphead]

Switched all the new cases to query error with a partial message match and an "Errors rather than returning wrong rows" comment.

[adriangb]

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.

[comphead]

Done. Each query now has its own comment with its expected result.

[adriangb]

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:

Suggested change

	2 NULL
	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;

[comphead]

Added both, thanks for running them.

[comphead]

Thanks @jayzhan211 and @adriangb.

@jayzhan211 the schema error with a filter between the aggregates is not new. Main hits the same error on the HAVING form of that query:

SELECT o.k, (SELECT count(*) FROM (SELECT count(*) AS n FROM nested_agg_inner AS i WHERE i.id = o.k HAVING count(*) >= 0))
FROM nested_agg_outer AS o ORDER BY o.k;
-- Schema error: Schema contains duplicate unqualified field name "count(Int64(1))"

The Aggregate branch kept rebuilding the plan after it set can_pull_up = false, and the rebuild failed with a schema error. It now stops the rewrite as soon as it reaches an aggregate over a pulled up scalar aggregate. The filter case and the HAVING form now both fail with Physical plan does not support logical expression ScalarSubquery, at the default max_passes. The filter case is added as a test.

@adriangb I changed "Closes #26508" to "Part of #26508".

[comphead]

Thanks @jayzhan211 and @adriangb for reviewing this, btw this is a second item for today which makes me think our test coverage is less than we would expect

[comphead]

@adriangb @jayzhan211 please have another look

[adriangb]

Thanks @comphead !

[comphead]

@jayzhan211 PTAL if you dont mind