fix: reject correlated IN/NOT IN over a groupless aggregate

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())?

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;