diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0c37f00b6435..71dd734250e7 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -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 @@ -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> { + 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 @@ -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)), } @@ -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)), } diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 7aa24d4c7fe3..79a9ea04dc43 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -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 () +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 () +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 () +02)--Subquery: +03)----Limit: skip=0, fetch=() +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 @@ -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 () +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