jayzhan211 commented on code in PR #26699:
URL: https://github.com/apache/datafusion/pull/26699#discussion_r4369529007
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -742,19 +741,49 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
// `GROUP BY` below returns a row for an outer row without
// a match, which the pulled up filter would drop.
if self.can_extend_partition_by &&
!self.pulled_up_scalar_agg {
- let extended =
- partition_window_by(window,
&local_correlated_cols)?;
- return self.pass_through_pulled_up_cols(extended);
+ partition_window_by(window, &local_correlated_cols)?
+ } else {
+ self.can_pull_up = false;
+ plan
}
- self.can_pull_up = false;
- }
- self.pass_through_pulled_up_cols(plan)
+ } else {
+ plan
+ };
+ let transformed = self
+ .pass_through_pulled_up_cols(plan,
input_count_expr_map.is_some())?;
+ (transformed, local_correlated_cols, None)
}
// `f_down` checked that the pulled up filters do not read an
// unnested column.
- LogicalPlan::Unnest(_) => self.pass_through_pulled_up_cols(plan),
- _ => Ok(Transformed::no(plan)),
+ LogicalPlan::Unnest(_)
Review Comment:
Recomputing the Join schema here lets the uncompensated EXISTS/IN fallback
(`rewrite(false)` in `build_join`) get past a Join that sits above a pulled-up
groupless aggregate. Before this PR that path failed with a schema error; now
it returns wrong rows:
```sql
CREATE TABLE t(a INT, b INT) AS VALUES (1,1),(2,2);
CREATE TABLE t2(a INT, b INT) AS VALUES (1,1),(3,30);
CREATE TABLE t3(c INT) AS VALUES (10),(20);
-- returns 1, expected 1, 2 (count(*) gives a row for every t row)
SELECT t.a FROM t WHERE EXISTS (SELECT 1 FROM (SELECT count(*) AS c FROM t2
WHERE t2.a = t.a) x CROSS JOIN t3);
-- returns nothing, expected 2
SELECT t.a FROM t WHERE EXISTS (SELECT 1 FROM (SELECT count(*) AS c FROM t2
WHERE t2.a = t.a) x CROSS JOIN t3 WHERE x.c = 0);
-- returns nothing, expected 2
SELECT t.a FROM t WHERE t.a IN (SELECT x.c + 2 FROM (SELECT count(*) AS c
FROM t2 WHERE t2.a = t.a) x CROSS JOIN t3);
```
Stopping at a Join above a pulled-up scalar aggregate puts these back to
"not supported" for me, and the subquery slt files still pass. Could you add
this arm and the three queries above as tests?
```rs
LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) if self.pulled_up_scalar_agg
=> {
self.can_pull_up = false;
return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop));
}
```
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -52,7 +52,13 @@ use datafusion_expr::{
#[derive(Debug)]
pub struct PullUpCorrelatedExpr {
pub join_filters: Vec<Expr>,
- /// mapping from the plan to its holding correlated columns
+ /// Correlated inner columns collected across the subquery, including their
+ /// re-qualified names after passing through
[`LogicalPlan::SubqueryAlias`].
+ pub correlated_subquery_cols: BTreeSet<Column>,
+ /// Deprecated: use [`Self::correlated_subquery_cols`] instead.
Review Comment:
nit: could the two kept maps use `#[deprecated(since = "56.0.0", note = "use
`correlated_subquery_cols` / `root_count_expr_map`")]` instead of a doc note,
so downstream users get a compiler warning? The internal writes would then need
`#[expect(deprecated)]`. The doc on `collected_count_expr_map` also still says
it maps each plan, but it now holds only the root entry.
##########
datafusion/sqllogictest/test_files/subquery.slt:
##########
@@ -4215,3 +4215,51 @@ SELECT * FROM (VALUES (1)) AS x(d)
WHERE x.d IN (SELECT t.c FROM (SELECT * FROM (VALUES (1)) AS y(c) WHERE y.c =
x.d) AS t WHERE x.d > 0);
----
1
+
+# Regression test for #26692: decorrelation must not lose correlated columns
+# when a subquery contains two structurally equal subplans.
+statement ok
+CREATE TABLE sq_dup_t(a INT);
+
+statement ok
+CREATE TABLE sq_dup_t2(a INT);
+
+statement ok
+INSERT INTO sq_dup_t VALUES (1), (2), (3);
+
+statement ok
+INSERT INTO sq_dup_t2 VALUES (1), (2);
+
+query I rowsort
+SELECT a FROM sq_dup_t WHERE EXISTS (
+ SELECT 1 FROM (SELECT s.a AS x FROM sq_dup_t2 AS s WHERE s.a = sq_dup_t.a)
+ CROSS JOIN (SELECT s.a AS y FROM sq_dup_t2 AS s)
+);
+----
+1
+2
+
+query I rowsort
+SELECT (
+ SELECT count(*) FROM (SELECT * FROM sq_dup_t2 WHERE sq_dup_t2.a =
sq_dup_t.a) AS s1
+ CROSS JOIN (SELECT * FROM sq_dup_t2) AS s2
+) FROM sq_dup_t;
+----
+0
+2
+2
+
+query I rowsort
+SELECT a FROM sq_dup_t WHERE EXISTS (
+ SELECT 1 FROM (SELECT sq_dup_t2.a FROM sq_dup_t2 WHERE sq_dup_t2.a =
sq_dup_t.a) AS x
+ CROSS JOIN LATERAL (SELECT sq_dup_t2.a FROM sq_dup_t2 WHERE sq_dup_t2.a =
x.a) AS l
Review Comment:
This LATERAL query already passes on main, so it doesn't cover the third
case from #26692. The wrong-rows case, where an outer `t2 AS s` shadows the
inner alias, isn't tested either. Could you add both from the issue as written?
It's fine to do in a follow-up if you'd rather.
```sql
SELECT t.a FROM t, t2 AS s
WHERE EXISTS (SELECT 1 FROM (SELECT * FROM t2 AS s WHERE s.a = t.a) x CROSS
JOIN t2 AS s);
-- expected: 1, 1
SELECT t.a FROM t
WHERE EXISTS (
SELECT 1 FROM (SELECT a FROM t2 WHERE t2.a = t.a) x
CROSS JOIN LATERAL (SELECT b FROM t2 WHERE t2.b = x.a) l
);
-- expected: 1
```
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -383,9 +408,14 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
}
fn f_up(&mut self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
- // Pop this node's own entry, pushed by the matching `f_down` call, so
- // `column_refs_above` holds only the nodes strictly above `plan`.
+ // Pop this node's own entries, pushed by the matching `f_down` call,
so
+ // `column_refs_above` holds only the nodes strictly above `plan` and
+ // `frame` holds the state pulled up from `plan`'s children.
self.column_refs_above.pop();
+ let PullUpFrame {
+ correlated_cols: local_correlated_cols,
+ count_expr_map: input_count_expr_map,
+ } = self.pulled_up_frames.pop().unwrap_or_default();
Review Comment:
nit: `unwrap_or_default()` hides a push/pop mismatch between `f_down` and
`f_up`, which is easy to introduce in an arm like the `Limit` one. Could this
return an internal error instead?
```rs
let Some(frame) = self.pulled_up_frames.pop() else {
return internal_err!("PullUpCorrelatedExpr: f_up without a matching
f_down frame");
};
```
--
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]
