jayzhan211 commented on code in PR #25391:
URL: https://github.com/apache/datafusion/pull/25391#discussion_r4324700818
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -677,31 +749,80 @@ impl PullUpCorrelatedExpr {
}
fn collect_missing_exprs(
- &self,
+ &mut self,
exprs: &[Expr],
correlated_subquery_cols: &BTreeSet<Column>,
+ subquery_schema: &DFSchemaRef,
) -> Result<Vec<Expr>> {
let mut missing_exprs = vec![];
for expr in exprs {
if !missing_exprs.contains(expr) {
missing_exprs.push(expr.clone())
}
}
+ // A correlated column compared bare (`t1.a = t2.b`) only ever
+ // matches one row per value, so grouping the subquery's own
+ // aggregate by that column is enough. But if the join filter
+ // wraps it in an expression (`t1.a = CAST(t2.b AS INT)`), two
+ // different column values can compare equal after the cast, and
+ // grouping by the bare column would compute the aggregate once
+ // per underlying value instead of once per outer row it actually
+ // joins against. Grouping by the wrapping expression instead
+ // keeps those together, matching what the join predicate itself
+ // treats as equal.
+ let join_filter_exprs =
+ collect_subquery_join_exprs(&self.join_filters, subquery_schema)?;
for col in correlated_subquery_cols.iter() {
- let col_expr = Expr::Column(col.clone());
- if !missing_exprs.contains(&col_expr) {
+ if let Some(existing_alias) = self.correlated_col_aliases.get(col)
{
+ if !collides_with_existing(&missing_exprs, existing_alias) {
+ missing_exprs.push(existing_alias.clone())
+ }
+ continue;
+ }
+ let wrapped = join_filter_exprs
Review Comment:
Thanks for dropping the alias rewrite. The second half of the earlier
suggestion, gating compensation on bare-column keys, is still needed though.
Without it, the LEFT JOIN matches both the `'1'` and `'01'` groups, so EXISTS
emits the outer row twice, and the two new dedup tests record those duplicates
as expected output:
```sql
-- DuckDB: 1, 2; main: 1; this PR: 1, 1, 2
SELECT o.k FROM dedup_outer o WHERE EXISTS (
SELECT count(*) FROM dedup_inner i WHERE CAST(i.k AS INT) = o.k
);
-- DuckDB: 1; main: 1, 2; this PR: 1, 1
SELECT o.k FROM dedup_outer o WHERE NOT EXISTS (
SELECT count(*) FROM dedup_inner i WHERE CAST(i.k AS INT) = o.k HAVING
count(*) = 0
);
```
Getting the right answer for wrapped keys is fine as a follow-up, but EXISTS
should never duplicate outer rows. Falling back to the old rewrite fixes that,
and with this change `subquery.slt` only differs in those two cases, which go
back to main's answers:
```rs
fn subquery_operands_are_bare_columns(
join_filters: &[Expr],
subquery_cols: &BTreeSet<Column>,
) -> bool {
join_filters.iter().all(|filter| {
let mut bare = true;
filter
.apply(|e| {
let refs = e.column_refs();
if !refs.is_empty() && refs.iter().all(|c|
subquery_cols.contains(*c)) {
bare = matches!(e, Expr::Column(_));
return Ok(if bare {
TreeNodeRecursion::Jump
} else {
TreeNodeRecursion::Stop
});
}
Ok(TreeNodeRecursion::Continue)
})
.expect("infallible");
bare
})
}
```
```diff
+ let subquery_cols: BTreeSet<Column> = pull_up
+ .correlated_subquery_cols_map
+ .values()
+ .flatten()
+ .cloned()
+ .collect();
+ let keys_are_bare = count_bug_compensation.is_none()
+ || subquery_operands_are_bare_columns(&pull_up.join_filters,
&subquery_cols);
let compensable = pull_up.can_pull_up
&& !reads_agg_output
+ && keys_are_bare
```
Could you also add a `TODO` with the follow-up issue link on the four dedup
cases, so their expected output isn't read as correct?
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -383,6 +454,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
// Correlated exist subquery, remove the limit(so that
correlated expressions can pull up)
(true, false) => Transformed::yes(match
limit.get_fetch_type()? {
FetchType::Literal(Some(0)) => {
+ self.forces_empty_result = true;
Review Comment:
I don't think this flag is needed, and it gives a wrong answer in the one
case where it's read. When the `LIMIT 0` sits above the aggregate, `f_down`
replaces it before the aggregate is visited, so there's no count map and the
plain semi/anti join already returns the right result. Both `LIMIT 0` slt cases
pass on main. The flag only matters when the `LIMIT 0` is under the aggregate,
and there `count(*)` still produces a row:
```sql
-- DuckDB: 11, 22, 33, 44; this PR: empty (plan has `Filter: Boolean(false)`)
SELECT t1_id FROM t1 WHERE EXISTS (
SELECT count(*) FROM (SELECT * FROM t2 WHERE t2.t2_int = t1.t1_int LIMIT
0) x
WHERE x.t2_id = t1.t1_id
);
```
Removing `forces_empty_result` and its `lit(false)` branch fixes this query
and keeps both `LIMIT 0` tests passing.
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -426,6 +433,50 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
None
};
+ if self.exists_sub_query
+ && self.pull_up_having_expr.is_some()
+ && let Some(expr) = conjunction(subquery_filters.clone())
+ {
+ let unqualified_expr = expr
Review Comment:
Dropping the qualifier here means the merged filter is resolved against the
joined schema, so it becomes ambiguous when the outer table has a column with
the same name. Main returned rows for this query; it now fails to plan:
```sql
CREATE TABLE o9(k INT, c INT) AS VALUES (1, 100), (3, 100);
CREATE TABLE i9(k INT) AS VALUES (1);
-- DuckDB: 1, 3; main: 1; this PR: Ambiguous reference to unqualified field c
SELECT k FROM o9 WHERE EXISTS (
SELECT * FROM (
SELECT count(*) AS c FROM i9 WHERE i9.k = o9.k HAVING count(*) < 10
) s WHERE s.c < 5
);
```
Every column in `pull_up_having_expr` is a subquery output, so requalifying
it with the alias at the top of `build_join_with_count_bug` fixes it (returns
1, 3):
```rs
let pull_up_having_expr = pull_up_having_expr
.map(|f| {
f.clone()
.transform_up(|e| match e {
Expr::Column(c) =>
Ok(Transformed::yes(Expr::Column(Column::new(
Some(alias),
c.name,
)))),
_ => Ok(Transformed::no(e)),
})
.data()
})
.transpose()?;
let pull_up_having_expr = pull_up_having_expr.as_ref();
```
--
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]
