Merging upstream version 11.0.1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 34 insertions, 0 deletions

diff --git a/sqlglot/optimizer/expand_laterals.py b/sqlglot/optimizer/expand_laterals.py
new file mode 100644
index 0000000..59f3fec
--- /dev/null
+++ b/sqlglot/optimizer/expand_laterals.py
@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+
+def expand_laterals(expression: exp.Expression) -> exp.Expression:
+    """
+    Expand lateral column alias references.
+
+    This assumes `qualify_columns` as already run.
+
+    Example:
+        >>> import sqlglot
+        >>> sql = "SELECT x.a + 1 AS b, b + 1 AS c FROM x"
+        >>> expression = sqlglot.parse_one(sql)
+        >>> expand_laterals(expression).sql()
+        'SELECT x.a + 1 AS b, x.a + 1 + 1 AS c FROM x'
+
+    Args:
+        expression: expression to optimize
+    Returns:
+        optimized expression
+    """
+    for select in expression.find_all(exp.Select):
+        alias_to_expression: t.Dict[str, exp.Expression] = {}
+        for projection in select.expressions:
+            for column in projection.find_all(exp.Column):
+                if not column.table and column.name in alias_to_expression:
+                    column.replace(alias_to_expression[column.name].copy())
+                if isinstance(projection, exp.Alias):
+                    alias_to_expression[projection.alias] = projection.this
+    return expression