fix bigquery -> psql transpilation of GENERATE_ARRAY

Author: alistairewj

src/mimic_utils/sqlglot_dialects/postgres.py

@@ -1,6 +1,7 @@
 import sqlglot
 import sqlglot.dialects.postgres
-from sqlglot import Expression, exp, select
+from sqlglot import Expression, exp
+from sqlglot.expressions import array, select
 
 # DATETIME: allow passing either a DATE directly, or multiple arguments
 # there isn't a class for the Datetime function, so we have to create it ourself,
@@ -94,18 +95,20 @@ def datetime_sql(self: Expression, expression: Expression):
 # https://cloud.google.com/bigquery/docs/reference/standard-sql/array_functions#generate_array
 # https://www.postgresql.org/docs/current/functions-srf.html
 def generate_array_sql(self: Expression, expression: Expression):
-    # first create a select statement which selects from generate_series
-    select_statement = select("*").from_(
+    # BigQuery's generate array returns an array data type,
+    # but PostgreSQL generate series returns a set of rows,
+    # so we wrap the output of generate series in an array
+    # constructor.
+    select_statement = array(select("*").from_(
         GenerateSeries(
             expressions=[
                 expression.expressions[0],
                 expression.expressions[1],
             ],
         )
-    )
+    ))
 
-    # now convert the select statement to an array
-    return f"ARRAY({self.sql(select_statement)})"
+    return self.generate(select_statement)
 sqlglot.dialects.postgres.Postgres.Generator.TRANSFORMS[GenerateArray] = generate_array_sql
 
 # we need to prevent the wrapping of the table alias in brackets for UNNEST

src/mimic_utils/transpile.py

@@ -6,8 +6,8 @@
 import sqlglot.dialects.bigquery
 import sqlglot.dialects.duckdb
 import sqlglot.dialects.postgres
-from sqlglot import Expression, exp, select
-from sqlglot.helper import seq_get
+from sqlglot import exp, select, alias
+from sqlglot.expressions import Array, array
 
 # Apply transformation monkey patches
 # these modules are imported for their side effects
@@ -47,6 +47,26 @@ def transpile_query(query: str, source_dialect: str="bigquery", destination_dial
                 quoted=False
             )
 
+    # HACK: sqlglot has a GenerateSeries transpilation in v25.13.0,
+    # which is inserted during the parse of BigQuery. However, it looks
+    # incorrect for postgres (at least), as it swaps GENERATE_ARRAY for GENERATE_SERIES.
+    # BigQuery's GENERATE_ARRAY outputs an array, but GENERATE_SERIES outputs exploded rows.
+    # We will manually replace the GENERATE_SERIES call with an anonymous function, so our
+    # custom transpile code can do the correct conversion for postgres.
+    if (source_dialect == 'bigquery') and (destination_dialect == 'postgres'):
+        for gs_function in sql_parsed.find_all(exp.GenerateSeries):
+            # rename to our anonymous generate array function, so the
+            # later loop will catch it
+            gs_function.replace(
+                exp.Anonymous(
+                    this='GENERATE_ARRAY',
+                    expressions=[
+                        gs_function.args['start'],
+                        gs_function.args['end']
+                    ]
+                )
+            )
+
     # BigQuery has a few functions which are not in sqlglot, so we have
     # created classes for them, and this loop replaces the anonymous functions
     # with the named functions