Fix(bigquery): unnest with structs closes #2125

tobymao · Aug 30, 2023 · d1ccb03 · d1ccb03
1 parent cc0a6e2
commit d1ccb03
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 6 deletions.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -94,14 +94,20 @@ def _unqualify_unnest(expression: exp.Expression) -> exp.Expression:
 
     These are added by the optimizer's qualify_column step.
     """
-    from sqlglot.optimizer.scope import Scope
+    from sqlglot.optimizer.scope import find_all_in_scope
 
     if isinstance(expression, exp.Select):
-        for unnest in expression.find_all(exp.Unnest):
-            if isinstance(unnest.parent, (exp.From, exp.Join)) and unnest.alias:
-                for column in Scope(expression).find_all(exp.Column):
-                    if column.table == unnest.alias:
-                        column.set("table", None)
+        unnest_aliases = {
+            unnest.alias
+            for unnest in find_all_in_scope(expression, exp.Unnest)
+            if isinstance(unnest.parent, (exp.From, exp.Join))
+        }
+        if unnest_aliases:
+            for column in expression.find_all(exp.Column):
+                if column.table in unnest_aliases:
+                    column.set("table", None)
+                elif column.db in unnest_aliases:
+                    column.set("db", None)
 
     return expression
 

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -6,8 +6,33 @@
 
 class TestBigQuery(Validator):
     dialect = "bigquery"
+    maxDiff = None
 
     def test_bigquery(self):
+        self.validate_all(
+            """SELECT
+  `u`.`harness_user_email` AS `harness_user_email`,
+  `d`.`harness_user_id` AS `harness_user_id`,
+  `harness_account_id` AS `harness_account_id`
+FROM `analytics_staging`.`stg_mongodb__users` AS `u`, UNNEST(`u`.`harness_cluster_details`) AS `d`, UNNEST(`d`.`harness_account_ids`) AS `harness_account_id`
+WHERE
+  NOT `harness_account_id` IS NULL""",
+            read={
+                "": """
+                SELECT
+                  "u"."harness_user_email" AS "harness_user_email",
+                  "_q_0"."d"."harness_user_id" AS "harness_user_id",
+                  "_q_1"."harness_account_id" AS "harness_account_id"
+                FROM
+                  "analytics_staging"."stg_mongodb__users" AS "u",
+                  UNNEST("u"."harness_cluster_details") AS "_q_0"("d"),
+                  UNNEST("_q_0"."d"."harness_account_ids") AS "_q_1"("harness_account_id")
+                WHERE
+                  NOT "_q_1"."harness_account_id" IS NULL
+                """
+            },
+            pretty=True,
+        )
         with self.assertRaises(TokenError):
             transpile("'\\'", read="bigquery")