Fix: select as struct transpilation closes #1788

tobymao · Jun 16, 2023 · 311380c · 311380c
1 parent 47db2cd
commit 311380c
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 5 deletions.
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
@@ -117,6 +117,9 @@ def get_start_end(token_type: TokenType) -> t.Tuple[t.Optional[str], t.Optional[
             "IDENTIFIER_ESCAPE": klass.tokenizer_class.IDENTIFIER_ESCAPES[0],
         }
 
+        if enum not in ("", "bigquery"):
+            dialect_properties["SELECT_KINDS"] = ()
+
         # Pass required dialect properties to the tokenizer, parser and generator classes
         for subclass in (klass.tokenizer_class, klass.parser_class, klass.generator_class):
             for name, value in dialect_properties.items():

diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
@@ -167,7 +167,7 @@ class Generator(generator.Generator):
             **generator.Generator.TRANSFORMS,
             exp.ApproxDistinct: approx_count_distinct_sql,
             exp.Array: lambda self, e: self.func("ARRAY", e.expressions[0])
-            if isinstance(seq_get(e.expressions, 0), exp.Select)
+            if e.expressions and e.expressions[0].find(exp.Select)
             else rename_func("LIST_VALUE")(self, e),
             exp.ArraySize: rename_func("ARRAY_LENGTH"),
             exp.ArraySort: _array_sort_sql,

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -2434,8 +2434,6 @@ class Select(Subqueryable):
         "expressions": False,
         "hint": False,
         "distinct": False,
-        "struct": False,  # https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#return_query_results_as_a_value_table
-        "value": False,
         "into": False,
         "from": False,
         **QUERY_MODIFIERS,

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -142,6 +142,9 @@ class Generator:
     # Whether or not comparing against booleans (e.g. x IS TRUE) is supported
     IS_BOOL_ALLOWED = True
 
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax
+    SELECT_KINDS: t.Tuple[str, ...] = ("STRUCT", "VALUE")
+
     TYPE_MAPPING = {
         exp.DataType.Type.NCHAR: "CHAR",
         exp.DataType.Type.NVARCHAR: "VARCHAR",
@@ -1565,9 +1568,30 @@ def select_sql(self, expression: exp.Select) -> str:
         hint = self.sql(expression, "hint")
         distinct = self.sql(expression, "distinct")
         distinct = f" {distinct}" if distinct else ""
-        kind = expression.args.get("kind")
-        kind = f" AS {kind}" if kind else ""
+        kind = self.sql(expression, "kind").upper()
         expressions = self.expressions(expression)
+
+        if kind:
+            if kind in self.SELECT_KINDS:
+                kind = f" AS {kind}"
+            else:
+                if kind == "STRUCT":
+                    expressions = self.expressions(
+                        sqls=[
+                            self.sql(
+                                exp.Struct(
+                                    expressions=[
+                                        exp.column(e.output_name).eq(
+                                            e.this if isinstance(e, exp.Alias) else e
+                                        )
+                                        for e in expression.expressions
+                                    ]
+                                )
+                            )
+                        ]
+                    )
+                kind = ""
+
         expressions = f"{self.sep()}{expressions}" if expressions else expressions
         sql = self.query_modifiers(
             expression,

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -23,6 +23,14 @@ def test_bigquery(self):
         self.validate_identity("SELECT b'abc'")
         self.validate_identity("""SELECT * FROM UNNEST(ARRAY<STRUCT<x INT64>>[1, 2])""")
         self.validate_identity("SELECT AS STRUCT 1 AS a, 2 AS b")
+        self.validate_all(
+            "SELECT AS STRUCT ARRAY(SELECT AS STRUCT b FROM x) AS y FROM z",
+            write={
+                "": "SELECT AS STRUCT ARRAY(SELECT AS STRUCT b FROM x) AS y FROM z",
+                "bigquery": "SELECT AS STRUCT ARRAY(SELECT AS STRUCT b FROM x) AS y FROM z",
+                "duckdb": "SELECT {'y': ARRAY(SELECT {'b': b} FROM x)} FROM z",
+            },
+        )
         self.validate_identity("SELECT DISTINCT AS STRUCT 1 AS a, 2 AS b")
         self.validate_identity("SELECT AS VALUE STRUCT(1 AS a, 2 AS b)")
         self.validate_identity("SELECT STRUCT<ARRAY<STRING>>(['2023-01-17'])")

diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
@@ -444,6 +444,7 @@ def test_duckdb(self):
 
     def test_array(self):
         self.validate_identity("ARRAY(SELECT id FROM t)")
+        self.validate_identity("ARRAY((SELECT id FROM t))")
 
     def test_cast(self):
         self.validate_identity("CAST(x AS REAL)")