scope Snowflake positional refs to pivot source

Author: treysp

sqlglot/optimizer/qualify_columns.py

@@ -670,14 +670,28 @@ def _qualify_columns(
                 and position.is_int
             ):
                 # Preserve references whose source shape cannot be resolved losslessly.
+                scope_pivot = next(
+                    (pivot for pivot in scope.pivots if pivot.alias == column_table), None
+                )
+                if not scope_pivot:
+                    scope_pivot = next(
+                        (
+                            pivot
+                            for pivot in reversed(scope.pivots)
+                            if pivot.parent and pivot.parent.alias_or_name == column_table
+                        ),
+                        None,
+                    )
                 if (
                     pivots
-                    or scope.pivots
+                    or scope_pivot
                     or (isinstance(column_source, exp.Table) and column_source.alias_column_names)
                     or (isinstance(column_source, Scope) and column_source.outer_columns)
                     or not source_columns
                     or "*" in source_columns
                 ):
+                    if scope_pivot:
+                        column.set("table", exp.to_identifier(scope_pivot.alias))
                     continue
 
                 positional_columns = resolver.get_source_columns(column_table, only_visible=True)

tests/test_optimizer.py

@@ -1200,6 +1200,30 @@ def test_qualify_snowflake_positional_columns_preserve_unresolved(self):
                 ("P.$2 AS _COL_0",),
                 {},
             ),
+            (
+                "aliasless PIVOT",
+                "WITH t AS (SELECT 1 AS id, 2 AS v, 'a' AS k) "
+                "SELECT t.$2 FROM t PIVOT(SUM(v) FOR k IN ('a' AS a))",
+                ("_0.$2 AS _COL_0",),
+                {},
+            ),
+            (
+                "chained aliasless PIVOT",
+                "WITH t AS (SELECT 1 AS id, 2 AS v, 'a' AS k) "
+                "SELECT t.$1 FROM t PIVOT(SUM(v) FOR k IN ('a' AS a)) "
+                "UNPIVOT(v2 FOR k2 IN (a))",
+                ("T.$1 AS _COL_0",),
+                {},
+            ),
+            (
+                "unrelated PIVOT",
+                "WITH t AS (SELECT 1 AS x), "
+                "s AS (SELECT 1 AS id, 2 AS v, 'a' AS k) "
+                "SELECT t.$1, p.$2 FROM t CROSS JOIN "
+                "s PIVOT(SUM(v) FOR k IN ('a' AS a)) AS p",
+                ("T.X AS X", "P.$2 AS _COL_1"),
+                {},
+            ),
             (
                 "unknown source",
                 "SELECT t.$1 FROM source AS t",