jayzhan211 commented on code in PR #26874:
URL: https://github.com/apache/datafusion/pull/26874#discussion_r4505021418
##########
datafusion/optimizer/src/decorrelate_lateral_join.rs:
##########
@@ -315,6 +330,40 @@ fn rewrite_internal(join: Join) ->
Result<Transformed<LogicalPlan>> {
Ok(Transformed::new(new_plan, true, TreeNodeRecursion::Jump))
}
+/// The `HAVING` of a subquery whose result is an ungrouped aggregate: a
+/// `Filter` between the subquery's output and that aggregate, with only
+/// projections and aliases around them. A filter below the aggregate only
+/// changes the aggregate's input, so it is not returned.
+fn having_over_ungrouped_aggregate(plan: &LogicalPlan) -> Option<&Filter> {
+ match plan {
+ LogicalPlan::Projection(projection) => {
+ having_over_ungrouped_aggregate(&projection.input)
+ }
+ LogicalPlan::SubqueryAlias(alias) => {
+ having_over_ungrouped_aggregate(&alias.input)
+ }
+ LogicalPlan::Filter(filter) if
feeds_ungrouped_aggregate(&filter.input) => {
+ Some(filter)
+ }
+ LogicalPlan::Filter(filter) =>
having_over_ungrouped_aggregate(&filter.input),
Review Comment:
Optional cleanup: the second `Filter` arm always returns `None`. If a filter
further down fed the aggregate, `feeds_ungrouped_aggregate(&filter.input)`
would already be true for this one, so the two arms can be one:
```suggestion
LogicalPlan::Filter(filter) => {
feeds_ungrouped_aggregate(&filter.input).then_some(filter)
}
```
##########
datafusion/optimizer/src/decorrelate_lateral_join.rs:
##########
@@ -88,6 +88,21 @@ fn rewrite_internal(join: Join) ->
Result<Transformed<LogicalPlan>> {
let subquery_plan = subquery.subquery.as_ref();
let original_join_filter = join.filter.clone();
+ // A HAVING over an ungrouped aggregate can remove the aggregate's only
+ // row. The rewrite below handles neither case correctly, so leave the join
+ // correlated and let the query error instead of returning wrong rows.
+ // <https://github.com/apache/datafusion/issues/21198>
+ // - An inner lateral join becomes an outer join that keeps every outer
row,
+ // so an outer row whose row the HAVING removes comes back with NULLs.
+ // - A HAVING that reads an outer column is moved into the join condition,
+ // where the count bug handling turns its unmatched rows into a count of
0,
+ // for a left lateral join too.
+ if let Some(having) = having_over_ungrouped_aggregate(subquery_plan)
Review Comment:
The gate also turns queries that work in the 55.2.0 release into errors.
Without a correlated `WHERE`, the aggregate isn't grouped by the correlation
columns and keeps one row per outer row, so a `HAVING` that reads an outer
column is just a join filter. Four such queries match DuckDB on 55.2.0 and
main, and error on this branch, for example:
```sql
SELECT t1.id, sub.cnt FROM t1, LATERAL (SELECT count(*) AS cnt FROM t2
HAVING count(*) > t1.id) AS sub;
-- 55.2.0, main and DuckDB: (1, 3), (2, 3); this PR: error
SELECT t1.id, sub.cnt FROM t1 LEFT JOIN LATERAL (SELECT count(*) AS cnt FROM
t2 HAVING count(*) > t1.id) AS sub ON true;
-- 55.2.0, main and DuckDB: (1, 3), (2, 3), (3, NULL); this PR: error
```
Gating on `pull_up.pulled_up_scalar_agg` (so after the pull-up) keeps all
your error cases and brings these back. `lateral_join.slt` passes with it and
with the tests below, and both tests fail on the current head. Could you apply
it and add the tests?
```diff
let subquery_plan = subquery.subquery.as_ref();
let original_join_filter = join.filter.clone();
+ // Walk the subquery plan bottom-up, extracting correlated filter
+ // predicates into join conditions and converting ungrouped aggregates
+ // into group-by aggregates keyed on the correlation columns.
+ let mut pull_up =
PullUpCorrelatedExpr::new().with_need_handle_count_bug(true);
+ let rewritten_subquery = subquery_plan.clone().rewrite(&mut
pull_up).data()?;
+ if !pull_up.can_pull_up {
+ return Ok(Transformed::no(LogicalPlan::Join(join)));
+ }
+
// A HAVING over an ungrouped aggregate can remove the aggregate's only
// row. The rewrite below handles neither case correctly, so leave the
join
// correlated and let the query error instead of returning wrong rows.
@@
- if let Some(having) = having_over_ungrouped_aggregate(subquery_plan)
+ if pull_up.pulled_up_scalar_agg
+ && let Some(having) = having_over_ungrouped_aggregate(subquery_plan)
&& (original_join_type == JoinType::Inner ||
having.predicate.contains_outer())
{
return Ok(Transformed::no(LogicalPlan::Join(join)));
}
- // Walk the subquery plan bottom-up, extracting correlated filter
- // predicates into join conditions and converting ungrouped aggregates
- // into group-by aggregates keyed on the correlation columns.
- let mut pull_up =
PullUpCorrelatedExpr::new().with_need_handle_count_bug(true);
- let rewritten_subquery = subquery_plan.clone().rewrite(&mut
pull_up).data()?;
- if !pull_up.can_pull_up {
- return Ok(Transformed::no(LogicalPlan::Join(join)));
- }
-
```
<details>
<summary>Tests for lateral_join.slt (both error on the current head, pass
with the change)</summary>
```
# Without a correlated WHERE the aggregate keeps its single row for every
# outer row, so a HAVING that reads an outer column is a plain join filter.
query II
SELECT t1.id, sub.cnt
FROM t1, LATERAL (
SELECT count(*) AS cnt FROM t2 HAVING count(*) > t1.id
) AS sub
ORDER BY t1.id;
----
1 3
2 3
query II
SELECT t1.id, sub.cnt
FROM t1 LEFT JOIN LATERAL (
SELECT count(*) AS cnt FROM t2 HAVING count(*) > t1.id
) AS sub ON true
ORDER BY t1.id;
----
1 3
2 3
3 NULL
```
The other two queries that changed: `max(t2.id) ... WHERE t2.id > 10 HAVING
max(t2.id) > t1.id * 10` (55.2.0 and main: (1, 30), (2, 30)) and `count(*) ...
WHERE t2.id > 100 HAVING count(*) >= t1.id - 1` (55.2.0 and main: (1, 0)). Both
match DuckDB on 55.2.0 and main, and error on this branch.
</details>
--
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]
