diff options
Diffstat (limited to 'sqlglot/optimizer/expand_laterals.py')
-rw-r--r-- | sqlglot/optimizer/expand_laterals.py | 34 |
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 |