Fix!: Make UNION column qualification recursive (#5508)

* Fix: Make UNION column qualification recursive

* Pre -> post

Author: VaggelisD

sqlglot/optimizer/qualify_columns.py

@@ -972,6 +972,44 @@ def all_columns(self) -> t.Set[str]:
             }
         return self._all_columns
 
+    def get_source_columns_from_set_op(self, expression: exp.Expression) -> t.List[str]:
+        if isinstance(expression, exp.Select):
+            return expression.named_selects
+        if isinstance(expression, exp.Subquery) and isinstance(expression.this, exp.SetOperation):
+            # Different types of SET modifiers can be chained together if they're explicitly grouped by nesting
+            return self.get_source_columns_from_set_op(expression.this)
+        if not isinstance(expression, exp.SetOperation):
+            raise OptimizeError(f"Unknown set operation: {expression}")
+
+        set_op = expression
+
+        # BigQuery specific set operations modifiers, e.g INNER UNION ALL BY NAME
+        on_column_list = set_op.args.get("on")
+
+        if on_column_list:
+            # The resulting columns are the columns in the ON clause:
+            # {INNER | LEFT | FULL} UNION ALL BY NAME ON (col1, col2, ...)
+            columns = [col.name for col in on_column_list]
+        elif set_op.side or set_op.kind:
+            side = set_op.side
+            kind = set_op.kind
+
+            # Visit the children UNIONs (if any) in a post-order traversal
+            left = self.get_source_columns_from_set_op(set_op.left)
+            right = self.get_source_columns_from_set_op(set_op.right)
+
+            # We use dict.fromkeys to deduplicate keys and maintain insertion order
+            if side == "LEFT":
+                columns = left
+            elif side == "FULL":
+                columns = list(dict.fromkeys(left + right))
+            elif kind == "INNER":
+                columns = list(dict.fromkeys(left).keys() & dict.fromkeys(right).keys())
+        else:
+            columns = set_op.named_selects
+
+        return columns
+
     def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequence[str]:
         """Resolve the source columns for a given source `name`."""
         cache_key = (name, only_visible)
@@ -996,31 +1034,8 @@ def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequenc
                         for k in source.expression.type.expressions:  # type: ignore
                             columns.append(k.name)
             elif isinstance(source, Scope) and isinstance(source.expression, exp.SetOperation):
-                set_op = source.expression
-
-                # BigQuery specific set operations modifiers, e.g INNER UNION ALL BY NAME
-                on_column_list = set_op.args.get("on")
-
-                if on_column_list:
-                    # The resulting columns are the columns in the ON clause:
-                    # {INNER | LEFT | FULL} UNION ALL BY NAME ON (col1, col2, ...)
-                    columns = [col.name for col in on_column_list]
-                elif set_op.side or set_op.kind:
-                    side = set_op.side
-                    kind = set_op.kind
-
-                    left = set_op.left.named_selects
-                    right = set_op.right.named_selects
-
-                    # We use dict.fromkeys to deduplicate keys and maintain insertion order
-                    if side == "LEFT":
-                        columns = left
-                    elif side == "FULL":
-                        columns = list(dict.fromkeys(left + right))
-                    elif kind == "INNER":
-                        columns = list(dict.fromkeys(left).keys() & dict.fromkeys(right).keys())
-                else:
-                    columns = set_op.named_selects
+                columns = self.get_source_columns_from_set_op(source.expression)
+
             else:
                 select = seq_get(source.expression.selects, 0)
 

tests/fixtures/optimizer/qualify_columns.sql

@@ -389,6 +389,25 @@ SELECT _q_0.foo AS foo, _q_0.qux AS qux FROM ((SELECT 1 AS foo, 2 AS bar LEFT UN
 SELECT * FROM (((SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) FULL UNION ALL BY NAME ON (foo, qux) SELECT 3 AS qux, 4 AS bar) INNER UNION ALL BY NAME ON (foo) SELECT 6 AS foo);
 SELECT _q_0.foo AS foo FROM (((SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) FULL UNION ALL BY NAME ON (foo, qux) SELECT 3 AS qux, 4 AS bar) INNER UNION ALL BY NAME ON (foo) SELECT 6 AS foo) AS _q_0;
 
+# Title: Nested set operations with modifiers
+# dialect: bigquery
+# execute: false
+WITH t1 AS (SELECT 1 AS a, 2 AS b), t2 AS (SELECT 2 AS b, 3 AS c), t3 AS (SELECT 2 AS c, 3 AS d), t4 AS (SELECT 2 AS e, 3 AS f) SELECT * FROM ((SELECT * FROM t1 FULL OUTER UNION ALL BY NAME (SELECT * FROM t2 FULL OUTER UNION ALL BY NAME (SELECT * FROM t3 FULL OUTER UNION ALL BY NAME SELECT * FROM t4))));
+WITH t1 AS (SELECT 1 AS a, 2 AS b), t2 AS (SELECT 2 AS b, 3 AS c), t3 AS (SELECT 2 AS c, 3 AS d), t4 AS (SELECT 2 AS e, 3 AS f) SELECT _q_0.a AS a, _q_0.b AS b, _q_0.c AS c, _q_0.d AS d, _q_0.e AS e, _q_0.f AS f FROM ((SELECT t1.a AS a, t1.b AS b FROM t1 AS t1 FULL OUTER UNION ALL BY NAME (SELECT t2.b AS b, t2.c AS c FROM t2 AS t2 FULL OUTER UNION ALL BY NAME (SELECT t3.c AS c, t3.d AS d FROM t3 AS t3 FULL OUTER UNION ALL BY NAME SELECT t4.e AS e, t4.f AS f FROM t4 AS t4))) AS _q_0);
+
+
+# Title: Nested set operations with different modifiers (FULL + INNER)
+# dialect: bigquery
+# execute: false
+WITH t1 AS (SELECT 1 AS a, 2 AS b), t2 AS (SELECT 2 AS b, 3 AS c), t3 AS (SELECT 2 AS c, 3 AS d), t4 AS (SELECT 2 AS e, 3 AS f) SELECT * FROM ((SELECT * FROM t1 FULL OUTER UNION ALL BY NAME (SELECT * FROM t2 INNER UNION ALL BY NAME (SELECT * FROM t3 FULL OUTER UNION ALL BY NAME SELECT * FROM t4))));
+WITH t1 AS (SELECT 1 AS a, 2 AS b), t2 AS (SELECT 2 AS b, 3 AS c), t3 AS (SELECT 2 AS c, 3 AS d), t4 AS (SELECT 2 AS e, 3 AS f) SELECT _q_0.a AS a, _q_0.b AS b, _q_0.c AS c FROM ((SELECT t1.a AS a, t1.b AS b FROM t1 AS t1 FULL OUTER UNION ALL BY NAME (SELECT t2.b AS b, t2.c AS c FROM t2 AS t2 INNER UNION ALL BY NAME (SELECT t3.c AS c, t3.d AS d FROM t3 AS t3 FULL OUTER UNION ALL BY NAME SELECT t4.e AS e, t4.f AS f FROM t4 AS t4))) AS _q_0);
+
+# Title: Nested set operations with different modifiers (FULL + LEFT)
+# dialect: bigquery
+# execute: false
+WITH t1 AS (SELECT 1 AS a, 2 AS b, 3 AS c, 4 AS d), t2 AS (SELECT 2 AS b, 3 AS c), t3 AS (SELECT 2 AS c, 3 AS d), t4 AS (SELECT 2 AS d, 3 AS e) SELECT * FROM ((SELECT * FROM t1 FULL OUTER UNION ALL BY NAME (SELECT * FROM t2 FULL UNION ALL BY NAME (SELECT * FROM t3 LEFT UNION ALL BY NAME SELECT * FROM t4))));
+WITH t1 AS (SELECT 1 AS a, 2 AS b, 3 AS c, 4 AS d), t2 AS (SELECT 2 AS b, 3 AS c), t3 AS (SELECT 2 AS c, 3 AS d), t4 AS (SELECT 2 AS d, 3 AS e) SELECT _q_0.a AS a, _q_0.b AS b, _q_0.c AS c, _q_0.d AS d FROM ((SELECT t1.a AS a, t1.b AS b, t1.c AS c, t1.d AS d FROM t1 AS t1 FULL OUTER UNION ALL BY NAME (SELECT t2.b AS b, t2.c AS c FROM t2 AS t2 FULL UNION ALL BY NAME (SELECT t3.c AS c, t3.d AS d FROM t3 AS t3 LEFT UNION ALL BY NAME SELECT t4.d AS d, t4.e AS e FROM t4 AS t4))) AS _q_0);
+
 --------------------------------------
 -- Subqueries
 --------------------------------------