fix(optimizer): don't merge ORDER BY into UNION (#3215)

tobymao · Mar 25, 2024 · ec4648f · ec4648f
1 parent af1b026
commit ec4648f
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 3 deletions.
diff --git a/sqlglot/optimizer/merge_subqueries.py b/sqlglot/optimizer/merge_subqueries.py
@@ -214,6 +214,7 @@ def _is_recursive():
         and not _outer_select_joins_on_inner_select_join()
         and not _is_a_window_expression_in_unmergable_operation()
         and not _is_recursive()
+        and not (inner_select.args.get("order") and outer_scope.is_union)
     )
 
 

diff --git a/tests/fixtures/optimizer/merge_subqueries.sql b/tests/fixtures/optimizer/merge_subqueries.sql
@@ -429,4 +429,20 @@ WHERE
       q.a AS a
     FROM q AS q
   );
-SELECT q.a AS a FROM x AS q WHERE q.a IN (SELECT y.b AS a FROM y AS y);
+SELECT q.a AS a FROM x AS q WHERE q.a IN (SELECT y.b AS a FROM y AS y);
+
+# title: dont merge when inner query has ORDER BY and outer query is UNION
+WITH q AS (
+  SELECT
+    x.a AS a
+  FROM x
+  ORDER BY x.a
+)
+SELECT
+  q.a AS a
+FROM q
+UNION ALL
+SELECT
+  1 AS a;
+WITH q AS (SELECT x.a AS a FROM x AS x ORDER BY x.a) SELECT q.a AS a FROM q AS q UNION ALL SELECT 1 AS a;
+
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -383,7 +383,8 @@ def test_simplify(self):
 
         self.assertIn("Anonymous.this expects a str or an Identifier, got 'int'.", str(e.exception))
 
-        sql = parse_one("""
+        sql = parse_one(
+            """
         WITH cte AS (select 1 union select 2), cte2 AS (
             SELECT ROW() OVER (PARTITION BY y) FROM (
                 (select 1) limit 10
@@ -395,7 +396,8 @@ def test_simplify(self):
           a div 1,
           filter("B", (x, y) -> x + y)
           FROM (z AS z CROSS JOIN z) AS f(a) LEFT JOIN a.b.c.d.e.f.g USING(n) ORDER BY 1
-        """)
+        """
+        )
         self.assertEqual(
             optimizer.simplify.gen(sql),
             """