Repository: tobymao/sqlglot
Pull request: fix(optimizer)!: Fix chained exp.SetOperation type annotation
Fixes #4261

Consider this repro:

```Python
>>> sql = """
WITH t AS (
SELECT NULL AS col
UNION ALL
SELECT 'a' AS col
UNION ALL
SELECT NULL AS col
)
SELECT * FROM t;
"""

>>> optimized = sqlglot.optimize(sql)
>>> optimized.selects[0].type
DataType(this=Type.NULL)
```

To annotate this query, `annotate_scope` will be called for each scope:

```
1. Scope<SELECT NULL AS "col">
2. Scope<SELECT 'a' AS "col">
3. Scope<SELECT NULL AS "col" UNION ALL SELECT 'a' AS "col">
4. Scope<SELECT NULL AS "col">
5. Scope<SELECT NULL AS "col" UNION ALL SELECT 'a' AS "col" UNION ALL SELECT NULL AS "col">
6. Scope<WITH "t" AS (SELECT NULL AS "col" UNION ALL SELECT 'a' AS "col" UNION ALL SELECT NULL AS "col") SELECT "t"."col" AS "col" FROM "t" AS "t">
```

The annotation fails because:
1. Set operations are processed only if there is least one source in the scope, so scopes (3) & (5) won't be traversed/annotated.

2. The top-level UNION (6) will be annotated, but the `left` and `right` properties of `exp.SetOperation` will only yield the very first & last `SELECT` statements of the chain thus inferring `col` as `NULL`.

This PR solves this ordering issue by:
- Iteratively visiting the child `exp.SetOperation` nodes & coercing their left - right counterparts into an intermediate result set
- For each intermediate result set, coerce it with a global result set to get the final column type mappings


Problem statement:
Optimizer relies on subquery's order to infer correct expression type with a multiple tables query
The Optimizer cannot infer the correct type for a given expression when analyzing a query with a `UNION ALL` clause with multiple tables (3 or more). It relies on the order of the `SELECT`.

### SQL  
The following code will result in a `DataType(this=Type.NULL)` for the `some_id_feature` expression.
```SQL
SELECT
  __source_table,
  some_id AS some_id_feature,
FROM (
  SELECT
    __source_table,
    some_id
  FROM (
    SELECT
      __source_table,
      NULL AS some_id
    FROM (
      SELECT
        'table_1' AS __source_table
      FROM table_1
    )
    UNION ALL
    SELECT
      __source_table,
      some_id
    FROM (
      SELECT
        table_2.some_id,
        'table_2' AS __source_table
      FROM table_2
    )
    UNION ALL
    SELECT
      __source_table,
      NULL AS some_id
    FROM (
      SELECT
        'table_3' AS __source_table
      FROM table_3
    )
  )
)
```

### Reproducible code  
```Python
from sqlglot import parse_one
from sqlglot.optimizer import optimize

node = parse_one(provided_sql, dialect='spark')

optimized = optimize(
    node,
    schema={
        'table_1': {
            'column_a': 'LONG', 'column_b': 'VARCHAR'}, 
        'table_2': {
            'some_id': 'VARCHAR'}, 
        'table_3': {
            'column_a': 'LONG', 'column_b': 'VARCHAR'}
    }
)

optimized.args["expressions"]
```

The  code outputs:
```
[Alias(
  this=Column(
    this=Identifier(this=__source_table, quoted=True),
    table=Identifier(this=_q_3, quoted=True),
    _type=DataType(this=Type.VARCHAR)),
  alias=Identifier(this=__source_table, quoted=True, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.VARCHAR)), Alias(
  this=Column(
    this=Identifier(this=some_id, quoted=True),
    table=Identifier(this=_q_3, quoted=True),
    _type=DataType(this=Type.NULL)),
  alias=Identifier(this=some_id_feature, quoted=True, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.NULL))]
```

However, if we change the order of the `SELECT` when using the `UNION ALL`, the correct `VARCHAR` will be inferred for `some_id_feature` as long as the `table_2`'s `SELECT` remains as the first or last one in the subquery:  

```SQL
SELECT
  __source_table,
  some_id AS some_id_feature,
FROM (
  SELECT
    __source_table,
    some_id
  FROM (
    SELECT
      __source_table,
      some_id
    FROM (
      SELECT
        table_2.some_id,
        'table_2' AS __source_table
      FROM table_2
    )
    UNION ALL
    SELECT
      __source_table,
      NULL AS some_id
    FROM (
      SELECT
        'table_1' AS __source_table
      FROM table_1
    )
    UNION ALL
    SELECT
      __source_table,
      NULL AS some_id
    FROM (
      SELECT
        'table_3' AS __source_table
      FROM table_3
    )
  )
)
```
It results:

```
[Alias(
  this=Column(
    this=Identifier(this=__source_table, quoted=True),
    table=Identifier(this=_q_3, quoted=True),
    _type=DataType(this=Type.VARCHAR)),
  alias=Identifier(this=__source_table, quoted=True, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.VARCHAR)), Alias(
  this=Column(
    this=Identifier(this=some_id, quoted=True),
    table=Identifier(this=_q_3, quoted=True),
    _type=DataType(this=Type.VARCHAR)),
  alias=Identifier(this=some_id_feature, quoted=True, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.VARCHAR))]
```
And, 

```SQL 
SELECT
  __source_table,
  some_id AS some_id_feature,
FROM (
  SELECT
    __source_table,
    some_id
  FROM (
    SELECT
      __source_table,
      NULL AS some_id
    FROM (
      SELECT
        'table_1' AS __source_table
      FROM table_1
    )
    UNION ALL
    SELECT
      __source_table,
      NULL AS some_id
    FROM (
      SELECT
        'table_3' AS __source_table
      FROM table_3
    )
    UNION ALL
    SELECT
      __source_table,
      some_id
    FROM (
      SELECT
        table_2.some_id,
        'table_2' AS __source_table
      FROM table_2
    )
  )
)
```

It results:

```
[Alias(
  this=Column(
    this=Identifier(this=__source_table, quoted=True),
    table=Identifier(this=_q_3, quoted=True),
    _type=DataType(this=Type.VARCHAR)),
  alias=Identifier(this=__source_table, quoted=True, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.VARCHAR)), Alias(
  this=Column(
    this=Identifier(this=some_id, quoted=True),
    table=Identifier(this=_q_3, quoted=True),
    _type=DataType(this=Type.VARCHAR)),
  alias=Identifier(this=some_id_feature, quoted=True, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.VARCHAR))]
```

Hints from issue discussion:
Appreciate the detailed report, we'll take a look shortly.

Commit message: Cache result of SetOperation subtrees

Patch to review:
diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py
index 052b74229c..be3e2e1240 100644
--- a/sqlglot/optimizer/annotate_types.py
+++ b/sqlglot/optimizer/annotate_types.py
@@ -192,6 +192,10 @@ def __init__(
         # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
         self._visited: t.Set[int] = set()
 
+        # Maps an exp.SetOperation's id (e.g. UNION) to it's annotated selected columns. This is done if the exp.SetOperation is
+        # the main expression of a source scope, as selecting from it multiple times would reprocess the entire subtree
+        self._setop_cols: t.Dict[int, t.Dict[str, exp.DataType.Type]] = dict()
+
     def _set_type(
         self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
     ) -> None:
@@ -231,22 +235,45 @@ def annotate_scope(self, scope: Scope) -> None:
             elif isinstance(expression, exp.SetOperation) and len(expression.left.selects) == len(
                 expression.right.selects
             ):
-                if expression.args.get("by_name"):
-                    r_type_by_select = {s.alias_or_name: s.type for s in expression.right.selects}
-                    selects[name] = {
-                        s.alias_or_name: self._maybe_coerce(
-                            t.cast(exp.DataType, s.type),
-                            r_type_by_select.get(s.alias_or_name) or exp.DataType.Type.UNKNOWN,
-                        )
-                        for s in expression.left.selects
-                    }
-                else:
-                    selects[name] = {
-                        ls.alias_or_name: self._maybe_coerce(
-                            t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
-                        )
-                        for ls, rs in zip(expression.left.selects, expression.right.selects)
-                    }
+                col_types: t.Dict[str, exp.DataType.Type] = self._setop_cols.get(id(expression), {})
+
+                if not col_types:
+                    # Process a chain / sub-tree of set operations
+                    for set_op in expression.walk(
+                        prune=lambda n: not isinstance(n, exp.SetOperation)
+                    ):
+                        if not isinstance(set_op, exp.SetOperation):
+                            continue
+
+                        if set_op.args.get("by_name"):
+                            r_type_by_select = {
+                                s.alias_or_name: s.type for s in set_op.right.selects
+                            }
+                            setop_cols = {
+                                s.alias_or_name: self._maybe_coerce(
+                                    t.cast(exp.DataType, s.type),
+                                    r_type_by_select.get(s.alias_or_name)
+                                    or exp.DataType.Type.UNKNOWN,
+                                )
+                                for s in set_op.left.selects
+                            }
+                        else:
+                            setop_cols = {
+                                ls.alias_or_name: self._maybe_coerce(
+                                    t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
+                                )
+                                for ls, rs in zip(set_op.left.selects, set_op.right.selects)
+                            }
+
+                        # Coerce intermediate results with the previously registered types, if they exist
+                        for col_name, col_type in setop_cols.items():
+                            col_types[col_name] = self._maybe_coerce(
+                                col_type, col_types.get(col_name, exp.DataType.Type.NULL)
+                            )
+
+                    self._setop_cols[id(expression)] = col_types
+
+                selects[name] = col_types
             else:
                 selects[name] = {s.alias_or_name: s.type for s in expression.selects}
 
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 93132859b7..b6aad8db9a 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1337,6 +1337,24 @@ def test_union_annotation(self):
         self.assertEqual(union_by_name.selects[0].type.this, exp.DataType.Type.BIGINT)
         self.assertEqual(union_by_name.selects[1].type.this, exp.DataType.Type.DOUBLE)
 
+        # Test chained UNIONs
+        sql = """
+            WITH t AS
+            (
+                SELECT NULL AS col
+                UNION
+                SELECT NULL AS col
+                UNION
+                SELECT 'a' AS col
+                UNION
+                SELECT NULL AS col
+                UNION
+                SELECT NULL AS col
+            )
+            SELECT col FROM t;
+        """
+        self.assertEqual(optimizer.optimize(sql).selects[0].type.this, exp.DataType.Type.VARCHAR)
+
     def test_recursive_cte(self):
         query = parse_one(
             """


Code review comments:
diff_hunk: @@ -192,6 +192,10 @@ def __init__(
         # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
         self._visited: t.Set[int] = set()
 
+        # Maps an exp.SetOperation's id (e.g. UNION) to it's annotated selected columns. This is done if the exp.SetOperation is
+        # the main expression of a source scope, as selecting from it multiple times would reprocess the entire subtree
+        self._setop_cols: t.Dict[int, t.Dict[str, exp.DataType.Type]] = dict()
path: sqlglot/optimizer/annotate_types.py
text: @user1:
```suggestion
        # Maps an exp.SetOperation's id (e.g. UNION) to its projection types. This is computed if the
        # exp.SetOperation is the expression of a scope source, as selecting from it multiple times
        # would reprocess the entire subtree to coerce the types of its operands' projections
        self._setop_column_types: t.Dict[int, t.Dict[str, exp.DataType.Type]] = {}
```

diff_hunk: @@ -231,22 +235,45 @@ def annotate_scope(self, scope: Scope) -> None:
             elif isinstance(expression, exp.SetOperation) and len(expression.left.selects) == len(
                 expression.right.selects
             ):
-                if expression.args.get("by_name"):
-                    r_type_by_select = {s.alias_or_name: s.type for s in expression.right.selects}
-                    selects[name] = {
-                        s.alias_or_name: self._maybe_coerce(
-                            t.cast(exp.DataType, s.type),
-                            r_type_by_select.get(s.alias_or_name) or exp.DataType.Type.UNKNOWN,
-                        )
-                        for s in expression.left.selects
-                    }
-                else:
-                    selects[name] = {
-                        ls.alias_or_name: self._maybe_coerce(
-                            t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
-                        )
-                        for ls, rs in zip(expression.left.selects, expression.right.selects)
-                    }
+                col_types: t.Dict[str, exp.DataType.Type] = self._setop_cols.get(id(expression), {})
path: sqlglot/optimizer/annotate_types.py
text: @user1:
Why do we need the annotation here? Also, we can do this instead:

```python
selects[name] = column_types = self._setop_column_types.setdefault(id(expression), {})
```

Then we'll be able to remove lines 274-276, because both dicts will be updated due to mutation.

@author:
That annotation was explicitly requested by mypy, probably due to the dict operations.