namanjain24-sudo opened a new pull request, #25529:
URL: https://github.com/apache/datafusion/pull/25529
## Which issue does this PR close?
- Closes #25519.
## Rationale for this change
A correlated subquery whose filter sits below an aggregate with a grouping
set
returns wrong results, with no error and no warning.
`PullUpCorrelatedExpr` moves the correlated filter above the aggregate and
adds
the correlated column to the aggregate's group expressions.
`LogicalPlanBuilder::aggregate` cross joins a plain group expression with the
sets a grouping set already holds, so `ROLLUP(i.k)`, which is
`GROUPING SETS ((i.k), ())`, becomes `GROUPING SETS ((i.k), (i.k, i.k))`. The
empty set is gone, and with it the grand total row the subquery returns for
every outer row, including the rows whose filter matches nothing. The join
that
replaces the filter cannot bring those rows back.
On `main` at `b4a8c824b4`, with `datafusion-cli`:
```sql
CREATE TABLE o(k INT) AS VALUES (1), (2), (NULL), (4), (5);
CREATE TABLE i(k INT) AS VALUES (1), (NULL), (5), (2);
SELECT o.k, EXISTS (SELECT 1 FROM i WHERE i.k = o.k GROUP BY ROLLUP(i.k)) AS
e FROM o ORDER BY o.k;
SELECT o.k, o.k IN (SELECT i.k FROM i WHERE i.k = o.k GROUP BY ROLLUP(i.k))
AS m FROM o ORDER BY o.k;
```
| `o.k` | `EXISTS` on `main` | `IN` on `main` | correct, per the issue |
| --- | --- | --- | --- |
| 1 | true | true | true / true |
| 2 | true | true | true / true |
| 4 | **false** | **false** | true / NULL |
| 5 | true | true | true / true |
| NULL | **false** | **false** | true / NULL |
The physical plan for the `EXISTS` query shows the loss directly. `ROLLUP(k)`
has two sets, and both of them now group by `k`:
```text
AggregateExec: group_by: (k), (k), mode: Partial
```
## What changes are included in this PR?
`datafusion/optimizer/src/decorrelate.rs`, in the `Aggregate` arm of
`PullUpCorrelatedExpr::f_up`:
- When the group expressions hold a grouping set, the pull up no longer adds
its
columns to them. If any set would be missing one, `can_pull_up` is set to
`false` and the subquery stays correlated, which is what the issue asks for
instead of a wrong result.
- `ROLLUP` and `CUBE` always contain the empty set, so they are only safe
when
the pull up has nothing to add. An explicit `GROUPING SETS` is checked set
by
set.
- When every set already groups by each column the pull up needs, the
aggregate
keeps its sets untouched and the subquery still decorrelates. Before this
PR
that case appended the column anyway and produced
`GROUPING SETS ((c, c), (c, b, c))`; it now stays `GROUPING SETS ((c), (c,
b))`.
- The columns checked are the ones `collect_missing_exprs` would append: the
correlated columns and the columns of a pulled up `HAVING`, minus the ones
the
group expressions already list on their own.
The rewrite still runs to completion in the unsupported case, the same way
the
existing `can_pull_over_aggregation` case does. The three callers
(`decorrelate_predicate_subquery`, `scalar_subquery_to_join`,
`decorrelate_lateral_join`) read `can_pull_up` only after the whole rewrite
has
finished, and the nodes above the aggregate still expect the pulled up
columns
in its output, so returning early there fails the rewrite with a schema error
rather than declining the transform.
## What is the testing strategy for this PR?
Five unit tests, in the two rules that reach this code path:
- `decorrelate_predicate_subquery.rs`: `ROLLUP` under `EXISTS`, `CUBE` under
`IN`, a `GROUPING SETS` whose second set groups by another column, and the
covering `GROUPING SETS` that must still decorrelate.
- `scalar_subquery_to_join.rs`: `ROLLUP` under a correlated scalar subquery.
Twelve queries in `datafusion/sqllogictest/test_files/subquery.slt` cover the
end-to-end behaviour: `ROLLUP`, `CUBE`, an explicit set list holding `()`, a
set
list whose sets group by different columns, a grouping set that does not
mention
the correlated column at all, and the `IN`, `NOT EXISTS` and correlated
scalar
forms; plus the four that must keep working, a covering `GROUPING SETS` with
one
and with two sets, a plain `GROUP BY`, and an uncorrelated `ROLLUP` subquery.
Ablation, with only `decorrelate.rs` reverted to `main` and every test kept
(the
tests live in other files, so the revert does not delete them): all five unit
tests fail, and seven of the eight `statement error` cases in `subquery.slt`
fail with "query is expected to fail, but actually succeed". The eighth, the
correlated scalar subquery, already errors on `main`.
On the branch: `cargo test -p datafusion-optimizer` passes, 901 + 26 + 5
tests;
the full `sqllogictest` suite passes, 521 files; `cargo clippy -p
datafusion-optimizer --all-targets -- -D warnings` and `cargo fmt --all --
--check` are clean.
## Are there any user-facing changes?
Yes, and it is worth a look before merging. The queries above stop returning
wrong rows, but they do not start returning right ones: leaving the subquery
correlated means the physical planner rejects it with
```text
This feature is not implemented: Physical plan does not support logical
expression Exists(...)
```
That is the behaviour #25519 asks for, and it matches how the rule already
declines a `Union`, `Sort` or `Extension` that holds an outer reference, and
a
`Limit` that holds one outside an `EXISTS`. Supporting these
plans properly needs the aggregate to keep its sets through decorrelation,
which
is a larger change than this one.
Queries whose grouping set already groups by the correlated column in every
set
are unaffected, other than no longer repeating that column inside each set.
--
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]
