Fix(optimizer): avoid merging window function nested under a projection (#4919)

Author: georgesittas

sqlglot/optimizer/merge_subqueries.py

@@ -129,8 +129,7 @@ def _mergeable(
     inner_select = inner_scope.expression.unnest()
 
     def _is_a_window_expression_in_unmergable_operation():
-        window_expressions = inner_select.find_all(exp.Window)
-        window_alias_names = {window.parent.alias_or_name for window in window_expressions}
+        window_aliases = {s.alias_or_name for s in inner_select.selects if s.find(exp.Window)}
         inner_select_name = from_or_join.alias_or_name
         unmergable_window_columns = [
             column
@@ -142,7 +141,7 @@ def _is_a_window_expression_in_unmergable_operation():
         window_expressions_in_unmergable = [
             column
             for column in unmergable_window_columns
-            if column.table == inner_select_name and column.name in window_alias_names
+            if column.table == inner_select_name and column.name in window_aliases
         ]
         return any(window_expressions_in_unmergable)
 

tests/fixtures/optimizer/merge_subqueries.sql

@@ -259,7 +259,7 @@ FROM
   t1;
 WITH t1 AS (SELECT x.a AS a, x.b AS b, ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.a) AS row_num FROM x AS x) SELECT SUM(t1.row_num) AS total_rows FROM t1 AS t1;
 
-# title: Test prevent merging of window if in group by func
+# title: Test prevent merging of window if in group by
 with t1 as (
   SELECT
     x.a,
@@ -277,7 +277,7 @@ GROUP BY t1.row_num
 ORDER BY t1.row_num;
 WITH t1 AS (SELECT x.a AS a, x.b AS b, ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.a) AS row_num FROM x AS x) SELECT t1.row_num AS row_num, SUM(t1.a) AS total FROM t1 AS t1 GROUP BY t1.row_num ORDER BY row_num;
 
-# title: Test prevent merging of window if in order by func
+# title: Test prevent merging of window if in order by
 with t1 as (
   SELECT
     x.a,
@@ -294,6 +294,23 @@ FROM
 ORDER BY t1.row_num, t1.a;
 WITH t1 AS (SELECT x.a AS a, x.b AS b, ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.a) AS row_num FROM x AS x) SELECT t1.row_num AS row_num, t1.a AS a FROM t1 AS t1 ORDER BY t1.row_num, t1.a;
 
+# title: Test preventing merging of window nested under complex projection if in order by
+WITH t1 AS (
+  SELECT
+    x.a,
+    x.b,
+    ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.a) - 1 AS row_num
+  FROM
+    x
+)
+SELECT
+  t1.row_num AS row_num,
+  t1.a AS a
+FROM
+  t1
+ORDER BY t1.row_num, t1.a;
+WITH t1 AS (SELECT x.a AS a, x.b AS b, ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.a) - 1 AS row_num FROM x AS x) SELECT t1.row_num AS row_num, t1.a AS a FROM t1 AS t1 ORDER BY t1.row_num, t1.a;
+
 # title: Test allow merging of window function
 with t1 as (
   SELECT