Skip to content

Commit 6cdc8e1

Browse files
VaggelisDgeorgesittas
authored andcommitted
feat(snowflake): Transpile BQ's ARRAY(SELECT AS STRUCT ...) (#5140)
* feat(snowflake): Transpile ARRAY(SELECT AS STRUCT ...) * Reuse existing selection
1 parent 052cf99 commit 6cdc8e1

File tree

2 files changed

+59
-1
lines changed

2 files changed

+59
-1
lines changed

sqlglot/dialects/snowflake.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1017,7 +1017,6 @@ class Generator(generator.Generator):
10171017
exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
10181018
exp.ArgMax: rename_func("MAX_BY"),
10191019
exp.ArgMin: rename_func("MIN_BY"),
1020-
exp.Array: inline_array_sql,
10211020
exp.ArrayConcat: lambda self, e: self.arrayconcat_sql(e, name="ARRAY_CAT"),
10221021
exp.ArrayContains: lambda self, e: self.func("ARRAY_CONTAINS", e.expression, e.this),
10231022
exp.AtTimeZone: lambda self, e: self.func(
@@ -1428,3 +1427,29 @@ def arrayagg_sql(self, expression: exp.ArrayAgg) -> str:
14281427
expr_sql = self.sql(exp.WithinGroup(this=expr_sql, expression=order))
14291428

14301429
return expr_sql
1430+
1431+
def array_sql(self, expression: exp.Array) -> str:
1432+
expressions = expression.expressions
1433+
1434+
first_expr = seq_get(expressions, 0)
1435+
if isinstance(first_expr, exp.Select):
1436+
# SELECT AS STRUCT foo AS alias_foo -> ARRAY_AGG(OBJECT_CONSTRUCT('alias_foo', foo))
1437+
if first_expr.text("kind").upper() == "STRUCT":
1438+
object_construct_args = []
1439+
for expr in first_expr.expressions:
1440+
# Alias case: SELECT AS STRUCT foo AS alias_foo -> OBJECT_CONSTRUCT('alias_foo', foo)
1441+
# Column case: SELECT AS STRUCT foo -> OBJECT_CONSTRUCT('foo', foo)
1442+
name = expr.this if isinstance(expr, exp.Alias) else expr
1443+
1444+
object_construct_args.extend([exp.Literal.string(expr.alias_or_name), name])
1445+
1446+
array_agg = exp.ArrayAgg(
1447+
this=_build_object_construct(args=object_construct_args)
1448+
)
1449+
1450+
first_expr.set("kind", None)
1451+
first_expr.set("expressions", [array_agg])
1452+
1453+
return self.sql(first_expr.subquery())
1454+
1455+
return inline_array_sql(self, expression)

tests/dialects/test_bigquery.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2576,3 +2576,36 @@ def test_array_agg(self):
25762576
"snowflake": "SELECT ARRAY_AGG(x) WITHIN GROUP (ORDER BY col1 ASC NULLS FIRST, col2 DESC NULLS LAST)",
25772577
},
25782578
)
2579+
2580+
def test_select_as_struct(self):
2581+
self.validate_all(
2582+
"SELECT ARRAY(SELECT AS STRUCT x1 AS x1, x2 AS x2 FROM t) AS array_col",
2583+
write={
2584+
"bigquery": "SELECT ARRAY(SELECT AS STRUCT x1 AS x1, x2 AS x2 FROM t) AS array_col",
2585+
"snowflake": "SELECT (SELECT ARRAY_AGG(OBJECT_CONSTRUCT('x1', x1, 'x2', x2)) FROM t) AS array_col",
2586+
},
2587+
)
2588+
2589+
self.validate_all(
2590+
"WITH t1 AS (SELECT ARRAY(SELECT AS STRUCT x1 AS alias_x1, x2 /* test */ FROM t2) AS array_col) SELECT array_col[0].alias_x1, array_col[0].x2 FROM t1",
2591+
write={
2592+
"bigquery": "WITH t1 AS (SELECT ARRAY(SELECT AS STRUCT x1 AS alias_x1, x2 /* test */ FROM t2) AS array_col) SELECT array_col[0].alias_x1, array_col[0].x2 FROM t1",
2593+
"snowflake": "WITH t1 AS (SELECT (SELECT ARRAY_AGG(OBJECT_CONSTRUCT('alias_x1', x1, 'x2', x2 /* test */)) FROM t2) AS array_col) SELECT array_col[0].alias_x1, array_col[0].x2 FROM t1",
2594+
},
2595+
)
2596+
2597+
self.validate_all(
2598+
"WITH t1 AS (SELECT ARRAY(SELECT AS STRUCT 1 AS a, 2 AS b) AS array_col) SELECT array_col[0].a, array_col[0].b FROM t1",
2599+
write={
2600+
"bigquery": "WITH t1 AS (SELECT ARRAY(SELECT AS STRUCT 1 AS a, 2 AS b) AS array_col) SELECT array_col[0].a, array_col[0].b FROM t1",
2601+
"snowflake": "WITH t1 AS (SELECT (SELECT ARRAY_AGG(OBJECT_CONSTRUCT('a', 1, 'b', 2))) AS array_col) SELECT array_col[0].a, array_col[0].b FROM t1",
2602+
},
2603+
)
2604+
2605+
self.validate_all(
2606+
"WITH t1 AS (SELECT ARRAY(SELECT AS STRUCT x1 AS alias_x1, x2 /* test */ FROM t2 WHERE x2 = 4) AS array_col) SELECT array_col[0].alias_x1, array_col[0].x2 FROM t1",
2607+
write={
2608+
"bigquery": "WITH t1 AS (SELECT ARRAY(SELECT AS STRUCT x1 AS alias_x1, x2 /* test */ FROM t2 WHERE x2 = 4) AS array_col) SELECT array_col[0].alias_x1, array_col[0].x2 FROM t1",
2609+
"snowflake": "WITH t1 AS (SELECT (SELECT ARRAY_AGG(OBJECT_CONSTRUCT('alias_x1', x1, 'x2', x2 /* test */)) FROM t2 WHERE x2 = 4) AS array_col) SELECT array_col[0].alias_x1, array_col[0].x2 FROM t1",
2610+
},
2611+
)

0 commit comments

Comments
 (0)