From 6f39c2fbcfa3972daaae275e0c2866786f4ce706 Mon Sep 17 00:00:00 2001 From: ryux1 Date: Mon, 7 Sep 2026 01:50:07 +0200 Subject: [PATCH 1/2] fix: preserve groupless aggregate subquery cardinality --- .../src/decorrelate_predicate_subquery.rs | 67 ++++++++++++ .../sqllogictest/test_files/issue_24960.slt | 102 ++++++++++++++++++ .../sqllogictest/test_files/subquery.slt | 10 +- 3 files changed, 170 insertions(+), 9 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/issue_24960.slt diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 5f623f1bef6f6..68a4acec71f83 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -95,6 +95,10 @@ impl OptimizerRule for DecorrelatePredicateSubquery { match extract_subquery_info(subquery_expr) { // The subquery expression is at the top level of the filter SubqueryPredicate::Top(subquery) => { + if let Some(expr) = single_row_subquery_predicate(&subquery) { + other_exprs.push(expr); + continue; + } match build_join_top(&subquery, &cur_input, config.alias_generator())? { Some(plan) => cur_input = plan, @@ -144,6 +148,24 @@ fn rewrite_inner_subqueries( let mut cur_input = outer; let alias = config.alias_generator(); let expr_without_subqueries = expr.transform(|e| match e { + Expr::Exists(Exists { subquery, negated }) + if correlated_guaranteed_single_row(subquery.subquery.as_ref()) => + { + Ok(Transformed::yes(lit(!negated))) + } + Expr::InSubquery(InSubquery { + expr, + subquery, + negated, + }) if correlated_guaranteed_single_row(subquery.subquery.as_ref()) => { + let scalar_subquery = Expr::ScalarSubquery(subquery); + let expr = if negated { + (*expr).not_eq(scalar_subquery) + } else { + (*expr).eq(scalar_subquery) + }; + Ok(Transformed::yes(expr)) + } Expr::Exists(Exists { subquery: Subquery { subquery, .. }, negated, @@ -225,6 +247,51 @@ fn has_subquery(expr: &Expr) -> bool { .unwrap() } +/// Returns whether the plan is guaranteed to produce one row. +/// +/// This deliberately recognizes only transparent wrappers around an +/// ungrouped aggregate. In particular, a `Filter` above the aggregate is a +/// `HAVING` clause and can remove the otherwise guaranteed row. +fn guaranteed_single_row(plan: &LogicalPlan) -> bool { + match plan { + LogicalPlan::Projection(projection) => { + guaranteed_single_row(projection.input.as_ref()) + } + LogicalPlan::SubqueryAlias(alias) => guaranteed_single_row(alias.input.as_ref()), + LogicalPlan::Aggregate(aggregate) => aggregate.group_expr.is_empty(), + _ => false, + } +} + +fn correlated_guaranteed_single_row(plan: &LogicalPlan) -> bool { + !plan.all_out_ref_exprs().is_empty() && guaranteed_single_row(plan) +} + +/// Rewrite predicates over a guaranteed-single-row subquery to scalar forms. +/// +/// An ungrouped aggregate without `HAVING` always produces one row, even when +/// its input is empty. `EXISTS` can therefore be folded by cardinality, while +/// `IN` is equivalent to comparison with a scalar subquery. The latter lets +/// `ScalarSubqueryToJoin` apply its existing empty-input aggregate +/// compensation. +fn single_row_subquery_predicate(query_info: &SubqueryInfo) -> Option { + if !correlated_guaranteed_single_row(query_info.query.subquery.as_ref()) { + return None; + } + + match &query_info.where_in_expr { + Some(expr) => { + let scalar_subquery = Expr::ScalarSubquery(query_info.query.clone()); + Some(if query_info.negated { + expr.clone().not_eq(scalar_subquery) + } else { + expr.clone().eq(scalar_subquery) + }) + } + None => Some(lit(!query_info.negated)), + } +} + /// Optimize the subquery to left-anti/left-semi join. /// If the subquery is a correlated subquery, we need extract the join predicate from the subquery. /// diff --git a/datafusion/sqllogictest/test_files/issue_24960.slt b/datafusion/sqllogictest/test_files/issue_24960.slt new file mode 100644 index 0000000000000..2e0e7fda664ce --- /dev/null +++ b/datafusion/sqllogictest/test_files/issue_24960.slt @@ -0,0 +1,102 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +statement ok +CREATE TABLE outer_values(value INT) AS VALUES (0), (1), (2), (3); + +statement ok +CREATE TABLE inner_values(value INT) AS VALUES (2), (2); + +# A groupless aggregate produces one row even when its input is empty. +query I rowsort +SELECT value +FROM outer_values o +WHERE EXISTS ( + SELECT count(*) FROM inner_values i WHERE o.value = i.value +); +---- +0 +1 +2 +3 + +query I +SELECT value +FROM outer_values o +WHERE NOT EXISTS ( + SELECT count(*) FROM inner_values i WHERE o.value = i.value +); +---- + +# IN must compare against the aggregate's empty-input value for outer rows +# that have no matching inner rows. +query I rowsort +SELECT value +FROM outer_values o +WHERE value IN ( + SELECT count(*) FROM inner_values i WHERE o.value = i.value +); +---- +0 +2 + +query I rowsort +SELECT value +FROM outer_values o +WHERE value NOT IN ( + SELECT count(*) FROM inner_values i WHERE o.value = i.value +); +---- +1 +3 + +# HAVING can remove the aggregate row, so these subqueries must not be folded +# based on the aggregate's cardinality. +query I +SELECT value +FROM outer_values o +WHERE EXISTS ( + SELECT count(*) + FROM inner_values i + WHERE o.value = i.value + HAVING count(*) > 0 +); +---- +2 + +# LIMIT 0 also removes the otherwise guaranteed aggregate row. +query I +SELECT value +FROM outer_values o +WHERE EXISTS ( + SELECT count(*) FROM inner_values i WHERE o.value = i.value LIMIT 0 +); +---- + +# The same cardinality rule applies when EXISTS is embedded in another +# predicate and would otherwise be implemented with a mark join. +query I rowsort +SELECT value +FROM outer_values o +WHERE value = 99 OR EXISTS ( + SELECT count(*) FROM inner_values i WHERE o.value = i.value +); +---- +0 +1 +2 +3 diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 9040d55f25260..cf1f5db428834 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -521,15 +521,7 @@ SELECT t1_id, (SELECT sum(t2_int) FROM t2 WHERE t2.t2_id = t1.t1_id group by t2_ query TT explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT sum(t1.t1_int + t2.t2_id) FROM t2 WHERE t1.t1_name = t2.t2_name) ---- -logical_plan -01)Projection: t1.t1_id, t1.t1_name -02)--Filter: EXISTS () -03)----Subquery: -04)------Projection: sum(outer_ref(t1.t1_int) + t2.t2_id) -05)--------Aggregate: groupBy=[[]], aggr=[[sum(CAST(outer_ref(t1.t1_int) + t2.t2_id AS Int64))]] -06)----------Filter: outer_ref(t1.t1_name) = t2.t2_name -07)------------TableScan: t2 -08)----TableScan: t1 projection=[t1_id, t1_name, t1_int] +logical_plan TableScan: t1 projection=[t1_id, t1_name] #support_agg_correlated_columns2 query TT From 19ff8418e3d226012a09e2a81ad9c4054723d717 Mon Sep 17 00:00:00 2001 From: ryux1 Date: Mon, 7 Sep 2026 02:20:06 +0200 Subject: [PATCH 2/2] test: cover single-row predicate rewrites --- .../src/decorrelate_predicate_subquery.rs | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 68a4acec71f83..87469da53b373 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -646,6 +646,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::expr_fn::count; macro_rules! assert_optimized_plan_equal { ( @@ -707,6 +708,110 @@ mod tests { optimizer.optimize(plan, &crate::OptimizerContext::new(), |_, _| {}) } + fn correlated_count_subquery() -> Result> { + Ok(Arc::new( + LogicalPlanBuilder::from(test_table_scan_with_name("sq")?) + .filter(col("sq.a").eq(out_ref_col(DataType::UInt32, "test.a")))? + .aggregate(Vec::::new(), vec![count(lit(1))])? + .build()?, + )) + } + + #[test] + fn single_row_cardinality_through_transparent_wrappers() -> Result<()> { + let aggregate = correlated_count_subquery()?; + assert!(guaranteed_single_row(aggregate.as_ref())); + + let projection = LogicalPlanBuilder::from(aggregate.as_ref().clone()) + .project(vec![lit(1)])? + .build()?; + assert!(guaranteed_single_row(&projection)); + + let alias = LogicalPlanBuilder::from(projection) + .alias("single")? + .build()?; + assert!(guaranteed_single_row(&alias)); + + let having = LogicalPlanBuilder::from(aggregate.as_ref().clone()) + .filter(lit(true))? + .build()?; + assert!(!guaranteed_single_row(&having)); + assert!(!guaranteed_single_row(&test_table_scan()?)); + + Ok(()) + } + + #[test] + fn correlated_single_row_top_level_predicates() -> Result<()> { + let predicates = [ + ( + exists(correlated_count_subquery()?), + "Filter: Boolean(true)", + ), + ( + not_exists(correlated_count_subquery()?), + "Filter: Boolean(false)", + ), + ( + in_subquery(col("test.b"), correlated_count_subquery()?), + "Filter: test.b = ()", + ), + ( + not_in_subquery(col("test.b"), correlated_count_subquery()?), + "Filter: test.b != ()", + ), + ]; + + for (predicate, expected_filter) in predicates { + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(predicate)? + .build()?; + let optimized = optimize_with_decorrelate(plan)?; + let display = optimized.display_indent().to_string(); + + assert!(display.starts_with(expected_filter), "{display}"); + assert!(!display.contains("LeftSemi Join"), "{display}"); + assert!(!display.contains("LeftAnti Join"), "{display}"); + } + + Ok(()) + } + + #[test] + fn correlated_single_row_embedded_predicates() -> Result<()> { + let predicates = [ + ( + exists(correlated_count_subquery()?), + "Filter: test.c = UInt32(1) OR Boolean(true)", + ), + ( + not_exists(correlated_count_subquery()?), + "Filter: test.c = UInt32(1) OR Boolean(false)", + ), + ( + in_subquery(col("test.b"), correlated_count_subquery()?), + "Filter: test.c = UInt32(1) OR test.b = ()", + ), + ( + not_in_subquery(col("test.b"), correlated_count_subquery()?), + "Filter: test.c = UInt32(1) OR test.b != ()", + ), + ]; + + for (predicate, expected_filter) in predicates { + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(col("test.c").eq(lit(1u32)).or(predicate))? + .build()?; + let optimized = optimize_with_decorrelate(plan)?; + let display = optimized.display_indent().to_string(); + + assert!(display.starts_with(expected_filter), "{display}"); + assert!(!display.contains("LeftMark Join"), "{display}"); + } + + Ok(()) + } + /// Test for several IN subquery expressions #[test] fn in_subquery_multiple() -> Result<()> {