fix(optimizer): Fix expansion of SELECT * REPLACE, RENAME (#3742)

* fix(optimizer): Fix expansion of SELECT * REPLACE, RENAME

* Refactor variable names

* PR Feedback 1

* Removing added empty line

Author: VaggelisD

sqlglot/optimizer/qualify_columns.py

@@ -513,7 +513,9 @@ def _expand_stars(
 
     new_selections = []
     except_columns: t.Dict[int, t.Set[str]] = {}
-    replace_columns: t.Dict[int, t.Dict[str, str]] = {}
+    replace_columns: t.Dict[int, t.Dict[str, exp.Alias]] = {}
+    rename_columns: t.Dict[int, t.Dict[str, str]] = {}
+
     coalesced_columns = set()
     dialect = resolver.schema.dialect
 
@@ -548,11 +550,13 @@ def _expand_stars(
             tables.extend(scope.selected_sources)
             _add_except_columns(expression, tables, except_columns)
             _add_replace_columns(expression, tables, replace_columns)
+            _add_rename_columns(expression, tables, rename_columns)
         elif expression.is_star:
             if not isinstance(expression, exp.Dot):
                 tables.append(expression.table)
                 _add_except_columns(expression.this, tables, except_columns)
                 _add_replace_columns(expression.this, tables, replace_columns)
+                _add_rename_columns(expression.this, tables, rename_columns)
             elif is_bigquery:
                 struct_fields = _expand_struct_stars(expression)
                 if struct_fields:
@@ -578,6 +582,8 @@ def _expand_stars(
 
             table_id = id(table)
             columns_to_exclude = except_columns.get(table_id) or set()
+            renamed_columns = rename_columns.get(table_id, {})
+            replaced_columns = replace_columns.get(table_id, {})
 
             if pivot:
                 if pivot_output_columns and pivot_exclude_columns:
@@ -606,10 +612,12 @@ def _expand_stars(
                         alias(exp.func("coalesce", *coalesce_args), alias=name, copy=False)
                     )
                 else:
-                    alias_ = replace_columns.get(table_id, {}).get(name, name)
-                    column = exp.column(name, table=table)
+                    alias_ = renamed_columns.get(name, name)
+                    selection_expr = replaced_columns.get(name) or exp.column(name, table=table)
                     new_selections.append(
-                        alias(column, alias_, copy=False) if alias_ != name else column
+                        alias(selection_expr, alias_, copy=False)
+                        if alias_ != name
+                        else selection_expr
                     )
 
     # Ensures we don't overwrite the initial selections with an empty list
@@ -631,15 +639,29 @@ def _add_except_columns(
         except_columns[id(table)] = columns
 
 
+def _add_rename_columns(
+    expression: exp.Expression, tables, rename_columns: t.Dict[int, t.Dict[str, str]]
+) -> None:
+    rename = expression.args.get("rename")
+
+    if not rename:
+        return
+
+    columns = {e.this.name: e.alias for e in rename}
+
+    for table in tables:
+        rename_columns[id(table)] = columns
+
+
 def _add_replace_columns(
-    expression: exp.Expression, tables, replace_columns: t.Dict[int, t.Dict[str, str]]
+    expression: exp.Expression, tables, replace_columns: t.Dict[int, t.Dict[str, exp.Alias]]
 ) -> None:
     replace = expression.args.get("replace")
 
     if not replace:
         return
 
-    columns = {e.this.name: e.alias for e in replace}
+    columns = {e.alias: e for e in replace}
 
     for table in tables:
         replace_columns[id(table)] = columns

sqlglot/optimizer/scope.py

@@ -284,6 +284,7 @@ def columns(self):
                             or column.name not in named_selects
                         )
                     )
+                    or (isinstance(ancestor, exp.Star) and not column.arg_key == "except")
                 ):
                     self._columns.append(column)
 

tests/fixtures/optimizer/qualify_columns.sql

@@ -385,14 +385,14 @@ WITH player AS (SELECT player.name, player.asset.info FROM players) SELECT * FRO
 WITH player AS (SELECT players.player.name AS name, players.player.asset.info AS info FROM players AS players) SELECT player.name AS name, player.info AS info FROM player AS player;
 
 --------------------------------------
--- Except and Replace
+-- Except, Replace, Rename
 --------------------------------------
 # execute: false
-SELECT * REPLACE(a AS d) FROM x;
+SELECT * RENAME(a AS d) FROM x;
 SELECT x.a AS d, x.b AS b FROM x AS x;
 
 # execute: false
-SELECT * EXCEPT(b) REPLACE(a AS d) FROM x;
+SELECT * EXCEPT(b) RENAME(a AS d) FROM x;
 SELECT x.a AS d FROM x AS x;
 
 SELECT x.* EXCEPT(a), y.* FROM x, y;
@@ -416,6 +416,30 @@ SELECT x.a AS a, x.b AS b, y.b AS b FROM x AS x LEFT JOIN x AS y ON x.a = y.a;
 SELECT COALESCE(CAST(t1.a AS VARCHAR), '') AS a, t2.* EXCEPT (a) FROM x AS t1, x AS t2;
 SELECT COALESCE(CAST(t1.a AS VARCHAR), '') AS a, t2.b AS b FROM x AS t1, x AS t2;
 
+# execute: false
+SELECT * REPLACE(2 AS a) FROM x;
+SELECT 2 AS a, x.b AS b FROM x AS x;
+
+# execute: false
+SELECT * EXCEPT (a, b) REPLACE (a AS a) FROM x;
+SELECT * EXCEPT (a, b) REPLACE (x.a AS a) FROM x AS x;
+
+# execute: false
+SELECT * REPLACE(COALESCE(b, a) AS a, a as b) FROM x;
+SELECT COALESCE(x.b, x.a) AS a, x.a AS b FROM x AS x;
+
+# execute: false
+SELECT * REPLACE(1 AS a) RENAME(b as alias_b) FROM x;
+SELECT 1 AS a, x.b AS alias_b FROM x AS x;
+
+# execute: false
+SELECT * EXCEPT(a) REPLACE(COALESCE(a, b) AS b) RENAME(b AS new_b) FROM x;
+SELECT COALESCE(x.a, x.b) AS new_b FROM x AS x;
+
+# execute: false
+SELECT * REPLACE(1 AS a, a AS b) RENAME(b AS new_b) FROM x;
+SELECT 1 AS a, x.a AS new_b FROM x AS x;
+
 --------------------------------------
 -- Using
 --------------------------------------