Repository: tobymao/sqlglot

Issue:
wrong qualify_columns on generated columns
**Before you file an issue**
- Make sure you specify the "read" dialect eg. `parse_one(sql, read="spark")`: Done
- Make sure you specify the "write" dialect eg. `ast.sql(dialect="duckdb")`: Done
- Check if the issue still exists on main: still exist

**Fully reproducible code snippet**

Prepare data:
```python
from sqlglot import parse_one, exp
from sqlglot.optimizer.qualify_columns import qualify_columns

sql ="""with data as (
  select 1 as id, 2 as my_id, 'a' as name, 'de' as full_name
)
select
  id as my_id,
  concat(id, name) as full_name
from data
group by my_id, full_name"""

schema = {
"catalog": {
      "schema1": {
          "customer": {
              "id": "INTEGER",
          },
      }
  }
}
```

Transpiler from bigquery to trino:
```python
smt = parse_one(sql, read="bigquery")
print(qualify_columns(smt, schema=schema).sql(dialect="trino", pretty=True))
```
Output:
```sql
WITH data AS (
  SELECT
    1 AS id,
    2 AS my_id,
    'a' AS name,
    'de' AS full_name
)
SELECT
  data.id AS my_id,
  CONCAT(CAST(data.id AS VARCHAR), CAST(data.name AS VARCHAR)) AS full_name
FROM data
GROUP BY
  data.my_id,
  data.full_name
```

The problem of the output is it can not be run on trino, because `data.id` and `data.name` is not an aggregate expression or appear in GROUP BY clause.

Additional test:
The output will expand correctly if we no not provide schema:
```python
smt = parse_one(sql, read="bigquery")
print(qualify_columns(smt, schema=None).sql(dialect="trino", pretty=True))

"""
WITH data AS (
  SELECT
    1 AS id,
    2 AS my_id,
    'a' AS name,
    'de' AS full_name
)
SELECT
  data.id AS my_id,
  CONCAT(CAST(data.id AS VARCHAR), CAST(data.name AS VARCHAR)) AS full_name
FROM data
GROUP BY
  data.id,
  CONCAT(CAST(data.id AS VARCHAR), CAST(data.name AS VARCHAR))
"""
```

The output will break again if I disable expand_alias_ref: 
```python
smt = parse_one(sql, read="bigquery")
print(qualify_columns(smt, schema=None, expand_alias_refs=False).sql(dialect="trino", pretty=True))
```

### Solution
I have made a patch to fix my case, but I think we can have better solution to fix this.

https://github.com/viplazylmht/sqlglot/blob/644bf67720442f964ffc4606693dcfc1b87ec6e6/sqlglot/optimizer/qualify_columns.py#L423-L431


**Official Documentation**
Please include links to official SQL documentation related to your issue.


Discussion:

Patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 4f557e7b..f21ba38c 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -221,6 +221,7 @@ class BigQuery(Dialect):
     SUPPORTS_SEMI_ANTI_JOIN = False
     LOG_BASE_FIRST = False
     HEX_LOWERCASE = True
+    FORCE_EARLY_ALIAS_REF_EXPANSION = True
 
     # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#case_sensitivity
     NORMALIZATION_STRATEGY = NormalizationStrategy.CASE_INSENSITIVE
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
index 209a2894..c0767753 100644
--- a/sqlglot/dialects/clickhouse.py
+++ b/sqlglot/dialects/clickhouse.py
@@ -102,6 +102,7 @@ class ClickHouse(Dialect):
     SUPPORTS_USER_DEFINED_TYPES = False
     SAFE_DIVISION = True
     LOG_BASE_FIRST: t.Optional[bool] = None
+    FORCE_EARLY_ALIAS_REF_EXPANSION = True
 
     UNESCAPED_SEQUENCES = {
         "\\0": "\0",
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 32d621a9..fca4d94c 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -325,6 +325,32 @@ class Dialect(metaclass=_Dialect):
     Whether COPY statement parameters are separated by comma or whitespace
     """
 
+    FORCE_EARLY_ALIAS_REF_EXPANSION = False
+    """
+    Whether alias reference expansion (_expand_alias_refs()) should run before column qualification (_qualify_columns()).
+
+    For example:
+        WITH data AS (
+        SELECT
+            1 AS id,
+            2 AS my_id
+        )
+        SELECT
+            id AS my_id
+        FROM
+            data
+        WHERE
+            my_id = 1
+        GROUP BY
+            my_id,
+        HAVING
+            my_id = 1
+
+    In most dialects "my_id" would refer to "data.my_id" (which is done in _qualify_columns()) across the query, except:
+        - BigQuery, which will forward the alias to GROUP BY + HAVING clauses i.e it resolves to "WHERE my_id = 1 GROUP BY id HAVING id = 1"
+        - Clickhouse, which will forward the alias across the query i.e it resolves to "WHERE id = 1 GROUP BY id HAVING id = 1"
+    """
+
     # --- Autofilled ---
 
     tokenizer_class = Tokenizer
diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py
index 6ef982f7..93b87c7a 100644
--- a/sqlglot/optimizer/qualify_columns.py
+++ b/sqlglot/optimizer/qualify_columns.py
@@ -60,8 +60,8 @@ def qualify_columns(
         _pop_table_column_aliases(scope.derived_tables)
         using_column_tables = _expand_using(scope, resolver)
 
-        if schema.empty and expand_alias_refs:
-            _expand_alias_refs(scope, resolver)
+        if (schema.empty or dialect.FORCE_EARLY_ALIAS_REF_EXPANSION) and expand_alias_refs:
+            _expand_alias_refs(scope, resolver, expand_only_groupby=(dialect == "bigquery"))
 
         _convert_columns_to_dots(scope, resolver)
         _qualify_columns(scope, resolver)
@@ -208,7 +208,7 @@ def _expand_using(scope: Scope, resolver: Resolver) -> t.Dict[str, t.Any]:
     return column_tables
 
 
-def _expand_alias_refs(scope: Scope, resolver: Resolver) -> None:
+def _expand_alias_refs(scope: Scope, resolver: Resolver, expand_only_groupby: bool = False) -> None:
     expression = scope.expression
 
     if not isinstance(expression, exp.Select):
@@ -219,7 +219,7 @@ def _expand_alias_refs(scope: Scope, resolver: Resolver) -> None:
     def replace_columns(
         node: t.Optional[exp.Expression], resolve_table: bool = False, literal_index: bool = False
     ) -> None:
-        if not node:
+        if not node or (expand_only_groupby and not isinstance(node, exp.Group)):
             return
 
         for column in walk_in_scope(node, prune=lambda node: node.is_star):


Test patch:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 81b9731b..b1d50942 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -473,15 +473,35 @@ SELECT :with,WITH :expressions,CTE :this,UNION :this,SELECT :expressions,1,:expr
             'SELECT "x"."a" + 1 AS "d", "x"."a" + 1 + 1 AS "e" FROM "x" AS "x" WHERE ("x"."a" + 2) > 1 GROUP BY "x"."a" + 1 + 1',
         )
 
+        unused_schema = {"l": {"c": "int"}}
         self.assertEqual(
             optimizer.qualify_columns.qualify_columns(
                 parse_one("SELECT CAST(x AS INT) AS y FROM z AS z"),
-                schema={"l": {"c": "int"}},
+                schema=unused_schema,
                 infer_schema=False,
             ).sql(),
             "SELECT CAST(x AS INT) AS y FROM z AS z",
         )
 
+        # BigQuery expands overlapping alias only for GROUP BY + HAVING
+        sql = "WITH data AS (SELECT 1 AS id, 2 AS my_id, 'a' AS name, 'b' AS full_name) SELECT id AS my_id, CONCAT(id, name) AS full_name FROM data WHERE my_id = 1 GROUP BY my_id, full_name HAVING my_id = 1"
+        self.assertEqual(
+            optimizer.qualify_columns.qualify_columns(
+                parse_one(sql, dialect="bigquery"),
+                schema=MappingSchema(schema=unused_schema, dialect="bigquery"),
+            ).sql(),
+            "WITH data AS (SELECT 1 AS id, 2 AS my_id, 'a' AS name, 'b' AS full_name) SELECT data.id AS my_id, CONCAT(data.id, data.name) AS full_name FROM data WHERE data.my_id = 1 GROUP BY data.id, CONCAT(data.id, data.name) HAVING data.id = 1",
+        )
+
+        # Clickhouse expands overlapping alias across the entire query
+        self.assertEqual(
+            optimizer.qualify_columns.qualify_columns(
+                parse_one(sql, dialect="clickhouse"),
+                schema=MappingSchema(schema=unused_schema, dialect="clickhouse"),
+            ).sql(),
+            "WITH data AS (SELECT 1 AS id, 2 AS my_id, 'a' AS name, 'b' AS full_name) SELECT data.id AS my_id, CONCAT(data.id, data.name) AS full_name FROM data WHERE data.id = 1 GROUP BY data.id, CONCAT(data.id, data.name) HAVING data.id = 1",
+        )
+
     def test_optimize_joins(self):
         self.check_file(
             "optimize_joins",