fix(optimizer): EXPLODE qualify and annotate (#7549)

* fix(optimizer): EXPLODE qualify and annotate

* change tests

* ref walrus

Author: geooo109

sqlglot/optimizer/annotate_types.py

@@ -300,14 +300,23 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]:
 
                     alias_column_names = expression.alias_column_names
 
-                    if (
-                        isinstance(expression, exp.Unnest)
-                        and expression.type
-                        and expression.type.is_type(exp.DType.STRUCT)
+                    if isinstance(expression, exp.Unnest):
+                        exp_type = expression.type
+                    elif isinstance(expression, exp.Lateral) and isinstance(
+                        expression.this, exp.Explode
                     ):
+                        exp_type = expression.this.type
+                    else:
+                        exp_type = None
+
+                    struct_type = (
+                        exp_type if exp_type and exp_type.is_type(exp.DType.STRUCT) else None
+                    )
+
+                    if struct_type:
                         selects[name] = {
                             col_def.name: t.cast(t.Union[exp.DataType, exp.DType], col_def.kind)
-                            for col_def in expression.type.expressions
+                            for col_def in struct_type.expressions
                             if isinstance(col_def, exp.ColumnDef) and col_def.kind
                         }
                     else:

sqlglot/optimizer/resolver.py

@@ -142,33 +142,34 @@ def get_source_columns(self, name: str, only_visible: bool = False) -> Sequence[
             if isinstance(source, exp.Table):
                 columns = self.schema.column_names(source, only_visible)
             elif isinstance(source, Scope) and isinstance(
-                source.expression, (exp.Values, exp.Unnest)
+                source_expr := source.expression, (exp.Values, exp.Unnest, exp.Lateral)
             ):
-                columns = source.expression.named_selects
+                columns = source_expr.named_selects
 
                 # 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 and isinstance(source.expression, exp.Unnest):
-                    unnest = source.expression
-
-                    # if type is not annotated yet, try to get it from the schema
-                    if not unnest.type or unnest.type.is_type(exp.DType.UNKNOWN):
-                        unnest_expr = seq_get(unnest.expressions, 0)
+                if self.dialect.UNNEST_COLUMN_ONLY and isinstance(source_expr, exp.Unnest):
+                    if not source_expr.type or source_expr.type.is_type(exp.DType.UNKNOWN):
+                        unnest_expr = seq_get(source_expr.expressions, 0)
                         if isinstance(unnest_expr, exp.Column) and self.scope.parent:
-                            col_type = self._get_unnest_column_type(unnest_expr)
-                            # extract element type if it's an ARRAY
+                            col_type = self._get_unnest_column_type(unnest_expr, self.scope.parent)
                             if col_type and col_type.is_type(exp.DType.ARRAY):
                                 element_types = col_type.expressions
                                 if element_types:
-                                    unnest.type = element_types[0].copy()
-                            else:
-                                if col_type:
-                                    unnest.type = col_type.copy()
-                    # check if the result type is a STRUCT - extract struct field names
-                    if unnest.is_type(exp.DType.STRUCT):
-                        for k in unnest.type.expressions:  # type: ignore
-                            columns.append(k.name)
+                                    source_expr.type = element_types[0].copy()
+                            elif col_type:
+                                source_expr.type = col_type.copy()
+
+                    columns.extend(self._struct_field_names(source_expr.type))
+                elif isinstance(source_expr, exp.Lateral) and isinstance(
+                    source_expr.this, exp.Explode
+                ):
+                    explode_col = source_expr.this.this
+
+                    if isinstance(explode_col, exp.Column) and source.parent:
+                        col_type = self._get_unnest_column_type(explode_col, source.parent)
+                        columns.extend(self._struct_field_names(col_type))
             elif isinstance(source, Scope) and isinstance(source.expression, exp.SetOperation):
                 columns = self.get_source_columns_from_set_op(source.expression)
             else:
@@ -338,19 +339,27 @@ def _get_unambiguous_columns(
 
         return unambiguous_columns
 
-    def _get_unnest_column_type(self, column: exp.Column) -> exp.DataType | None:
+    def _struct_field_names(self, col_type: exp.DataType | None) -> list[str]:
+        if col_type and col_type.is_type(exp.DType.ARRAY):
+            col_type = seq_get(col_type.expressions, 0)
+
+        return (
+            [k.name for k in col_type.expressions]
+            if col_type and col_type.is_type(exp.DType.STRUCT)
+            else []
+        )
+
+    def _get_unnest_column_type(self, column: exp.Column, scope: Scope) -> exp.DataType | None:
         """
-        Get the type of a column being unnested, tracing through CTEs/subqueries to find the base table.
+        Get the type of a column being unnested/exploded, tracing through CTEs/subqueries to find the base table.
 
         Args:
-            column: The column expression being unnested.
+            column: The column expression being unnested/exploded.
+            scope: The scope to resolve the column in.
 
         Returns:
             The DataType of the column, or None if not found.
         """
-        scope = self.scope.parent
-        assert scope
-
         # if column is qualified, use that table, otherwise disambiguate using the resolver
         if column.table:
             table_name = column.table

tests/test_optimizer.py

@@ -651,6 +651,45 @@ def test_validate_columns(self):
             )
             optimizer.qualify_columns.validate_qualify_columns(qualified)
 
+        schema = {"my_table": {"items": "ARRAY<STRUCT<name STRING, age INT>>"}}
+        expression = annotate_types(
+            optimizer.qualify.qualify(
+                parse_one(
+                    "SELECT ci.name, ci.age FROM my_table LATERAL VIEW EXPLODE(items) ci AS ci",
+                    read="spark",
+                ),
+                schema=schema,
+                dialect="spark",
+            ),
+            schema=schema,
+            dialect="spark",
+        )
+        self.assertEqual(
+            expression.sql(dialect="spark"),
+            "SELECT `ci`.`name` AS `name`, `ci`.`age` AS `age` FROM `my_table` AS `my_table` LATERAL VIEW EXPLODE(`my_table`.`items`) ci AS `ci`",
+        )
+        self.assertEqual(expression.selects[0].type, exp.DataType.build("STRING", dialect="spark"))
+        self.assertEqual(expression.selects[1].type, exp.DataType.build("INT", dialect="spark"))
+
+        schema = {"my_table": {"items": "ARRAY<STRUCT<amount FLOAT, type STRING>>"}}
+        expression = annotate_types(
+            optimizer.qualify.qualify(
+                parse_one(
+                    "SELECT (SELECT SUM(ci.amount) FROM my_table LATERAL VIEW EXPLODE(items) ci AS ci WHERE ci.type = 'promotion') AS total FROM my_table",
+                    read="spark",
+                ),
+                schema=schema,
+                dialect="spark",
+            ),
+            schema=schema,
+            dialect="spark",
+        )
+        self.assertEqual(
+            expression.sql(dialect="spark"),
+            "SELECT (SELECT SUM(`ci`.`amount`) AS `_col_0` FROM `my_table` AS `my_table` LATERAL VIEW EXPLODE(`my_table`.`items`) ci AS `ci` WHERE `ci`.`type` = 'promotion') AS `total` FROM `my_table` AS `my_table`",
+        )
+        self.assertEqual(expression.selects[0].type, exp.DataType.build("DOUBLE", dialect="spark"))
+
     def test_qualify_columns__with_invisible(self):
         schema = MappingSchema(self.schema, {"x": {"a"}, "y": {"b"}, "z": {"b"}})
         self.check_file("qualify_columns__with_invisible", qualify_columns, schema=schema)