fix(bigquery): Inline type-annotated ARRAY literals (#4671)

tobymao · Jan 29, 2025 · c45f174 · c45f174
1 parent 1a91913
commit c45f174
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -1229,3 +1229,11 @@ def contains_sql(self, expression: exp.Contains) -> str:
                 expr = expr.this
 
             return self.func("CONTAINS_SUBSTR", this, expr)
+
+        def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) -> str:
+            this = expression.this
+
+            if isinstance(this, exp.Array):
+                return f"{self.sql(expression, 'to')}{self.sql(this)}"
+
+            return super().cast_sql(expression, safe_prefix=safe_prefix)
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -2035,10 +2035,10 @@ def test_inline_constructor(self):
         )
 
         self.validate_all(
-            "SELECT ARRAY<INT>[1, 2, 3]",
+            "SELECT ARRAY<FLOAT64>[1, 2, 3]",
             write={
-                "bigquery": "SELECT CAST([1, 2, 3] AS ARRAY<INT64>)",
-                "duckdb": "SELECT CAST([1, 2, 3] AS INT[])",
+                "bigquery": "SELECT ARRAY<FLOAT64>[1, 2, 3]",
+                "duckdb": "SELECT CAST([1, 2, 3] AS DOUBLE[])",
             },
         )
         self.validate_all(
@@ -2051,14 +2051,14 @@ def test_inline_constructor(self):
         self.validate_all(
             "SELECT * FROM UNNEST(ARRAY<STRUCT<x INT64>>[])",
             write={
-                "bigquery": "SELECT * FROM UNNEST(CAST([] AS ARRAY<STRUCT<x INT64>>))",
+                "bigquery": "SELECT * FROM UNNEST(ARRAY<STRUCT<x INT64>>[])",
                 "duckdb": "SELECT * FROM (SELECT UNNEST(CAST([] AS STRUCT(x BIGINT)[]), max_depth => 2))",
             },
         )
         self.validate_all(
             "SELECT * FROM UNNEST(ARRAY<STRUCT<device_id INT64, time DATETIME, signal INT64, state STRING>>[STRUCT(1, DATETIME '2023-11-01 09:34:01', 74, 'INACTIVE'),STRUCT(4, DATETIME '2023-11-01 09:38:01', 80, 'ACTIVE')])",
             write={
-                "bigquery": "SELECT * FROM UNNEST(CAST([STRUCT(1, CAST('2023-11-01 09:34:01' AS DATETIME), 74, 'INACTIVE'), STRUCT(4, CAST('2023-11-01 09:38:01' AS DATETIME), 80, 'ACTIVE')] AS ARRAY<STRUCT<device_id INT64, time DATETIME, signal INT64, state STRING>>))",
+                "bigquery": "SELECT * FROM UNNEST(ARRAY<STRUCT<device_id INT64, time DATETIME, signal INT64, state STRING>>[STRUCT(1, CAST('2023-11-01 09:34:01' AS DATETIME), 74, 'INACTIVE'), STRUCT(4, CAST('2023-11-01 09:38:01' AS DATETIME), 80, 'ACTIVE')])",
                 "duckdb": "SELECT * FROM (SELECT UNNEST(CAST([ROW(1, CAST('2023-11-01 09:34:01' AS TIMESTAMP), 74, 'INACTIVE'), ROW(4, CAST('2023-11-01 09:38:01' AS TIMESTAMP), 80, 'ACTIVE')] AS STRUCT(device_id BIGINT, time TIMESTAMP, signal BIGINT, state TEXT)[]), max_depth => 2))",
             },
         )

diff --git a/tests/dialects/test_presto.py b/tests/dialects/test_presto.py
@@ -104,7 +104,7 @@ def test_cast(self):
         self.validate_all(
             "CAST(ARRAY[1, 2] AS ARRAY(BIGINT))",
             write={
-                "bigquery": "CAST([1, 2] AS ARRAY<INT64>)",
+                "bigquery": "ARRAY<INT64>[1, 2]",
                 "duckdb": "CAST([1, 2] AS BIGINT[])",
                 "presto": "CAST(ARRAY[1, 2] AS ARRAY(BIGINT))",
                 "spark": "CAST(ARRAY(1, 2) AS ARRAY<BIGINT>)",