From c666d06b6ee95bc6a255ba04bd72b5a7395be67e Mon Sep 17 00:00:00 2001 From: Fredrik Fornwall Date: Mon, 24 Aug 2026 22:56:32 +0200 Subject: [PATCH] fix: keep the OFFSET of a correlated EXISTS subquery A correlated subquery refers to columns of the query it is nested in, so in principle it has to be evaluated once per outer row: ```sql SELECT k FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2.v = t1.k) ``` Re-running the subquery for every outer row would be slow, and DataFusion has no operator that does so. Instead the optimizer rewrites it into a join, which is how analytical engines execute correlated subqueries efficiently: the condition that mentions the outer column (`t2.v = t1.k`) is pulled up out of the subquery and becomes the join condition of a semi join between `t1` and `t2`. Any operator that sits between that condition and the top of the subquery is in the way of the pull up, and a `LIMIT` is one such operator. For an `EXISTS` subquery a limit was simply deleted to clear the way, because `EXISTS` only asks whether the subquery returns any row, and a `LIMIT n` with `n > 0` cannot change that (a `LIMIT 0` was turned into an empty relation). This overlooked that the same plan node also carries the `OFFSET`, and an offset does change the answer: it skips rows, so a subquery that would return one row returns none after `OFFSET 1`. The offset was dropped along with the fetch and the query answered wrongly. With two rows of `v = 1` and one row of `v = 3` in `t2` ```sql SELECT k FROM t1 WHERE EXISTS (SELECT * FROM t2 WHERE t2.v = t1.k OFFSET 1) ``` returned both `1` and `3`, although only `k = 1` has a second row to skip past. Only remove the limit when it cannot change whether the subquery is empty, which is a zero offset with a literal fetch. A zero fetch still becomes an empty relation. Anything else, a positive offset or an offset or fetch that is not a literal, marks the subquery as one that cannot be pulled up, so it stays a correlated subquery in the plan and is reported as unsupported rather than answered wrongly. The decision to remove a limit was also made from the wrong information. It keyed on whether any correlated condition had been collected so far in the subquery, not on whether one sat below this particular limit. In a join inside the `EXISTS` the correlated side is visited first, so a `LIMIT` on an unrelated, uncorrelated sibling branch was deleted too and changed the result. Decide per limit instead, in `f_down`, from the outer references of that limit's own subtree, which is how `IN` subqueries were already handled. It has to happen in `f_down`: by `f_up` the filters below have been rewritten and their outer references are gone. The limit is rewritten right there in `f_down`, and the unsupported shapes now bail in `f_down` like the other unsupported shapes in this rewriter. Signed-off-by: Fredrik Fornwall --- datafusion/optimizer/src/decorrelate.rs | 74 +++++----- .../sqllogictest/test_files/subquery.slt | 139 ++++++++++++++++++ 2 files changed, 180 insertions(+), 33 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0c37f00b64355..71dd734250e72 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 7aa24d4c7fe37..79a9ea04dc431 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