fix(optimizer)!: postgres qualify GENERATE_SERIES and table projection (#6373)

Author: geooo109

sqlglot/dialects/dialect.py

@@ -671,6 +671,14 @@ class Dialect(metaclass=_Dialect):
     Whether JSON_EXTRACT_SCALAR returns null if a non-scalar value is selected.
     """
 
+    DEFAULT_FUNCTIONS_COLUMN_NAMES: t.Dict[t.Type[exp.Func], t.Union[str, t.Tuple[str, ...]]] = {}
+    """
+    Maps function expressions to their default output column name(s).
+
+    For example, in Postgres, generate_series function outputs a column named "generate_series" by default,
+    so we map the ExplodingGenerateSeries expression to "generate_series" string.
+    """
+
     # --- Autofilled ---
 
     tokenizer_class = Tokenizer

sqlglot/dialects/postgres.py

@@ -297,6 +297,11 @@ class Postgres(Dialect):
     NULL_ORDERING = "nulls_are_large"
     TIME_FORMAT = "'YYYY-MM-DD HH24:MI:SS'"
     TABLESAMPLE_SIZE_IS_PERCENT = True
+    TABLES_REFERENCEABLE_AS_COLUMNS = True
+
+    DEFAULT_FUNCTIONS_COLUMN_NAMES = {
+        exp.ExplodingGenerateSeries: "generate_series",
+    }
 
     TIME_MAPPING = {
         "d": "%u",  # 1-based day of week

sqlglot/optimizer/qualify_columns.py

@@ -598,7 +598,7 @@ def _qualify_columns(
                 and len(column.parts) == 1
                 and column_name in scope.selected_sources
             ):
-                # BigQuery allows tables to be referenced as columns, treating them as structs
+                # BigQuery and Postgres allow tables to be referenced as columns, treating them as structs/records
                 scope.replace(column, exp.TableColumn(this=column.this))
 
     for pivot in scope.pivots:

sqlglot/optimizer/qualify_tables.py

@@ -4,7 +4,7 @@
 
 from sqlglot import exp
 from sqlglot.dialects.dialect import Dialect, DialectType
-from sqlglot.helper import name_sequence, seq_get
+from sqlglot.helper import name_sequence, seq_get, ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.scope import Scope, traverse_scope
 
@@ -79,6 +79,7 @@ def _set_alias(
         target_alias: t.Optional[str] = None,
         scope: t.Optional[Scope] = None,
         normalize: bool = False,
+        columns: t.Optional[t.List[t.Union[str, exp.Identifier]]] = None,
     ) -> None:
         alias = expression.args.get("alias") or exp.TableAlias()
 
@@ -96,6 +97,10 @@ def _set_alias(
         quoted = True if canonicalize_table_aliases or not target_alias else None
 
         alias.set("this", exp.to_identifier(new_alias_name, quoted=quoted))
+
+        if columns:
+            alias.set("columns", [exp.to_identifier(c) for c in columns])
+
         expression.set("alias", alias)
 
         if scope:
@@ -132,11 +137,27 @@ def _set_alias(
                 if pivot := seq_get(source.args.get("pivots") or [], 0):
                     name = source.name
 
+                table_this = source.this
+                table_alias = source.args.get("alias")
+                function_columns: t.List[t.Union[str, exp.Identifier]] = []
+                if isinstance(table_this, exp.Func):
+                    if not table_alias:
+                        function_columns = ensure_list(
+                            dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES.get(type(table_this))
+                        )
+                    elif columns := table_alias.columns:
+                        function_columns = columns
+                    elif type(table_this) in dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES:
+                        function_columns = ensure_list(source.alias_or_name)
+                        source.set("alias", None)
+                        name = None
+
                 _set_alias(
                     source,
                     canonical_aliases,
                     target_alias=name or source.name or None,
                     normalize=True,
+                    columns=function_columns,
                 )
 
                 source_fqn = ".".join(p.name for p in source.parts)

tests/fixtures/optimizer/qualify_columns.sql

@@ -287,6 +287,11 @@ SELECT DATE_TRUNC(t.col1, WEEK(MONDAY)) AS _col_0, t.col2 AS col2 FROM t AS t;
 SELECT first, second FROM (SELECT 'val' AS col, STACK(2, 1, 2, 3) AS (first, second)) AS tbl;
 SELECT tbl.first AS first, tbl.second AS second FROM (SELECT 'val' AS col, STACK(2, 1, 2, 3) AS (first, second)) AS tbl;
 
+# execute: false
+# dialect: postgres
+WITH t AS (SELECT 1 AS c) SELECT t FROM t;
+WITH t AS (SELECT 1 AS c) SELECT t AS _col_0 FROM t AS t;
+
 --------------------------------------
 -- Derived tables
 --------------------------------------

tests/fixtures/optimizer/qualify_tables.sql

@@ -260,3 +260,24 @@ SELECT * FROM c.db.x AS "_1" WHERE "_1".a = (SELECT SUM("_0".c) AS c FROM c.db.y
 # canonicalize_table_aliases: true
 SELECT t.foo FROM t AS t, (SELECT t.bar FROM t AS t);
 SELECT "_2".foo FROM c.db.t AS "_2", (SELECT "_0".bar FROM c.db.t AS "_0") AS "_1";
+
+# title: Qualify GENERATE_SERIES with its default column generate_series
+# dialect: postgres
+SELECT generate_series FROM GENERATE_SERIES(1,2);
+SELECT generate_series FROM GENERATE_SERIES(1, 2) AS "_0"(generate_series);
+
+# title: Qualify GENERATE_SERIES with alias by wrapping it
+# dialect: postgres
+SELECT g FROM GENERATE_SERIES(1,2) AS g;
+SELECT g FROM GENERATE_SERIES(1, 2) AS "_0"(g);
+
+# title: Qualify GENERATE_SERIES with alias on table and columns
+# dialect: postgres
+SELECT g FROM GENERATE_SERIES(1,2) AS t(g);
+SELECT g FROM GENERATE_SERIES(1, 2) AS t(g);
+
+# title: Qualify GENERATE_SERIES with explicit column and canonicalize_table_aliases
+# dialect: postgres
+# canonicalize_table_aliases: true
+SELECT g FROM GENERATE_SERIES(1,2) AS t(g);
+SELECT g FROM GENERATE_SERIES(1, 2) AS "_0"(g);