Skip to content

Commit

Permalink
Fix: workaround bigquery grouped alias with order
Browse files Browse the repository at this point in the history
  • Loading branch information
tobymao committed Jun 28, 2023
1 parent 3600240 commit eaed790
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 4 deletions.
35 changes: 34 additions & 1 deletion sqlglot/dialects/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,34 @@ def _unqualify_unnest(expression: exp.Expression) -> exp.Expression:
return expression


# https://issuetracker.google.com/issues/162294746
# workaround for bigquery bug when grouping by an expression and then ordering
# WITH x AS (SELECT 1 y)
# SELECT y + 1 z
# FROM x
# GROUP BY x + 1
# ORDER by z
def _alias_ordered_group(expression: exp.Expression) -> exp.Expression:
if isinstance(expression, exp.Select):
group = expression.args.get("group")
order = expression.args.get("order")

if group and order:
aliases = {
select.this: select.args["alias"]
for select in expression.selects
if isinstance(select, exp.Alias)
}

for e in group.expressions:
alias = aliases.get(e)

if alias:
e.replace(exp.column(alias))

return expression


def _pushdown_cte_column_names(expression: exp.Expression) -> exp.Expression:
"""BigQuery doesn't allow column names when defining a CTE, so we try to push them down."""
if isinstance(expression, exp.CTE) and expression.alias_column_names:
Expand Down Expand Up @@ -360,7 +388,12 @@ class Generator(generator.Generator):
),
exp.RegexpLike: rename_func("REGEXP_CONTAINS"),
exp.Select: transforms.preprocess(
[transforms.explode_to_unnest, _unqualify_unnest, transforms.eliminate_distinct_on]
[
transforms.explode_to_unnest,
_unqualify_unnest,
transforms.eliminate_distinct_on,
_alias_ordered_group,
]
),
exp.StrToDate: lambda self, e: f"PARSE_DATE({self.format_time(e)}, {self.sql(e, 'this')})",
exp.StrToTime: lambda self, e: f"PARSE_TIMESTAMP({self.format_time(e)}, {self.sql(e, 'this')})",
Expand Down
10 changes: 10 additions & 0 deletions tests/dialects/test_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -517,6 +517,16 @@ def test_bigquery(self):
},
)

self.validate_identity(
"SELECT y + 1 z FROM x GROUP BY y + 1 ORDER BY z",
"SELECT y + 1 AS z FROM x GROUP BY z ORDER BY z",
)
self.validate_identity(
"SELECT y + 1 z FROM x GROUP BY y + 1",
"SELECT y + 1 AS z FROM x GROUP BY y + 1",
)
self.validate_identity("SELECT y + 1 FROM x GROUP BY y + 1 ORDER BY 1")

def test_user_defined_functions(self):
self.validate_identity(
"CREATE TEMPORARY FUNCTION a(x FLOAT64, y FLOAT64) RETURNS FLOAT64 NOT DETERMINISTIC LANGUAGE js AS 'return x*y;'"
Expand Down
6 changes: 3 additions & 3 deletions tests/fixtures/optimizer/optimizer.sql
Original file line number Diff line number Diff line change
Expand Up @@ -676,9 +676,9 @@ WHERE
`top_terms`.`rank` = 1
AND CAST(`top_terms`.`refresh_date` AS DATE) >= DATE_SUB(CURRENT_DATE, INTERVAL 2 WEEK)
GROUP BY
`top_terms`.`refresh_date`,
`top_terms`.`term`,
`top_terms`.`rank`
`day`,
`top_term`,
`rank`
ORDER BY
`day` DESC;

Expand Down

0 comments on commit eaed790

Please sign in to comment.