From 43e8ba9e171111cede67009db7a93c2b355ca6dd Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Tue, 15 Sep 2026 10:46:45 -0400 Subject: [PATCH 01/10] fix: count-bug compensation for groupless aggregates --- datafusion/optimizer/src/decorrelate.rs | 85 +++++++-- .../src/decorrelate_predicate_subquery.rs | 174 +++++++++++++++++- .../sqllogictest/test_files/subquery.slt | 150 +++++++++++++++ 3 files changed, 395 insertions(+), 14 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0a1db9b2128c6..117d6fef2a61d 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -426,6 +426,51 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { None }; + if self.pull_up_having_expr.is_some() + && let Some(expr) = conjunction(subquery_filters.clone()) + { + let unqualified_expr = expr + .transform_up(|e| { + if let Expr::Column(Column { name, .. }) = &e { + Ok(Transformed::yes(Expr::Column( + Column::new_unqualified(name), + ))) + } else { + Ok(Transformed::no(e)) + } + }) + .data()?; + let combined = match self.pull_up_having_expr.take() { + Some(existing) => existing.and(unqualified_expr), + None => unqualified_expr, + }; + self.pull_up_having_expr = Some(combined); + let new_plan = + LogicalPlanBuilder::from((*plan_filter.input).clone()).build()?; + // This filter has no correlated column of its own, but + // the aggregate underneath it does. + let mut carried_correlated_cols = correlated_subquery_cols; + if let Some(existing) = + self.correlated_subquery_cols_map.get(&*plan_filter.input) + { + carried_correlated_cols.extend(existing.iter().cloned()); + } + self.correlated_subquery_cols_map + .insert(new_plan.clone(), carried_correlated_cols); + if !expr_result_map_for_count_bug.is_empty() { + self.collected_count_expr_map + .insert(new_plan.clone(), expr_result_map_for_count_bug); + } else if let Some(input_map) = self + .collected_count_expr_map + .get(&*plan_filter.input) + .cloned() + { + self.collected_count_expr_map + .insert(new_plan.clone(), input_map); + } + return Ok(Transformed::yes(new_plan)); + } + match (&pull_up_expr_opt, &self.pull_up_having_expr) { (Some(_), Some(_)) => { // Error path @@ -472,7 +517,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { self.collected_count_expr_map.get(&*projection.input) { proj_exprs_evaluation_result_on_empty_batch( - &projection.expr, + &missing_exprs, projection.input.schema(), expr_result_map, &mut expr_result_map_for_count_bug, @@ -813,19 +858,22 @@ impl PullUpCorrelatedExpr { } for col in correlated_subquery_cols.iter() { let col_expr = Expr::Column(col.clone()); - if !missing_exprs.contains(&col_expr) { + if !collides_with_existing(&missing_exprs, &col_expr) { missing_exprs.push(col_expr) } } if let Some(pull_up_having) = &self.pull_up_having_expr { - let filter_apply_columns = pull_up_having.column_refs(); - for col in filter_apply_columns { - // add to missing_exprs if not already there - let contains = missing_exprs + for col in pull_up_having.column_refs() { + let col_expr = Expr::Column(col.clone()); + // `col` is unqualified but the projection may already have + // it as `agg.c`. Check by bare name so that doesn't get + // added twice. + let already_present = missing_exprs .iter() - .any(|expr| matches!(expr, Expr::Column(c) if c == col)); - if !contains { - missing_exprs.push(Expr::Column(col.clone())) + .any(|expr| matches!(expr, Expr::Column(c) if c.name == col.name)) + || collides_with_existing(&missing_exprs, &col_expr); + if !already_present { + missing_exprs.push(col_expr) } } } @@ -925,6 +973,16 @@ fn partitions_by_all(expr: &Expr, cols: &BTreeSet) -> bool { }) } +/// True if `candidate` collides with an existing entry's schema name, +/// e.g. a bare column already wrapped in an equally-named `CAST`. +/// Colliding means `LogicalPlanBuilder::project` would reject it as a duplicate. +fn collides_with_existing(exprs: &[Expr], candidate: &Expr) -> bool { + let candidate_name = candidate.schema_name().to_string(); + exprs + .iter() + .any(|expr| expr.schema_name().to_string() == candidate_name) +} + fn can_pullup_over_aggregation(expr: &Expr) -> bool { if let Expr::BinaryExpr(BinaryExpr { left, @@ -1105,9 +1163,14 @@ fn filter_exprs_evaluation_result_on_empty_batch( let simplifier = ExprSimplifier::new(info); let result_expr = simplifier.simplify(result_expr)?; match &result_expr { - // evaluate to false or null on empty batch, no need to pull up Expr::Literal(ScalarValue::Null, _) - | Expr::Literal(ScalarValue::Boolean(Some(false)), _) => None, + | Expr::Literal(ScalarValue::Boolean(None), _) + | Expr::Literal(ScalarValue::Boolean(Some(false)), _) => { + for (name, exprs) in input_expr_result_map_for_count_bug { + expr_result_map_for_count_bug.insert(name.clone(), exprs.clone()); + } + None + } // evaluate to true on empty batch, need to pull up the expr Expr::Literal(ScalarValue::Boolean(Some(true)), _) => { for (name, exprs) in input_expr_result_map_for_count_bug { diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 99b57a5494aa8..acfc4f9a6ad26 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -20,7 +20,8 @@ use std::collections::BTreeSet; use std::ops::Deref; use std::sync::Arc; -use crate::decorrelate::PullUpCorrelatedExpr; +use crate::analyzer::type_coercion::TypeCoercionRewriter; +use crate::decorrelate::{PullUpCorrelatedExpr, UN_MATCHED_ROW_INDICATOR}; use crate::extract_equijoin_predicate::split_eq_and_noneq_join_predicate; use crate::optimizer::ApplyOrder; use crate::utils::replace_qualified_name; @@ -779,13 +780,16 @@ fn build_join( ) -> Result> { let mut pull_up = PullUpCorrelatedExpr::new() .with_in_predicate_opt(in_predicate_opt.cloned()) - .with_exists_sub_query(in_predicate_opt.is_none()); + .with_exists_sub_query(in_predicate_opt.is_none()) + .with_need_handle_count_bug(true); let new_plan = subquery.clone().rewrite(&mut pull_up).data()?; if !pull_up.can_pull_up { return Ok(None); } + let count_bug_compensation = pull_up.collected_count_expr_map.get(&new_plan).cloned(); + let sub_query_alias = LogicalPlanBuilder::from(new_plan) .alias(alias.to_string())? .build()?; @@ -801,6 +805,27 @@ fn build_join( replace_qualified_name(filter, &all_correlated_cols, alias).map(Some) })?; + if let Some(expr_map) = count_bug_compensation + && !expr_map.is_empty() + { + return build_join_with_count_bug( + left, + sub_query_alias, + join_filter_opt, + in_predicate_opt, + join_type, + alias, + &expr_map, + pull_up.pull_up_having_expr.as_ref(), + ) + .map(|plan| { + Some(BuiltJoin { + plan, + mark_is_three_valued_exact: false, + }) + }); + } + // The two sides of the `IN` predicate: the value from the outer plan and // the subquery output it is compared with, renamed to the alias. A // multi-column `(a, b) IN (SELECT x, y ...)` has one pair per tuple @@ -1059,6 +1084,116 @@ fn build_join( })) } +/// Builds the join for a correlated `EXISTS` subquery whose groupless +/// aggregate requires count-bug compensation. +#[expect(clippy::too_many_arguments)] +fn build_join_with_count_bug( + left: &LogicalPlan, + sub_query_alias: LogicalPlan, + join_filter_opt: Option, + in_predicate_opt: Option<&Expr>, + join_type: JoinType, + alias: &str, + expr_map: &crate::decorrelate::ExprResultMap, + pull_up_having_expr: Option<&Expr>, +) -> Result { + if in_predicate_opt.is_some() { + return not_impl_err!( + "build_join_with_count_bug: `IN`/`NOT IN` count-bug compensation is not implemented" + ); + } + + let joined = LogicalPlanBuilder::from(left.clone()) + .join_on(sub_query_alias, JoinType::Left, join_filter_opt)? + .build()?; + + let left_projection: Vec = left + .schema() + .columns() + .into_iter() + .map(Expr::from) + .collect(); + + let indicator_col = Expr::Column(Column::new(Some(alias), UN_MATCHED_ROW_INDICATOR)); + + let having_arm = pull_up_having_expr.map(|f| f.clone().is_not_true()); + + let mut expr_rewrite = TypeCoercionRewriter::new(joined.schema()); + + // The subquery's HAVING clause, evaluated against an unmatched row's + // default aggregate values (e.g. count(*) defaults to 0, sum(x) to + // NULL), is usually `true`, but not once a second filter is combined + // into it, so it cannot be assumed. + let unmatched_having_default = match pull_up_having_expr { + Some(f) => f + .clone() + .transform_up(|e| { + if let Expr::Column(Column { name, .. }) = &e + && let Some(default_value) = expr_map.get(name) + { + return Ok(Transformed::yes(default_value.clone())); + } + Ok(Transformed::no(e)) + }) + .data() + .and_then(|simplified| simplified.rewrite(&mut expr_rewrite).data())?, + None => lit(true), + }; + + // EXISTS is true by default (the groupless aggregate always + // produces a row), unless either the row joined against a group + // whose HAVING predicate failed, or an unmatched row's own + // default aggregate values fail that same HAVING clause. + let exists_expr = match having_arm { + Some(when_expr) => { + when(indicator_col.clone().is_null(), unmatched_having_default) + .when(when_expr, lit(false)) + .otherwise(lit(true))? + .rewrite(&mut expr_rewrite) + .data()? + } + None => lit(true), + }; + + // The correlation condition can match more than one subquery row per + // outer row (e.g. two rows that compare equal only after a `CAST`). + // A plain `LEFT JOIN` keeps every one of those matches, so `distinct` + // collapses them back down to the one row `LeftSemi`/`LeftAnti`/`LeftMark` + // would have produced without count-bug compensation. + let new_plan = match join_type { + JoinType::LeftMark | JoinType::RightMark => { + let mut proj_exprs = left_projection; + proj_exprs.push(exists_expr.alias_qualified(Some(alias), "mark")); + LogicalPlanBuilder::from(joined) + .project(proj_exprs)? + .distinct()? + .build()? + } + JoinType::LeftAnti => LogicalPlanBuilder::from(joined) + .filter(not(exists_expr))? + .project(left_projection)? + .distinct()? + .build()?, + JoinType::LeftSemi => LogicalPlanBuilder::from(joined) + .filter(exists_expr)? + .project(left_projection)? + .distinct()? + .build()?, + _ => { + return internal_err!( + "build_join_with_count_bug: unsupported join type {join_type:?}" + ); + } + }; + + debug!( + "predicate subquery (count bug) optimized:\n{}", + new_plan.display_indent() + ); + + Ok(new_plan) +} + #[derive(Debug)] struct SubqueryInfo { query: Subquery, @@ -1111,6 +1246,7 @@ mod tests { TableScanBuilder, and, binary_expr, col, cube, grouping_set, out_ref_col, rollup, table_scan, }; + use datafusion_functions_aggregate::count::count_udaf; macro_rules! assert_optimized_plan_equal { ( @@ -1438,7 +1574,7 @@ mod tests { SubqueryAlias: __correlated_sq_2 [o_custkey:Int64] Projection: orders.o_custkey [o_custkey:Int64] TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] - " + " ) } @@ -2855,6 +2991,38 @@ mod tests { ) } + /// Test for correlated exists subquery filter with disjunction and count bug + #[test] + fn exists_subquery_disjunction_with_count_bug() -> Result<()> { + let sq = Arc::new( + LogicalPlanBuilder::from(test_table_scan_with_name("sq")?) + .filter(out_ref_col(DataType::UInt32, "test.a").eq(col("sq.a")))? + .aggregate(Vec::::new(), vec![count_udaf().call(vec![])])? + .build()?, + ); + + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(exists(sq).or(col("test.c").eq(lit(1))))? + .project(vec![col("test.b")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.b [b:UInt32] + Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32] + Filter: __correlated_sq_1.mark OR test.c = Int32(1) [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] + Distinct: [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] + Projection: test.a, test.b, test.c, Boolean(true) AS mark [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] + Left Join: Filter: test.a = __correlated_sq_1.a [a:UInt32, b:UInt32, c:UInt32, a:UInt32;N, __always_true:Boolean;N, count():Int64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + SubqueryAlias: __correlated_sq_1 [a:UInt32, __always_true:Boolean, count():Int64] + Aggregate: groupBy=[[sq.a, Boolean(true) AS __always_true]], aggr=[[count()]] [a:UInt32, __always_true:Boolean, count():Int64] + TableScan: sq [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + /// Test for correlated EXISTS subquery filter #[test] fn exists_subquery_correlated() -> Result<()> { diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 39b3961c82451..7ab4b125dc872 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1857,6 +1857,156 @@ item1 1970-01-01T00:00:03 75 statement ok drop table source_table; + + +# correlated_exists_groupless_count_agg_count_bug +query II rowsort +SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- +11 1 +22 2 +33 3 +44 4 + +# correlated_exists_groupless_count_agg_count_bug_having +query TT +explain SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) +---- +logical_plan +01)Projection: t1.t1_id, t1.t1_int +02)--Aggregate: groupBy=[[t1.t1_id, t1.t1_name, t1.t1_int]], aggr=[[]] +03)----Projection: t1.t1_id, t1.t1_name, t1.t1_int +04)------Filter: __correlated_sq_1.__always_true IS NULL OR __correlated_sq_1.__always_true IS NOT NULL AND NOT __correlated_sq_1.count(Int64(1)) = Int64(0) IS NOT TRUE +05)--------Projection: t1.t1_id, t1.t1_name, t1.t1_int, __correlated_sq_1.count(Int64(1)), __correlated_sq_1.__always_true +06)----------Left Join: t1.t1_id = __correlated_sq_1.t2_id +07)------------TableScan: t1 projection=[t1_id, t1_name, t1_int] +08)------------SubqueryAlias: __correlated_sq_1 +09)--------------Projection: t2.t2_id, count(Int64(1)), Boolean(true) AS __always_true +10)----------------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1))]] +11)------------------TableScan: t2 projection=[t2_id] + +query II +SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) +---- +33 3 + +# correlated_exists_groupless_compound_having_select_dummy_column +query II +SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT 1 FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0 AND count(*) < 100) +---- +33 3 + +# correlated_not_exists_groupless_compound_having_select_dummy_column +query II rowsort +SELECT t1_id, t1_int FROM t1 WHERE NOT EXISTS (SELECT 1 FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0 AND count(*) < 100) +---- +11 1 +22 2 +44 4 + +# correlated_exists_groupless_having_with_stacked_filter_via_cte +query II +SELECT t1_id, t1_int FROM t1 WHERE EXISTS ( + WITH agg AS (SELECT count(*) AS c FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) + SELECT c FROM agg WHERE c >= 0 +) +---- +33 3 + +# correlated_not_exists_groupless_having_with_stacked_filter_via_cte +query II rowsort +SELECT t1_id, t1_int FROM t1 WHERE NOT EXISTS ( + WITH agg AS (SELECT count(*) AS c FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) + SELECT c FROM agg WHERE c >= 0 +) +---- +11 1 +22 2 +44 4 + +# correlated_exists_groupless_having_with_stacked_filter_on_different_aggregate +query II +SELECT t1_id, t1_int FROM t1 WHERE EXISTS ( + WITH agg AS (SELECT count(*) AS c, sum(t2_int) AS s FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) + SELECT c FROM agg WHERE s > 5 +) +---- + +# correlated_not_exists_groupless_count_agg_count_bug +query II +SELECT t1_id, t1_int FROM t1 WHERE NOT EXISTS (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- + +# correlated_exists_groupless_null_default_agg_always_true +query II rowsort +SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- +11 1 +22 2 +33 3 +44 4 + +# correlated_not_exists_groupless_null_default_agg_always_true +query II +SELECT t1_id, t1_int FROM t1 WHERE NOT EXISTS (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- + +# correlated_exists_groupless_count_agg_count_bug_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or exists (select count(*) from t2 where t1.t1_id = t2.t2_id having count(*) = 0) +---- +33 +44 + +# correlated_not_exists_groupless_count_agg_count_bug_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or not exists (select count(*) from t2 where t1.t1_id = t2.t2_id having count(*) = 0) +---- +11 +22 +44 + +# correlated_exists_groupless_null_default_agg_count_bug_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or exists (select sum(t2_int) from t2 where t1.t1_id = t2.t2_id having sum(t2_int) is null) +---- +33 +44 + +# correlated_not_exists_groupless_null_default_agg_count_bug_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or not exists (select sum(t2_int) from t2 where t1.t1_id = t2.t2_id having sum(t2_int) is null) +---- +11 +22 +44 + +# correlated_exists_groupless_count_agg_duplicate_match_dedup +statement ok +CREATE TABLE dedup_outer(k INT) AS VALUES (1); + +statement ok +CREATE TABLE dedup_inner(k TEXT) AS VALUES ('1'), ('01'); + +query I +SELECT o.k +FROM dedup_outer o +WHERE EXISTS ( + SELECT count(*) FROM dedup_inner i WHERE CAST(i.k AS INT) = o.k +) +---- +1 + +statement ok +DROP TABLE dedup_outer; + +statement ok +DROP TABLE dedup_inner; + statement count 0 drop table t1; From fa4e3fb734072d078f5be33740bede4c7dacf091 Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Tue, 15 Sep 2026 16:17:57 -0400 Subject: [PATCH 02/10] fix: dedup by grouping on the join comparison --- datafusion/expr/src/utils.rs | 32 +++++++ datafusion/optimizer/src/decorrelate.rs | 87 ++++++++++++++++--- .../src/decorrelate_predicate_subquery.rs | 63 ++++++-------- .../sqllogictest/test_files/subquery.slt | 49 ++++++++--- 4 files changed, 175 insertions(+), 56 deletions(-) diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 777e8350782a8..bf7976ca8999b 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -1574,6 +1574,38 @@ pub fn collect_subquery_cols( }) } +/// Determine the subquery-side expressions a join filter compares, +/// stopping at the first subquery-referencing sub-expression rather than +/// descending to the bare columns inside of it. +/// +/// `t1.a = t2.b` yields `t2.b`, matching [`collect_subquery_cols`]. But +/// `t1.a = CAST(t2.b AS INT)` yields `CAST(t2.b AS INT)` whole, not `t2.b`. +/// The two differ whenever the correlated predicate wraps the subquery's +/// column in an expression, in which case we need to evaluate the full +/// expression before grouping, so that two column values which compare +/// equal after the wrapping expression end up in the same group. +pub fn collect_subquery_join_exprs( + exprs: &[Expr], + subquery_schema: &DFSchema, +) -> Result> { + let mut result: Vec = vec![]; + for expr in exprs { + expr.apply(|e| { + if e.column_refs() + .into_iter() + .all(|col| subquery_schema.has_column(col)) + && !e.column_refs().is_empty() + && !result.contains(e) + { + result.push(e.clone()); + return Ok(TreeNodeRecursion::Jump); + } + Ok(TreeNodeRecursion::Continue) + })?; + } + Ok(result) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 117d6fef2a61d..02a68ab304f0e 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -33,7 +33,8 @@ use datafusion_expr::expr::{Alias, GroupingSet}; use datafusion_expr::logical_plan::{Join, JoinType}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ - collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, + collect_subquery_cols, collect_subquery_join_exprs, conjunction, find_join_exprs, + split_conjunction, }; use datafusion_expr::{ Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, Extension, @@ -100,6 +101,14 @@ pub struct PullUpCorrelatedExpr { /// grouping set that omits a correlated column can tell whether /// something above it reads the column the pull up would fill in. column_refs_above: Vec>, + /// A correlated column wrapped in an expression (e.g. `CAST(t2.b AS INT)`) + /// gets grouped by that entire expression instead of the bare column, + /// aliased to a generated name, once, in the `Aggregate` this column belongs to. Every + /// later reference to that column, in a `Projection` above it for + /// instance, needs to use the same alias instead of the now + /// unresolvable bare column. This records the mapping the first + /// time it's made. + correlated_col_aliases: HashMap, } impl Default for PullUpCorrelatedExpr { @@ -123,6 +132,7 @@ impl PullUpCorrelatedExpr { pulled_up_scalar_agg: false, correlated_filters: Vec::new(), column_refs_above: Vec::new(), + correlated_col_aliases: HashMap::new(), } } @@ -447,8 +457,6 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { self.pull_up_having_expr = Some(combined); let new_plan = LogicalPlanBuilder::from((*plan_filter.input).clone()).build()?; - // This filter has no correlated column of its own, but - // the aggregate underneath it does. let mut carried_correlated_cols = correlated_subquery_cols; if let Some(existing) = self.correlated_subquery_cols_map.get(&*plan_filter.input) @@ -509,8 +517,11 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { &mut local_correlated_cols, ); // add missing columns to Projection - let mut missing_exprs = - self.collect_missing_exprs(&projection.expr, &local_correlated_cols)?; + let mut missing_exprs = self.collect_missing_exprs( + &projection.expr, + &local_correlated_cols, + projection.input.schema(), + )?; let mut expr_result_map_for_count_bug = HashMap::new(); if let Some(expr_result_map) = @@ -556,7 +567,6 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { &self.correlated_subquery_cols_map, &mut local_correlated_cols, ); - // A grouping set that omits a column the pull up adds either // keeps its sets as they are (every set already groups by the // column), or gains the column in every set, `ROLLUP(i.k)` @@ -589,6 +599,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { self.collect_missing_exprs( &aggregate.group_expr, &local_correlated_cols, + aggregate.input.schema(), )? } GroupingSetsPullUp::Unsafe => { @@ -603,6 +614,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { self.collect_missing_exprs( &aggregate.group_expr, &local_correlated_cols, + aggregate.input.schema(), )? } } @@ -611,6 +623,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { self.collect_missing_exprs( &aggregate.group_expr, &local_correlated_cols, + aggregate.input.schema(), )? }; @@ -655,8 +668,16 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { ); let mut new_correlated_cols = BTreeSet::new(); for col in local_correlated_cols.iter() { - new_correlated_cols - .insert(Column::new(Some(alias.alias.clone()), col.name.clone())); + let requalified = + Column::new(Some(alias.alias.clone()), col.name.clone()); + // A column already folded into a group-by alias (see + // `collect_missing_exprs`) needs that mapping carried + // forward under its requalified name. + if let Some(existing_alias) = self.correlated_col_aliases.get(col) { + self.correlated_col_aliases + .insert(requalified.clone(), existing_alias.clone()); + } + new_correlated_cols.insert(requalified); } let new_plan = if alias.input.schema().fields().len() @@ -845,9 +866,10 @@ impl PullUpCorrelatedExpr { } fn collect_missing_exprs( - &self, + &mut self, exprs: &[Expr], correlated_subquery_cols: &BTreeSet, + subquery_schema: &DFSchemaRef, ) -> Result> { let mut missing_exprs = Vec::with_capacity(exprs.len() + correlated_subquery_cols.len()); @@ -856,8 +878,53 @@ impl PullUpCorrelatedExpr { 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 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 + .iter() + .find(|e| e.column_refs().len() == 1 && e.column_refs().contains(col)); + let col_expr = match wrapped { + Some(e) if !matches!(e, Expr::Column(_)) => { + let alias_name = format!("__correlated_group_expr_{}", col.name); + let aliased = e.clone().alias(alias_name.clone()); + let reference = Expr::Column(Column::new_unqualified(&alias_name)); + self.join_filters = self + .join_filters + .iter() + .map(|f| { + f.clone() + .transform_up(|node| { + if &node == e { + Ok(Transformed::yes(reference.clone())) + } else { + Ok(Transformed::no(node)) + } + }) + .data() + }) + .collect::>>()?; + self.correlated_col_aliases.insert(col.clone(), reference); + aliased + } + _ => Expr::Column(col.clone()), + }; if !collides_with_existing(&missing_exprs, &col_expr) { missing_exprs.push(col_expr) } diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index acfc4f9a6ad26..f7f85752269cf 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -1155,29 +1155,21 @@ fn build_join_with_count_bug( None => lit(true), }; - // The correlation condition can match more than one subquery row per - // outer row (e.g. two rows that compare equal only after a `CAST`). - // A plain `LEFT JOIN` keeps every one of those matches, so `distinct` - // collapses them back down to the one row `LeftSemi`/`LeftAnti`/`LeftMark` - // would have produced without count-bug compensation. let new_plan = match join_type { JoinType::LeftMark | JoinType::RightMark => { let mut proj_exprs = left_projection; proj_exprs.push(exists_expr.alias_qualified(Some(alias), "mark")); LogicalPlanBuilder::from(joined) .project(proj_exprs)? - .distinct()? .build()? } JoinType::LeftAnti => LogicalPlanBuilder::from(joined) .filter(not(exists_expr))? .project(left_projection)? - .distinct()? .build()?, JoinType::LeftSemi => LogicalPlanBuilder::from(joined) .filter(exists_expr)? .project(left_projection)? - .distinct()? .build()?, _ => { return internal_err!( @@ -1804,10 +1796,10 @@ mod tests { plan, @r" Projection: customer.c_custkey [c_custkey:Int64] - LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey AND (customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.o_orderkey = Int32(1)) [c_custkey:Int64, c_name:Utf8] + LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey AND (customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.__correlated_group_expr_o_orderkey) [c_custkey:Int64, c_name:Utf8] TableScan: customer [c_custkey:Int64, c_name:Utf8] - SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, o_orderkey:Int64] - Projection: orders.o_custkey, orders.o_orderkey [o_custkey:Int64, o_orderkey:Int64] + SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] + Projection: orders.o_custkey, orders.o_orderkey = Int32(1) AS __correlated_group_expr_o_orderkey [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] " ) @@ -2860,10 +2852,10 @@ mod tests { plan, @r" Projection: customer.c_custkey [c_custkey:Int64] - LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.o_orderkey = Int32(1) [c_custkey:Int64, c_name:Utf8] + LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.__correlated_group_expr_o_orderkey [c_custkey:Int64, c_name:Utf8] TableScan: customer [c_custkey:Int64, c_name:Utf8] - SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, o_orderkey:Int64] - Projection: orders.o_custkey, orders.o_orderkey [o_custkey:Int64, o_orderkey:Int64] + SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] + Projection: orders.o_custkey, orders.o_orderkey = Int32(1) AS __correlated_group_expr_o_orderkey [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] " ) @@ -3012,13 +3004,12 @@ mod tests { Projection: test.b [b:UInt32] Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32] Filter: __correlated_sq_1.mark OR test.c = Int32(1) [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] - Distinct: [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] - Projection: test.a, test.b, test.c, Boolean(true) AS mark [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] - Left Join: Filter: test.a = __correlated_sq_1.a [a:UInt32, b:UInt32, c:UInt32, a:UInt32;N, __always_true:Boolean;N, count():Int64;N] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [a:UInt32, __always_true:Boolean, count():Int64] - Aggregate: groupBy=[[sq.a, Boolean(true) AS __always_true]], aggr=[[count()]] [a:UInt32, __always_true:Boolean, count():Int64] - TableScan: sq [a:UInt32, b:UInt32, c:UInt32] + Projection: test.a, test.b, test.c, Boolean(true) AS mark [a:UInt32, b:UInt32, c:UInt32, mark:Boolean] + Left Join: Filter: test.a = __correlated_sq_1.a [a:UInt32, b:UInt32, c:UInt32, a:UInt32;N, __always_true:Boolean;N, count():Int64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + SubqueryAlias: __correlated_sq_1 [a:UInt32, __always_true:Boolean, count():Int64] + Aggregate: groupBy=[[sq.a, Boolean(true) AS __always_true]], aggr=[[count()]] [a:UInt32, __always_true:Boolean, count():Int64] + TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) } @@ -3157,10 +3148,10 @@ mod tests { plan, @r" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, a:UInt32] - Projection: UInt32(1), sq.a [UInt32(1):UInt32, a:UInt32] + SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, __correlated_group_expr_a:UInt32] + Projection: UInt32(1), UInt32(1) + sq.a AS __correlated_group_expr_a [UInt32(1):UInt32, __correlated_group_expr_a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) @@ -3216,11 +3207,11 @@ mod tests { plan, @r" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [c:UInt32, a:UInt32] - Distinct: [c:UInt32, a:UInt32] - Projection: sq.c, sq.a [c:UInt32, a:UInt32] + SubqueryAlias: __correlated_sq_1 [c:UInt32, __correlated_group_expr_a:UInt32] + Distinct: [c:UInt32, __correlated_group_expr_a:UInt32] + Projection: sq.c, UInt32(1) + sq.a AS __correlated_group_expr_a [c:UInt32, __correlated_group_expr_a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) @@ -3247,11 +3238,11 @@ mod tests { plan, @r" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [sq.b + sq.c:UInt32, a:UInt32] - Distinct: [sq.b + sq.c:UInt32, a:UInt32] - Projection: sq.b + sq.c, sq.a [sq.b + sq.c:UInt32, a:UInt32] + SubqueryAlias: __correlated_sq_1 [sq.b + sq.c:UInt32, __correlated_group_expr_a:UInt32] + Distinct: [sq.b + sq.c:UInt32, __correlated_group_expr_a:UInt32] + Projection: sq.b + sq.c, UInt32(1) + sq.a AS __correlated_group_expr_a [sq.b + sq.c:UInt32, __correlated_group_expr_a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) @@ -3278,11 +3269,11 @@ mod tests { plan, @r" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, c:UInt32, a:UInt32] - Distinct: [UInt32(1):UInt32, c:UInt32, a:UInt32] - Projection: UInt32(1), sq.c, sq.a [UInt32(1):UInt32, c:UInt32, a:UInt32] + SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, c:UInt32, __correlated_group_expr_a:UInt32] + Distinct: [UInt32(1):UInt32, c:UInt32, __correlated_group_expr_a:UInt32] + Projection: UInt32(1), sq.c, UInt32(1) + sq.a AS __correlated_group_expr_a [UInt32(1):UInt32, c:UInt32, __correlated_group_expr_a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 7ab4b125dc872..7a235d6138417 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1874,16 +1874,14 @@ explain SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE ---- logical_plan 01)Projection: t1.t1_id, t1.t1_int -02)--Aggregate: groupBy=[[t1.t1_id, t1.t1_name, t1.t1_int]], aggr=[[]] -03)----Projection: t1.t1_id, t1.t1_name, t1.t1_int -04)------Filter: __correlated_sq_1.__always_true IS NULL OR __correlated_sq_1.__always_true IS NOT NULL AND NOT __correlated_sq_1.count(Int64(1)) = Int64(0) IS NOT TRUE -05)--------Projection: t1.t1_id, t1.t1_name, t1.t1_int, __correlated_sq_1.count(Int64(1)), __correlated_sq_1.__always_true -06)----------Left Join: t1.t1_id = __correlated_sq_1.t2_id -07)------------TableScan: t1 projection=[t1_id, t1_name, t1_int] -08)------------SubqueryAlias: __correlated_sq_1 -09)--------------Projection: t2.t2_id, count(Int64(1)), Boolean(true) AS __always_true -10)----------------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1))]] -11)------------------TableScan: t2 projection=[t2_id] +02)--Filter: __correlated_sq_1.__always_true IS NULL OR __correlated_sq_1.__always_true IS NOT NULL AND NOT __correlated_sq_1.count(Int64(1)) = Int64(0) IS NOT TRUE +03)----Projection: t1.t1_id, t1.t1_int, __correlated_sq_1.count(Int64(1)), __correlated_sq_1.__always_true +04)------Left Join: t1.t1_id = __correlated_sq_1.t2_id +05)--------TableScan: t1 projection=[t1_id, t1_int] +06)--------SubqueryAlias: __correlated_sq_1 +07)----------Projection: t2.t2_id, count(Int64(1)), Boolean(true) AS __always_true +08)------------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1))]] +09)--------------TableScan: t2 projection=[t2_id] query II SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) @@ -2007,6 +2005,37 @@ DROP TABLE dedup_outer; statement ok DROP TABLE dedup_inner; +statement ok +CREATE TABLE cte_dedup_outer(k INT) AS VALUES (1), (2), (3); + +statement ok +CREATE TABLE cte_dedup_inner(k TEXT, v INT) AS VALUES ('1', 10), ('01', 20), ('2', 30); + +# correlated_exists_groupless_count_agg_duplicate_match_dedup_via_cte +query I +SELECT o.k FROM cte_dedup_outer o WHERE EXISTS ( + WITH agg AS (SELECT count(*) AS c FROM cte_dedup_inner i WHERE CAST(i.k AS INT) = o.k) + SELECT c FROM agg WHERE c = 2 +) +---- +1 + +# correlated_not_exists_groupless_count_agg_duplicate_match_dedup_via_cte +query I rowsort +SELECT o.k FROM cte_dedup_outer o WHERE NOT EXISTS ( + WITH agg AS (SELECT count(*) AS c FROM cte_dedup_inner i WHERE CAST(i.k AS INT) = o.k) + SELECT c FROM agg WHERE c = 2 +) +---- +2 +3 + +statement ok +DROP TABLE cte_dedup_outer; + +statement ok +DROP TABLE cte_dedup_inner; + statement count 0 drop table t1; From d64fc1ef4c756dae564be18ff74a3aeb86faca62 Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Tue, 15 Sep 2026 16:34:31 -0400 Subject: [PATCH 03/10] test: align dedup test data with checked-in shapes --- .../sqllogictest/test_files/subquery.slt | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 7a235d6138417..35fd9c88fbeaa 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1983,21 +1983,28 @@ where t1_id > 40 or not exists (select sum(t2_int) from t2 where t1.t1_id = t2.t 22 44 -# correlated_exists_groupless_count_agg_duplicate_match_dedup statement ok -CREATE TABLE dedup_outer(k INT) AS VALUES (1); +CREATE TABLE dedup_outer(k INT) AS VALUES (1), (2); statement ok -CREATE TABLE dedup_inner(k TEXT) AS VALUES ('1'), ('01'); +CREATE TABLE dedup_inner(k TEXT, v INT) AS VALUES ('1', 10), ('1', 20), ('01', 30), ('01', 40); +# correlated_exists_groupless_count_agg_duplicate_match_dedup query I -SELECT o.k -FROM dedup_outer o -WHERE EXISTS ( +SELECT o.k FROM dedup_outer o WHERE EXISTS ( SELECT count(*) FROM dedup_inner i WHERE CAST(i.k AS INT) = o.k ) ---- 1 +2 + +# correlated_not_exists_groupless_count_agg_duplicate_match_dedup +query I +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 +) +---- +1 statement ok DROP TABLE dedup_outer; From afdb738344c4fa89837311664abfea93d1f3b92a Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Wed, 16 Sep 2026 10:38:35 -0400 Subject: [PATCH 04/10] fix: correlated exists/not exists over zero-row subquery limit --- datafusion/optimizer/src/decorrelate.rs | 7 +++++ .../src/decorrelate_predicate_subquery.rs | 31 ++++++++++++------- .../sqllogictest/test_files/subquery.slt | 27 ++++++++++++++++ 3 files changed, 54 insertions(+), 11 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 02a68ab304f0e..0ab51aaeb3eac 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -109,6 +109,11 @@ pub struct PullUpCorrelatedExpr { /// unresolvable bare column. This records the mapping the first /// time it's made. correlated_col_aliases: HashMap, + /// `LIMIT 0` forces the subquery to zero rows unconditionally, regardless + /// of whether the correlation matched, collapsing it to an `EmptyRelation`. + /// Join-compensation needs this flag to + /// distinguish "empty" from its usual "matched" default. + pub forces_empty_result: bool, } impl Default for PullUpCorrelatedExpr { @@ -133,6 +138,7 @@ impl PullUpCorrelatedExpr { correlated_filters: Vec::new(), column_refs_above: Vec::new(), correlated_col_aliases: HashMap::new(), + forces_empty_result: false, } } @@ -295,6 +301,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { // with an empty relation. let fetch = limit.get_fetch_type()?; if matches!(fetch, FetchType::Literal(Some(0))) { + self.forces_empty_result = true; return Ok(Transformed::yes(LogicalPlan::EmptyRelation( EmptyRelation { produce_one_row: false, diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index f7f85752269cf..cc45ea1e6fc12 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -817,6 +817,7 @@ fn build_join( alias, &expr_map, pull_up.pull_up_having_expr.as_ref(), + pull_up.forces_empty_result, ) .map(|plan| { Some(BuiltJoin { @@ -1096,6 +1097,7 @@ fn build_join_with_count_bug( alias: &str, expr_map: &crate::decorrelate::ExprResultMap, pull_up_having_expr: Option<&Expr>, + forces_empty_result: bool, ) -> Result { if in_predicate_opt.is_some() { return not_impl_err!( @@ -1141,18 +1143,25 @@ fn build_join_with_count_bug( }; // EXISTS is true by default (the groupless aggregate always - // produces a row), unless either the row joined against a group - // whose HAVING predicate failed, or an unmatched row's own - // default aggregate values fail that same HAVING clause. - let exists_expr = match having_arm { - Some(when_expr) => { - when(indicator_col.clone().is_null(), unmatched_having_default) - .when(when_expr, lit(false)) - .otherwise(lit(true))? - .rewrite(&mut expr_rewrite) - .data()? + // produces a row), unless one of three cases holds: (1) the outer + // row joined against an inner group whose HAVING predicate failed, + // (2) an unmatched outer row's own default aggregate values fail + // that same HAVING clause, or (3) the subquery was truncated to + // zero rows unconditionally (LIMIT 0), which is false for every + // outer row regardless of whether it matched. + let exists_expr = if forces_empty_result { + lit(false) + } else { + match having_arm { + Some(when_expr) => { + when(indicator_col.clone().is_null(), unmatched_having_default) + .when(when_expr, lit(false)) + .otherwise(lit(true))? + .rewrite(&mut expr_rewrite) + .data()? + } + None => lit(true), } - None => lit(true), }; let new_plan = match join_type { diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 35fd9c88fbeaa..f9680ff6c64a8 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2043,6 +2043,33 @@ DROP TABLE cte_dedup_outer; statement ok DROP TABLE cte_dedup_inner; +statement ok +CREATE TABLE limit_zero_outer(k INT) AS VALUES (1); + +statement ok +CREATE TABLE limit_zero_inner(k INT, v INT) AS VALUES (1, NULL); + +# correlated_exists_groupless_count_agg_limit_zero +query I +SELECT o.k FROM limit_zero_outer o WHERE EXISTS ( + SELECT count(*) FROM limit_zero_inner i WHERE i.k = o.k LIMIT 0 +) +---- + +# correlated_not_exists_groupless_count_agg_limit_zero +query I +SELECT o.k FROM limit_zero_outer o WHERE NOT EXISTS ( + SELECT count(*) FROM limit_zero_inner i WHERE i.k = o.k LIMIT 0 +) +---- +1 + +statement ok +DROP TABLE limit_zero_outer; + +statement ok +DROP TABLE limit_zero_inner; + statement count 0 drop table t1; From ba10f146b6866ec90c7c06045ec02e5c1db14576 Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Wed, 16 Sep 2026 13:00:59 -0400 Subject: [PATCH 05/10] test: update stale query plan expectations --- datafusion/sqllogictest/test_files/joins.slt | 51 +++++++++++--------- 1 file changed, 27 insertions(+), 24 deletions(-) diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 6f498a9925df9..aa9bb02d70841 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -2206,10 +2206,11 @@ WHERE EXISTS ( ) ---- logical_plan -01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----TableScan: join_t2 projection=[t2_id] +04)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +05)------TableScan: join_t2 projection=[t2_id] statement ok set datafusion.optimizer.repartition_joins = false; @@ -2258,10 +2259,10 @@ WHERE EXISTS ( ) ---- logical_plan -01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: join_t2.t2_id +04)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id 05)------Filter: join_t2.t2_int < UInt32(3) 06)--------TableScan: join_t2 projection=[t2_id, t2_int] @@ -2306,11 +2307,12 @@ WHERE EXISTS ( ) ---- logical_plan -01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--Filter: join_t1.t1_int < UInt32(3) 03)----TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 04)--SubqueryAlias: __correlated_sq_1 -05)----TableScan: join_t2 projection=[t2_id] +05)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +06)------TableScan: join_t2 projection=[t2_id] statement ok set datafusion.optimizer.repartition_joins = false; @@ -2353,10 +2355,11 @@ WHERE NOT EXISTS ( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----TableScan: join_t2 projection=[t2_id] +04)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +05)------TableScan: join_t2 projection=[t2_id] statement ok set datafusion.optimizer.repartition_joins = false; @@ -2399,12 +2402,12 @@ WHERE NOT EXISTS ( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: join_t2.t2_id -05)------Aggregate: groupBy=[[join_t2.t2_int, join_t2.t2_id]], aggr=[[]] -06)--------Projection: join_t2.t2_int, join_t2.t2_id +04)----Projection: __correlated_group_expr_t2_id +05)------Aggregate: groupBy=[[join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] +06)--------Projection: join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] statement ok @@ -2448,12 +2451,12 @@ WHERE NOT EXISTS( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: join_t2.t2_id -05)------Aggregate: groupBy=[[join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, join_t2.t2_id]], aggr=[[]] -06)--------Projection: join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, join_t2.t2_id +04)----Projection: __correlated_group_expr_t2_id +05)------Aggregate: groupBy=[[join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] +06)--------Projection: join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] statement ok @@ -2499,12 +2502,12 @@ WHERE NOT EXISTS( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: join_t2.t2_id -05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, join_t2.t2_id]], aggr=[[]] -06)--------Projection: Int64(1), join_t2.t2_int, join_t2.t2_id +04)----Projection: __correlated_group_expr_t2_id +05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] +06)--------Projection: Int64(1), join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] query ITI @@ -2530,12 +2533,12 @@ WHERE NOT EXISTS( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: join_t2.t2_id -05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, join_t2.t2_id]], aggr=[[]] -06)--------Projection: Int64(1), join_t2.t2_int, join_t2.t2_id +04)----Projection: __correlated_group_expr_t2_id +05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] +06)--------Projection: Int64(1), join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] query ITI From 589d646d1e909024b66f421d1b12ef2ec3e435bb Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Wed, 16 Sep 2026 18:51:54 -0400 Subject: [PATCH 06/10] chore: some cleanup --- datafusion/optimizer/src/decorrelate.rs | 10 +++++----- .../src/decorrelate_predicate_subquery.rs | 14 ++++++++------ 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0ab51aaeb3eac..a2f29911271d5 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -103,11 +103,11 @@ pub struct PullUpCorrelatedExpr { column_refs_above: Vec>, /// A correlated column wrapped in an expression (e.g. `CAST(t2.b AS INT)`) /// gets grouped by that entire expression instead of the bare column, - /// aliased to a generated name, once, in the `Aggregate` this column belongs to. Every - /// later reference to that column, in a `Projection` above it for - /// instance, needs to use the same alias instead of the now - /// unresolvable bare column. This records the mapping the first - /// time it's made. + /// aliased to a generated name, once, in the `Aggregate` this column + /// belongs to. Every later reference to that column, in a + /// `Projection` above it for instance, needs to use the same alias + /// instead of the now unresolvable bare column. This records the + /// mapping the first time it's made. correlated_col_aliases: HashMap, /// `LIMIT 0` forces the subquery to zero rows unconditionally, regardless /// of whether the correlation matched, collapsing it to an `EmptyRelation`. diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index cc45ea1e6fc12..67cefeb55b0b7 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -1101,7 +1101,7 @@ fn build_join_with_count_bug( ) -> Result { if in_predicate_opt.is_some() { return not_impl_err!( - "build_join_with_count_bug: `IN`/`NOT IN` count-bug compensation is not implemented" + "IN/NOT IN count-bug compensation for groupless aggregates is not implemented" ); } @@ -1143,12 +1143,14 @@ fn build_join_with_count_bug( }; // EXISTS is true by default (the groupless aggregate always - // produces a row), unless one of three cases holds: (1) the outer - // row joined against an inner group whose HAVING predicate failed, + // produces a row), unless one of three cases holds: + // (1) the outer row joined against an inner group whose HAVING + // predicate failed, // (2) an unmatched outer row's own default aggregate values fail - // that same HAVING clause, or (3) the subquery was truncated to - // zero rows unconditionally (LIMIT 0), which is false for every - // outer row regardless of whether it matched. + // that same HAVING clause, + // (3) the subquery was truncated to zero rows unconditionally + // (LIMIT 0), which is false for every outer row regardless of + // whether it matched. let exists_expr = if forces_empty_result { lit(false) } else { From 9a7059214ca23f1f86d89c0e4a6eab2873c471ae Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Wed, 16 Sep 2026 20:47:43 -0400 Subject: [PATCH 07/10] chore: cleanup error messages --- datafusion/optimizer/src/decorrelate.rs | 2 +- datafusion/optimizer/src/decorrelate_predicate_subquery.rs | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index a2f29911271d5..af202d12d18ab 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -1111,7 +1111,7 @@ fn remove_duplicated_filter( Expr::BinaryExpr(b) => b.op.swap() == Some(b.op), _ => true, }, - "remove_duplicated_filter: in_predicate must use a commutative operator" + "in_predicate must use a commutative operator" ); Ok(filters diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 67cefeb55b0b7..59a7205ddca23 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -1183,9 +1183,7 @@ fn build_join_with_count_bug( .project(left_projection)? .build()?, _ => { - return internal_err!( - "build_join_with_count_bug: unsupported join type {join_type:?}" - ); + return internal_err!("unsupported join type {join_type:?}"); } }; From 38d8380fd0014c122844cf03cbdb3aa9ef027eed Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Fri, 18 Sep 2026 20:36:53 -0400 Subject: [PATCH 08/10] chore: update test post-rebase --- datafusion/sqllogictest/test_files/subquery.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index f9680ff6c64a8..41a3e473d4b9f 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1874,7 +1874,7 @@ explain SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE ---- logical_plan 01)Projection: t1.t1_id, t1.t1_int -02)--Filter: __correlated_sq_1.__always_true IS NULL OR __correlated_sq_1.__always_true IS NOT NULL AND NOT __correlated_sq_1.count(Int64(1)) = Int64(0) IS NOT TRUE +02)--Filter: CASE WHEN __correlated_sq_1.__always_true IS NULL THEN Boolean(true) WHEN __correlated_sq_1.count(Int64(1)) = Int64(0) IS NOT TRUE THEN Boolean(false) ELSE Boolean(true) END 03)----Projection: t1.t1_id, t1.t1_int, __correlated_sq_1.count(Int64(1)), __correlated_sq_1.__always_true 04)------Left Join: t1.t1_id = __correlated_sq_1.t2_id 05)--------TableScan: t1 projection=[t1_id, t1_int] From d5d7fd6116327744a3049a9f2a978d29a142137e Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi <81446894+mohammadnaqvi04@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:58:57 -0400 Subject: [PATCH 09/10] chore: address feedback --- datafusion/optimizer/src/decorrelate.rs | 76 ++--------- .../src/decorrelate_predicate_subquery.rs | 118 ++++++++++-------- datafusion/sqllogictest/test_files/joins.slt | 51 ++++---- .../sqllogictest/test_files/subquery.slt | 73 ++++++++++- 4 files changed, 169 insertions(+), 149 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index af202d12d18ab..3b193ce5f1f9e 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -33,8 +33,7 @@ use datafusion_expr::expr::{Alias, GroupingSet}; use datafusion_expr::logical_plan::{Join, JoinType}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ - collect_subquery_cols, collect_subquery_join_exprs, conjunction, find_join_exprs, - split_conjunction, + collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, }; use datafusion_expr::{ Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, Extension, @@ -101,14 +100,6 @@ pub struct PullUpCorrelatedExpr { /// grouping set that omits a correlated column can tell whether /// something above it reads the column the pull up would fill in. column_refs_above: Vec>, - /// A correlated column wrapped in an expression (e.g. `CAST(t2.b AS INT)`) - /// gets grouped by that entire expression instead of the bare column, - /// aliased to a generated name, once, in the `Aggregate` this column - /// belongs to. Every later reference to that column, in a - /// `Projection` above it for instance, needs to use the same alias - /// instead of the now unresolvable bare column. This records the - /// mapping the first time it's made. - correlated_col_aliases: HashMap, /// `LIMIT 0` forces the subquery to zero rows unconditionally, regardless /// of whether the correlation matched, collapsing it to an `EmptyRelation`. /// Join-compensation needs this flag to @@ -137,7 +128,6 @@ impl PullUpCorrelatedExpr { pulled_up_scalar_agg: false, correlated_filters: Vec::new(), column_refs_above: Vec::new(), - correlated_col_aliases: HashMap::new(), forces_empty_result: false, } } @@ -443,7 +433,8 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { None }; - if self.pull_up_having_expr.is_some() + if self.exists_sub_query + && self.pull_up_having_expr.is_some() && let Some(expr) = conjunction(subquery_filters.clone()) { let unqualified_expr = expr @@ -675,16 +666,8 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { ); let mut new_correlated_cols = BTreeSet::new(); for col in local_correlated_cols.iter() { - let requalified = - Column::new(Some(alias.alias.clone()), col.name.clone()); - // A column already folded into a group-by alias (see - // `collect_missing_exprs`) needs that mapping carried - // forward under its requalified name. - if let Some(existing_alias) = self.correlated_col_aliases.get(col) { - self.correlated_col_aliases - .insert(requalified.clone(), existing_alias.clone()); - } - new_correlated_cols.insert(requalified); + new_correlated_cols + .insert(Column::new(Some(alias.alias.clone()), col.name.clone())); } let new_plan = if alias.input.schema().fields().len() @@ -876,7 +859,7 @@ impl PullUpCorrelatedExpr { &mut self, exprs: &[Expr], correlated_subquery_cols: &BTreeSet, - subquery_schema: &DFSchemaRef, + _subquery_schema: &DFSchemaRef, ) -> Result> { let mut missing_exprs = Vec::with_capacity(exprs.len() + correlated_subquery_cols.len()); @@ -885,53 +868,8 @@ impl PullUpCorrelatedExpr { 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() { - 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 - .iter() - .find(|e| e.column_refs().len() == 1 && e.column_refs().contains(col)); - let col_expr = match wrapped { - Some(e) if !matches!(e, Expr::Column(_)) => { - let alias_name = format!("__correlated_group_expr_{}", col.name); - let aliased = e.clone().alias(alias_name.clone()); - let reference = Expr::Column(Column::new_unqualified(&alias_name)); - self.join_filters = self - .join_filters - .iter() - .map(|f| { - f.clone() - .transform_up(|node| { - if &node == e { - Ok(Transformed::yes(reference.clone())) - } else { - Ok(Transformed::no(node)) - } - }) - .data() - }) - .collect::>>()?; - self.correlated_col_aliases.insert(col.clone(), reference); - aliased - } - _ => Expr::Column(col.clone()), - }; + let col_expr = Expr::Column(col.clone()); if !collides_with_existing(&missing_exprs, &col_expr) { missing_exprs.push(col_expr) } diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 59a7205ddca23..d02094bda04e5 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -778,18 +778,43 @@ fn build_join( alias: &str, needs_null_aware_mark: bool, ) -> Result> { - let mut pull_up = PullUpCorrelatedExpr::new() - .with_in_predicate_opt(in_predicate_opt.cloned()) - .with_exists_sub_query(in_predicate_opt.is_none()) - .with_need_handle_count_bug(true); - - let new_plan = subquery.clone().rewrite(&mut pull_up).data()?; + let exists_sub_query = in_predicate_opt.is_none(); + let rewrite = + |handle_count_bug: bool| -> Result<(LogicalPlan, PullUpCorrelatedExpr)> { + let mut pull_up = PullUpCorrelatedExpr::new() + .with_in_predicate_opt(in_predicate_opt.cloned()) + .with_exists_sub_query(exists_sub_query) + .with_need_handle_count_bug(handle_count_bug); + let new_plan = subquery.clone().rewrite(&mut pull_up).data()?; + Ok((new_plan, pull_up)) + }; + let (mut new_plan, mut pull_up) = rewrite(exists_sub_query)?; + + let mut count_bug_compensation = pull_up + .collected_count_expr_map + .get(&new_plan) + .filter(|m| !m.is_empty()) + .cloned(); + // A HAVING-failure miss looks identical to an absent group once compensated, so bail if the join filter reads the aggregate's own output. + let reads_agg_output = count_bug_compensation.as_ref().is_some_and(|expr_map| { + pull_up + .join_filters + .iter() + .flat_map(|f| f.column_refs()) + .any(|c| expr_map.contains_key(&c.name)) + }); + // Fall back to the uncompensated rewrite whenever compensation can't safely apply, so a HAVING never gets silently dropped. + let compensable = pull_up.can_pull_up + && !reads_agg_output + && (count_bug_compensation.is_some() || pull_up.pull_up_having_expr.is_none()); + if exists_sub_query && !compensable { + (new_plan, pull_up) = rewrite(false)?; + count_bug_compensation = None; + } if !pull_up.can_pull_up { return Ok(None); } - let count_bug_compensation = pull_up.collected_count_expr_map.get(&new_plan).cloned(); - let sub_query_alias = LogicalPlanBuilder::from(new_plan) .alias(alias.to_string())? .build()?; @@ -812,7 +837,6 @@ fn build_join( left, sub_query_alias, join_filter_opt, - in_predicate_opt, join_type, alias, &expr_map, @@ -1092,19 +1116,12 @@ fn build_join_with_count_bug( left: &LogicalPlan, sub_query_alias: LogicalPlan, join_filter_opt: Option, - in_predicate_opt: Option<&Expr>, join_type: JoinType, alias: &str, expr_map: &crate::decorrelate::ExprResultMap, pull_up_having_expr: Option<&Expr>, forces_empty_result: bool, ) -> Result { - if in_predicate_opt.is_some() { - return not_impl_err!( - "IN/NOT IN count-bug compensation for groupless aggregates is not implemented" - ); - } - let joined = LogicalPlanBuilder::from(left.clone()) .join_on(sub_query_alias, JoinType::Left, join_filter_opt)? .build()?; @@ -1155,13 +1172,14 @@ fn build_join_with_count_bug( lit(false) } else { match having_arm { - Some(when_expr) => { - when(indicator_col.clone().is_null(), unmatched_having_default) - .when(when_expr, lit(false)) - .otherwise(lit(true))? - .rewrite(&mut expr_rewrite) - .data()? - } + Some(when_expr) => when( + indicator_col.clone().is_null(), + unmatched_having_default.is_true(), + ) + .when(when_expr, lit(false)) + .otherwise(lit(true))? + .rewrite(&mut expr_rewrite) + .data()?, None => lit(true), } }; @@ -1803,12 +1821,12 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: customer.c_custkey [c_custkey:Int64] - LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey AND (customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.__correlated_group_expr_o_orderkey) [c_custkey:Int64, c_name:Utf8] + LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey AND (customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.o_orderkey = Int32(1)) [c_custkey:Int64, c_name:Utf8] TableScan: customer [c_custkey:Int64, c_name:Utf8] - SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] - Projection: orders.o_custkey, orders.o_orderkey = Int32(1) AS __correlated_group_expr_o_orderkey [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] + SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, o_orderkey:Int64] + Projection: orders.o_custkey, orders.o_orderkey [o_custkey:Int64, o_orderkey:Int64] TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] " ) @@ -2859,12 +2877,12 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: customer.c_custkey [c_custkey:Int64] - LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.__correlated_group_expr_o_orderkey [c_custkey:Int64, c_name:Utf8] + LeftSemi Join: Filter: customer.c_custkey = __correlated_sq_1.o_custkey OR __correlated_sq_1.o_orderkey = Int32(1) [c_custkey:Int64, c_name:Utf8] TableScan: customer [c_custkey:Int64, c_name:Utf8] - SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] - Projection: orders.o_custkey, orders.o_orderkey = Int32(1) AS __correlated_group_expr_o_orderkey [o_custkey:Int64, __correlated_group_expr_o_orderkey:Boolean] + SubqueryAlias: __correlated_sq_1 [o_custkey:Int64, o_orderkey:Int64] + Projection: orders.o_custkey, orders.o_orderkey [o_custkey:Int64, o_orderkey:Int64] TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] " ) @@ -3155,12 +3173,12 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, __correlated_group_expr_a:UInt32] - Projection: UInt32(1), UInt32(1) + sq.a AS __correlated_group_expr_a [UInt32(1):UInt32, __correlated_group_expr_a:UInt32] + SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, a:UInt32] + Projection: UInt32(1), sq.a [UInt32(1):UInt32, a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) @@ -3214,13 +3232,13 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [c:UInt32, __correlated_group_expr_a:UInt32] - Distinct: [c:UInt32, __correlated_group_expr_a:UInt32] - Projection: sq.c, UInt32(1) + sq.a AS __correlated_group_expr_a [c:UInt32, __correlated_group_expr_a:UInt32] + SubqueryAlias: __correlated_sq_1 [c:UInt32, a:UInt32] + Distinct: [c:UInt32, a:UInt32] + Projection: sq.c, sq.a [c:UInt32, a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) @@ -3245,13 +3263,13 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [sq.b + sq.c:UInt32, __correlated_group_expr_a:UInt32] - Distinct: [sq.b + sq.c:UInt32, __correlated_group_expr_a:UInt32] - Projection: sq.b + sq.c, UInt32(1) + sq.a AS __correlated_group_expr_a [sq.b + sq.c:UInt32, __correlated_group_expr_a:UInt32] + SubqueryAlias: __correlated_sq_1 [sq.b + sq.c:UInt32, a:UInt32] + Distinct: [sq.b + sq.c:UInt32, a:UInt32] + Projection: sq.b + sq.c, sq.a [sq.b + sq.c:UInt32, a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) @@ -3276,13 +3294,13 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: test.b [b:UInt32] - LeftSemi Join: Filter: __correlated_sq_1.__correlated_group_expr_a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] + LeftSemi Join: Filter: UInt32(1) + __correlated_sq_1.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32] TableScan: test [a:UInt32, b:UInt32, c:UInt32] - SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, c:UInt32, __correlated_group_expr_a:UInt32] - Distinct: [UInt32(1):UInt32, c:UInt32, __correlated_group_expr_a:UInt32] - Projection: UInt32(1), sq.c, UInt32(1) + sq.a AS __correlated_group_expr_a [UInt32(1):UInt32, c:UInt32, __correlated_group_expr_a:UInt32] + SubqueryAlias: __correlated_sq_1 [UInt32(1):UInt32, c:UInt32, a:UInt32] + Distinct: [UInt32(1):UInt32, c:UInt32, a:UInt32] + Projection: UInt32(1), sq.c, sq.a [UInt32(1):UInt32, c:UInt32, a:UInt32] TableScan: sq [a:UInt32, b:UInt32, c:UInt32] " ) diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index aa9bb02d70841..6f498a9925df9 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -2206,11 +2206,10 @@ WHERE EXISTS ( ) ---- logical_plan -01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id -05)------TableScan: join_t2 projection=[t2_id] +04)----TableScan: join_t2 projection=[t2_id] statement ok set datafusion.optimizer.repartition_joins = false; @@ -2259,10 +2258,10 @@ WHERE EXISTS ( ) ---- logical_plan -01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +04)----Projection: join_t2.t2_id 05)------Filter: join_t2.t2_int < UInt32(3) 06)--------TableScan: join_t2 projection=[t2_id, t2_int] @@ -2307,12 +2306,11 @@ WHERE EXISTS ( ) ---- logical_plan -01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftSemi Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--Filter: join_t1.t1_int < UInt32(3) 03)----TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 04)--SubqueryAlias: __correlated_sq_1 -05)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id -06)------TableScan: join_t2 projection=[t2_id] +05)----TableScan: join_t2 projection=[t2_id] statement ok set datafusion.optimizer.repartition_joins = false; @@ -2355,11 +2353,10 @@ WHERE NOT EXISTS ( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id -05)------TableScan: join_t2 projection=[t2_id] +04)----TableScan: join_t2 projection=[t2_id] statement ok set datafusion.optimizer.repartition_joins = false; @@ -2402,12 +2399,12 @@ WHERE NOT EXISTS ( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: __correlated_group_expr_t2_id -05)------Aggregate: groupBy=[[join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] -06)--------Projection: join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +04)----Projection: join_t2.t2_id +05)------Aggregate: groupBy=[[join_t2.t2_int, join_t2.t2_id]], aggr=[[]] +06)--------Projection: join_t2.t2_int, join_t2.t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] statement ok @@ -2451,12 +2448,12 @@ WHERE NOT EXISTS( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: __correlated_group_expr_t2_id -05)------Aggregate: groupBy=[[join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] -06)--------Projection: join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +04)----Projection: join_t2.t2_id +05)------Aggregate: groupBy=[[join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, join_t2.t2_id]], aggr=[[]] +06)--------Projection: join_t2.t2_id + join_t2.t2_int, join_t2.t2_int, join_t2.t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] statement ok @@ -2502,12 +2499,12 @@ WHERE NOT EXISTS( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: __correlated_group_expr_t2_id -05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] -06)--------Projection: Int64(1), join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +04)----Projection: join_t2.t2_id +05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, join_t2.t2_id]], aggr=[[]] +06)--------Projection: Int64(1), join_t2.t2_int, join_t2.t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] query ITI @@ -2533,12 +2530,12 @@ WHERE NOT EXISTS( ) ---- logical_plan -01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > __correlated_sq_1.__correlated_group_expr_t2_id +01)LeftAnti Join: Filter: CAST(join_t1.t1_id AS Int64) + Int64(1) > CAST(__correlated_sq_1.t2_id AS Int64) * Int64(2) 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 -04)----Projection: __correlated_group_expr_t2_id -05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, __correlated_group_expr_t2_id]], aggr=[[]] -06)--------Projection: Int64(1), join_t2.t2_int, CAST(join_t2.t2_id AS Int64) * Int64(2) AS __correlated_group_expr_t2_id +04)----Projection: join_t2.t2_id +05)------Aggregate: groupBy=[[Int64(1), join_t2.t2_int, join_t2.t2_id]], aggr=[[]] +06)--------Projection: Int64(1), join_t2.t2_int, join_t2.t2_id 07)----------TableScan: join_t2 projection=[t2_id, t2_int] query ITI diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 41a3e473d4b9f..3542c5a76acbe 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1294,6 +1294,19 @@ select t1.t1_int from t1 where ( 2 4 +# correlated_scalar_subquery_having_default_not_coerced_to_match_stacked_filter +query II rowsort +SELECT t1_id, ( + SELECT c FROM ( + SELECT count(*) AS c FROM t2 WHERE t2.t2_id = t1.t1_id HAVING count(*) < 10 + ) s WHERE c > 5 +) AS v FROM t1 +---- +11 NULL +22 NULL +33 NULL +44 NULL + #correlated_scalar_subquery_count_agg_in_case_when query TT explain @@ -1902,6 +1915,14 @@ SELECT t1_id, t1_int FROM t1 WHERE NOT EXISTS (SELECT 1 FROM t2 WHERE t1.t1_id = 22 2 44 4 +# correlated_exists_groupless_having_reads_outer_column_falls_back +query I +SELECT t1_id FROM t1 WHERE EXISTS ( + SELECT count(*) FROM t2 WHERE t2.t2_id = t1.t1_id HAVING count(*) = t1.t1_int +) +---- +11 + # correlated_exists_groupless_having_with_stacked_filter_via_cte query II SELECT t1_id, t1_int FROM t1 WHERE EXISTS ( @@ -1949,6 +1970,34 @@ query II SELECT t1_id, t1_int FROM t1 WHERE NOT EXISTS (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id) ---- +# correlated_not_exists_groupless_null_default_having_is_null_not_false +query I rowsort +SELECT t1_id FROM t1 WHERE NOT EXISTS ( + SELECT c FROM ( + SELECT count(*) AS c, sum(t2_int) AS s FROM t2 + WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0 + ) agg WHERE s > 5 +) +---- +11 +22 +33 +44 + +# correlated_not_exists_groupless_null_default_having_is_null_not_false_embedded_in_or +query I rowsort +SELECT t1_id FROM t1 WHERE t1_id > 100 OR NOT EXISTS ( + SELECT c FROM ( + SELECT count(*) AS c, sum(t2_int) AS s FROM t2 + WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0 + ) agg WHERE s > 5 +) +---- +11 +22 +33 +44 + # correlated_exists_groupless_count_agg_count_bug_embedded_in_or query I rowsort select t1_id from t1 @@ -1990,21 +2039,23 @@ statement ok CREATE TABLE dedup_inner(k TEXT, v INT) AS VALUES ('1', 10), ('1', 20), ('01', 30), ('01', 40); # correlated_exists_groupless_count_agg_duplicate_match_dedup -query I +query I rowsort SELECT o.k FROM dedup_outer o WHERE EXISTS ( SELECT count(*) FROM dedup_inner i WHERE CAST(i.k AS INT) = o.k ) ---- 1 +1 2 # correlated_not_exists_groupless_count_agg_duplicate_match_dedup -query I +query I rowsort 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 ) ---- 1 +1 statement ok DROP TABLE dedup_outer; @@ -2025,7 +2076,6 @@ SELECT o.k FROM cte_dedup_outer o WHERE EXISTS ( SELECT c FROM agg WHERE c = 2 ) ---- -1 # correlated_not_exists_groupless_count_agg_duplicate_match_dedup_via_cte query I rowsort @@ -2034,6 +2084,7 @@ SELECT o.k FROM cte_dedup_outer o WHERE NOT EXISTS ( SELECT c FROM agg WHERE c = 2 ) ---- +1 2 3 @@ -2064,6 +2115,22 @@ SELECT o.k FROM limit_zero_outer o WHERE NOT EXISTS ( ---- 1 +# correlated_in_groupless_max_agg_is_not_count_bug_compensated +query I +SELECT t1_id FROM t1 +WHERE t1_int + 2 IN (SELECT max(t2_int) FROM t2 WHERE t2.t2_id = t1.t1_id) +---- +11 + +# correlated_not_in_groupless_count_agg_is_not_count_bug_compensated +query I rowsort +SELECT t1_id FROM t1 +WHERE t1_int NOT IN (SELECT count(*) FROM t2 WHERE t2.t2_id = t1.t1_id) +---- +22 +33 +44 + statement ok DROP TABLE limit_zero_outer; From d96dfb91ac02c773ae351413bd8532f307cbc023 Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi <81446894+mohammadnaqvi04@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:22:00 -0400 Subject: [PATCH 10/10] fix: gate count-bug compensation on bare keys --- datafusion/optimizer/src/decorrelate.rs | 7 -- .../src/decorrelate_predicate_subquery.rs | 84 ++++++++++++++----- .../sqllogictest/test_files/subquery.slt | 8 +- 3 files changed, 67 insertions(+), 32 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 3b193ce5f1f9e..7dcf8fde935e5 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -100,11 +100,6 @@ pub struct PullUpCorrelatedExpr { /// grouping set that omits a correlated column can tell whether /// something above it reads the column the pull up would fill in. column_refs_above: Vec>, - /// `LIMIT 0` forces the subquery to zero rows unconditionally, regardless - /// of whether the correlation matched, collapsing it to an `EmptyRelation`. - /// Join-compensation needs this flag to - /// distinguish "empty" from its usual "matched" default. - pub forces_empty_result: bool, } impl Default for PullUpCorrelatedExpr { @@ -128,7 +123,6 @@ impl PullUpCorrelatedExpr { pulled_up_scalar_agg: false, correlated_filters: Vec::new(), column_refs_above: Vec::new(), - forces_empty_result: false, } } @@ -291,7 +285,6 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { // with an empty relation. let fetch = limit.get_fetch_type()?; if matches!(fetch, FetchType::Literal(Some(0))) { - self.forces_empty_result = true; return Ok(Transformed::yes(LogicalPlan::EmptyRelation( EmptyRelation { produce_one_row: false, diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index d02094bda04e5..9c20108da261c 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -28,7 +28,9 @@ use crate::utils::replace_qualified_name; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::alias::AliasGenerator; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_common::{ Column, DFSchema, NullEquality, Result, ScalarValue, assert_or_internal_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, @@ -761,6 +763,30 @@ fn with_null_aware_value_keys( )) } +fn subquery_operands_are_bare_columns( + join_filters: &[Expr], + subquery_cols: &BTreeSet, +) -> 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 + }) +} + /// The outcome of [`build_join`]. struct BuiltJoin { /// The outer plan with the subquery joined into it. @@ -803,9 +829,18 @@ fn build_join( .flat_map(|f| f.column_refs()) .any(|c| expr_map.contains_key(&c.name)) }); + let subquery_cols: BTreeSet = 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); // Fall back to the uncompensated rewrite whenever compensation can't safely apply, so a HAVING never gets silently dropped. let compensable = pull_up.can_pull_up && !reads_agg_output + && keys_are_bare && (count_bug_compensation.is_some() || pull_up.pull_up_having_expr.is_none()); if exists_sub_query && !compensable { (new_plan, pull_up) = rewrite(false)?; @@ -841,7 +876,6 @@ fn build_join( alias, &expr_map, pull_up.pull_up_having_expr.as_ref(), - pull_up.forces_empty_result, ) .map(|plan| { Some(BuiltJoin { @@ -1111,7 +1145,6 @@ fn build_join( /// Builds the join for a correlated `EXISTS` subquery whose groupless /// aggregate requires count-bug compensation. -#[expect(clippy::too_many_arguments)] fn build_join_with_count_bug( left: &LogicalPlan, sub_query_alias: LogicalPlan, @@ -1120,8 +1153,22 @@ fn build_join_with_count_bug( alias: &str, expr_map: &crate::decorrelate::ExprResultMap, pull_up_having_expr: Option<&Expr>, - forces_empty_result: bool, ) -> Result { + 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(); + let joined = LogicalPlanBuilder::from(left.clone()) .join_on(sub_query_alias, JoinType::Left, join_filter_opt)? .build()?; @@ -1164,24 +1211,17 @@ fn build_join_with_count_bug( // (1) the outer row joined against an inner group whose HAVING // predicate failed, // (2) an unmatched outer row's own default aggregate values fail - // that same HAVING clause, - // (3) the subquery was truncated to zero rows unconditionally - // (LIMIT 0), which is false for every outer row regardless of - // whether it matched. - let exists_expr = if forces_empty_result { - lit(false) - } else { - match having_arm { - Some(when_expr) => when( - indicator_col.clone().is_null(), - unmatched_having_default.is_true(), - ) - .when(when_expr, lit(false)) - .otherwise(lit(true))? - .rewrite(&mut expr_rewrite) - .data()?, - None => lit(true), - } + // that same HAVING clause. + let exists_expr = match having_arm { + Some(when_expr) => when( + indicator_col.clone().is_null(), + unmatched_having_default.is_true(), + ) + .when(when_expr, lit(false)) + .otherwise(lit(true))? + .rewrite(&mut expr_rewrite) + .data()?, + None => lit(true), }; let new_plan = match join_type { diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 3542c5a76acbe..3cd3bfbf193bb 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2038,6 +2038,7 @@ CREATE TABLE dedup_outer(k INT) AS VALUES (1), (2); statement ok CREATE TABLE dedup_inner(k TEXT, v INT) AS VALUES ('1', 10), ('1', 20), ('01', 30), ('01', 40); +# TODO(#26014): should match DuckDB (1, 2), not main's answer. # correlated_exists_groupless_count_agg_duplicate_match_dedup query I rowsort SELECT o.k FROM dedup_outer o WHERE EXISTS ( @@ -2045,9 +2046,8 @@ SELECT o.k FROM dedup_outer o WHERE EXISTS ( ) ---- 1 -1 -2 +# TODO(#26014): should match DuckDB, not main's answer. # correlated_not_exists_groupless_count_agg_duplicate_match_dedup query I rowsort SELECT o.k FROM dedup_outer o WHERE NOT EXISTS ( @@ -2055,7 +2055,7 @@ SELECT o.k FROM dedup_outer o WHERE NOT EXISTS ( ) ---- 1 -1 +2 statement ok DROP TABLE dedup_outer; @@ -2069,6 +2069,7 @@ CREATE TABLE cte_dedup_outer(k INT) AS VALUES (1), (2), (3); statement ok CREATE TABLE cte_dedup_inner(k TEXT, v INT) AS VALUES ('1', 10), ('01', 20), ('2', 30); +# TODO(#26014): should match DuckDB, not main's answer. # correlated_exists_groupless_count_agg_duplicate_match_dedup_via_cte query I SELECT o.k FROM cte_dedup_outer o WHERE EXISTS ( @@ -2077,6 +2078,7 @@ SELECT o.k FROM cte_dedup_outer o WHERE EXISTS ( ) ---- +# TODO(#26014): should match DuckDB, not main's answer. # correlated_not_exists_groupless_count_agg_duplicate_match_dedup_via_cte query I rowsort SELECT o.k FROM cte_dedup_outer o WHERE NOT EXISTS (