Fix aggregate FILTER subquery scheduling

Author: liudy1

core/translate/group_by.rs

@@ -363,16 +363,25 @@ fn collect_agg_leaf_columns(aggregates: &[Aggregate], plan: &SelectPlan) -> Resu
                 }
                 Ok(WalkControl::SkipChildren)
             }
-            ast::Expr::SubqueryResult { subquery_id, .. } => {
+            ast::Expr::SubqueryResult {
+                subquery_id,
+                query_type,
+                ..
+            } => {
                 let is_correlated = plan
                     .non_from_clause_subqueries
                     .iter()
                     .find(|s| s.internal_id == *subquery_id)
                     .is_some_and(|s| s.correlated);
                 if is_correlated && !leaf_columns.iter().any(|e| exprs_are_equivalent(e, expr)) {
                     leaf_columns.push(expr.clone());
+                    return Ok(WalkControl::SkipChildren);
+                }
+                if matches!(query_type, ast::SubqueryType::In { .. }) {
+                    Ok(WalkControl::Continue)
+                } else {
+                    Ok(WalkControl::SkipChildren)
                 }
-                Ok(WalkControl::SkipChildren)
             }
             _ => Ok(WalkControl::Continue),
         }

core/translate/subquery.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use rustc_hash::FxHashMap as HashMap;
-use turso_parser::ast::{self, SortOrder, SubqueryType};
+use turso_parser::ast::{self, SortOrder, SubqueryType, TableInternalId};
 
 use crate::{
     alloc::TursoIteratorExt,
@@ -15,7 +15,8 @@ use crate::{
             emit_program_for_select_with_resolver, emit_query,
         },
         expr::{
-            compare_affinity, get_expr_affinity_info, unwrap_parens, walk_expr_mut, WalkControl,
+            compare_affinity, expr_references_subquery_id, get_expr_affinity_info, unwrap_parens,
+            walk_expr_mut, WalkControl,
         },
         optimizer::optimize_select_plan,
         plan::{
@@ -1987,13 +1988,40 @@ fn assign_select_subquery_eval_phases(plan: &mut SelectPlan) {
         .group_by
         .as_ref()
         .is_some_and(|group_by| !group_by.exprs.is_empty());
+    let aggregate_step_subquery_ids: Vec<TableInternalId> = plan
+        .non_from_clause_subqueries
+        .iter()
+        .filter_map(|subquery| {
+            aggregate_step_references_subquery(&plan.aggregates, subquery.internal_id)
+                .then_some(subquery.internal_id)
+        })
+        .collect();
 
     for subquery in plan.non_from_clause_subqueries.iter_mut() {
+        let referenced_by_aggregate_step =
+            aggregate_step_subquery_ids.contains(&subquery.internal_id);
         subquery.eval_phase = match subquery.origin {
-            SubqueryOrigin::SelectHaving | SubqueryOrigin::SelectOrderBy if has_grouped_output => {
+            SubqueryOrigin::SelectHaving | SubqueryOrigin::SelectOrderBy
+                if has_grouped_output && !referenced_by_aggregate_step =>
+            {
                 SubqueryEvalPhase::GroupedOutput
             }
             _ => subquery.origin.phase_floor(),
         };
     }
 }
+
+fn aggregate_step_references_subquery(
+    aggregates: &[Aggregate],
+    subquery_id: TableInternalId,
+) -> bool {
+    aggregates.iter().any(|agg| {
+        agg.args
+            .iter()
+            .any(|arg| expr_references_subquery_id(arg, subquery_id))
+            || agg
+                .filter_expr
+                .as_ref()
+                .is_some_and(|expr| expr_references_subquery_id(expr, subquery_id))
+    })
+}

testing/sqltests/tests/agg-functions/memory.sqltest

@@ -113,6 +113,33 @@ expect {
     1|null||
 }
 
+@cross-check-integrity
+test grouped-having-aggregate-filter-constant-in-subquery {
+    SELECT 1 GROUP BY 1 HAVING COUNT(*) FILTER (WHERE 1 IN (SELECT 1)) = 1;
+}
+expect {
+    1
+}
+
+@cross-check-integrity
+test grouped-having-aggregate-filter-empty-in-subquery {
+    SELECT 1 GROUP BY 1 HAVING COUNT(*) FILTER (WHERE 1 IN (SELECT 2)) = 1;
+}
+expect {
+}
+
+@cross-check-integrity
+test grouped-having-aggregate-filter-column-in-subquery {
+    CREATE TABLE t (x INTEGER, y INTEGER);
+    CREATE TABLE u (z INTEGER);
+    INSERT INTO t VALUES (1, 1), (2, 2);
+    INSERT INTO u VALUES (1);
+    SELECT y, SUM(x) FROM t GROUP BY y HAVING AVG(x) FILTER (WHERE y IN (SELECT z FROM u ORDER BY z LIMIT 1)) > 0;
+}
+expect {
+    1|1
+}
+
 @cross-check-integrity
 test case-when-aggregate-returns-column-value {
     CREATE TABLE t1(a INTEGER, b TEXT);