diff --git a/sqlglot/optimizer/merge_subqueries.py b/sqlglot/optimizer/merge_subqueries.py index 1e4a8307ef..0e438285a6 100644 --- a/sqlglot/optimizer/merge_subqueries.py +++ b/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) diff --git a/tests/fixtures/optimizer/merge_subqueries.sql b/tests/fixtures/optimizer/merge_subqueries.sql index 27434fd736..0eeffd43bb 100644 --- a/tests/fixtures/optimizer/merge_subqueries.sql +++ b/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;