Skip to content

fix: keep the OFFSET of a correlated EXISTS subquery #25284

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fornwall wants to merge 1 commit into
base: main
Choose a base branch
from
+180 −33
Open

fix: keep the OFFSET of a correlated EXISTS subquery #25284

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 41 additions & 33 deletions datafusion/optimizer/src/decorrelate.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use datafusion_expr::utils::{
};
use datafusion_expr::{
BinaryExpr, Cast, EmptyRelation, Expr, ExprSchemable, FetchType, LogicalPlan,
LogicalPlanBuilder, Operator, expr, lit,
LogicalPlanBuilder, Operator, SkipType, expr, lit,
};

/// This struct rewrite the sub query plan by pull up the correlated
Expand Down Expand Up @@ -117,6 +117,13 @@ impl PullUpCorrelatedExpr {
self.exists_sub_query = exists_sub_query;
self
}

/// Mark the plan as one whose correlated expressions cannot be pulled up
/// and stop descending into it
fn unsupported(&mut self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
self.can_pull_up = false;
Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
}
}

/// Used to indicate the unmatched rows from the inner(subquery) table after the left out Join
Expand Down Expand Up @@ -145,28 +152,44 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
LogicalPlan::Union(_) | LogicalPlan::Sort(_) | LogicalPlan::Extension(_) => {
let plan_hold_outer = !plan.all_out_ref_exprs().is_empty();
if plan_hold_outer {
// the unsupported case
self.can_pull_up = false;
Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
self.unsupported(plan)
} else {
Ok(Transformed::no(plan))
}
}
LogicalPlan::Limit(_) => {
let plan_hold_outer = !plan.all_out_ref_exprs().is_empty();
match (self.exists_sub_query, plan_hold_outer) {
(false, true) => {
// the unsupported case
self.can_pull_up = false;
Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
LogicalPlan::Limit(ref limit) => {
if plan.all_out_ref_exprs().is_empty() {
return Ok(Transformed::no(plan));
}
if !self.exists_sub_query {
return self.unsupported(plan);
}
// Only emptiness matters for EXISTS, so remove a limit that
// cannot make its input empty and replace one that always does
// with an empty relation.
let fetch = limit.get_fetch_type()?;
if matches!(fetch, FetchType::Literal(Some(0))) {
return Ok(Transformed::yes(LogicalPlan::EmptyRelation(
EmptyRelation {
produce_one_row: false,
schema: Arc::clone(limit.input.schema()),
},
)));
}
match (limit.get_skip_type()?, fetch) {
(SkipType::Literal(0), FetchType::Literal(_)) => {
// The rewriter does not call `f_down` on the returned
// node, so do it here
let mut t = self.f_down((*limit.input).clone())?;
t.transformed = true;
Ok(t)
}
_ => Ok(Transformed::no(plan)),
_ => self.unsupported(plan),
}
}
_ if plan.contains_outer_reference() => {
// the unsupported cases, the plan expressions contain out reference columns(like window expressions)
self.can_pull_up = false;
Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
self.unsupported(plan)
}
_ => Ok(Transformed::no(plan)),
}
Expand Down Expand Up @@ -375,28 +398,13 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
}
}
LogicalPlan::Limit(limit) => {
let input_expr_map =
self.collected_count_expr_map.get(&*limit.input).cloned();
// handling the limit clause in the subquery
let new_plan = match (self.exists_sub_query, self.join_filters.is_empty())
if let Some(input_map) =
self.collected_count_expr_map.get(&*limit.input).cloned()
{
// 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)) => {
LogicalPlan::EmptyRelation(EmptyRelation {
produce_one_row: false,
schema: Arc::clone(limit.input.schema()),
})
}
_ => LogicalPlanBuilder::from((*limit.input).clone()).build()?,
}),
_ => Transformed::no(plan),
};
if let Some(input_map) = input_expr_map {
self.collected_count_expr_map
.insert(new_plan.data.clone(), input_map);
.insert(plan.clone(), input_map);
}
Ok(new_plan)
Ok(Transformed::no(plan))
}
_ => Ok(Transformed::no(plan)),
}
Expand Down
139 changes: 139 additions & 0 deletions datafusion/sqllogictest/test_files/subquery.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -656,6 +656,125 @@ SELECT t1_id, t1_name FROM t1 WHERE NOT EXISTS (SELECT * FROM t2 WHERE t2_id = t
33 c
44 d

#exists_subquery_with_offset0
#de-correlated, limit is removed
query TT
explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id limit 1 offset 0)
----
logical_plan
01)LeftSemi Join: t1.t1_id = __correlated_sq_1.t2_id
02)--TableScan: t1 projection=[t1_id, t1_name]
03)--SubqueryAlias: __correlated_sq_1
04)----TableScan: t2 projection=[t2_id]

query IT rowsort
SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id limit 1 offset 0)
----
11 a
22 b
44 d

#exists_subquery_with_offset
#not de-correlated, the offset could make the subquery empty
query TT
explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id offset 1)
----
logical_plan
01)Filter: EXISTS (<subquery>)
02)--Subquery:
03)----Limit: skip=1, fetch=None
04)------Projection: t2.t2_id, t2.t2_name, t2.t2_int
05)--------Filter: t2.t2_id = outer_ref(t1.t1_id)
06)----------TableScan: t2
07)--TableScan: t1 projection=[t1_id, t1_name]

# errors rather than returning wrong rows
query error Physical plan does not support logical expression Exists
SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id offset 1)

#not_exists_subquery_with_offset
#not de-correlated, the offset could make the subquery empty
query TT
explain SELECT t1_id, t1_name FROM t1 WHERE NOT EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id offset 1)
----
logical_plan
01)Filter: NOT EXISTS (<subquery>)
02)--Subquery:
03)----Limit: skip=1, fetch=None
04)------Projection: t2.t2_id, t2.t2_name, t2.t2_int
05)--------Filter: t2.t2_id = outer_ref(t1.t1_id)
06)----------TableScan: t2
07)--TableScan: t1 projection=[t1_id, t1_name]

# errors rather than returning wrong rows
query error Physical plan does not support logical expression Exists
SELECT t1_id, t1_name FROM t1 WHERE NOT EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id offset 1)

#exists_subquery_with_offset_in_disjunction
#not de-correlated, errors rather than returning wrong rows
query error Physical plan does not support logical expression Exists
SELECT t1_id, t1_name FROM t1 WHERE t1_id > 40 OR EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id offset 1)

#exists_subquery_with_limit0_and_offset
#de-correlated, limit 0 is empty whatever the offset
query TT
explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id limit 0 offset 1)
----
logical_plan EmptyRelation: rows=0

query IT rowsort
SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id limit 0 offset 1)
----

#exists_subquery_with_non_literal_limit
#not de-correlated, the fetch could evaluate to 0 and make the subquery empty
query TT
explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id limit (SELECT count(*) FROM t2))
----
logical_plan
01)Filter: EXISTS (<subquery>)
02)--Subquery:
03)----Limit: skip=0, fetch=(<subquery>)
04)------Subquery:
05)--------Projection: count(Int64(1)) AS count(*)
06)----------Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]
07)------------TableScan: t2
08)------Projection: t2.t2_id, t2.t2_name, t2.t2_int
09)--------Filter: t2.t2_id = outer_ref(t1.t1_id)
10)----------TableScan: t2
11)--TableScan: t1 projection=[t1_id, t1_name]

#exists_subquery_with_limit_on_uncorrelated_branch
#de-correlated, the limit is on an uncorrelated branch so it stays
query TT
explain SELECT t1_id FROM t1 WHERE EXISTS (SELECT 1 FROM (SELECT * FROM t2 WHERE t2_id = t1_id) a JOIN (SELECT * FROM t3 ORDER BY t3_id LIMIT 1) b ON a.t2_int = b.t3_int)
----
logical_plan
01)LeftSemi Join: t1.t1_id = __correlated_sq_1.t2_id
02)--TableScan: t1 projection=[t1_id]
03)--SubqueryAlias: __correlated_sq_1
04)----Projection: a.t2_id
05)------LeftSemi Join: a.t2_int = b.t3_int
06)--------SubqueryAlias: a
07)----------TableScan: t2 projection=[t2_id, t2_int]
08)--------SubqueryAlias: b
09)----------Projection: t3.t3_int
10)------------Sort: t3.t3_id ASC NULLS LAST, fetch=1
11)--------------TableScan: t3 projection=[t3_id, t3_int]

query I rowsort
SELECT t1_id FROM t1 WHERE EXISTS (SELECT 1 FROM (SELECT * FROM t2 WHERE t2_id = t1_id) a JOIN (SELECT * FROM t3 ORDER BY t3_id LIMIT 1) b ON a.t2_int = b.t3_int)
----
11
44

#exists_subquery_with_offset_on_uncorrelated_branch
#de-correlated, the offset is on an uncorrelated branch so it stays
query I rowsort
SELECT t1_id FROM t1 WHERE EXISTS (SELECT 1 FROM (SELECT * FROM t2 WHERE t2_id = t1_id) a JOIN (SELECT * FROM t3 ORDER BY t3_id LIMIT 1 OFFSET 1) b ON a.t2_int = b.t3_int)
----
22

#in_correlated_subquery_with_limit
#not de-correlated
query TT
Expand Down Expand Up @@ -717,6 +836,26 @@ logical_plan
09)----------TableScan: t2
10)--TableScan: t1 projection=[t1_id, t1_name]

#exists_subquery_with_limit_over_union
#not de-correlated, the union below the removed limit still holds outer references
query TT
explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2_id = t1_id UNION ALL SELECT * FROM t2 WHERE upper(t2_name) = upper(t1.t1_name) LIMIT 1)
----
logical_plan
01)Filter: EXISTS (<subquery>)
02)--Subquery:
03)----Limit: skip=0, fetch=1
04)------Union
05)--------Projection: t2.t2_id, t2.t2_name, t2.t2_int
06)----------Limit: skip=0, fetch=1
07)------------Filter: t2.t2_id = outer_ref(t1.t1_id)
08)--------------TableScan: t2
09)--------Projection: t2.t2_id, t2.t2_name, t2.t2_int
10)----------Limit: skip=0, fetch=1
11)------------Filter: upper(t2.t2_name) = upper(outer_ref(t1.t1_name))
12)--------------TableScan: t2
13)--TableScan: t1 projection=[t1_id, t1_name]

#simple_uncorrelated_scalar_subquery
query TT
explain select (select count(*) from t1) as b
Expand Down