[comphead]

Which issue does this PR close?

• Closes 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. 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 and LATERAL forms, 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. 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?

[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;