feat(bigquery): Add support for side & kind on set operators (#4959)

* feat(bigquery): Support set operations

* PR Feedback 1

Author: VaggelisD

sqlglot/expressions.py

@@ -3370,6 +3370,9 @@ class SetOperation(Query):
         "expression": True,
         "distinct": False,
         "by_name": False,
+        "side": False,
+        "kind": False,
+        "on": False,
         **QUERY_MODIFIERS,
     }
 
@@ -3408,6 +3411,14 @@ def left(self) -> Query:
     def right(self) -> Query:
         return self.expression
 
+    @property
+    def kind(self) -> str:
+        return self.text("kind").upper()
+
+    @property
+    def side(self) -> str:
+        return self.text("side").upper()
+
 
 class Union(SetOperation):
     pass

sqlglot/generator.py

@@ -1445,12 +1445,18 @@ def set_operation(self, expression: exp.SetOperation) -> str:
                 self.unsupported(f"{op_name} requires DISTINCT or ALL to be specified")
 
         if distinct is default_distinct:
-            kind = ""
+            distinct_or_all = ""
         else:
-            kind = " DISTINCT" if distinct else " ALL"
+            distinct_or_all = " DISTINCT" if distinct else " ALL"
+
+        side_kind = " ".join(filter(None, [expression.side, expression.kind]))
+        side_kind = f"{side_kind} " if side_kind else ""
 
         by_name = " BY NAME" if expression.args.get("by_name") else ""
-        return f"{op_name}{kind}{by_name}"
+        on = self.expressions(expression, key="on", flat=True)
+        on = f" ON ({on})" if on else ""
+
+        return f"{side_kind}{op_name}{distinct_or_all}{by_name}{on}"
 
     def set_operations(self, expression: exp.SetOperation) -> str:
         if not self.SET_OP_MODIFIERS:

sqlglot/optimizer/pushdown_projections.py

@@ -49,25 +49,31 @@ def pushdown_projections(expression, schema=None, remove_unused_selections=True)
             parent_selections = {SELECT_ALL}
 
         if isinstance(scope.expression, exp.SetOperation):
-            left, right = scope.union_scopes
-            if len(left.expression.selects) != len(right.expression.selects):
-                scope_sql = scope.expression.sql()
-                raise OptimizeError(f"Invalid set operation due to column mismatch: {scope_sql}.")
-
-            referenced_columns[left] = parent_selections
-
-            if any(select.is_star for select in right.expression.selects):
-                referenced_columns[right] = parent_selections
-            elif not any(select.is_star for select in left.expression.selects):
-                if scope.expression.args.get("by_name"):
-                    referenced_columns[right] = referenced_columns[left]
-                else:
-                    referenced_columns[right] = [
-                        right.expression.selects[i].alias_or_name
-                        for i, select in enumerate(left.expression.selects)
-                        if SELECT_ALL in parent_selections
-                        or select.alias_or_name in parent_selections
-                    ]
+            set_op = scope.expression
+            if not (set_op.kind or set_op.side):
+                # Do not optimize this set operation if it's using the BigQuery specific
+                # kind / side syntax (e.g INNER UNION ALL BY NAME) which changes the semantics of the operation
+                left, right = scope.union_scopes
+                if len(left.expression.selects) != len(right.expression.selects):
+                    scope_sql = scope.expression.sql()
+                    raise OptimizeError(
+                        f"Invalid set operation due to column mismatch: {scope_sql}."
+                    )
+
+                referenced_columns[left] = parent_selections
+
+                if any(select.is_star for select in right.expression.selects):
+                    referenced_columns[right] = parent_selections
+                elif not any(select.is_star for select in left.expression.selects):
+                    if scope.expression.args.get("by_name"):
+                        referenced_columns[right] = referenced_columns[left]
+                    else:
+                        referenced_columns[right] = [
+                            right.expression.selects[i].alias_or_name
+                            for i, select in enumerate(left.expression.selects)
+                            if SELECT_ALL in parent_selections
+                            or select.alias_or_name in parent_selections
+                        ]
 
         if isinstance(scope.expression, exp.Select):
             if remove_unused_selections:

sqlglot/optimizer/qualify_columns.py

@@ -916,6 +916,32 @@ def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequenc
                     if source.expression.is_type(exp.DataType.Type.STRUCT):
                         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
             else:
                 columns = source.expression.named_selects
 

sqlglot/parser.py

@@ -4596,39 +4596,69 @@ def _parse_locks(self) -> t.List[exp.Lock]:
 
         return locks
 
-    def _parse_set_operations(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
-        while this and self._match_set(self.SET_OPERATIONS):
-            token_type = self._prev.token_type
+    def parse_set_operation(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
+        start = self._index
+        _, side_token, kind_token = self._parse_join_parts()
 
-            if token_type == TokenType.UNION:
-                operation: t.Type[exp.SetOperation] = exp.Union
-            elif token_type == TokenType.EXCEPT:
-                operation = exp.Except
-            else:
-                operation = exp.Intersect
+        side = side_token.text if side_token else None
+        kind = kind_token.text if kind_token else None
 
-            comments = self._prev.comments
+        if not self._match_set(self.SET_OPERATIONS):
+            self._retreat(start)
+            return None
 
-            if self._match(TokenType.DISTINCT):
-                distinct: t.Optional[bool] = True
-            elif self._match(TokenType.ALL):
-                distinct = False
-            else:
-                distinct = self.dialect.SET_OP_DISTINCT_BY_DEFAULT[operation]
-                if distinct is None:
-                    self.raise_error(f"Expected DISTINCT or ALL for {operation.__name__}")
+        token_type = self._prev.token_type
 
-            by_name = self._match_text_seq("BY", "NAME")
-            expression = self._parse_select(nested=True, parse_set_operation=False)
+        if token_type == TokenType.UNION:
+            operation: t.Type[exp.SetOperation] = exp.Union
+        elif token_type == TokenType.EXCEPT:
+            operation = exp.Except
+        else:
+            operation = exp.Intersect
 
-            this = self.expression(
-                operation,
-                comments=comments,
-                this=this,
-                distinct=distinct,
-                by_name=by_name,
-                expression=expression,
-            )
+        comments = self._prev.comments
+
+        if self._match(TokenType.DISTINCT):
+            distinct: t.Optional[bool] = True
+        elif self._match(TokenType.ALL):
+            distinct = False
+        else:
+            distinct = self.dialect.SET_OP_DISTINCT_BY_DEFAULT[operation]
+            if distinct is None:
+                self.raise_error(f"Expected DISTINCT or ALL for {operation.__name__}")
+
+        by_name = self._match_text_seq("BY", "NAME") or self._match_text_seq(
+            "STRICT", "CORRESPONDING"
+        )
+        if self._match_text_seq("CORRESPONDING"):
+            by_name = True
+            if not side and not kind:
+                kind = "INNER"
+
+        on_column_list = None
+        if by_name and self._match_texts(("ON", "BY")):
+            on_column_list = self._parse_wrapped_csv(self._parse_column)
+
+        expression = self._parse_select(nested=True, parse_set_operation=False)
+
+        return self.expression(
+            operation,
+            comments=comments,
+            this=this,
+            distinct=distinct,
+            by_name=by_name,
+            expression=expression,
+            side=side,
+            kind=kind,
+            on=on_column_list,
+        )
+
+    def _parse_set_operations(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
+        while True:
+            setop = self.parse_set_operation(this)
+            if not setop:
+                break
+            this = setop
 
         if isinstance(this, exp.SetOperation) and self.MODIFIERS_ATTACHED_TO_SET_OP:
             expression = this.expression

tests/dialects/test_bigquery.py

@@ -2382,3 +2382,43 @@ def test_annotate_timestamps(self):
 
         for select in annotated.selects:
             self.assertEqual(select.type.sql("bigquery"), "TIMESTAMP")
+
+    def test_set_operations(self):
+        self.validate_identity("SELECT 1 AS foo INNER UNION ALL SELECT 3 AS foo, 4 AS bar")
+
+        for side in ("", " LEFT", " FULL"):
+            for kind in ("", " OUTER"):
+                for name in (
+                    "",
+                    " BY NAME",
+                    " BY NAME ON (foo, bar)",
+                ):
+                    with self.subTest(f"Testing {side} {kind} {name} in test_set_operations"):
+                        self.validate_identity(
+                            f"SELECT 1 AS foo{side}{kind} UNION ALL{name} SELECT 3 AS foo, 4 AS bar",
+                        )
+
+        self.validate_identity(
+            "SELECT 1 AS x UNION ALL CORRESPONDING SELECT 2 AS x",
+            "SELECT 1 AS x INNER UNION ALL BY NAME SELECT 2 AS x",
+        )
+
+        self.validate_identity(
+            "SELECT 1 AS x UNION ALL CORRESPONDING BY (foo, bar) SELECT 2 AS x",
+            "SELECT 1 AS x INNER UNION ALL BY NAME ON (foo, bar) SELECT 2 AS x",
+        )
+
+        self.validate_identity(
+            "SELECT 1 AS x LEFT UNION ALL CORRESPONDING SELECT 2 AS x",
+            "SELECT 1 AS x LEFT UNION ALL BY NAME SELECT 2 AS x",
+        )
+
+        self.validate_identity(
+            "SELECT 1 AS x UNION ALL STRICT CORRESPONDING SELECT 2 AS x",
+            "SELECT 1 AS x UNION ALL BY NAME SELECT 2 AS x",
+        )
+
+        self.validate_identity(
+            "SELECT 1 AS x UNION ALL STRICT CORRESPONDING BY (foo, bar) SELECT 2 AS x",
+            "SELECT 1 AS x UNION ALL BY NAME ON (foo, bar) SELECT 2 AS x",
+        )

tests/fixtures/optimizer/qualify_columns.sql

@@ -325,6 +325,65 @@ SELECT _q_0.a AS a FROM (SELECT x.a AS a FROM x AS x UNION SELECT x.a AS a FROM
 ((select a from x where a < 1)) UNION ((select a from x where a > 2));
 ((SELECT x.a AS a FROM x AS x WHERE x.a < 1)) UNION ((SELECT x.a AS a FROM x AS x WHERE x.a > 2));
 
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar INNER UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.bar AS bar FROM (SELECT 1 AS foo, 2 AS bar INNER UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar UNION ALL CORRESPONDING SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.bar AS bar FROM (SELECT 1 AS foo, 2 AS bar INNER UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.foo AS foo, _q_0.bar AS bar FROM (SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.foo AS foo, _q_0.bar AS bar, _q_0.baz AS baz FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar LEFT UNION ALL CORRESPONDING SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.foo AS foo, _q_0.bar AS bar FROM (SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL CORRESPONDING SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.foo AS foo, _q_0.bar AS bar, _q_0.baz AS baz FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL CORRESPONDING BY (foo, bar) SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.foo AS foo, _q_0.bar AS bar FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL BY NAME ON (foo, bar) SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL BY NAME ON (foo, bar) SELECT 3 AS bar, 4 AS baz);
+SELECT _q_0.foo AS foo, _q_0.bar AS bar FROM (SELECT 1 AS foo, 2 AS bar FULL UNION ALL BY NAME ON (foo, bar) SELECT 3 AS bar, 4 AS baz) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+SELECT * FROM ((SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) LEFT UNION ALL BY NAME ON (bar) SELECT 3 AS foo, 4 AS bar);
+SELECT _q_0.bar AS bar FROM ((SELECT 1 AS foo, 2 AS bar LEFT UNION ALL BY NAME SELECT 3 AS bar, 4 AS baz) LEFT UNION ALL BY NAME ON (bar) SELECT 3 AS foo, 4 AS bar) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+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);
+SELECT _q_0.foo AS foo, _q_0.qux AS qux 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) AS _q_0;
+
+# dialect: bigquery
+# execute: false
+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;
+
 --------------------------------------
 -- Subqueries
 --------------------------------------