fix(optimizer): EXISTS over a scalar aggregate is always true [CLAUDE]

sqlglot/optimizer/unnest_subqueries.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from sqlglot import exp
 from sqlglot.helper import name_sequence
-from sqlglot.optimizer.scope import ScopeType, find_in_scope, traverse_scope
+from sqlglot.optimizer.scope import ScopeType, find_all_in_scope, find_in_scope, traverse_scope
 from sqlglot._typing import E
 
 
@@ -31,7 +31,9 @@ def unnest_subqueries(expression: E) -> E:
         if not parent:
             continue
         if scope.external_columns:
-            decorrelate(select, parent, scope.external_columns, next_alias_name)
+            # a correlated set operation branch can't be hoisted out on its own
+            if scope.scope_type != ScopeType.SET_OPERATION:
+                decorrelate(select, parent, scope.external_columns, next_alias_name)
         elif scope.scope_type == ScopeType.SUBQUERY:
             unnest(select, parent, next_alias_name)
 
@@ -148,9 +150,23 @@ def unnest(select, parent_select, next_alias_name):
 def decorrelate(select, parent_select, external_columns, next_alias_name):
     where = select.args.get("where")
 
-    if not where or where.find(exp.Or) or select.find(exp.Limit, exp.Offset):
+    if not where or where.find(exp.Or) or select.find(exp.Limit, exp.Offset, exp.Fetch):
         return
 
+    parent_predicate = select.find_ancestor(exp.Predicate)
+
+    # find_ancestor crosses query boundaries, so the predicate can belong to another query
+    if parent_predicate is not None and parent_predicate.parent_select is not parent_select:
+        return
+
+    if isinstance(parent_predicate, exp.Exists) and not select.args.get("group"):
+        if select.args.get("having") or select.args.get("qualify"):
+            return
+
+        if _has_aggregate_projection(select):
+            _replace(parent_predicate, exp.true())
+            return
+
     table_alias = next_alias_name()
     keys = []
 
@@ -207,8 +223,6 @@ def decorrelate(select, parent_select, external_columns, next_alias_name):
             if isinstance(predicate, exp.EQ) and key not in group_by:
                 group_by.append(key)
 
-    parent_predicate = select.find_ancestor(exp.Predicate)
-
     # When the subquery is embedded inside a function (e.g. COALESCE, TRIM) in the SELECT list,
     # the ancestor chain contains no Predicate node AND the subquery is not a direct projection.
     if parent_predicate is None and not is_subquery_projection:
@@ -338,6 +352,32 @@ def _replace(expression: exp.Expr, condition: exp.ExpOrStr) -> exp.Expr:
     return expression.replace(exp.condition(condition))
 
 
+def _is_windowed(agg: exp.Expr) -> bool:
+    # a window applies to exactly one function, its `this`; an aggregate anywhere else groups
+    node = agg
+    parent = node.parent
+
+    # parens, FILTER and IGNORE NULLS wrap that function without changing which one it is
+    while parent is not None and parent.this is node:
+        if isinstance(parent, exp.Window):
+            return True
+
+        if isinstance(parent, exp.Func):
+            return False
+
+        node, parent = parent, parent.parent
+
+    return False
+
+
+def _has_aggregate_projection(select: exp.Select) -> bool:
+    return any(
+        not _is_windowed(agg)
+        for projection in select.selects
+        for agg in find_all_in_scope(projection, exp.AggFunc)
+    )
+
+
 def _other_operand(expression: object) -> exp.Expr | None:
     if isinstance(expression, exp.In):
         return expression.this

tests/fixtures/optimizer/unnest_subqueries.sql

@@ -35,6 +35,73 @@ SELECT * FROM x LEFT JOIN (SELECT SUM(y.a) AS a, y.a AS _u_1, ARRAY_AGG(y.b) AS
 SELECT * FROM x WHERE EXISTS (SELECT y.a AS a, y.b AS b FROM y WHERE x.a = y.a);
 SELECT * FROM x LEFT JOIN (SELECT y.a AS a FROM y WHERE TRUE GROUP BY y.a) AS _u_0 ON x.a = _u_0.a WHERE NOT _u_0.a IS NULL;
 
+# title: EXISTS over a scalar aggregate always matches, it returns exactly one row
+SELECT * FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE y.a = x.a);
+SELECT * FROM x WHERE TRUE;
+
+# title: NOT EXISTS over a scalar aggregate never matches
+SELECT * FROM x WHERE NOT EXISTS (SELECT SUM(y.b) FROM y WHERE y.a = x.a);
+SELECT * FROM x WHERE NOT TRUE;
+
+# title: EXISTS over a scalar aggregate with a HAVING is not rewritten
+SELECT * FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE y.a = x.a HAVING COUNT(*) = 0);
+SELECT * FROM x WHERE EXISTS(SELECT COUNT(*) FROM y WHERE y.a = x.a HAVING COUNT(*) = 0);
+
+# title: EXISTS over a scalar aggregate with a FETCH is not rewritten, it can return no rows
+SELECT * FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE y.a = x.a FETCH FIRST 0 ROWS ONLY);
+SELECT * FROM x WHERE EXISTS(SELECT COUNT(*) FROM y WHERE y.a = x.a FETCH FIRST 0 ROWS ONLY);
+
+# title: EXISTS over a windowed aggregate is not a scalar aggregate
+SELECT * FROM x WHERE EXISTS (SELECT COUNT(*) OVER () FROM y WHERE y.a = x.a);
+SELECT * FROM x LEFT JOIN (SELECT y.a AS _u_1 FROM y WHERE TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a WHERE NOT _u_0._u_1 IS NULL;
+
+# title: EXISTS is not folded when the aggregate belongs to a derived table inside it
+SELECT * FROM x WHERE EXISTS (SELECT * FROM (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) AS t WHERE t.c > 5);
+SELECT * FROM x WHERE EXISTS(SELECT * FROM (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) AS t WHERE t.c > 5);
+
+# title: EXISTS is not folded when the aggregate belongs to a CTE inside it
+SELECT * FROM x WHERE EXISTS (WITH t AS (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) SELECT t.c AS c FROM t WHERE t.c > 5);
+SELECT * FROM x WHERE EXISTS(WITH t AS (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) SELECT t.c AS c FROM t WHERE t.c > 5);
+
+# title: EXISTS is not folded when the aggregate is only one branch of a set operation
+SELECT * FROM x WHERE EXISTS (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a INTERSECT SELECT z.a AS a FROM z);
+SELECT * FROM x WHERE EXISTS(SELECT COUNT(*) AS c FROM y WHERE y.a = x.a INTERSECT SELECT z.a AS a FROM z);
+
+# title: a correlated branch of a set operation is not decorrelated, it can't be hoisted out
+SELECT * FROM x WHERE EXISTS (SELECT y.a AS a FROM y WHERE y.a = x.a INTERSECT SELECT z.a AS a FROM z);
+SELECT * FROM x WHERE EXISTS(SELECT y.a AS a FROM y WHERE y.a = x.a INTERSECT SELECT z.a AS a FROM z);
+
+# title: a parenthesized correlated branch of a set operation is not decorrelated either
+SELECT * FROM x WHERE EXISTS ((SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) INTERSECT (SELECT z.a AS a FROM z));
+SELECT * FROM x WHERE EXISTS((SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) INTERSECT (SELECT z.a AS a FROM z));
+
+SELECT * FROM x WHERE EXISTS ((SELECT y.a AS a FROM y WHERE y.a = x.a) EXCEPT (SELECT z.a AS a FROM z));
+SELECT * FROM x WHERE EXISTS((SELECT y.a AS a FROM y WHERE y.a = x.a) EXCEPT (SELECT z.a AS a FROM z));
+
+# title: EXISTS over a scalar aggregate with a QUALIFY is not rewritten
+SELECT * FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE y.a = x.a QUALIFY ROW_NUMBER() OVER () = 2);
+SELECT * FROM x WHERE EXISTS(SELECT COUNT(*) FROM y WHERE y.a = x.a QUALIFY ROW_NUMBER() OVER () = 2);
+
+# title: an aggregate in a window spec still groups the subquery into a single row
+SELECT * FROM x WHERE EXISTS (SELECT RANK() OVER (ORDER BY SUM(y.b)) FROM y WHERE y.a = x.a);
+SELECT * FROM x WHERE TRUE;
+
+# title: a parenthesized aggregate in a window spec still groups the subquery
+SELECT * FROM x WHERE EXISTS (SELECT RANK() OVER (ORDER BY (SUM(y.b))) FROM y WHERE y.a = x.a);
+SELECT * FROM x WHERE TRUE;
+
+# title: an aggregate in the arguments of a windowed function still groups the subquery
+SELECT * FROM x WHERE EXISTS (SELECT LAG(SUM(y.b)) OVER (ORDER BY 1) FROM y WHERE y.a = x.a);
+SELECT * FROM x WHERE TRUE;
+
+# title: a FILTER between the window and the aggregate leaves it windowed
+SELECT * FROM x WHERE EXISTS (SELECT SUM(y.b) FILTER(WHERE y.b > 1) OVER () FROM y WHERE y.a = x.a);
+SELECT * FROM x LEFT JOIN (SELECT y.a AS _u_1 FROM y WHERE TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a WHERE NOT _u_0._u_1 IS NULL;
+
+# title: EXISTS over a scalar aggregate is folded inside an outer window function
+SELECT COUNT(CASE WHEN EXISTS(SELECT COUNT(*) FROM y WHERE y.a = x.a) THEN 1 END) OVER () FROM x;
+SELECT COUNT(CASE WHEN TRUE THEN 1 END) OVER () FROM x;
+
 SELECT * FROM x WHERE x.a IN (SELECT y.a AS a FROM y LIMIT 10);
 SELECT * FROM x WHERE x.a IN (SELECT y.a AS a FROM y LIMIT 10);
 

tests/test_executor.py

@@ -862,6 +862,110 @@ def test_correlated_count(self):
             ],
         )
 
+    def test_correlated_exists_over_scalar_aggregate(self):
+        tables = {"x": [{"a": 1}, {"a": 2}, {"a": None}], "y": [{"b": 2}, {"b": 3}]}
+        schema = {"x": {"a": "int"}, "y": {"b": "int"}}
+        all_rows = [1, 2, None]
+
+        for sql, expected in (
+            ("SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE b = x.a)", all_rows),
+            ("SELECT a FROM x WHERE NOT EXISTS (SELECT COUNT(*) FROM y WHERE b = x.a)", []),
+            ("SELECT a FROM x WHERE EXISTS (SELECT SUM(b) FROM y WHERE b = x.a)", all_rows),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT MAX(b) FROM y WHERE b = x.a AND 1 = 2)",
+                all_rows,
+            ),
+            ("SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE b > x.a)", all_rows),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT DISTINCT COUNT(*) FROM y WHERE b = x.a)",
+                all_rows,
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT (SELECT COUNT(*) FROM y) FROM y WHERE b = x.a)",
+                [2],
+            ),
+            ("SELECT a FROM x WHERE EXISTS (SELECT 1, COUNT(*) FROM y WHERE b = x.a)", all_rows),
+            ("SELECT a FROM x WHERE EXISTS (SELECT 1, 2 FROM y WHERE b = x.a)", [2]),
+            ("SELECT a FROM x WHERE NOT EXISTS (SELECT 1, 2 FROM y WHERE b = x.a)", [1, None]),
+            ("SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) OVER () FROM y WHERE b = x.a)", [2]),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT RANK() OVER (ORDER BY SUM(b)) FROM y WHERE b = x.a)",
+                all_rows,
+            ),
+            ("SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE b = x.a GROUP BY b)", [2]),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT RANK() OVER (ORDER BY (SUM(b))) FROM y WHERE b = x.a)",
+                all_rows,
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT LAG(SUM(b)) OVER (ORDER BY 1) FROM y WHERE b = x.a)",
+                all_rows,
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT SUM(b) FILTER (WHERE b > 1) OVER () FROM y WHERE b = x.a)",
+                [2],
+            ),
+            # a HAVING is declined, so the executor evaluates the subquery per outer row
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE b = x.a HAVING COUNT(*) = 0)",
+                [1, None],
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) FROM y WHERE b = x.a HAVING COUNT(*) > 0)",
+                [2],
+            ),
+        ):
+            with self.subTest(sql):
+                self.assertEqual(
+                    sorted(
+                        (row[0] for row in execute(sql, tables=tables, schema=schema).rows), key=str
+                    ),
+                    sorted(expected, key=str),
+                )
+
+    def test_correlated_exists_over_a_nested_scalar_aggregate(self):
+        # the aggregate belongs to a nested query, so the EXISTS is still conditional
+        tables = {
+            "x": [{"a": 1}, {"a": 2}, {"a": None}],
+            "y": [{"a": 2, "b": 20}, {"a": 3, "b": 30}],
+            "z": [{"a": 1}],
+        }
+        schema = {"x": {"a": "int"}, "y": {"a": "int", "b": "int"}, "z": {"a": "int"}}
+
+        for sql, expected in (
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT * FROM (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) AS t WHERE t.c > 5)",
+                [],
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (WITH t AS (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) SELECT t.c AS c FROM t WHERE t.c > 5)",
+                [],
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT COUNT(*) AS c FROM y WHERE y.a = x.a INTERSECT SELECT z.a AS a FROM z)",
+                [2],
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS (SELECT y.a AS a FROM y WHERE y.a = x.a INTERSECT SELECT z.a AS a FROM z)",
+                [],
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS ((SELECT COUNT(*) AS c FROM y WHERE y.a = x.a) INTERSECT (SELECT z.a AS a FROM z))",
+                [2],
+            ),
+            (
+                "SELECT a FROM x WHERE EXISTS ((SELECT y.a AS a FROM y WHERE y.a = x.a) EXCEPT (SELECT z.a AS a FROM z))",
+                [2],
+            ),
+        ):
+            with self.subTest(sql):
+                self.assertEqual(
+                    sorted(
+                        (row[0] for row in execute(sql, tables=tables, schema=schema).rows), key=str
+                    ),
+                    sorted(expected, key=str),
+                )
+
     def test_table_depth_mismatch(self):
         tables = {"table": []}
         schema = {"db": {"table": {"col": "VARCHAR"}}}