Repository: tobymao/sqlglot
Pull request: Feat(optimizer): expand join constructs into SELECT * from subqueries
Fixes #1554

<details>
<summary>Click this to see the execution of all added test cases in both trino and postgres.</summary>

Trino
-----

```
trino> WITH tbl AS (SELECT 1) SELECT * FROM (tbl AS tbl) AS _q_0; -- Only valid in Trino, postgres doesn't allow this
 _col0 
-------
     1 
(1 row)

trino> WITH tbl AS (SELECT 1) SELECT * FROM ((tbl AS tbl)) AS _q_0;
 _col0 
-------
     1 
(1 row)

trino> WITH tbl AS (SELECT 1) SELECT * FROM (((tbl AS tbl))) AS _q_0;
 _col0 
-------
     1 
(1 row)

trino> WITH tbl AS (SELECT 1) SELECT * FROM (SELECT * FROM tbl AS tbl) AS _q_0; -- optimized
 _col0 
-------
     1 
(1 row)

trino> WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0;
 id1 | id2 | id3 
-----+-----+-----
   1 |   1 |   1 
(1 row)

trino> WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM ((tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3)) AS _q_0;
 id1 | id2 | id3 
-----+-----+-----
   1 |   1 |   1 
(1 row)

trino> WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0; -- optimized
 id1 | id2 | id3 
-----+-----+-----
   1 |   1 |   1 
(1 row)

trino> WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (tbl1 AS tbl1 JOIN (tbl2 AS tbl2 JOIN tbl3 AS tbl3 ON id2 = id3) AS _q_0 ON id1 = id3) AS _q_1;;
 id1 | id2 | id3 
-----+-----+-----
   1 |   1 |   1 
(1 row)

trino> WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN (SELECT * FROM tbl2 AS tbl2 JOIN tbl3 AS tbl3 ON id2 = id3) AS _q_0 ON id1 = id3) AS _q_1; -- optimized
 id1 | id2 | id3 
-----+-----+-----
   1 |   1 |   1 
(1 row)
```

Postgres
-----

```
georgesittas=# WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0;
 id1 | id2 | id3
-----+-----+-----
   1 |   1 |   1
(1 row)

georgesittas=# WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM ((tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3)) AS _q_0;
 id1 | id2 | id3
-----+-----+-----
   1 |   1 |   1
(1 row)

georgesittas=# WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0; -- optimized
 id1 | id2 | id3
-----+-----+-----
   1 |   1 |   1
(1 row)

georgesittas=# WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (tbl1 AS tbl1 JOIN (tbl2 AS tbl2 JOIN tbl3 AS tbl3 ON id2 = id3) AS _q_0 ON id1 = id3) AS _q_1;
 id1 | id2 | id3
-----+-----+-----
   1 |   1 |   1
(1 row)

georgesittas=# WITH tbl1(id1) AS (SELECT 1), tbl2(id2) AS (SELECT 1), tbl3(id3) AS (SELECT 1) SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN (SELECT * FROM tbl2 AS tbl2 JOIN tbl3 AS tbl3 ON id2 = id3) AS _q_0 ON id1 = id3) AS _q_1; -- optimized
 id1 | id2 | id3
-----+-----+-----
   1 |   1 |   1
(1 row)
```

</details>

Problem statement:
build_scope and optimize fail on query
Hi,

dialect = trino

It is not possible to build the scope or optimize this sql statement,

```
sql = """
WITH
v_select_44724 AS (SELECT field_4705 Name, id id, field_4707 Active, field_4729 num, field_4730 num1, field_4731 num2 FROM askai_34.public.database_table_271),
v_query_28404 AS (SELECT Company.Name nameCompany, Company.Active activeCompany,  count(Company.Active) as Count FROM (v_select_44724 AS Company) GROUP BY Company.Name, Company.Active)
select * from v_query_28404
"""
ast = sqlglot.parse(sql)
optimized_ast = sqlglot.optimizer.optimize(ast[0].root())
scope = build_scope(optimized_ast)
```

We get the following error,

```
raise OptimizeError(f"Unexpected expression type: {type(scope.expression)}")
sqlglot.errors.OptimizeError: Unexpected expression type: <class 'sqlglot.expressions.Table'>
```

It looks like this snippet seems to fail, 
```
(v_select_44724 AS Company) 
```

If we remove the enclosing parenthesis it works fine.
```
v_select_44724 AS Company
```

I think adding the following conditional to `_traverse_scope()` should solve it,
```
elif isinstance(scope.expression, exp.Table):
    pass
```

Can this be fixed?

Hints from issue discussion:
Will take a look soon, thanks for reporting.

Commit message: Remove unnecessary copying

Patch to review:
diff --git a/sqlglot/optimizer/expand_join_constructs.py b/sqlglot/optimizer/expand_join_constructs.py
new file mode 100644
index 0000000000..649a794847
--- /dev/null
+++ b/sqlglot/optimizer/expand_join_constructs.py
@@ -0,0 +1,33 @@
+import typing as t
+
+from sqlglot import exp
+from sqlglot.optimizer.scope import traverse_scope
+
+
+def expand_join_constructs(expression: exp.Expression) -> exp.Expression:
+    """
+    Replace "join constructs" (*) by equivalent SELECT * subqueries.
+
+    Example:
+    >>> import sqlglot
+    >>> expression = sqlglot.parse_one("SELECT * FROM (tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2) AS tbl")
+    >>> expand_join_constructs(expression).sql()
+    'SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2) AS tbl'
+
+    (*) See section 7.2.1.2 in https://www.postgresql.org/docs/current/queries-table-expressions.html
+    """
+    for scope in traverse_scope(expression):
+        if isinstance(scope.expression, exp.Table):
+            for join in scope.expression.args.get("joins", []):
+                if isinstance(join.this, exp.Subquery):
+                    expand_join_constructs(join.this.unnest())
+
+            outermost_subquery = t.cast(exp.Expression, scope.expression.parent)
+            while isinstance(outermost_subquery.parent, exp.Subquery):
+                outermost_subquery = outermost_subquery.parent
+
+            outermost_subquery.this.replace(
+                exp.select("*").from_(scope.expression.copy(), copy=False)
+            )
+
+    return expression
diff --git a/sqlglot/optimizer/optimizer.py b/sqlglot/optimizer/optimizer.py
index c165ffe23b..5fa10b99f2 100644
--- a/sqlglot/optimizer/optimizer.py
+++ b/sqlglot/optimizer/optimizer.py
@@ -10,6 +10,7 @@
 from sqlglot.optimizer.eliminate_ctes import eliminate_ctes
 from sqlglot.optimizer.eliminate_joins import eliminate_joins
 from sqlglot.optimizer.eliminate_subqueries import eliminate_subqueries
+from sqlglot.optimizer.expand_join_constructs import expand_join_constructs
 from sqlglot.optimizer.expand_multi_table_selects import expand_multi_table_selects
 from sqlglot.optimizer.isolate_table_selects import isolate_table_selects
 from sqlglot.optimizer.lower_identities import lower_identities
@@ -27,6 +28,7 @@
 RULES = (
     lower_identities,
     qualify_tables,
+    expand_join_constructs,
     isolate_table_selects,
     qualify_columns,
     pushdown_projections,
diff --git a/sqlglot/optimizer/scope.py b/sqlglot/optimizer/scope.py
index b582eb00f7..115b0073d4 100644
--- a/sqlglot/optimizer/scope.py
+++ b/sqlglot/optimizer/scope.py
@@ -510,7 +510,7 @@ def _traverse_scope(scope):
         yield from _traverse_union(scope)
     elif isinstance(scope.expression, exp.Subquery):
         yield from _traverse_subqueries(scope)
-    elif isinstance(scope.expression, exp.UDTF):
+    elif isinstance(scope.expression, (exp.UDTF, exp.Table)):
         pass
     else:
         raise OptimizeError(f"Unexpected expression type: {type(scope.expression)}")
diff --git a/tests/fixtures/optimizer/expand_join_constructs.sql b/tests/fixtures/optimizer/expand_join_constructs.sql
new file mode 100644
index 0000000000..cdf5962b91
--- /dev/null
+++ b/tests/fixtures/optimizer/expand_join_constructs.sql
@@ -0,0 +1,18 @@
+-- This is valid in Trino, so we treat the (tbl AS tbl) as a "join construct" per postgres' terminology.
+SELECT * FROM (tbl AS tbl) AS _q_0;
+SELECT * FROM (SELECT * FROM tbl AS tbl) AS _q_0;
+
+SELECT * FROM ((tbl AS tbl)) AS _q_0;
+SELECT * FROM (SELECT * FROM tbl AS tbl) AS _q_0;
+
+SELECT * FROM (((tbl AS tbl))) AS _q_0;
+SELECT * FROM (SELECT * FROM tbl AS tbl) AS _q_0;
+
+SELECT * FROM (tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0;
+SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0;
+
+SELECT * FROM ((tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3)) AS _q_0;
+SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2 JOIN tbl3 AS tbl3 ON id1 = id3) AS _q_0;
+
+SELECT * FROM (tbl1 AS tbl1 JOIN (tbl2 AS tbl2 JOIN tbl3 AS tbl3 ON id2 = id3) AS _q_0 ON id1 = id3) AS _q_1;
+SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN (SELECT * FROM tbl2 AS tbl2 JOIN tbl3 AS tbl3 ON id2 = id3) AS _q_0 ON id1 = id3) AS _q_1;
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index fec35b4a0e..2a18074f71 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -235,6 +235,11 @@ def test_expand_laterals(self):
             execute=True,
         )
 
+    def test_expand_join_constructs(self):
+        self.check_file(
+            "expand_join_constructs", optimizer.expand_join_constructs.expand_join_constructs
+        )
+
     def test_expand_multi_table_selects(self):
         self.check_file(
             "expand_multi_table_selects",


Code review comments:
diff_hunk: @@ -0,0 +1,33 @@
+import typing as t
+
+from sqlglot import exp
+from sqlglot.optimizer.scope import traverse_scope
+
+
+def expand_join_constructs(expression: exp.Expression) -> exp.Expression:
+    """
+    Replace "join constructs" (*) by equivalent SELECT * subqueries.
+
+    Example:
+    >>> import sqlglot
+    >>> expression = sqlglot.parse_one("SELECT * FROM (tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2) AS tbl")
+    >>> expand_join_constructs(expression).sql()
+    'SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2) AS tbl'
+
+    (*) See section 7.2.1.2 in https://www.postgresql.org/docs/current/queries-table-expressions.html
+    """
+    for scope in traverse_scope(expression):
path: sqlglot/optimizer/expand_join_constructs.py
text: @user1:
Do you need scope here? I.e. what does it add beyond expression.walk?

@author:
Hmm, that's a good question. It seemed like a natural choice to me at first because it captures the `(tbl as tbl)` expressions nicely (i.e we capture subqueries as "scopes" and we can just check if their expression is an `exp.Table`).

But now that you mention it, we can probably walk the tree and look for `exp.Subquery` expressions for which the `.this.unnest()` returns an `exp.Table` and do the same thing? Does that make sense?

diff_hunk: @@ -0,0 +1,33 @@
+import typing as t
+
+from sqlglot import exp
+from sqlglot.optimizer.scope import traverse_scope
+
+
+def expand_join_constructs(expression: exp.Expression) -> exp.Expression:
+    """
+    Replace "join constructs" (*) by equivalent SELECT * subqueries.
+
+    Example:
+    >>> import sqlglot
+    >>> expression = sqlglot.parse_one("SELECT * FROM (tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2) AS tbl")
+    >>> expand_join_constructs(expression).sql()
+    'SELECT * FROM (SELECT * FROM tbl1 AS tbl1 JOIN tbl2 AS tbl2 ON id1 = id2) AS tbl'
+
+    (*) See section 7.2.1.2 in https://www.postgresql.org/docs/current/queries-table-expressions.html
+    """
+    for scope in traverse_scope(expression):
+        if isinstance(scope.expression, exp.Table):
+            for join in scope.expression.args.get("joins", []):
+                if isinstance(join.this, exp.Subquery):
+                    expand_join_constructs(join.this.unnest())
+
+            outermost_subquery = t.cast(exp.Expression, scope.expression.parent)
+            while isinstance(outermost_subquery.parent, exp.Subquery):
path: sqlglot/optimizer/expand_join_constructs.py
text: @user1:
Kinda strange to recurse, but then crawl back up the tree and do a replacement.

Can we do the replacement at the top of the tree, where we initiate the recursion?

@author:
That's a good point, this does feel a bit iffy. I'll think about your suggestion!

(This might also be somewhat costly due to repeatedly calling `traverse_scope`...)