sqlglot.generators.bigquery

  1from __future__ import annotations
  2
  3import logging
  4import re
  5import typing as t
  6from decimal import Decimal
  7from fractions import Fraction
  8
  9from sqlglot import exp, generator, transforms
 10from sqlglot.dialects.dialect import (
 11    arg_max_or_min_no_count,
 12    array_sort_order,
 13    date_add_interval_sql,
 14    datestrtodate_sql,
 15    filter_array_using_unnest,
 16    generate_series_sql,
 17    if_sql,
 18    inline_array_unless_query,
 19    max_or_greatest,
 20    min_or_least,
 21    no_filter_sql,
 22    no_ilike_sql,
 23    regexp_replace_sql,
 24    rename_func,
 25    sha256_sql,
 26    timestrtotime_sql,
 27    ts_or_ds_add_cast,
 28    unit_to_var,
 29    strposition_sql,
 30    groupconcat_sql,
 31    sha2_digest_sql,
 32)
 33from sqlglot.generator import unsupported_args
 34from sqlglot.helper import seq_get
 35
 36logger = logging.getLogger("sqlglot")
 37
 38JSON_EXTRACT_TYPE = t.Union[exp.JSONExtract, exp.JSONExtractScalar, exp.JSONExtractArray]
 39
 40DQUOTES_ESCAPING_JSON_FUNCTIONS = ("JSON_QUERY", "JSON_VALUE", "JSON_QUERY_ARRAY")
 41
 42
 43def _derived_table_values_to_unnest(self: BigQueryGenerator, expression: exp.Values) -> str:
 44    if not expression.find_ancestor(exp.From, exp.Join):
 45        return self.values_sql(expression)
 46
 47    structs = []
 48    alias = expression.args.get("alias")
 49    for tup in expression.find_all(exp.Tuple):
 50        field_aliases = (
 51            alias.columns
 52            if alias and alias.columns
 53            else (f"_c{i}" for i in range(len(tup.expressions)))
 54        )
 55        expressions = [
 56            exp.PropertyEQ(this=exp.to_identifier(name), expression=fld)
 57            for name, fld in zip(field_aliases, tup.expressions)
 58        ]
 59        structs.append(exp.Struct(expressions=expressions))
 60
 61    # Due to `UNNEST_COLUMN_ONLY`, it is expected that the table alias be contained in the columns expression
 62    alias_name_only = exp.TableAlias(columns=[alias.this]) if alias else None
 63    return self.unnest_sql(
 64        exp.Unnest(expressions=[exp.array(*structs, copy=False)], alias=alias_name_only)
 65    )
 66
 67
 68def _returnsproperty_sql(self: BigQueryGenerator, expression: exp.ReturnsProperty) -> str:
 69    this = expression.this
 70    if isinstance(this, exp.Schema):
 71        this = f"{self.sql(this, 'this')} <{self.expressions(this)}>"
 72    else:
 73        this = self.sql(this)
 74    return f"RETURNS {this}"
 75
 76
 77def _create_sql(self: BigQueryGenerator, expression: exp.Create) -> str:
 78    returns = expression.find(exp.ReturnsProperty)
 79    if expression.kind == "FUNCTION" and returns and returns.args.get("is_table"):
 80        expression.set("kind", "TABLE FUNCTION")
 81
 82        if isinstance(expression.expression, (exp.Subquery, exp.Literal)):
 83            expression.set("expression", expression.expression.this)
 84
 85    return self.create_sql(expression)
 86
 87
 88# https://issuetracker.google.com/issues/162294746
 89# workaround for bigquery bug when grouping by an expression and then ordering
 90# WITH x AS (SELECT 1 y)
 91# SELECT y + 1 z
 92# FROM x
 93# GROUP BY x + 1
 94# ORDER by z
 95def _alias_ordered_group(expression: exp.Expr) -> exp.Expr:
 96    if isinstance(expression, exp.Select):
 97        group = expression.args.get("group")
 98        order = expression.args.get("order")
 99
100        if group and order:
101            aliases = {
102                select.this: select.args["alias"]
103                for select in expression.selects
104                if isinstance(select, exp.Alias)
105            }
106
107            for grouped in group.expressions:
108                if grouped.is_int:
109                    continue
110                alias = aliases.get(grouped)
111                if alias:
112                    grouped.replace(exp.column(alias))
113
114    return expression
115
116
117def _pushdown_cte_column_names(expression: exp.Expr) -> exp.Expr:
118    """BigQuery doesn't allow column names when defining a CTE, so we try to push them down."""
119    if isinstance(expression, exp.CTE) and expression.alias_column_names:
120        cte_query = expression.this
121
122        if cte_query.is_star:
123            logger.warning(
124                "Can't push down CTE column names for star queries. Run the query through"
125                " the optimizer or use 'qualify' to expand the star projections first."
126            )
127            return expression
128
129        column_names = expression.alias_column_names
130        expression.args["alias"].set("columns", None)
131
132        for name, select in zip(column_names, cte_query.selects):
133            to_replace = select
134
135            if isinstance(select, exp.Alias):
136                select = select.this
137
138            # Inner aliases are shadowed by the CTE column names
139            to_replace.replace(exp.alias_(select, name))
140
141    return expression
142
143
144def _unnest_explode_generate_series(expression: exp.Expr) -> exp.Expr:
145    """
146    Rewrites exploding GENERATE_SERIES projections into table references, e.g.
147
148        SELECT GENERATE_SERIES(1, 2) AS x           -> SELECT x FROM GENERATE_SERIES(1, 2) AS x
149        SELECT y, GENERATE_SERIES(1, 2) AS x FROM t -> SELECT y, x FROM t CROSS JOIN GENERATE_SERIES(1, 2) AS x
150
151    since BigQuery can't explode in the projection and must unnest it in the FROM clause instead.
152    The resulting table reference is unnested downstream by `transforms.unnest_generate_series`.
153    """
154    if isinstance(expression, exp.Select):
155        for projection in expression.selects:
156            if isinstance(series := projection.unalias(), exp.ExplodingGenerateSeries):
157                column_name = projection.output_name or "_gen_series_value"
158
159                projection.replace(exp.column(column_name))
160                table = exp.Table(
161                    this=series, alias=exp.TableAlias(this=exp.to_identifier(column_name))
162                )
163
164                if expression.args.get("from_"):
165                    expression.join(table, copy=False, join_type="CROSS")
166                else:
167                    expression.set("from_", exp.From(this=table))
168
169    return expression
170
171
172def _array_contains_sql(self: BigQueryGenerator, expression: exp.ArrayContains) -> str:
173    return self.sql(
174        exp.Exists(
175            this=exp.select("1")
176            .from_(exp.Unnest(expressions=[expression.left]).as_("_unnest", table=["_col"]))
177            .where(exp.column("_col").eq(expression.right))
178        )
179    )
180
181
182def _widen_decimal(expression: exp.Expr) -> exp.Expr:
183    # NUMERIC holds at most 29 integer and 9 fractional digits, e.g. DECIMAL(38, 0) needs BIGNUMERIC
184    for node in expression.find_all(exp.DataType):
185        if node.is_type(exp.DType.DECIMAL) and node.expressions:
186            precision, *rest = (param.this for param in node.expressions)
187            scale = rest[0] if rest else exp.Literal.number(0)
188            if precision.is_int and scale.is_int:
189                scale_digits = int(scale.name)
190                if int(precision.name) - scale_digits > 29 or scale_digits > 9:
191                    node.set("this", exp.DType.BIGDECIMAL)
192    return expression
193
194
195def _extract_part(part: str) -> t.Callable[[BigQueryGenerator, exp.Func], str]:
196    return lambda self, e: self.sql(exp.func("EXTRACT", exp.var(part), e.this, copy=False))
197
198
199def _day_of_week_iso_sql(self: BigQueryGenerator, expression: exp.DayOfWeekIso) -> str:
200    # DAYOFWEEK is Sunday=1..Saturday=7, ISO is Monday=1..Sunday=7
201    template: exp.Expr = exp.maybe_parse("(MOD(EXTRACT(DAYOFWEEK FROM :this) + 5, 7) + 1)")
202    return self.sql(exp.replace_placeholders(template, this=expression.this))
203
204
205def _array_agg_fold_sql(
206    self: BigQueryGenerator, expression: exp.ArrayDistinct | exp.ArraySort
207) -> str:
208    distinct = False
209    ignore_nulls = False
210    sort: exp.ArraySort | None = None
211    agg: exp.Expr = expression
212    while isinstance(agg, (exp.ArrayDistinct, exp.ArraySort)) and not agg.expression:
213        if isinstance(agg, exp.ArrayDistinct):
214            distinct = True
215            ignore_nulls |= agg.args.get("check_null") is False
216        elif sort is None:
217            sort = agg
218        agg = agg.this
219
220    order = array_sort_order(sort) if sort else None
221
222    # An ARRAY_AGG with its own ORDER BY or LIMIT can't take another
223    if (
224        not isinstance(agg, exp.ArrayAgg)
225        or isinstance(agg.this, (exp.Order, exp.Limit))
226        or (sort is not None and order is None)
227    ):
228        self.unsupported(
229            f"{expression.sql_name()} is only supported on an ARRAY_AGG without ORDER BY or LIMIT and with a literal order"
230        )
231        return self.func(expression.sql_name(), expression.this, expression.expression)
232
233    value = agg.this
234    if isinstance(value, exp.Distinct):
235        distinct, value = True, value.expressions[0]
236
237    folded: exp.Expr = exp.Distinct(expressions=[value]) if distinct else value
238    if order is not None:
239        asc, nulls_first = order
240        keys = [exp.Ordered(this=value.copy(), desc=not asc or None, nulls_first=asc)]
241
242        # Aggregate ORDER BY only puts NULLs first when ascending, so other placements need a key
243        if not ignore_nulls and nulls_first is not None and nulls_first != asc:
244            if distinct:
245                self.unsupported("ARRAY_AGG(DISTINCT) only puts NULLs first when ascending")
246            else:
247                is_null = value.copy().is_(exp.null())
248                keys.insert(
249                    0,
250                    exp.Ordered(
251                        this=is_null, desc=nulls_first or None, nulls_first=not nulls_first
252                    ),
253                )
254        folded = exp.Order(this=folded, expressions=keys)
255    agg.set("this", folded)
256    if ignore_nulls:
257        # All-NULL input yields [], but an empty input still yields NULL.
258        agg = exp.If(
259            this=exp.Count(this=exp.Star()).eq(0),
260            true=exp.null(),
261            false=exp.func("COALESCE", exp.IgnoreNulls(this=agg), exp.array()),
262        )
263    return self.sql(agg)
264
265
266def _ts_or_ds_add_sql(self: BigQueryGenerator, expression: exp.TsOrDsAdd) -> str:
267    return date_add_interval_sql("DATE", "ADD")(self, ts_or_ds_add_cast(expression))
268
269
270def _ts_or_ds_diff_sql(self: BigQueryGenerator, expression: exp.TsOrDsDiff) -> str:
271    expression.this.replace(exp.cast(expression.this, exp.DType.TIMESTAMP))
272    expression.expression.replace(exp.cast(expression.expression, exp.DType.TIMESTAMP))
273    unit = unit_to_var(expression)
274    return self.func("DATE_DIFF", expression.this, expression.expression, unit)
275
276
277def _unix_to_time_sql(self: BigQueryGenerator, expression: exp.UnixToTime) -> str:
278    scale = expression.args.get("scale")
279    timestamp = expression.this
280
281    if scale in (None, exp.UnixToTime.SECONDS):
282        return self.func("TIMESTAMP_SECONDS", timestamp)
283    if scale == exp.UnixToTime.MILLIS:
284        return self.func("TIMESTAMP_MILLIS", timestamp)
285    if scale == exp.UnixToTime.MICROS:
286        return self.func("TIMESTAMP_MICROS", timestamp)
287
288    unix_seconds = exp.cast(
289        exp.Div(this=timestamp, expression=exp.func("POW", 10, scale)), exp.DType.BIGINT
290    )
291    return self.func("TIMESTAMP_SECONDS", unix_seconds)
292
293
294def _str_to_datetime_sql(self: BigQueryGenerator, expression: exp.StrToDate | exp.StrToTime) -> str:
295    this = self.sql(expression, "this")
296    dtype = "DATE" if isinstance(expression, exp.StrToDate) else "TIMESTAMP"
297
298    if expression.args.get("safe"):
299        fmt = self.format_time(
300            expression,
301            self.dialect.INVERSE_FORMAT_MAPPING,
302            self.dialect.INVERSE_FORMAT_TRIE,
303        )
304        return f"SAFE_CAST({this} AS {dtype} FORMAT {fmt})"
305
306    fmt = self.format_time(expression)
307    return self.func(f"PARSE_{dtype}", fmt, this, expression.args.get("zone"))
308
309
310@unsupported_args("ins_cost", "del_cost", "sub_cost")
311def _levenshtein_sql(self: BigQueryGenerator, expression: exp.Levenshtein) -> str:
312    max_dist = expression.args.get("max_dist")
313    if max_dist:
314        max_dist = exp.Kwarg(this=exp.var("max_distance"), expression=max_dist)
315
316    return self.func("EDIT_DISTANCE", expression.this, expression.expression, max_dist)
317
318
319def _json_extract_sql(self: BigQueryGenerator, expression: JSON_EXTRACT_TYPE) -> str:
320    upper = self.function_name(expression).upper()
321
322    dquote_escaping = upper in DQUOTES_ESCAPING_JSON_FUNCTIONS
323
324    if dquote_escaping:
325        self._quote_json_path_key_using_brackets = False
326
327    sql = rename_func(upper)(self, expression)
328
329    if dquote_escaping:
330        self._quote_json_path_key_using_brackets = True
331
332    return sql
333
334
335class BigQueryGenerator(generator.Generator):
336    TRY_SUPPORTED = False
337    SUPPORTS_UESCAPE = False
338    SUPPORTS_DECODE_CASE = False
339    INTERVAL_ALLOWS_PLURAL_FORM = False
340    JOIN_HINTS = False
341    QUERY_HINTS = False
342    TABLE_HINTS = False
343    LIMIT_FETCH = "LIMIT"
344    RENAME_TABLE_WITH_DB = False
345    NVL2_SUPPORTED = False
346    UNNEST_WITH_ORDINALITY = False
347    CAST_TO_INT_ROUNDS = True
348    COLLATE_IS_FUNC = True
349    LIMIT_ONLY_LITERALS = True
350    SUPPORTS_TABLE_ALIAS_COLUMNS = False
351    SUPPORTS_NAMED_CTE_COLUMNS = False
352    UNPIVOT_ALIASES_ARE_IDENTIFIERS = False
353    JSON_KEY_VALUE_PAIR_SEP = ","
354    NULL_ORDERING_SUPPORTED: bool | None = False
355    IGNORE_NULLS_IN_FUNC = True
356    JSON_PATH_SINGLE_QUOTE_ESCAPE = True
357    CAN_IMPLEMENT_ARRAY_ANY = True
358    SUPPORTS_TO_NUMBER = False
359    NAMED_PLACEHOLDER_TOKEN = "@"
360    HEX_FUNC = "TO_HEX"
361    WITH_PROPERTIES_PREFIX = "OPTIONS"
362    SUPPORTS_EXPLODING_PROJECTIONS = False
363    EXCEPT_INTERSECT_SUPPORT_ALL_CLAUSE = False
364    SUPPORTS_UNIX_SECONDS = True
365    DECLARE_DEFAULT_ASSIGNMENT = "DEFAULT"
366
367    SAFE_JSON_PATH_KEY_RE = re.compile(r"^[\-\w]*$")
368
369    WINDOW_FUNCS_WITH_NULL_ORDERING = (
370        exp.CumeDist,
371        exp.DenseRank,
372        exp.FirstValue,
373        exp.Lag,
374        exp.LastValue,
375        exp.Lead,
376        exp.NthValue,
377        exp.Ntile,
378        exp.PercentRank,
379        exp.Rank,
380        exp.RowNumber,
381    )
382
383    TS_OR_DS_TYPES: t.ClassVar = (
384        exp.TsOrDsToDatetime,
385        exp.TsOrDsToTimestamp,
386        exp.TsOrDsToTime,
387        exp.TsOrDsToDate,
388    )
389
390    TRANSFORMS = {
391        **generator.Generator.TRANSFORMS,
392        exp.AIEmbed: rename_func("EMBED"),
393        exp.AIGenerate: rename_func("GENERATE"),
394        exp.AISimilarity: rename_func("SIMILARITY"),
395        exp.ApproxTopK: rename_func("APPROX_TOP_COUNT"),
396        exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
397        exp.ArgMax: arg_max_or_min_no_count("MAX_BY"),
398        exp.ArgMin: arg_max_or_min_no_count("MIN_BY"),
399        exp.Array: inline_array_unless_query,
400        exp.ArrayContains: _array_contains_sql,
401        exp.ArrayDistinct: _array_agg_fold_sql,
402        exp.ArraySort: _array_agg_fold_sql,
403        exp.ArrayFilter: filter_array_using_unnest,
404        exp.ArrayRemove: filter_array_using_unnest,
405        exp.BitwiseAndAgg: rename_func("BIT_AND"),
406        exp.BitwiseOrAgg: rename_func("BIT_OR"),
407        exp.BitwiseXorAgg: rename_func("BIT_XOR"),
408        exp.BitwiseCount: rename_func("BIT_COUNT"),
409        exp.ByteLength: rename_func("BYTE_LENGTH"),
410        exp.Cast: transforms.preprocess(
411            [_widen_decimal, transforms.remove_precision_parameterized_types]
412        ),
413        exp.CollateProperty: lambda self, e: (
414            f"DEFAULT COLLATE {self.sql(e, 'this')}"
415            if e.args.get("default")
416            else f"COLLATE {self.sql(e, 'this')}"
417        ),
418        exp.Commit: lambda *_: "COMMIT TRANSACTION",
419        exp.CountIf: rename_func("COUNTIF"),
420        exp.Create: _create_sql,
421        exp.CTE: transforms.preprocess([_pushdown_cte_column_names]),
422        exp.DateAdd: date_add_interval_sql("DATE", "ADD"),
423        exp.DateDiff: lambda self, e: self.func("DATE_DIFF", e.this, e.expression, unit_to_var(e)),
424        exp.DateFromParts: rename_func("DATE"),
425        exp.DateStrToDate: datestrtodate_sql,
426        exp.DateSub: date_add_interval_sql("DATE", "SUB"),
427        exp.DatetimeAdd: date_add_interval_sql("DATETIME", "ADD"),
428        exp.DatetimeSub: date_add_interval_sql("DATETIME", "SUB"),
429        exp.DateFromUnixDate: rename_func("DATE_FROM_UNIX_DATE"),
430        exp.Day: _extract_part("DAY"),
431        exp.DayOfMonth: _extract_part("DAY"),
432        exp.DayOfWeekIso: _day_of_week_iso_sql,
433        exp.DayOfYear: _extract_part("DAYOFYEAR"),
434        exp.Filter: no_filter_sql,
435        exp.FromTimeZone: lambda self, e: self.func(
436            "DATETIME", self.func("TIMESTAMP", e.this, e.args.get("zone")), "'UTC'"
437        ),
438        exp.GenerateSeries: generate_series_sql("GENERATE_ARRAY"),
439        exp.GroupConcat: lambda self, e: groupconcat_sql(
440            self, e, func_name="STRING_AGG", within_group=False, sep=None
441        ),
442        exp.Hex: lambda self, e: self.func("UPPER", self.func("TO_HEX", self.sql(e, "this"))),
443        exp.HexString: lambda self, e: self.hexstring_sql(e, binary_function_repr="FROM_HEX"),
444        exp.Hour: _extract_part("HOUR"),
445        exp.If: if_sql(false_value="NULL"),
446        exp.ILike: no_ilike_sql,
447        exp.IntDiv: rename_func("DIV"),
448        exp.Int64: rename_func("INT64"),
449        exp.JSONBool: rename_func("BOOL"),
450        exp.JSONExtract: _json_extract_sql,
451        exp.JSONExtractArray: _json_extract_sql,
452        exp.JSONExtractScalar: _json_extract_sql,
453        exp.JSONFormat: lambda self, e: self.func(
454            "TO_JSON" if e.args.get("to_json") else "TO_JSON_STRING",
455            e.this,
456            e.args.get("options"),
457        ),
458        exp.JSONKeysAtDepth: rename_func("JSON_KEYS"),
459        exp.JSONValueArray: rename_func("JSON_VALUE_ARRAY"),
460        exp.Levenshtein: _levenshtein_sql,
461        exp.Max: max_or_greatest,
462        exp.MD5: lambda self, e: self.func("TO_HEX", self.func("MD5", e.this)),
463        exp.MD5Digest: rename_func("MD5"),
464        exp.Min: min_or_least,
465        exp.Month: _extract_part("MONTH"),
466        exp.Normalize: lambda self, e: self.func(
467            "NORMALIZE_AND_CASEFOLD" if e.args.get("is_casefold") else "NORMALIZE",
468            e.this,
469            e.args.get("form"),
470        ),
471        exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}",
472        exp.Quarter: _extract_part("QUARTER"),
473        exp.RegexpExtract: lambda self, e: self.func(
474            "REGEXP_EXTRACT",
475            e.this,
476            e.expression,
477            e.args.get("position"),
478            e.args.get("occurrence"),
479        ),
480        exp.RegexpExtractAll: lambda self, e: self.func("REGEXP_EXTRACT_ALL", e.this, e.expression),
481        exp.RegexpReplace: regexp_replace_sql,
482        exp.RegexpLike: rename_func("REGEXP_CONTAINS"),
483        exp.ReturnsProperty: _returnsproperty_sql,
484        exp.Rollback: lambda *_: "ROLLBACK TRANSACTION",
485        exp.ParseTime: lambda self, e: self.func("PARSE_TIME", self.format_time(e), e.this),
486        exp.ParseDatetime: lambda self, e: self.func("PARSE_DATETIME", self.format_time(e), e.this),
487        exp.Select: transforms.preprocess(
488            [
489                _unnest_explode_generate_series,
490                transforms.explode_projection_to_unnest(),
491                transforms.unqualify_unnest,
492                transforms.eliminate_distinct_on,
493                _alias_ordered_group,
494                transforms.eliminate_semi_and_anti_joins,
495            ]
496        ),
497        exp.SHA: rename_func("SHA1"),
498        exp.SHA2: sha256_sql,
499        exp.SHA1Digest: rename_func("SHA1"),
500        exp.SHA2Digest: sha2_digest_sql,
501        exp.StabilityProperty: lambda self, e: (
502            "DETERMINISTIC" if e.name == "IMMUTABLE" else "NOT DETERMINISTIC"
503        ),
504        exp.String: rename_func("STRING"),
505        exp.StrPosition: lambda self, e: strposition_sql(
506            self, e, func_name="INSTR", supports_position=True, supports_occurrence=True
507        ),
508        exp.StrToDate: _str_to_datetime_sql,
509        exp.StrToTime: _str_to_datetime_sql,
510        exp.SessionUser: lambda *_: "SESSION_USER()",
511        exp.Table: transforms.preprocess([transforms.unnest_generate_series]),
512        exp.TimeAdd: date_add_interval_sql("TIME", "ADD"),
513        exp.TimeFromParts: rename_func("TIME"),
514        exp.TimestampFromParts: rename_func("DATETIME"),
515        exp.TimeSub: date_add_interval_sql("TIME", "SUB"),
516        exp.TimestampAdd: date_add_interval_sql("TIMESTAMP", "ADD"),
517        exp.TimestampDiff: rename_func("TIMESTAMP_DIFF"),
518        exp.TimestampSub: date_add_interval_sql("TIMESTAMP", "SUB"),
519        exp.TimeStrToTime: timestrtotime_sql,
520        exp.Transaction: lambda *_: "BEGIN TRANSACTION",
521        exp.TsOrDsAdd: _ts_or_ds_add_sql,
522        exp.TsOrDsDiff: _ts_or_ds_diff_sql,
523        exp.TsOrDsToTime: rename_func("TIME"),
524        exp.TsOrDsToDatetime: rename_func("DATETIME"),
525        exp.TsOrDsToTimestamp: rename_func("TIMESTAMP"),
526        exp.Unhex: rename_func("FROM_HEX"),
527        exp.UnixDate: rename_func("UNIX_DATE"),
528        exp.UnixToTime: _unix_to_time_sql,
529        exp.Uuid: lambda *_: "GENERATE_UUID()",
530        exp.Values: _derived_table_values_to_unnest,
531        exp.VariancePop: rename_func("VAR_POP"),
532        exp.WeekOfYear: _extract_part("ISOWEEK"),
533        exp.Year: _extract_part("YEAR"),
534        exp.SafeDivide: rename_func("SAFE_DIVIDE"),
535    }
536
537    SUPPORTED_JSON_PATH_PARTS = {
538        exp.JSONPathKey,
539        exp.JSONPathRoot,
540        exp.JSONPathSubscript,
541    }
542
543    TYPE_MAPPING = {
544        **generator.Generator.TYPE_MAPPING,
545        exp.DType.BIGDECIMAL: "BIGNUMERIC",
546        exp.DType.BIGINT: "INT64",
547        exp.DType.BINARY: "BYTES",
548        exp.DType.BLOB: "BYTES",
549        exp.DType.BOOLEAN: "BOOL",
550        exp.DType.CHAR: "STRING",
551        exp.DType.DECIMAL: "NUMERIC",
552        exp.DType.DOUBLE: "FLOAT64",
553        exp.DType.FLOAT: "FLOAT64",
554        exp.DType.INT: "INT64",
555        exp.DType.NCHAR: "STRING",
556        exp.DType.NVARCHAR: "STRING",
557        exp.DType.SMALLINT: "INT64",
558        exp.DType.TEXT: "STRING",
559        exp.DType.TIMESTAMP: "DATETIME",
560        exp.DType.TIMESTAMPNTZ: "DATETIME",
561        exp.DType.TIMESTAMPTZ: "TIMESTAMP",
562        exp.DType.TIMESTAMPLTZ: "TIMESTAMP",
563        exp.DType.TINYINT: "INT64",
564        exp.DType.ROWVERSION: "BYTES",
565        exp.DType.UUID: "STRING",
566        exp.DType.VARBINARY: "BYTES",
567        exp.DType.VARCHAR: "STRING",
568        exp.DType.VARIANT: "ANY TYPE",
569    }
570
571    PROPERTIES_LOCATION = {
572        **generator.Generator.PROPERTIES_LOCATION,
573        exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
574        exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
575    }
576
577    # WINDOW comes after QUALIFY
578    # https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#window_clause
579    # BigQuery requires QUALIFY before WINDOW
580    AFTER_HAVING_MODIFIER_TRANSFORMS = {
581        "qualify": generator.AFTER_HAVING_MODIFIER_TRANSFORMS["qualify"],
582        "windows": generator.AFTER_HAVING_MODIFIER_TRANSFORMS["windows"],
583    }
584
585    # from: https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#reserved_keywords
586    RESERVED_KEYWORDS = {
587        "all",
588        "and",
589        "any",
590        "array",
591        "as",
592        "asc",
593        "assert_rows_modified",
594        "at",
595        "between",
596        "by",
597        "case",
598        "cast",
599        "collate",
600        "contains",
601        "create",
602        "cross",
603        "cube",
604        "current",
605        "default",
606        "define",
607        "desc",
608        "distinct",
609        "else",
610        "end",
611        "enum",
612        "escape",
613        "except",
614        "exclude",
615        "exists",
616        "extract",
617        "false",
618        "fetch",
619        "following",
620        "for",
621        "from",
622        "full",
623        "group",
624        "grouping",
625        "groups",
626        "hash",
627        "having",
628        "if",
629        "ignore",
630        "in",
631        "inner",
632        "intersect",
633        "interval",
634        "into",
635        "is",
636        "join",
637        "lateral",
638        "left",
639        "like",
640        "limit",
641        "lookup",
642        "merge",
643        "natural",
644        "new",
645        "no",
646        "not",
647        "null",
648        "nulls",
649        "of",
650        "on",
651        "or",
652        "order",
653        "outer",
654        "over",
655        "partition",
656        "preceding",
657        "proto",
658        "qualify",
659        "range",
660        "recursive",
661        "respect",
662        "right",
663        "rollup",
664        "rows",
665        "select",
666        "set",
667        "some",
668        "struct",
669        "tablesample",
670        "then",
671        "to",
672        "treat",
673        "true",
674        "unbounded",
675        "union",
676        "unnest",
677        "using",
678        "when",
679        "where",
680        "window",
681        "with",
682        "within",
683    }
684
685    def frombase_sql(self, expression: exp.FromBase) -> str:
686        # Hex is the only base BigQuery can cast from, and the sign goes before the prefix, e.g. -0xff
687        if expression.expression.name != "16":
688            self.unsupported("FROM_BASE is only supported for base 16")
689            return self.function_fallback_sql(expression)
690
691        prefixed = exp.func(
692            "REGEXP_REPLACE",
693            expression.this,
694            exp.Literal.string("^([+-]?)"),
695            exp.Literal.string("\\10x"),
696            copy=False,
697        )
698        return self.sql(exp.cast(prefixed, exp.DType.BIGINT, copy=False))
699
700    def approxquantile_sql(self, expression: exp.ApproxQuantile) -> str:
701        node: exp.Expr = expression
702        while node.parent and not isinstance(node.parent, (exp.Window, exp.AggFunc, exp.Select)):
703            node = node.parent
704        if isinstance(node.parent, exp.Window) and node.arg_key == "this":
705            self.unsupported("APPROX_QUANTILES does not support window functions")
706            return self.function_fallback_sql(expression)
707
708        # The k/n quantile is APPROX_QUANTILES(x, n)[k], e.g. 0.95 -> n = 20, k = 19. n is at most
709        # 100,000, so the quantile has at most 5 decimal places, and Decimal comparisons are exact
710        quantile = expression.args["quantile"]
711        weight = expression.args.get("weight")
712        value = Decimal(quantile.to_py()) if quantile.is_number else None
713        if (
714            value is None
715            or not 0 <= value <= 1
716            or round(value, 5) != value
717            or (weight and (not weight.is_number or weight.to_py() != 1))
718        ):
719            self.unsupported("APPROX_QUANTILE needs a literal quantile and no weight")
720            return self.function_fallback_sql(expression)
721
722        fraction = Fraction(value)
723        quantiles = exp.func("APPROX_QUANTILES", expression.this, fraction.denominator, copy=False)
724        return self.sql(quantiles[fraction.numerator])
725
726    def weekstart_sql(self, expression: exp.WeekStart) -> str:
727        if expression.this.name.upper() == "SUNDAY":
728            # BigQuery specific optimization since WEEK(SUNDAY) == WEEK
729            return "WEEK"
730
731        return self.func("WEEK", expression.this)
732
733    def datetrunc_sql(self, expression: exp.DateTrunc) -> str:
734        unit = expression.unit
735        unit_sql = unit.name if unit.is_string else self.sql(unit)
736        return self.func("DATE_TRUNC", expression.this, unit_sql, expression.args.get("zone"))
737
738    def mod_sql(self, expression: exp.Mod) -> str:
739        this = expression.this
740        expr = expression.expression
741        return self.func(
742            "MOD",
743            this.unnest() if isinstance(this, exp.Paren) else this,
744            expr.unnest() if isinstance(expr, exp.Paren) else expr,
745        )
746
747    def column_parts(self, expression: exp.Column) -> str:
748        if expression.meta_get("quoted_column"):
749            # If a column reference is of the form `dataset.table`.name, we need
750            # to preserve the quoted table path, otherwise the reference breaks
751            table_parts = ".".join(p.name for p in expression.parts[:-1])
752            table_path = self.sql(exp.Identifier(this=table_parts, quoted=True))
753            return f"{table_path}.{self.sql(expression, 'this')}"
754
755        return super().column_parts(expression)
756
757    def table_parts(self, expression: exp.Table) -> str:
758        # Depending on the context, `x.y` may not resolve to the same data source as `x`.`y`, so
759        # we need to make sure the correct quoting is used in each case.
760        #
761        # For example, if there is a CTE x that clashes with a schema name, then the former will
762        # return the table y in that schema, whereas the latter will return the CTE's y column:
763        #
764        # - WITH x AS (SELECT [1, 2] AS y) SELECT * FROM x, `x.y`   -> cross join
765        # - WITH x AS (SELECT [1, 2] AS y) SELECT * FROM x, `x`.`y` -> implicit unnest
766        if expression.meta_get("quoted_table"):
767            table_parts = ".".join(p.name for p in expression.parts)
768            return self.sql(exp.Identifier(this=table_parts, quoted=True))
769
770        return super().table_parts(expression)
771
772    def timetostr_sql(self, expression: exp.TimeToStr) -> str:
773        this = expression.this
774        if isinstance(this, exp.TsOrDsToDatetime):
775            func_name = "FORMAT_DATETIME"
776        elif isinstance(this, exp.TsOrDsToTimestamp):
777            func_name = "FORMAT_TIMESTAMP"
778        elif isinstance(this, exp.TsOrDsToTime):
779            func_name = "FORMAT_TIME"
780        else:
781            func_name = "FORMAT_DATE"
782
783        time_expr = this if isinstance(this, self.TS_OR_DS_TYPES) else expression
784        return self.func(
785            func_name, self.format_time(expression), time_expr.this, expression.args.get("zone")
786        )
787
788    def eq_sql(self, expression: exp.EQ) -> str:
789        # Operands of = cannot be NULL in BigQuery
790        if isinstance(expression.left, exp.Null) or isinstance(expression.right, exp.Null):
791            if not isinstance(expression.parent, exp.Update):
792                return "NULL"
793
794        return self.binary(expression, "=")
795
796    def attimezone_sql(self, expression: exp.AtTimeZone) -> str:
797        parent = expression.parent
798
799        # BigQuery allows CAST(.. AS {STRING|TIMESTAMP} [FORMAT <fmt> [AT TIME ZONE <tz>]]).
800        # Only the TIMESTAMP one should use the below conversion, when AT TIME ZONE is included.
801        if not isinstance(parent, exp.Cast) or not parent.to.is_type("text"):
802            return self.func(
803                "TIMESTAMP", self.func("DATETIME", expression.this, expression.args.get("zone"))
804            )
805
806        return super().attimezone_sql(expression)
807
808    def trycast_sql(self, expression: exp.TryCast) -> str:
809        return self.cast_sql(expression, safe_prefix="SAFE_")
810
811    def bracket_sql(self, expression: exp.Bracket) -> str:
812        this = expression.this
813        expressions = expression.expressions
814
815        if len(expressions) == 1 and this and this.is_type(exp.DType.STRUCT):
816            arg = expressions[0]
817            if arg.type is None:
818                from sqlglot.optimizer.annotate_types import annotate_types
819
820                arg = annotate_types(arg, dialect=self.dialect)
821
822            if arg.type and arg.type.this in exp.DataType.TEXT_TYPES:
823                # BQ doesn't support bracket syntax with string values for structs
824                return f"{self.sql(this)}.{arg.name}"
825
826        expressions_sql = self.expressions(expression, flat=True)
827        offset = expression.args.get("offset")
828
829        if offset == 0:
830            expressions_sql = f"OFFSET({expressions_sql})"
831        elif offset == 1:
832            expressions_sql = f"ORDINAL({expressions_sql})"
833        elif offset is not None:
834            self.unsupported(f"Unsupported array offset: {offset}")
835
836        if expression.args.get("safe"):
837            expressions_sql = f"SAFE_{expressions_sql}"
838
839        return f"{self.sql(this)}[{expressions_sql}]"
840
841    def in_unnest_op(self, expression: exp.Unnest) -> str:
842        return self.sql(expression)
843
844    def version_sql(self, expression: exp.Version) -> str:
845        if expression.name == "TIMESTAMP":
846            expression.set("this", "SYSTEM_TIME")
847        return super().version_sql(expression)
848
849    def contains_sql(self, expression: exp.Contains) -> str:
850        this = expression.this
851        expr = expression.expression
852
853        if isinstance(this, exp.Lower) and isinstance(expr, exp.Lower):
854            this = this.this
855            expr = expr.this
856
857        return self.func("CONTAINS_SUBSTR", this, expr, expression.args.get("json_scope"))
858
859    def parsejson_sql(self, expression: exp.ParseJSON) -> str:
860        if expression.args.get("is_literal"):
861            return f"JSON {self.sql(expression, 'this')}"
862        return super().parsejson_sql(expression)
863
864    def cast_sql(self, expression: exp.Cast, safe_prefix: str | None = None) -> str:
865        this = expression.this
866
867        # This ensures that inline type-annotated ARRAY literals like ARRAY<INT64>[1, 2, 3]
868        # are roundtripped unaffected. The inner check excludes ARRAY(SELECT ...) expressions,
869        # because they aren't literals and so the above syntax is invalid BigQuery.
870        if isinstance(this, exp.Array):
871            elem = seq_get(this.expressions, 0)
872            if not (elem and elem.find(exp.Query)):
873                return f"{self.sql(expression, 'to')}{self.sql(this)}"
874
875        return super().cast_sql(expression, safe_prefix=safe_prefix)
876
877    def clusterproperty_sql(self, expression: exp.ClusterProperty) -> str:
878        if expression.this:
879            self.unsupported(f"Unsupported CLUSTER BY {self.sql(expression, 'this')}")
880            return ""
881        return self.op_expressions("CLUSTER BY", expression)