fix(optimizer): unpivot annotate types (#7543)

* fix(optimizer): unpivot annotate types for bq

* fix varchar

* fix

* simplify

* refactor

* remove copy

* revert copy

* refactor impl

* varchar default

* more robust impl

* remove typing

* simplify

* correlated subquery case

* simplify logic

* reverse types

* feedback 1 george (partial)

* feedback 1 (george - full resolved)

* refactor

* ref

* ref 2

* remove dup tests

* doc str update

Author: geooo109

sqlglot-integration-tests

@@ -8,6 +8,7 @@
 from sqlglot.helper import trait, ensure_list
 from sqlglot.expressions.core import (
     Aliases,
+    Column,
     Condition,
     Distinct,
     Dot,
@@ -1762,6 +1763,60 @@ def unpivot(self) -> bool:
     def fields(self) -> list[Expr]:
         return self.args.get("fields", [])
 
+    def output_columns(self, pre_pivot_columns: t.Iterable[str]) -> list[str]:
+        """
+        Returns the columns produced by this (UN)PIVOT, in order.
+
+        Example:
+            >>> from sqlglot import parse_one, exp
+            >>> piv = parse_one("SELECT * FROM t UNPIVOT(val FOR name IN (a, b))").find(exp.Pivot)
+            >>> piv.output_columns(["a", "b", "c"])
+            ['c', 'name', 'val']
+
+        AST shape:
+            PIVOT(SUM(val) FOR name IN ('a', 'b')):
+                expressions: aggregate(s), e.g. [Sum(this=Column(val))]
+                fields:      [In(this=Column(name), expressions=[Literal('a'), Literal('b')])]
+                columns:     optional explicit output identifiers (e.g. set by Snowflake)
+
+            UNPIVOT(val FOR name IN (a, b)):
+                expressions: value Identifier(s), or Tuple(Identifiers) for multi-value
+                fields:      [In(this=Identifier(name), expressions=[Column(a), Column(b)])]
+                             For literal-aliased entries (`a AS 'x'`) the IN expressions
+                             are wrapped in PivotAlias(this=Column, alias=Literal).
+
+        Args:
+            pre_pivot_columns: Columns visible to the operator before it runs
+                (e.g. the source table or subquery's projections).
+        """
+        if self.unpivot:
+            excluded: set[str] = set()
+            name_columns: list[Identifier] = []
+            for field in self.fields:
+                if not isinstance(field, In):
+                    continue
+                if isinstance(field.this, Identifier):
+                    name_columns.append(field.this)
+                for e in field.expressions:
+                    excluded.update(c.output_name for c in e.find_all(Column))
+            value_columns = [
+                ident
+                for e in self.expressions
+                for ident in (e.expressions if isinstance(e, Tuple) else [e])
+                if isinstance(ident, Identifier)
+            ]
+            outputs = [i.name for i in name_columns + value_columns]
+        else:
+            excluded = {c.output_name for c in self.find_all(Column)}
+            outputs = [c.output_name for c in self.args.get("columns") or []]
+            if not outputs:
+                outputs = [c.alias_or_name for c in self.expressions]
+
+        if not excluded or not outputs:
+            return []
+
+        return [c for c in pre_pivot_columns if c not in excluded] + outputs
+
 
 class UnpivotColumns(Expression):
     arg_types = {"this": True, "expressions": True}

sqlglot/expressions/query.py

@@ -330,6 +330,10 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]:
                 elif isinstance(expression, exp.Selectable):
                     selects[name] = {s.alias_or_name: s.type for s in expression.selects if s.type}
 
+            for pivot in scope.pivots:
+                if pivot.unpivot and (col_types := self._get_unpivot_column_types(pivot, selects)):
+                    selects[pivot.alias] = col_types
+
             self._scope_selects[scope] = selects
 
         return self._scope_selects[scope]
@@ -419,7 +423,15 @@ def _annotate_expression(
                         source_scope = source_scope.parent
 
                 if isinstance(source, exp.Table):
-                    self._set_type(expr, self.schema.get_column_type(source, expr))
+                    table_col_type = self.schema.get_column_type(source, expr)
+                    if table_col_type.is_type(exp.DType.UNKNOWN) and source.args.get("pivots"):
+                        table_col_type = (
+                            self._get_scope_selects(source_scope or scope)
+                            .get(expr.table, {})
+                            .get(expr.name)
+                            or table_col_type
+                        )
+                    self._set_type(expr, table_col_type)
                 elif source and source_scope:
                     col_type = (
                         self._get_scope_selects(source_scope).get(expr.table, {}).get(expr.name)
@@ -430,6 +442,16 @@ def _annotate_expression(
                         self._set_type(expr, source.expression.type)
                     else:
                         self._set_type(expr, exp.DType.UNKNOWN)
+                elif (
+                    not source
+                    and scope.pivots
+                    and (
+                        pivot_type := self._get_scope_selects(scope)
+                        .get(expr.table, {})
+                        .get(expr.name)
+                    )
+                ):
+                    self._set_type(expr, pivot_type)
                 else:
                     self._set_type(expr, exp.DType.UNKNOWN)
 
@@ -591,6 +613,39 @@ def _get_setop_column_types(
         self._setop_column_types[setop_id] = col_types
         return col_types
 
+    def _get_unpivot_column_types(
+        self, pivot: exp.Pivot, selects: dict[str, dict[str, t.Any]]
+    ) -> dict[str, t.Any]:
+        src_types = selects.get(parent.alias_or_name, {}) if (parent := pivot.parent) else {}
+        new_types: dict[str, t.Any] = {}
+
+        for field in pivot.fields:
+            field_col = field.this
+            first = seq_get(field.expressions, 0)
+
+            if isinstance(first, exp.PivotAlias) and (alias_node := first.args.get("alias")):
+                new_types[field_col.name] = alias_node.type
+                in_src = first.this
+            else:
+                new_types[field_col.name] = exp.DType.VARCHAR.into_expr()
+                in_src = first
+
+            in_cols = in_src.expressions if isinstance(in_src, exp.Tuple) else [in_src]
+            val_expr = seq_get(pivot.expressions, 0)
+            val_cols = val_expr.expressions if isinstance(val_expr, exp.Tuple) else [val_expr]
+            for val_col, in_col in zip(val_cols, in_cols):
+                new_types[val_col.output_name] = (
+                    src_types.get(in_col.output_name)
+                    if in_col.is_type(exp.DType.UNKNOWN)
+                    else in_col.type
+                )
+
+        return {
+            name: type_
+            for name in pivot.output_columns(src_types)
+            if (type_ := new_types.get(name) or src_types.get(name))
+        }
+
     def _annotate_binary(self, expression: B) -> B:
         left, right = expression.left, expression.right
         if not left or not right:

sqlglot/optimizer/annotate_types.py

@@ -15,7 +15,7 @@
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
-    from collections.abc import Iterator, Iterable, Sequence
+    from collections.abc import Iterable
 
 
 def qualify_columns(
@@ -181,49 +181,6 @@ def _separate_pseudocolumns(scope: Scope, pseudocolumns: set[str]) -> None:
         scope.clear_cache()
 
 
-def _pivot_output_columns(pivot: exp.Pivot, pre_pivot_columns: Sequence[str]) -> list[str]:
-    """Compute the columns exposed after a (UN)PIVOT, given its pre-pivot source columns.
-
-    Returns an empty list for degenerate pivots (no IN-list or no output names) so callers
-    can fall through to their non-pivot handling.
-    """
-    if pivot.unpivot:
-        excluded = {
-            c.output_name
-            for field in pivot.fields
-            if isinstance(field, exp.In)
-            for e in field.expressions
-            for c in e.find_all(exp.Column)
-        }
-        outputs = [i.name for i in _unpivot_columns(pivot)]
-    else:
-        excluded = {c.output_name for c in pivot.find_all(exp.Column)}
-        outputs = [c.output_name for c in pivot.args.get("columns") or []]
-        if not outputs:
-            outputs = [c.alias_or_name for c in pivot.expressions]
-
-    if not excluded or not outputs:
-        return []
-
-    return [c for c in pre_pivot_columns if c not in excluded] + outputs
-
-
-def _unpivot_columns(unpivot: exp.Pivot) -> Iterator[exp.Identifier]:
-    name_columns = [
-        field.this
-        for field in unpivot.fields
-        if isinstance(field, exp.In) and isinstance(field.this, exp.Identifier)
-    ]
-    value_columns = (
-        ident
-        for e in unpivot.expressions
-        for ident in (e.expressions if isinstance(e, exp.Tuple) else [e])
-        if isinstance(ident, exp.Identifier)
-    )
-
-    return itertools.chain(name_columns, value_columns)
-
-
 def _pop_table_column_aliases(derived_tables: Iterable[exp.Expr]) -> None:
     """
     Remove table column aliases.
@@ -638,7 +595,7 @@ def _qualify_columns(
             if isinstance(column_source, exp.Table) and (
                 pivots := column_source.args.get("pivots")
             ):
-                source_columns = _pivot_output_columns(pivots[0], source_columns)
+                source_columns = pivots[0].output_columns(source_columns)
             if (
                 not allow_partial_qualification
                 and source_columns
@@ -879,7 +836,7 @@ def _expand_stars(
             replaced_columns = replace_columns.get(table_id, {})
 
             if pivot:
-                pivot_columns = pivot.alias_column_names or _pivot_output_columns(pivot, columns)
+                pivot_columns = pivot.alias_column_names or pivot.output_columns(columns)
 
                 if pivot_columns:
                     new_selections.extend(