Repository: tobymao/sqlglot
Pull request: Fix(optimizer)!: resolve table "columns" in bigquery that produce structs
Fixes #5207

I've described the issue in this comment: https://github.com/tobymao/sqlglot/issues/5207#issuecomment-2961924796. Let me know what you think.

Problem statement:
Optimizer fails in some specific cases
Sqlglot version: 26.25.3

Reading dialect: BigQuery
Writing dialect: N/A

**Fully reproducible code snippet**

```python
from loguru import logger
from sqlglot import Expression
from sqlglot.dialects.dialect import Dialect
from sqlglot.errors import OptimizeError
from sqlglot.optimizer.qualify_columns import (
  pushdown_cte_alias_columns,
  qualify_columns,
  quote_identifiers,
  validate_qualify_columns,
)
from sqlglot.optimizer.qualify_tables import qualify_tables
from sqlglot.schema import ensure_schema


def qualify(expression: Expression, dialect="bigquery"):
  schema = ensure_schema(None, dialect=dialect)
  expression = qualify_tables(expression, schema=schema)

  if Dialect.get_or_raise(dialect).PREFER_CTE_ALIAS_COLUMN:
    expression = pushdown_cte_alias_columns(expression)

  try:
    expression = qualify_columns(
      expression,
      schema,
      expand_alias_refs=True,
      expand_stars=True,
      infer_schema=None,
    )
  except OptimizeError as e:
    logger.error(f"Exception while qualifying query: [{e}].")
    raise

  expression = quote_identifiers(expression, dialect=dialect, identify=True)
  validate_qualify_columns(expression)

  return expression

ast = qualify(sqlglot.parse_one(sql, dialect="bigquery"))
```

This sql code fails:

```sql
WITH some_cte AS (
    SELECT 
      foo,
      bar
    FROM `database.schema.table`
)

SELECT
  *,
  TO_JSON_STRING(some_cte) AS serialized_row
FROM some_cte;
```

With the error message
```
File "(...)/sqlglot/optimizer/qualify_columns.py", line 126, in validate_qualify_columns
  sqlglot.errors.OptimizeError: Column '"some_cte"' could not be resolved
```

And here is a slightly different example that does work:

```sql
WITH some_cte AS (
    SELECT 
      *
    FROM `database.schema.table`
)

SELECT
  *,
  TO_JSON_STRING(some_cte) AS serialized_row
FROM some_cte;
```

Both examples are valid in bigquery.

Hints from issue discussion:
Thanks, this is a legit bug. An easier way to see what's happening is:

```python
from sqlglot import parse_one
from sqlglot.optimizer.qualify import qualify

ast = parse_one("with t as (select 1) select to_json_string(t) from t", dialect="bigquery")
qualify(ast, dialect="bigquery")
# sqlglot.errors.OptimizeError: Column '"t"' could not be resolved
```

Both this and your example use a table as a column, which is valid in BigQuery because it treats it as a struct with the table's columns for its fields. Sqlglot's optimizer doesn't fully support this behavior today, but we'll take a look.

Commit message: Refactor

Patch to review:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 21db7f8fba..8521d8cb76 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -7045,6 +7045,12 @@ class Semicolon(Expression):
     arg_types = {}
 
 
+# BigQuery allows SELECT t FROM t and treats the projection as a struct value. This expression
+# type is intended to be constructed by qualify so that we can properly annotate its type later
+class TableAsStruct(Expression):
+    pass
+
+
 def _norm_arg(arg):
     return arg.lower() if type(arg) is str else arg
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 458256eb5b..4b036da41e 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -201,6 +201,7 @@ class Generator(metaclass=_Generator):
         exp.StreamingTableProperty: lambda *_: "STREAMING",
         exp.StrictProperty: lambda *_: "STRICT",
         exp.SwapTable: lambda self, e: f"SWAP WITH {self.sql(e, 'this')}",
+        exp.TableAsStruct: lambda self, e: self.sql(e.this),
         exp.Tags: lambda self, e: f"TAG ({self.expressions(e, flat=True)})",
         exp.TemporaryProperty: lambda *_: "TEMPORARY",
         exp.TitleColumnConstraint: lambda self, e: f"TITLE {self.sql(e, 'this')}",
diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py
index d460b1cc0f..f3bc1f1aaf 100644
--- a/sqlglot/optimizer/annotate_types.py
+++ b/sqlglot/optimizer/annotate_types.py
@@ -12,7 +12,7 @@
     seq_get,
 )
 from sqlglot.optimizer.scope import Scope, traverse_scope
-from sqlglot.schema import Schema, ensure_schema
+from sqlglot.schema import MappingSchema, Schema, ensure_schema
 from sqlglot.dialects.dialect import Dialect
 
 if t.TYPE_CHECKING:
@@ -290,9 +290,48 @@ def annotate_scope(self, scope: Scope) -> None:
                 elif isinstance(source.expression, exp.Unnest):
                     self._set_type(col, source.expression.type)
 
+        if isinstance(self.schema, MappingSchema):
+            for pseudocolumn in scope.pseudocolumns:
+                if isinstance(pseudocolumn, exp.TableAsStruct):
+                    source = scope.sources.get(pseudocolumn.name)
+
+                    if isinstance(source, exp.Table):
+                        schema = self.schema.find(
+                            source, raise_on_missing=False, ensure_data_types=True
+                        )
+                        if not isinstance(schema, dict):
+                            continue
+
+                        struct_type = exp.DataType(
+                            this=exp.DataType.Type.STRUCT,
+                            expressions=[
+                                exp.ColumnDef(this=exp.to_identifier(c), kind=kind)
+                                for c, kind in schema.items()
+                            ],
+                            nested=True,
+                        )
+                        self._set_type(pseudocolumn, struct_type)
+                    elif (
+                        isinstance(source, Scope)
+                        and isinstance(source.expression, exp.Query)
+                        and source.expression.is_type(exp.DataType.Type.STRUCT)
+                    ):
+                        self._set_type(pseudocolumn, source.expression.type)
+
         # Then (possibly) annotate the remaining expressions in the scope
         self._maybe_annotate(scope.expression)
 
+        if self.schema.dialect == "bigquery" and isinstance(scope.expression, exp.Query):
+            struct_type = exp.DataType(
+                this=exp.DataType.Type.STRUCT,
+                expressions=[
+                    exp.ColumnDef(this=exp.to_identifier(select.output_name), kind=select.type)
+                    for select in scope.expression.selects
+                ],
+                nested=True,
+            )
+            self._set_type(scope.expression, struct_type)
+
     def _maybe_annotate(self, expression: E) -> E:
         if id(expression) in self._visited:
             return expression  # We've already inferred the expression's type
diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py
index dedd973d8f..85a8e0ba07 100644
--- a/sqlglot/optimizer/qualify_columns.py
+++ b/sqlglot/optimizer/qualify_columns.py
@@ -529,6 +529,13 @@ def _qualify_columns(scope: Scope, resolver: Resolver, allow_partial_qualificati
             column_table = resolver.get_table(column_name)
             if column_table:
                 column.set("table", column_table)
+            elif (
+                resolver.schema.dialect == "bigquery"
+                and len(column.parts) == 1
+                and column_name in scope.selected_sources
+            ):
+                # BigQuery allows tables to be referenced as columns, treating them as structs
+                scope.replace(column, exp.TableAsStruct(this=column.this))
 
     for pivot in scope.pivots:
         for column in pivot.find_all(exp.Column):
diff --git a/sqlglot/optimizer/scope.py b/sqlglot/optimizer/scope.py
index df6072c4b1..ae54d23ea7 100644
--- a/sqlglot/optimizer/scope.py
+++ b/sqlglot/optimizer/scope.py
@@ -88,6 +88,7 @@ def __init__(
     def clear_cache(self):
         self._collected = False
         self._raw_columns = None
+        self._pseudocolumns = None
         self._stars = None
         self._derived_tables = None
         self._udtfs = None
@@ -125,6 +126,7 @@ def _collect(self):
         self._derived_tables = []
         self._udtfs = []
         self._raw_columns = []
+        self._pseudocolumns = []
         self._stars = []
         self._join_hints = []
         self._semi_anti_join_tables = set()
@@ -156,6 +158,8 @@ def _collect(self):
                 self._derived_tables.append(node)
             elif isinstance(node, exp.UNWRAPPED_QUERIES):
                 self._subqueries.append(node)
+            elif isinstance(node, exp.TableAsStruct):
+                self._pseudocolumns.append(node)
 
         self._collected = True
 
@@ -309,6 +313,13 @@ def columns(self):
 
         return self._columns
 
+    @property
+    def pseudocolumns(self):
+        if self._columns is None:
+            self._ensure_collected()
+
+        return self._pseudocolumns
+
     @property
     def selected_sources(self):
         """
diff --git a/tests/fixtures/optimizer/qualify_columns.sql b/tests/fixtures/optimizer/qualify_columns.sql
index aaf7150faf..8dfc80786c 100644
--- a/tests/fixtures/optimizer/qualify_columns.sql
+++ b/tests/fixtures/optimizer/qualify_columns.sql
@@ -273,6 +273,11 @@ SELECT * FROM QUARTERLY_SALES AS QUARTERLY_SALES PIVOT(SUM(QUARTERLY_SALES.AMOUN
 SELECT PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY x) AS x FROM t;
 SELECT PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY t.x) AS x FROM t AS t;
 
+# execute: false
+# dialect: bigquery
+WITH t AS (SELECT 1 AS c) SELECT TO_JSON_STRING(t) FROM t;
+WITH t AS (SELECT 1 AS c) SELECT TO_JSON_STRING(t) AS _col_0 FROM t AS t;
+
 --------------------------------------
 -- Derived tables
 --------------------------------------
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index ac5e0f1f53..1b7e7d050c 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1580,3 +1580,31 @@ def test_manually_annotate_snowflake(self):
 
         self.assertTrue(annotated.selects[0].is_type("INT"))
         self.assertTrue(annotated.selects[1].is_type("VARCHAR"))
+
+    def test_annotate_table_as_struct_bigquery(self):
+        dialect = "bigquery"
+        schema = {"d": {"s": {"t": {"c1": "int64", "c2": "struct<f1 int64, f2 string>"}}}}
+
+        def _annotate(query: str) -> exp.Expression:
+            expression = parse_one(example_query, dialect=dialect)
+            qual = optimizer.qualify.qualify(expression, schema=schema, dialect=dialect)
+            return optimizer.annotate_types.annotate_types(qual, schema=schema, dialect=dialect)
+
+        example_query = "SELECT t FROM d.s.t"
+        annotated = _annotate(example_query)
+
+        self.assertTrue(
+            annotated.selects[0].is_type("STRUCT<c1 BIGINT, c2 STRUCT<f1 BIGINT, f2 TEXT>>")
+        )
+
+        example_query = "SELECT subq FROM (SELECT * from d.s.t) subq"
+        annotated = _annotate(example_query)
+
+        self.assertTrue(
+            annotated.selects[0].is_type("STRUCT<c1 BIGINT, c2 STRUCT<f1 BIGINT, f2 TEXT>>")
+        )
+
+        example_query = "WITH t AS (SELECT 1 AS c) SELECT t FROM t"
+        annotated = _annotate(example_query)
+
+        self.assertTrue(annotated.selects[0].is_type("STRUCT<c INt>"))


Code review comments:
diff_hunk: @@ -290,9 +290,48 @@ def annotate_scope(self, scope: Scope) -> None:
                 elif isinstance(source.expression, exp.Unnest):
                     self._set_type(col, source.expression.type)
 
+        if isinstance(self.schema, MappingSchema):
+            for pseudocolumn in scope.pseudocolumns:
+                if isinstance(pseudocolumn, exp.TableAsStruct):
+                    source = scope.sources.get(pseudocolumn.name)
+
+                    if isinstance(source, exp.Table):
+                        schema = self.schema.find(
+                            source, raise_on_missing=False, ensure_data_types=True
+                        )
+                        if not isinstance(schema, dict):
+                            continue
+
+                        struct_type = exp.DataType(
+                            this=exp.DataType.Type.STRUCT,
+                            expressions=[
+                                exp.ColumnDef(this=exp.to_identifier(c), kind=kind)
+                                for c, kind in schema.items()
+                            ],
+                            nested=True,
+                        )
+                        self._set_type(pseudocolumn, struct_type)
+                    elif (
+                        isinstance(source, Scope)
+                        and isinstance(source.expression, exp.Query)
+                        and source.expression.is_type(exp.DataType.Type.STRUCT)
+                    ):
+                        self._set_type(pseudocolumn, source.expression.type)
+
         # Then (possibly) annotate the remaining expressions in the scope
         self._maybe_annotate(scope.expression)
 
+        if self.schema.dialect == "bigquery" and isinstance(scope.expression, exp.Query):
+            struct_type = exp.DataType(
+                this=exp.DataType.Type.STRUCT,
+                expressions=[
+                    exp.ColumnDef(this=exp.to_identifier(select.output_name), kind=select.type)
+                    for select in scope.expression.selects
+                ],
+                nested=True,
+            )
+            self._set_type(scope.expression, struct_type)
path: sqlglot/optimizer/annotate_types.py
text: @user1:
What does this do? Could we have other projections in that `selects` list that do not belong to the table/struct?

@author:
BigQuery allows `SELECT t FROM t`, and `t` could be anything. It could be a physical table, a CTE, a subquery, etc. It can also depend on other nested objects, so we need to recursively infer the types of each query and propagate them "upwards".

TL;DR, this annotates the `Select`'s type in `SELECT 1 AS x, 2 AS y` to `STRUCT<x INT, y INT>` only for BigQuery so we can reuse this type elsewhere.

> Could we have other projections in that selects list that do not belong to the table/struct?

Not sure what you mean here.

@author:
I'll make it so if there's at least one unknown, the whole thing will be an unknown.

@author:
Done^.

diff_hunk: @@ -156,6 +158,8 @@ def _collect(self):
                 self._derived_tables.append(node)
             elif isinstance(node, exp.UNWRAPPED_QUERIES):
                 self._subqueries.append(node)
+            elif isinstance(node, exp.TableAsStruct):
+                self._pseudocolumns.append(node)
path: sqlglot/optimizer/scope.py
text: @user1:
Should we rename these as `table_as_columns` or something equivalent? I'd believe pseudocolumns would hint towards dialect peculiarities, INFORMATION_SCHEMA columns etc

@author:
Sure, we can narrow it down.

@author:
Renamed.