diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index febbd3a6c84b5..858c6baac60d3 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -854,6 +854,35 @@ fn build_join( if !pull_up.can_pull_up { return Ok(None); } + // TODO(#24960): stopgap until count-bug compensation covers `IN`/`NOT IN`. + // The subquery stays correlated unless an outer row with no match gets the + // right answer anyway. + if !exists_sub_query && pull_up.pulled_up_scalar_agg { + let null_reads_as_false = match join_type { + JoinType::LeftSemi => true, + JoinType::LeftMark => !needs_null_aware_mark, + _ => false, + }; + let count_bug_is_harmless = match rewrite(true) { + Ok((plan, pull_up)) => pull_up + .collected_count_expr_map + .get(&plan) + .filter(|m| !m.is_empty()) + .is_none_or(|empty_input_values| { + null_reads_as_false + && empty_input_values + .get(subquery.schema().field(0).name()) + .is_some_and(|value| { + matches!(value, Expr::Literal(v, _) if v.is_null()) + }) + }), + // Count-bug handling errors on some shapes it cannot compensate. + Err(_) => false, + }; + if !count_bug_is_harmless { + return Ok(None); + } + } let sub_query_alias = LogicalPlanBuilder::from(new_plan) .alias(alias.to_string())? diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index d4a42cca0c801..565d3c481c014 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2124,14 +2124,44 @@ 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 +# correlated_in_groupless_max_agg_embedded_in_or_is_not_count_bug_compensated +query I SELECT t1_id FROM t1 -WHERE t1_int NOT IN (SELECT count(*) FROM t2 WHERE t2.t2_id = t1.t1_id) +WHERE t1_int + 2 IN (SELECT max(t2_int) FROM t2 WHERE t2.t2_id = t1.t1_id) OR t1_id = 33 +ORDER BY t1_id ---- -22 +11 33 -44 + +# correlated_in_groupless_count_agg_stays_correlated +# Errors rather than returning wrong rows. Expected `33`. +query error Physical plan does not support logical expression InSubquery +SELECT t1_id FROM t1 +WHERE 0 IN (SELECT count(*) FROM t2 WHERE t2.t2_id = t1.t1_id) + +# correlated_not_in_groupless_count_agg_stays_correlated +# Errors rather than returning wrong rows. Expected `11, 22, 44`. +query error Physical plan does not support logical expression InSubquery +SELECT t1_id FROM t1 +WHERE 0 NOT IN (SELECT count(*) FROM t2 WHERE t2.t2_id = t1.t1_id) + +# correlated_in_groupless_count_agg_in_projection_stays_correlated +# Errors rather than returning wrong rows. +# Expected `(11, false), (22, false), (33, true), (44, false)`. +query error Physical plan does not support logical expression InSubquery +SELECT t1_id, 0 IN (SELECT count(*) FROM t2 WHERE t2.t2_id = t1.t1_id) FROM t1 + +# correlated_not_in_groupless_max_agg_stays_correlated +# Errors rather than returning wrong rows. Expected `22, 44`. +query error Physical plan does not support logical expression InSubquery +SELECT t1_id FROM t1 +WHERE t1_int + 2 NOT IN (SELECT max(t2_int) FROM t2 WHERE t2.t2_id = t1.t1_id) + +# correlated_in_groupless_max_agg_in_projection_stays_correlated +# Errors rather than returning wrong rows. +# Expected `(11, true), (22, false), (33, NULL), (44, false)`. +query error Physical plan does not support logical expression InSubquery +SELECT t1_id, t1_int + 2 IN (SELECT max(t2_int) FROM t2 WHERE t2.t2_id = t1.t1_id) FROM t1 statement ok DROP TABLE limit_zero_outer;