fix(optimizer)!: merge window issue

sqlglot/optimizer/merge_subqueries.py

@@ -136,21 +136,27 @@ def _mergeable(
     """
     inner_select = inner_scope.expression.unnest()
 
-    def _is_a_window_expression_in_unmergable_operation():
+    def _window_projection_blocks_merge():
+        """A window function's result depends on the full row set it sees, so merging the
+        subquery into the outer query is unsafe when:
+          - the outer query filters or joins (WHERE/JOIN), which changes that row set, or
+          - a window column is referenced in an operation that isn't pushed down
+            (GROUP BY, ORDER BY, HAVING, aggregate).
+        """
         window_aliases = {s.alias_or_name for s in inner_select.selects if s.find(exp.Window)}
         if not window_aliases:
             return False
+
+        outer = outer_scope.expression
+        if outer.args.get("where") or outer.args.get("joins"):
+            return True
+
         inner_select_name = from_or_join.alias_or_name
-        unmergable_window_columns = [
-            column
-            for column in outer_scope.columns
-            if column.find_ancestor(
-                exp.Where, exp.Group, exp.Order, exp.Join, exp.Having, exp.AggFunc
-            )
-        ]
         return any(
-            column.table == inner_select_name and column.name in window_aliases
-            for column in unmergable_window_columns
+            column.table == inner_select_name
+            and column.name in window_aliases
+            and column.find_ancestor(exp.Group, exp.Order, exp.Having, exp.AggFunc)
+            for column in outer_scope.columns
         )
 
     def _outer_select_joins_on_inner_select_join():
@@ -223,7 +229,7 @@ def _is_recursive():
             )
         )
         and not _outer_select_joins_on_inner_select_join()
-        and not _is_a_window_expression_in_unmergable_operation()
+        and not _window_projection_blocks_merge()
         and not _is_recursive()
         and not (inner_select.args.get("order") and outer_scope.is_union)
         and not isinstance(seq_get(inner_select.expressions, 0), exp.QueryTransform)

tests/fixtures/optimizer/merge_subqueries.sql

@@ -543,3 +543,11 @@ WITH t0 AS (SELECT 5 AS id), t1 AS (SELECT 1 AS id, 'US' AS cid), t2 AS (SELECT
 # title: Dont replace GROUP and ORDER BY if expression is literal
 WITH t1 AS (SELECT 1 AS col) SELECT a, SUM(b) AS b FROM (SELECT 6 AS a, col AS b FROM t1) AS t GROUP BY a ORDER BY a;
 WITH t1 AS (SELECT 1 AS col) SELECT 6 AS a, SUM(t1.col) AS b FROM t1 AS t1 GROUP BY a ORDER BY a;
+
+# title: Window function is not merged when the outer query filters its input rows
+SELECT s.w FROM (SELECT a, ROW_NUMBER() OVER (ORDER BY a) AS w FROM x) AS s WHERE s.a > 5;
+SELECT s.w AS w FROM (SELECT x.a AS a, ROW_NUMBER() OVER (ORDER BY x.a) AS w FROM x AS x) AS s WHERE s.a > 5;
+
+# title: Window function is not merged when the outer query joins, changing its input rows
+SELECT s.w FROM (SELECT b, ROW_NUMBER() OVER (ORDER BY b) AS w FROM x) AS s JOIN y ON s.b = y.b;
+SELECT s.w AS w FROM (SELECT x.b AS b, ROW_NUMBER() OVER (ORDER BY x.b) AS w FROM x AS x) AS s JOIN y AS y ON s.b = y.b;