fix(optimizer)!: query schema directly when type annotation fails for processing UNNEST source

Author: fivetran-BradfordPaskewitz

sqlglot/optimizer/resolver.py

@@ -144,8 +144,22 @@ def get_source_columns(self, name: str, only_visible: bool = False) -> t.Sequenc
                 # in bigquery, unnest structs are automatically scoped as tables, so you can
                 # directly select a struct field in a query.
                 # this handles the case where the unnest is statically defined.
-                if self.dialect.UNNEST_COLUMN_ONLY:
-                    if source.expression.is_type(exp.DataType.Type.STRUCT):
+                if self.dialect.UNNEST_COLUMN_ONLY and isinstance(source.expression, exp.Unnest):
+                    unnest_type = source.expression.type
+
+                    # if type is not annotated yet, try to get it from the schema
+                    if not unnest_type or unnest_type.is_type(exp.DataType.Type.UNKNOWN):
+                        unnest_expr = seq_get(source.expression.expressions, 0)
+                        if isinstance(unnest_expr, exp.Column) and self.scope.parent:
+                            unnest_type = self._get_unnest_column_type(unnest_expr)
+
+                    # check if unnesting an ARRAY of STRUCTs - extract struct field names
+                    if unnest_type and unnest_type.is_type(exp.DataType.Type.ARRAY):
+                        element_types = unnest_type.expressions
+                        if element_types and element_types[0].is_type(exp.DataType.Type.STRUCT):
+                            for field in element_types[0].expressions:  # type: ignore
+                                columns.append(field.name)
+                    elif 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):
@@ -299,3 +313,66 @@ def _get_unambiguous_columns(
                 unambiguous_columns[column] = table
 
         return unambiguous_columns
+
+    def _get_unnest_column_type(self, column: exp.Column) -> t.Optional[exp.DataType]:
+        """
+        Get the type of a column being unnested, tracing through CTEs/subqueries to find the base table.
+
+        Args:
+            column: The column expression being unnested.
+
+        Returns:
+            The DataType of the column, or None if not found.
+        """
+        # start from parent scope and trace through sources to find the actual table
+        scope = self.scope.parent
+        if not scope:
+            return None
+
+        # try each source in the parent scope to find which one contains this column
+        for source_name in scope.sources:
+            source = scope.sources[source_name]
+            col_type: t.Optional[exp.DataType]
+
+            if isinstance(source, exp.Table):
+                # found a base table - get the column type from schema
+                col_type = self.schema.get_column_type(source, column)
+                if col_type and not col_type.is_type(exp.DataType.Type.UNKNOWN):
+                    return col_type
+            elif isinstance(source, Scope):
+                # CTE or subquery - recursively check its sources
+                col_type = self._get_column_type_from_scope(source, column.name)
+                if col_type:
+                    return col_type
+
+        return None
+
+    def _get_column_type_from_scope(self, scope: Scope, col_name: str) -> t.Optional[exp.DataType]:
+        """
+        Recursively find a column's type by tracing through nested scopes to the base table.
+
+        Args:
+            scope: The scope to search.
+            col_name: The column name to find.
+
+        Returns:
+            The DataType of the column, or None if not found.
+        """
+        for source_name in scope.sources:
+            source = scope.sources[source_name]
+            col_type: t.Optional[exp.DataType]
+
+            if isinstance(source, exp.Table):
+                # found a base table - try to get the column type
+                col_type = self.schema.get_column_type(
+                    source, exp.Column(this=exp.to_identifier(col_name))
+                )
+                if col_type and not col_type.is_type(exp.DataType.Type.UNKNOWN):
+                    return col_type
+            elif isinstance(source, Scope):
+                # nested scope - recurse
+                col_type = self._get_column_type_from_scope(source, col_name)
+                if col_type:
+                    return col_type
+
+        return None

tests/test_optimizer.py

@@ -143,6 +143,9 @@ def setUp(self):
             "t_bool": {
                 "a": "BOOLEAN",
             },
+            "table1": {
+                "credits": "ARRAY<STRUCT<name STRING, amount FLOAT64, type STRING>>",
+            },
         }
 
     def check_file(
@@ -516,6 +519,48 @@ def test_qualify_columns(self, logger):
             "SELECT a.b_id AS b_id FROM a AS a JOIN b AS b ON a.b_id = b.b_id JOIN c AS c ON b.b_id = c.b_id JOIN d AS d ON b.d_id = d.d_id",
         )
 
+        self.assertEqual(
+            optimizer.qualify.qualify(
+                parse_one(
+                    """
+                    SELECT
+                      (SELECT SUM(c.amount)
+                       FROM UNNEST(credits) AS c
+                       WHERE type != 'promotion') as total
+                    FROM billing
+                    """,
+                    read="bigquery",
+                ),
+                schema={"billing": {"credits": "ARRAY<STRUCT<amount FLOAT64, type STRING>>"}},
+                dialect="bigquery",
+            ).sql(dialect="bigquery"),
+            "SELECT (SELECT SUM(`c`.`amount`) AS `_col_0` FROM UNNEST(`billing`.`credits`) AS `c` WHERE `type` <> 'promotion') AS `total` FROM `billing` AS `billing`",
+        )
+
+        self.assertEqual(
+            optimizer.qualify.qualify(
+                parse_one(
+                    """
+                    WITH cte AS (SELECT * FROM base_table)
+                    SELECT
+                      (SELECT SUM(item.price)
+                       FROM UNNEST(items) AS item
+                       WHERE category = 'electronics') as electronics_total
+                    FROM cte
+                    """,
+                    read="bigquery",
+                ),
+                schema={
+                    "base_table": {
+                        "id": "INT64",
+                        "items": "ARRAY<STRUCT<price FLOAT64, category STRING>>",
+                    }
+                },
+                dialect="bigquery",
+            ).sql(dialect="bigquery"),
+            "WITH `cte` AS (SELECT `base_table`.`id` AS `id`, `base_table`.`items` AS `items` FROM `base_table` AS `base_table`) SELECT (SELECT SUM(`item`.`price`) AS `_col_0` FROM UNNEST(`cte`.`items`) AS `item` WHERE `category` = 'electronics') AS `electronics_total` FROM `cte` AS `cte`",
+        )
+
         self.check_file(
             "qualify_columns",
             qualify_columns,