From 6bac38abb36ba30b692bfb0d8fd1c13ec7413f8d Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Sun, 6 Sep 2026 22:57:52 -0400 Subject: [PATCH 1/4] fix: EXISTS/IN subqueries hit the count bug for groupless aggregates --- .../src/decorrelate_predicate_subquery.rs | 269 +++++++++++++++--- .../sqllogictest/test_files/subquery.slt | 108 +++++++ 2 files changed, 339 insertions(+), 38 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 5f623f1bef6f6..5446304772dc2 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -20,17 +20,18 @@ 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; +use crate::utils::{evaluates_to_null, replace_qualified_name}; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::alias::AliasGenerator; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{ - Column, DFSchemaRef, ExprSchema, NullEquality, Result, assert_or_internal_err, - plan_err, + Column, DFSchemaRef, ExprSchema, NullEquality, Result, ScalarValue, + assert_or_internal_err, internal_err, plan_err, }; use datafusion_expr::expr::{Exists, InSubquery}; use datafusion_expr::expr_rewriter::create_col_from_scalar_expr; @@ -38,7 +39,7 @@ use datafusion_expr::logical_plan::{JoinType, Subquery}; use datafusion_expr::utils::{conjunction, expr_to_columns, split_conjunction_owned}; use datafusion_expr::{ BinaryExpr, Expr, Filter, LogicalPlan, LogicalPlanBuilder, Operator, exists, - in_subquery, lit, not, not_exists, not_in_subquery, + in_subquery, lit, not, not_exists, not_in_subquery, when, }; use log::debug; @@ -285,7 +286,7 @@ fn build_join_top( subquery, in_predicate_opt.as_ref(), join_type, - subquery_alias, + &subquery_alias, ) } @@ -317,7 +318,7 @@ fn mark_join( let exists_expr = if negated { !exists_col } else { exists_col }; Ok( - build_join(left, subquery, in_predicate_opt, JoinType::LeftMark, alias)? + build_join(left, subquery, in_predicate_opt, JoinType::LeftMark, &alias)? .map(|plan| (plan, exists_expr)), ) } @@ -355,22 +356,62 @@ fn join_keys_may_be_null( Ok(false) } +/// Combines the join filter derived from the subquery's correlated +/// predicates with an `IN`-subquery's equality predicate, folding +/// the latter into an equality against the right-side (subquery) column +/// where needed. +fn build_join_filter_with_in_predicate( + join_filter_opt: Option, + in_predicate_opt: Option<&Expr>, + alias: &str, +) -> Result> { + // Validate and extract the equality operands from the `IN` predicate. + let eq_operands = in_predicate_opt + .map(|in_predicate| match in_predicate { + Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::Eq, + right, + }) => Ok((left.deref().clone(), right.deref().clone())), + _ => internal_err!("expected a binary equality IN predicate"), + }) + .transpose()?; + + let join_filter = match (join_filter_opt, eq_operands) { + (Some(join_filter), Some((left, right))) => { + let right_col = create_col_from_scalar_expr(&right, alias.to_string())?; + let in_predicate = Expr::eq(left, Expr::Column(right_col)); + Some(in_predicate.and(join_filter)) + } + (Some(join_filter), None) => Some(join_filter), + (None, Some((left, right))) => { + let right_col = create_col_from_scalar_expr(&right, alias.to_string())?; + Some(Expr::eq(left, Expr::Column(right_col))) + } + (None, None) => None, + }; + Ok(join_filter) +} + fn build_join( left: &LogicalPlan, subquery: &LogicalPlan, in_predicate_opt: Option<&Expr>, join_type: JoinType, - alias: String, + alias: &str, ) -> 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()?; @@ -383,37 +424,33 @@ fn build_join( // alias the join filter let join_filter_opt = conjunction(pull_up.join_filters) .map_or(Ok(None), |filter| { - replace_qualified_name(filter, &all_correlated_cols, &alias).map(Some) + replace_qualified_name(filter, &all_correlated_cols, alias).map(Some) })?; - let join_filter = match (join_filter_opt, in_predicate_opt.cloned()) { - ( - Some(join_filter), - Some(Expr::BinaryExpr(BinaryExpr { - left, - op: Operator::Eq, - right, - })), - ) => { - let right_col = create_col_from_scalar_expr(&right, alias)?; - let in_predicate = Expr::eq(left.deref().clone(), Expr::Column(right_col)); - in_predicate.and(join_filter) - } - (Some(join_filter), _) => join_filter, - ( - _, - Some(Expr::BinaryExpr(BinaryExpr { - left, - op: Operator::Eq, - right, - })), - ) => { - let right_col = create_col_from_scalar_expr(&right, alias)?; + 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(Some); + } - Expr::eq(left.deref().clone(), Expr::Column(right_col)) - } - (None, None) => lit(true), - _ => return Ok(None), + let join_filter = match build_join_filter_with_in_predicate( + join_filter_opt, + in_predicate_opt, + alias, + )? { + Some(join_filter) => join_filter, + None if in_predicate_opt.is_none() => lit(true), + None => return Ok(None), }; if matches!(join_type, JoinType::LeftMark | JoinType::RightMark) { @@ -465,6 +502,9 @@ fn build_join( // nullable mark column when the predicate can be implemented by hash keys. // Non-equality correlated filters stay on the legacy path because hash join // execution cannot mark UNKNOWN candidates for residual predicates. + // Only handles LeftMark, no caller produces RightMark today. + // The join below has no join keys set yet, so this flag has no effect + // unless a later pass sets them from the filter before execution. let null_aware = join_type == JoinType::LeftMark && in_predicate_opt.is_some() && mark_filter_is_hashable_only @@ -501,6 +541,9 @@ fn build_join( // // Additionally, if the join keys are non-nullable on both sides, we don't need // null-aware semantics because NULLs cannot exist in the data. + // Only handles LeftAnti, no caller produces RightAnti today. + // The join below has no join keys set yet, so this flag has no effect + // unless a later pass sets them from the filter before execution. let null_aware = join_type == JoinType::LeftAnti && in_predicate_opt.is_some() && join_keys_may_be_null(&join_filter, left.schema(), sub_query_alias.schema())?; @@ -530,6 +573,124 @@ fn build_join( Ok(Some(new_plan)) } +/// Builds the join for a correlated `EXISTS`/`IN` 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 { + // The `IN` equality must not be folded into the join filter here. This is a + // `LEFT JOIN` on the correlation condition alone, and the value comparison + // (with count-bug compensation) is applied afterward, in `exists_expr` below. + // Folding it in would make the join itself conditional on the compared value, + // causing rows with a real match to be NULL-padded and wrongly compensated. + 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| not(f.clone())); + + let mut expr_rewrite = TypeCoercionRewriter { + schema: joined.schema(), + }; + + let exists_expr = if let Some(in_predicate) = in_predicate_opt { + let Expr::BinaryExpr(BinaryExpr { + left: in_left, + op: Operator::Eq, + right: in_right, + }) = in_predicate + else { + return internal_err!("expected a binary equality IN predicate"); + }; + let value_name = create_col_from_scalar_expr(in_right, alias.to_string())?.name; + let compensation = expr_map.get(&value_name).cloned(); + let value_col = Expr::Column(Column::new(Some(alias), value_name)); + let compensated_value = match compensation { + // Only substitute when the empty-input value isn't NULL itself. + // If it is NULL, the unmatched row's raw NULL from the LEFT JOIN + // is already the correct value, so no substitution is needed. + Some(empty_batch_result) + if !evaluates_to_null( + empty_batch_result.clone(), + empty_batch_result.column_refs(), + )? => + { + let mut builder = + when(indicator_col.clone().is_null(), empty_batch_result); + if let Some(when_expr) = &having_arm { + builder = builder.when(when_expr.clone(), lit(ScalarValue::Null)); + } + builder + .otherwise(value_col)? + .rewrite(&mut expr_rewrite) + .data()? + } + _ => value_col, + }; + in_left.deref().clone().eq(compensated_value) + } else { + // EXISTS is true by default (the groupless aggregate always + // produces a row), unless the row has an actual join match whose + // HAVING predicate fails. + match having_arm { + Some(when_expr) => when(indicator_col.clone().is_null(), lit(true)) + .when(when_expr, lit(false)) + .otherwise(lit(true))? + .rewrite(&mut expr_rewrite) + .data()?, + None => lit(true), + } + }; + + let new_plan = match join_type { + JoinType::LeftMark | JoinType::RightMark => { + let mut proj_exprs = left_projection; + // "mark" must match the column name mark_join expects back. + proj_exprs.push(exists_expr.alias_qualified(Some(alias), "mark")); + LogicalPlanBuilder::from(joined) + .project(proj_exprs)? + .build()? + } + JoinType::LeftAnti => LogicalPlanBuilder::from(joined) + .filter(not(exists_expr))? + .project(left_projection)? + .build()?, + JoinType::LeftSemi => LogicalPlanBuilder::from(joined) + .filter(exists_expr)? + .project(left_projection)? + .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, @@ -579,6 +740,7 @@ mod tests { use arrow::datatypes::{DataType, Field, Schema}; use datafusion_expr::builder::table_source; use datafusion_expr::{and, binary_expr, col, out_ref_col, table_scan}; + use datafusion_functions_aggregate::count::count_udaf; macro_rules! assert_optimized_plan_equal { ( @@ -765,7 +927,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] - " + " ) } @@ -1881,6 +2043,37 @@ 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] + 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 9040d55f25260..878f399edfe43 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1138,6 +1138,65 @@ select t1.t1_int from t1 where ( 2 4 +# 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)--Filter: __correlated_sq_1.__always_true IS NULL OR __correlated_sq_1.__always_true IS NOT NULL AND (__correlated_sq_1.count(Int64(1)) != Int64(0)) IS DISTINCT FROM Boolean(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) +---- +33 3 + +# 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_in_groupless_count_agg_count_bug +query II +SELECT t1_id, t1_int FROM t1 WHERE 0 IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- +33 3 + +# 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_in_groupless_null_default_agg_wrong_zero_substitution +query II +SELECT t1_id, t1_int FROM t1 WHERE t1_id = 33 AND 0 IN (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- + #correlated_scalar_subquery_count_agg_in_case_when query TT explain @@ -1370,6 +1429,55 @@ where t1_id > 40 or not exists (select 1 from t2 where t2.t2_int > t1.t1_int) 33 44 +# 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_in_groupless_count_agg_count_bug_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or 0 in (select count(*) from t2 where t1.t1_id = t2.t2_id) +---- +33 +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_in_groupless_null_default_agg_wrong_zero_substitution_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or 0 in (select sum(t2_int) from t2 where t1.t1_id = t2.t2_id) +---- +44 + ########## # Regression for https://github.com/apache/datafusion/issues/23010: # a projection that selects / reorders a subset of columns over a mark join. From b52f9168b8d4eacc4b7b31ae79602ef4d2d26b7d Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Mon, 7 Sep 2026 16:29:54 -0400 Subject: [PATCH 2/4] test: add not in coverage --- .../sqllogictest/test_files/subquery.slt | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 878f399edfe43..90a809299281f 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1197,6 +1197,24 @@ query II SELECT t1_id, t1_int FROM t1 WHERE t1_id = 33 AND 0 IN (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id) ---- +# correlated_not_in_groupless_count_agg_count_bug +query II rowsort +SELECT t1_id, t1_int FROM t1 WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- +11 1 +22 2 +44 4 + +# correlated_not_in_groupless_count_agg_count_bug_having +query II +SELECT t1_id, t1_int FROM t1 WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) +---- + +# correlated_not_in_groupless_null_default_agg_wrong_zero_substitution +query II +SELECT t1_id, t1_int FROM t1 WHERE t1_id = 33 AND 0 NOT IN (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id) +---- + #correlated_scalar_subquery_count_agg_in_case_when query TT explain @@ -1478,6 +1496,22 @@ where t1_id > 40 or 0 in (select sum(t2_int) from t2 where t1.t1_id = t2.t2_id) ---- 44 +# correlated_not_in_groupless_count_agg_count_bug_embedded_in_or +query I rowsort +select t1_id from t1 +where t1_id > 40 or 0 not in (select count(*) from t2 where t1.t1_id = t2.t2_id having count(*) = 0) +---- +44 + +# correlated_not_in_groupless_null_default_agg_embedded_in_or_plans_successfully +query I rowsort +select t1_id from t1 +where t1_id > 40 or 0 not in (select sum(t2_int) from t2 where t1.t1_id = t2.t2_id) +---- +11 +22 +44 + ########## # Regression for https://github.com/apache/datafusion/issues/23010: # a projection that selects / reorders a subset of columns over a mark join. From a8b0fb4fa76670650b11f18a7a9c4c110bb188c4 Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Sun, 13 Sep 2026 14:18:56 -0400 Subject: [PATCH 3/4] fix: correct having-failed rows for not-in --- .../src/decorrelate_predicate_subquery.rs | 40 +++++++++++-------- .../sqllogictest/test_files/subquery.slt | 26 +++++++++++- 2 files changed, 48 insertions(+), 18 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 5446304772dc2..b98528868fa7d 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -30,8 +30,8 @@ use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::alias::AliasGenerator; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{ - Column, DFSchemaRef, ExprSchema, NullEquality, Result, ScalarValue, - assert_or_internal_err, internal_err, plan_err, + Column, DFSchemaRef, ExprSchema, NullEquality, Result, assert_or_internal_err, + internal_err, plan_err, }; use datafusion_expr::expr::{Exists, InSubquery}; use datafusion_expr::expr_rewriter::create_col_from_scalar_expr; @@ -622,29 +622,35 @@ fn build_join_with_count_bug( let value_name = create_col_from_scalar_expr(in_right, alias.to_string())?.name; let compensation = expr_map.get(&value_name).cloned(); let value_col = Expr::Column(Column::new(Some(alias), value_name)); - let compensated_value = match compensation { - // Only substitute when the empty-input value isn't NULL itself. - // If it is NULL, the unmatched row's raw NULL from the LEFT JOIN - // is already the correct value, so no substitution is needed. + // A HAVING-failed row means zero subquery rows for that outer row, so + // for `IN`-truth purposes it's `false` (no value to match). + let having_failed_result = lit(false); + // Only substitute the empty-input value when it isn't NULL itself. + // If it is NULL, the unmatched row's raw NULL from the LEFT JOIN is + // already the correct value, so no substitution is needed. + let unmatched_result = match &compensation { Some(empty_batch_result) if !evaluates_to_null( empty_batch_result.clone(), empty_batch_result.column_refs(), )? => { - let mut builder = - when(indicator_col.clone().is_null(), empty_batch_result); - if let Some(when_expr) = &having_arm { - builder = builder.when(when_expr.clone(), lit(ScalarValue::Null)); - } - builder - .otherwise(value_col)? - .rewrite(&mut expr_rewrite) - .data()? + in_left.deref().clone().eq(empty_batch_result.clone()) } - _ => value_col, + _ => in_left.deref().clone().eq(value_col.clone()), }; - in_left.deref().clone().eq(compensated_value) + let matched_result = in_left.deref().clone().eq(value_col); + match &having_arm { + Some(when_expr) => when(indicator_col.clone().is_null(), unmatched_result) + .when(when_expr.clone(), having_failed_result) + .otherwise(matched_result)? + .rewrite(&mut expr_rewrite) + .data()?, + None => when(indicator_col.clone().is_null(), unmatched_result) + .otherwise(matched_result)? + .rewrite(&mut expr_rewrite) + .data()?, + } } else { // EXISTS is true by default (the groupless aggregate always // produces a row), unless the row has an actual join match whose diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 1bca92b4fd215..8aa9cb34d3682 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1178,6 +1178,12 @@ SELECT t1_id, t1_int FROM t1 WHERE 0 IN (SELECT count(*) FROM t2 WHERE t1.t1_id ---- 33 3 +# correlated_in_groupless_count_agg_count_bug_having +query II +SELECT t1_id, t1_int FROM t1 WHERE 0 IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) +---- +33 3 + # 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) @@ -1206,9 +1212,25 @@ SELECT t1_id, t1_int FROM t1 WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t1.t1 44 4 # correlated_not_in_groupless_count_agg_count_bug_having -query II +query TT +explain SELECT t1_id, t1_int FROM t1 WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) +---- +logical_plan +01)LeftSemi Join: t1.t1_id = __correlated_sq_1.t2_id +02)--TableScan: t1 projection=[t1_id, t1_int] +03)--SubqueryAlias: __correlated_sq_1 +04)----Projection: t2.t2_id +05)------Filter: __common_expr_2 IS NOT DISTINCT FROM Boolean(true) OR __common_expr_2 +06)--------Projection: count(Int64(1)) != Int64(0) AS __common_expr_2, t2.t2_id +07)----------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1))]] +08)------------TableScan: t2 projection=[t2_id] + +query II rowsort SELECT t1_id, t1_int FROM t1 WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0) ---- +11 1 +22 2 +44 4 # correlated_not_in_groupless_null_default_agg_wrong_zero_substitution query II @@ -1501,6 +1523,8 @@ query I rowsort select t1_id from t1 where t1_id > 40 or 0 not in (select count(*) from t2 where t1.t1_id = t2.t2_id having count(*) = 0) ---- +11 +22 44 # correlated_not_in_groupless_null_default_agg_embedded_in_or_plans_successfully From e39d0b40cfaaa9775b90d8221be462e9cb6c9b38 Mon Sep 17 00:00:00 2001 From: Mohammad Naqvi Date: Mon, 14 Sep 2026 18:14:24 -0400 Subject: [PATCH 4/4] fix: correctly propagate compensation through stacked filters --- datafusion/optimizer/src/decorrelate.rs | 85 ++++++++++++++++--- .../src/decorrelate_predicate_subquery.rs | 45 ++++++++-- .../sqllogictest/test_files/subquery.slt | 64 ++++++++++++-- 3 files changed, 169 insertions(+), 25 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0c37f00b64355..da08339c29571 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -214,6 +214,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 @@ -260,7 +305,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, @@ -417,19 +462,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) } } } @@ -437,6 +485,16 @@ impl PullUpCorrelatedExpr { } } +/// 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, @@ -617,9 +675,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 b98528868fa7d..8cd2584f2bf53 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -604,12 +604,31 @@ fn build_join_with_count_bug( let indicator_col = Expr::Column(Column::new(Some(alias), UN_MATCHED_ROW_INDICATOR)); - let having_arm = pull_up_having_expr.map(|f| not(f.clone())); + let having_arm = pull_up_having_expr.map(|f| f.clone().is_not_true()); let mut expr_rewrite = TypeCoercionRewriter { schema: joined.schema(), }; + // The subquery's HAVING clause, evaluated against an unmatched row's + // default aggregate values, 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), + }; + let exists_expr = if let Some(in_predicate) = in_predicate_opt { let Expr::BinaryExpr(BinaryExpr { left: in_left, @@ -628,7 +647,7 @@ fn build_join_with_count_bug( // Only substitute the empty-input value when it isn't NULL itself. // If it is NULL, the unmatched row's raw NULL from the LEFT JOIN is // already the correct value, so no substitution is needed. - let unmatched_result = match &compensation { + let unmatched_value_result = match &compensation { Some(empty_batch_result) if !evaluates_to_null( empty_batch_result.clone(), @@ -639,6 +658,11 @@ fn build_join_with_count_bug( } _ => in_left.deref().clone().eq(value_col.clone()), }; + let unmatched_result = when( + unmatched_having_default.clone().is_not_true(), + having_failed_result.clone(), + ) + .otherwise(unmatched_value_result)?; let matched_result = in_left.deref().clone().eq(value_col); match &having_arm { Some(when_expr) => when(indicator_col.clone().is_null(), unmatched_result) @@ -653,14 +677,17 @@ fn build_join_with_count_bug( } } else { // EXISTS is true by default (the groupless aggregate always - // produces a row), unless the row has an actual join match whose - // HAVING predicate fails. + // 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. match having_arm { - Some(when_expr) => when(indicator_col.clone().is_null(), lit(true)) - .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) + .when(when_expr, lit(false)) + .otherwise(lit(true))? + .rewrite(&mut expr_rewrite) + .data()? + } None => lit(true), } }; diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 8aa9cb34d3682..83e93471b41a2 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1153,7 +1153,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 (__correlated_sq_1.count(Int64(1)) != Int64(0)) IS DISTINCT FROM Boolean(true) +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] @@ -1167,6 +1167,56 @@ SELECT t1_id, t1_int FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t1.t1_i ---- 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_in_groupless_having_with_stacked_filter_via_cte +query II +SELECT t1_id, t1_int FROM t1 WHERE t1_int IN ( + 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 +) +---- + +# 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) @@ -1184,6 +1234,11 @@ SELECT t1_id, t1_int FROM t1 WHERE 0 IN (SELECT count(*) FROM t2 WHERE t1.t1_id ---- 33 3 +# correlated_in_groupless_null_default_agg_having_with_type_coercion_plans_successfully +query I +SELECT t1_id FROM t1 WHERE 5.0 IN (SELECT sum(t2_int) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0 OR sum(t2_int) IS NULL) +---- + # 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) @@ -1220,10 +1275,9 @@ logical_plan 02)--TableScan: t1 projection=[t1_id, t1_int] 03)--SubqueryAlias: __correlated_sq_1 04)----Projection: t2.t2_id -05)------Filter: __common_expr_2 IS NOT DISTINCT FROM Boolean(true) OR __common_expr_2 -06)--------Projection: count(Int64(1)) != Int64(0) AS __common_expr_2, t2.t2_id -07)----------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1))]] -08)------------TableScan: t2 projection=[t2_id] +05)------Filter: count(Int64(1)) = Int64(0) IS NOT TRUE OR count(Int64(1)) != Int64(0) +06)--------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1))]] +07)----------TableScan: t2 projection=[t2_id] query II rowsort SELECT t1_id, t1_int FROM t1 WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t1.t1_id = t2.t2_id HAVING count(*) = 0)