sqlglot.parsers.bigquery

  1from __future__ import annotations
  2
  3import re
  4import typing as t
  5
  6from sqlglot import exp, parser
  7from sqlglot.dialects.dialect import (
  8    binary_from_function,
  9    build_date_delta_with_interval,
 10    build_formatted_time,
 11)
 12from sqlglot.helper import seq_get, split_num_words
 13from sqlglot.tokens import TokenType
 14
 15if t.TYPE_CHECKING:
 16    from sqlglot._typing import E
 17
 18
 19def _build_contains_substring(args: list) -> exp.Contains:
 20    this = exp.Lower(this=seq_get(args, 0))
 21    expr = exp.Lower(this=seq_get(args, 1))
 22    return exp.Contains(this=this, expression=expr, json_scope=seq_get(args, 2))
 23
 24
 25def _build_date(args: list) -> exp.Date | exp.DateFromParts:
 26    expr_type = exp.DateFromParts if len(args) == 3 else exp.Date
 27    return expr_type.from_arg_list(args)
 28
 29
 30def build_date_diff(args: list) -> exp.Expr:
 31    expr = exp.DateDiff(
 32        this=seq_get(args, 0),
 33        expression=seq_get(args, 1),
 34        unit=seq_get(args, 2),
 35        date_part_boundary=True,
 36    )
 37
 38    unit = expr.args.get("unit")
 39    if isinstance(unit, exp.Var) and unit.name.upper() == "WEEK":
 40        expr.set("unit", exp.WeekStart(this=exp.var("SUNDAY")))
 41
 42    return expr
 43
 44
 45def _build_datetime(args: list) -> exp.Func:
 46    if len(args) == 1:
 47        return exp.TsOrDsToDatetime.from_arg_list(args)
 48    if len(args) == 2:
 49        return exp.Datetime.from_arg_list(args)
 50    return exp.TimestampFromParts.from_arg_list(args)
 51
 52
 53def _build_extract_json_with_default_path(
 54    expr_type: type[E],
 55) -> t.Callable:
 56    def _builder(args: list, dialect: t.Any) -> E:
 57        if len(args) == 1:
 58            args.append(exp.Literal.string("$"))
 59        return parser.build_extract_json_with_path(expr_type)(args, dialect)
 60
 61    return _builder
 62
 63
 64def _build_format_time(expr_type: type[exp.Expr]) -> t.Callable[[list], exp.TimeToStr]:
 65    def _builder(args: list) -> exp.TimeToStr:
 66        formatted_time = build_formatted_time(exp.TimeToStr, "bigquery")(
 67            [expr_type(this=seq_get(args, 1)), seq_get(args, 0)]
 68        )
 69        formatted_time.set("zone", seq_get(args, 2))
 70        return formatted_time
 71
 72    return _builder
 73
 74
 75def _build_json_strip_nulls(args: list) -> exp.JSONStripNulls:
 76    expression = exp.JSONStripNulls(this=seq_get(args, 0))
 77    for arg in args[1:]:
 78        if isinstance(arg, exp.Kwarg):
 79            expression.set(arg.this.name.lower(), arg)
 80        else:
 81            expression.set("expression", arg)
 82    return expression
 83
 84
 85def _build_levenshtein(args: list) -> exp.Levenshtein:
 86    max_dist = seq_get(args, 2)
 87    return exp.Levenshtein(
 88        this=seq_get(args, 0),
 89        expression=seq_get(args, 1),
 90        max_dist=max_dist.expression if max_dist else None,
 91    )
 92
 93
 94def _build_parse_timestamp(args: list) -> exp.StrToTime:
 95    this = build_formatted_time(exp.StrToTime, "bigquery")([seq_get(args, 1), seq_get(args, 0)])
 96    this.set("zone", seq_get(args, 2))
 97    return this
 98
 99
100def _build_regexp_extract(expr_type: type[E], default_group: exp.Expr | None = None) -> t.Callable:
101    def _builder(args: list, dialect: t.Any) -> E:
102        try:
103            group = re.compile(args[1].name).groups == 1
104        except re.error:
105            group = False
106
107        return expr_type(
108            this=seq_get(args, 0),
109            expression=seq_get(args, 1),
110            position=seq_get(args, 2),
111            occurrence=seq_get(args, 3),
112            group=exp.Literal.number(1) if group else default_group,
113            **(
114                {"null_if_pos_overflow": dialect.REGEXP_EXTRACT_POSITION_OVERFLOW_RETURNS_NULL}
115                if expr_type is exp.RegexpExtract
116                else {}
117            ),
118        )
119
120    return _builder
121
122
123def _build_time(args: list) -> exp.Func:
124    if len(args) == 1:
125        return exp.TsOrDsToTime(this=args[0])
126    if len(args) == 2:
127        return exp.Time.from_arg_list(args)
128    return exp.TimeFromParts.from_arg_list(args)
129
130
131def _build_timestamp(args: list) -> exp.Timestamp:
132    timestamp = exp.Timestamp.from_arg_list(args)
133    timestamp.set("with_tz", True)
134    return timestamp
135
136
137def _build_to_hex(args: list) -> exp.Hex | exp.MD5:
138    arg = seq_get(args, 0)
139    return exp.MD5(this=arg.this) if isinstance(arg, exp.MD5Digest) else exp.LowerHex(this=arg)
140
141
142MAKE_INTERVAL_KWARGS = ["year", "month", "day", "hour", "minute", "second"]
143
144
145class BigQueryParser(parser.Parser):
146    PREFIXED_PIVOT_COLUMNS: t.ClassVar = True
147    LOG_DEFAULTS_TO_LN: t.ClassVar = True
148    SUPPORTS_IMPLICIT_UNNEST: t.ClassVar = True
149    JOINS_HAVE_EQUAL_PRECEDENCE: t.ClassVar = True
150
151    # BigQuery does not allow ASC/DESC to be used as an identifier, allows GRANT as an identifier
152    ID_VAR_TOKENS: t.ClassVar = {
153        *parser.Parser.ID_VAR_TOKENS,
154        TokenType.GRANT,
155    } - {TokenType.ASC, TokenType.DESC}
156
157    ALIAS_TOKENS: t.ClassVar = {
158        *parser.Parser.ALIAS_TOKENS,
159        TokenType.GRANT,
160    } - {TokenType.ASC, TokenType.DESC}
161
162    TABLE_ALIAS_TOKENS: t.ClassVar = {
163        *parser.Parser.TABLE_ALIAS_TOKENS,
164        TokenType.ANTI,
165        TokenType.GRANT,
166        TokenType.SEMI,
167    } - {TokenType.ASC, TokenType.DESC}
168
169    COMMENT_TABLE_ALIAS_TOKENS: t.ClassVar = {
170        *parser.Parser.COMMENT_TABLE_ALIAS_TOKENS,
171        TokenType.GRANT,
172    } - {TokenType.ASC, TokenType.DESC}
173
174    UPDATE_ALIAS_TOKENS: t.ClassVar = {
175        *parser.Parser.UPDATE_ALIAS_TOKENS,
176        TokenType.GRANT,
177    } - {TokenType.ASC, TokenType.DESC}
178
179    FUNCTIONS: t.ClassVar[dict[str, t.Callable]] = {
180        **{k: v for k, v in parser.Parser.FUNCTIONS.items() if k != "SEARCH"},
181        "APPROX_TOP_COUNT": exp.ApproxTopK.from_arg_list,
182        "BIT_AND": exp.BitwiseAndAgg.from_arg_list,
183        "BIT_OR": exp.BitwiseOrAgg.from_arg_list,
184        "BIT_XOR": exp.BitwiseXorAgg.from_arg_list,
185        "BIT_COUNT": exp.BitwiseCount.from_arg_list,
186        "BOOL": exp.JSONBool.from_arg_list,
187        "CONTAINS_SUBSTR": _build_contains_substring,
188        "DATE": _build_date,
189        "DATE_ADD": build_date_delta_with_interval(exp.DateAdd),
190        "DATE_DIFF": build_date_diff,
191        "DATE_SUB": build_date_delta_with_interval(exp.DateSub),
192        "DATE_TRUNC": lambda args: exp.DateTrunc(
193            unit=seq_get(args, 1),
194            this=seq_get(args, 0),
195            zone=seq_get(args, 2),
196        ),
197        "DATETIME": _build_datetime,
198        "DATETIME_ADD": build_date_delta_with_interval(exp.DatetimeAdd),
199        "DATETIME_SUB": build_date_delta_with_interval(exp.DatetimeSub),
200        "DIV": binary_from_function(exp.IntDiv),
201        "EDIT_DISTANCE": _build_levenshtein,
202        "EMBED": exp.AIEmbed.from_arg_list,
203        "FORMAT_DATE": _build_format_time(exp.TsOrDsToDate),
204        "GENERATE": exp.AIGenerate.from_arg_list,
205        "GENERATE_ARRAY": exp.GenerateSeries.from_arg_list,
206        "JSON_EXTRACT_SCALAR": _build_extract_json_with_default_path(exp.JSONExtractScalar),
207        "JSON_EXTRACT_ARRAY": _build_extract_json_with_default_path(exp.JSONExtractArray),
208        "JSON_EXTRACT_STRING_ARRAY": _build_extract_json_with_default_path(exp.JSONValueArray),
209        "JSON_KEYS": exp.JSONKeysAtDepth.from_arg_list,
210        "JSON_QUERY": parser.build_extract_json_with_path(exp.JSONExtract),
211        "JSON_QUERY_ARRAY": _build_extract_json_with_default_path(exp.JSONExtractArray),
212        "JSON_STRIP_NULLS": _build_json_strip_nulls,
213        "JSON_VALUE": _build_extract_json_with_default_path(exp.JSONExtractScalar),
214        "JSON_VALUE_ARRAY": _build_extract_json_with_default_path(exp.JSONValueArray),
215        "LENGTH": lambda args: exp.Length(this=seq_get(args, 0), binary=True),
216        "MD5": exp.MD5Digest.from_arg_list,
217        "SHA1": exp.SHA1Digest.from_arg_list,
218        "NORMALIZE_AND_CASEFOLD": lambda args: exp.Normalize(
219            this=seq_get(args, 0), form=seq_get(args, 1), is_casefold=True
220        ),
221        "OCTET_LENGTH": exp.ByteLength.from_arg_list,
222        "TO_HEX": _build_to_hex,
223        "PARSE_DATE": lambda args: build_formatted_time(exp.StrToDate, "bigquery")(
224            [seq_get(args, 1), seq_get(args, 0)]
225        ),
226        "PARSE_TIME": lambda args: build_formatted_time(exp.ParseTime, "bigquery")(
227            [seq_get(args, 1), seq_get(args, 0)]
228        ),
229        "PARSE_TIMESTAMP": _build_parse_timestamp,
230        "PARSE_DATETIME": lambda args: build_formatted_time(exp.ParseDatetime, "bigquery")(
231            [seq_get(args, 1), seq_get(args, 0)]
232        ),
233        "REGEXP_CONTAINS": exp.RegexpLike.from_arg_list,
234        "REGEXP_EXTRACT": _build_regexp_extract(exp.RegexpExtract),
235        "REGEXP_SUBSTR": _build_regexp_extract(exp.RegexpExtract),
236        "REGEXP_EXTRACT_ALL": _build_regexp_extract(
237            exp.RegexpExtractAll, default_group=exp.Literal.number(0)
238        ),
239        "SHA256": lambda args: exp.SHA2Digest(
240            this=seq_get(args, 0), length=exp.Literal.number(256)
241        ),
242        "SHA512": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(512)),
243        "SIMILARITY": exp.AISimilarity.from_arg_list,
244        "SPLIT": lambda args: exp.Split(
245            # https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#split
246            this=seq_get(args, 0),
247            expression=seq_get(args, 1) or exp.Literal.string(","),
248        ),
249        "STRPOS": exp.StrPosition.from_arg_list,
250        "TIME": _build_time,
251        "TIME_ADD": build_date_delta_with_interval(exp.TimeAdd),
252        "TIME_SUB": build_date_delta_with_interval(exp.TimeSub),
253        "TIMESTAMP": _build_timestamp,
254        "TIMESTAMP_ADD": build_date_delta_with_interval(exp.TimestampAdd),
255        "TIMESTAMP_SUB": build_date_delta_with_interval(exp.TimestampSub),
256        "TIMESTAMP_MICROS": lambda args: exp.UnixToTime(
257            this=seq_get(args, 0), scale=exp.UnixToTime.MICROS
258        ),
259        "TIMESTAMP_MILLIS": lambda args: exp.UnixToTime(
260            this=seq_get(args, 0), scale=exp.UnixToTime.MILLIS
261        ),
262        "TIMESTAMP_SECONDS": lambda args: exp.UnixToTime(this=seq_get(args, 0)),
263        "TO_JSON": lambda args: exp.JSONFormat(
264            this=seq_get(args, 0), options=seq_get(args, 1), to_json=True
265        ),
266        "TO_JSON_STRING": exp.JSONFormat.from_arg_list,
267        "FORMAT_DATETIME": _build_format_time(exp.TsOrDsToDatetime),
268        "FORMAT_TIMESTAMP": _build_format_time(exp.TsOrDsToTimestamp),
269        "FORMAT_TIME": _build_format_time(exp.TsOrDsToTime),
270        "FROM_HEX": exp.Unhex.from_arg_list,
271        "WEEK": lambda args: exp.WeekStart(this=exp.var(seq_get(args, 0))),
272    }
273
274    FUNCTION_PARSERS = {
275        **{k: v for k, v in parser.Parser.FUNCTION_PARSERS.items() if k != "TRIM"},
276        "ARRAY": lambda self: self.expression(
277            exp.Array(expressions=[self._parse_statement()], struct_name_inheritance=True)
278        ),
279        "JSON_ARRAY": lambda self: self.expression(
280            exp.JSONArray(expressions=self._parse_csv(self._parse_bitwise))
281        ),
282        "MAKE_INTERVAL": lambda self: self._parse_make_interval(),
283        "PREDICT": lambda self: self._parse_ml(exp.Predict),
284        "TRANSLATE": lambda self: self._parse_translate(),
285        "FEATURES_AT_TIME": lambda self: self._parse_features_at_time(),
286        "GENERATE_EMBEDDING": lambda self: self._parse_ml(exp.GenerateEmbedding),
287        "GENERATE_TEXT_EMBEDDING": lambda self: self._parse_ml(exp.GenerateEmbedding, is_text=True),
288        "GENERATE_TEXT": lambda self: self._parse_generate(exp.GenerateText),
289        "GENERATE_TABLE": lambda self: self._parse_generate(exp.GenerateTable),
290        "GENERATE_BOOL": lambda self: self._parse_generate(exp.GenerateBool),
291        "GENERATE_INT": lambda self: self._parse_generate(exp.GenerateInt),
292        "GENERATE_DOUBLE": lambda self: self._parse_generate(exp.GenerateDouble),
293        "VECTOR_SEARCH": lambda self: self._parse_vector_search(),
294        "FORECAST": lambda self: self._parse_forecast(),
295    }
296
297    NO_PAREN_FUNCTIONS: t.ClassVar = {
298        **parser.Parser.NO_PAREN_FUNCTIONS,
299        TokenType.CURRENT_DATETIME: exp.CurrentDatetime,
300    }
301
302    NESTED_TYPE_TOKENS: t.ClassVar = {
303        *parser.Parser.NESTED_TYPE_TOKENS,
304        TokenType.TABLE,
305    }
306
307    PROPERTY_PARSERS: t.ClassVar = {
308        **parser.Parser.PROPERTY_PARSERS,
309        "NOT DETERMINISTIC": lambda self: self.expression(
310            exp.StabilityProperty(this=exp.Literal.string("VOLATILE"))
311        ),
312        "OPTIONS": lambda self: self._parse_with_property(),
313    }
314
315    CONSTRAINT_PARSERS: t.ClassVar = {
316        **parser.Parser.CONSTRAINT_PARSERS,
317        "OPTIONS": lambda self: exp.Properties(expressions=self._parse_with_property()),
318    }
319
320    RANGE_PARSERS: t.ClassVar = {
321        k: v for k, v in parser.Parser.RANGE_PARSERS.items() if k != TokenType.OVERLAPS
322    }
323
324    DASHED_TABLE_PART_FOLLOW_TOKENS: t.ClassVar = {
325        TokenType.DOT,
326        TokenType.L_PAREN,
327        TokenType.R_PAREN,
328    }
329
330    STATEMENT_PARSERS: t.ClassVar = {
331        **parser.Parser.STATEMENT_PARSERS,
332        TokenType.ELSE: lambda self: self._parse_as_command(self._prev),
333        TokenType.END: lambda self: self._parse_as_command(self._prev),
334        TokenType.FOR: lambda self: self._parse_for_in(),
335        TokenType.EXPORT: lambda self: self._parse_export_data(),
336        TokenType.DECLARE: lambda self: self._parse_declare(),
337    }
338
339    BRACKET_OFFSETS: t.ClassVar = {
340        "OFFSET": (0, False),
341        "ORDINAL": (1, False),
342        "SAFE_OFFSET": (0, True),
343        "SAFE_ORDINAL": (1, True),
344    }
345
346    def _parse_for_in(self) -> exp.ForIn | exp.Command:
347        index = self._index
348        this = self._parse_range()
349        self._match_text_seq("DO")
350        if self._match(TokenType.COMMAND):
351            self._retreat(index)
352            return self._parse_as_command(self._prev)
353        return self.expression(exp.ForIn(this=this, expression=self._parse_statement()))
354
355    def _parse_table_part(self, schema: bool = False) -> exp.Expr | None:
356        this = super()._parse_table_part(schema=schema) or self._parse_number()
357
358        # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#table_names
359        if isinstance(this, exp.Identifier):
360            table_name = this.name
361            while self._match(TokenType.DASH, advance=False) and self._next:
362                start = self._curr
363                while self._is_connected() and not self._match_set(
364                    self.DASHED_TABLE_PART_FOLLOW_TOKENS, advance=False
365                ):
366                    self._advance()
367
368                if start == self._curr:
369                    break
370
371                table_name += self._find_sql(start, self._prev)
372
373            this = exp.Identifier(this=table_name, quoted=this.args.get("quoted")).update_positions(
374                this
375            )
376        elif isinstance(this, exp.Literal):
377            table_name = this.name
378
379            if self._is_connected() and self._parse_var(any_token=True):
380                table_name += self._prev.text
381
382            this = exp.Identifier(this=table_name, quoted=True).update_positions(this)
383
384        return this
385
386    def _parse_table_parts(
387        self,
388        schema: bool = False,
389        is_db_reference: bool = False,
390        wildcard: bool = False,
391        fast: bool = False,
392    ) -> exp.Table | exp.Dot | None:
393        table = super()._parse_table_parts(
394            schema=schema, is_db_reference=is_db_reference, wildcard=True, fast=fast
395        )
396
397        if not isinstance(table, exp.Table):
398            return table
399
400        # proj-1.db.tbl -- `1.` is tokenized as a float so we need to unravel it here
401        if not table.catalog:
402            if table.db:
403                previous_db = table.args["db"]
404                parts = table.db.split(".")
405                if len(parts) == 2 and not table.args["db"].quoted:
406                    table.set(
407                        "catalog", exp.Identifier(this=parts[0]).update_positions(previous_db)
408                    )
409                    table.set("db", exp.Identifier(this=parts[1]).update_positions(previous_db))
410            else:
411                previous_this = table.this
412                parts = table.name.split(".")
413                if len(parts) == 2 and not table.this.quoted:
414                    table.set("db", exp.Identifier(this=parts[0]).update_positions(previous_this))
415                    table.set("this", exp.Identifier(this=parts[1]).update_positions(previous_this))
416
417        if isinstance(table.this, exp.Identifier) and any("." in p.name for p in table.parts):
418            alias = table.this
419            catalog, db, this_id, *rest = (
420                exp.to_identifier(p, quoted=True)
421                for p in split_num_words(".".join(p.name for p in table.parts), ".", 3)
422            )
423
424            for part in (catalog, db, this_id):
425                if part:
426                    part.update_positions(table.this)
427
428            this: exp.Expr | None = this_id
429            if rest and this:
430                this = exp.Dot.build([this, *rest])  # type: ignore[list-item]
431
432            table = exp.Table(this=this, db=db, catalog=catalog, pivots=table.args.get("pivots"))
433            table.meta["quoted_table"] = True
434        else:
435            alias = None
436
437        # The `INFORMATION_SCHEMA` views in BigQuery need to be qualified by a region or
438        # dataset, so if the project identifier is omitted we need to fix the ast so that
439        # the `INFORMATION_SCHEMA.X` bit is represented as a single (quoted) Identifier.
440        # Otherwise, we wouldn't correctly qualify a `Table` node that references these
441        # views, because it would seem like the "catalog" part is set, when it'd actually
442        # be the region/dataset. Merging the two identifiers into a single one is done to
443        # avoid producing a 4-part Table reference, which would cause issues in the schema
444        # module, when there are 3-part table names mixed with information schema views.
445        #
446        # See: https://cloud.google.com/bigquery/docs/information-schema-intro#syntax
447        table_parts = table.parts
448        if len(table_parts) > 1 and table_parts[-2].name.upper() == "INFORMATION_SCHEMA":
449            # We need to alias the table here to avoid breaking existing qualified columns.
450            # This is expected to be safe, because if there's an actual alias coming up in
451            # the token stream, it will overwrite this one. If there isn't one, we are only
452            # exposing the name that can be used to reference the view explicitly (a no-op).
453            exp.alias_(
454                table,
455                t.cast(exp.Identifier, alias or table_parts[-1]),
456                table=True,
457                copy=False,
458            )
459
460            info_schema_view = f"{table_parts[-2].name}.{table_parts[-1].name}"
461            new_this = exp.Identifier(this=info_schema_view, quoted=True).update_positions(
462                line=table_parts[-2].meta.get("line"),
463                col=table_parts[-1].meta.get("col"),
464                start=table_parts[-2].meta.get("start"),
465                end=table_parts[-1].meta.get("end"),
466            )
467            table.set("this", new_this)
468            table.set("db", seq_get(table_parts, -3))
469            table.set("catalog", seq_get(table_parts, -4))
470
471        return table
472
473    def _parse_column(self) -> exp.Expr | None:
474        column = super()._parse_column()
475        if isinstance(column, exp.Column):
476            parts = column.parts
477            if any("." in p.name for p in parts):
478                catalog, db, table, this, *rest = (
479                    exp.to_identifier(p, quoted=True)
480                    for p in split_num_words(".".join(p.name for p in parts), ".", 4)
481                )
482
483                if rest and this:
484                    this = exp.Dot.build([this, *rest])  # type: ignore
485
486                column = exp.Column(this=this, table=table, db=db, catalog=catalog)
487                column.meta["quoted_column"] = True
488
489        return column
490
491    @t.overload
492    def _parse_json_object(self, agg: t.Literal[False]) -> exp.JSONObject: ...
493
494    @t.overload
495    def _parse_json_object(self, agg: t.Literal[True]) -> exp.JSONObjectAgg: ...
496
497    def _parse_json_object(self, agg=False):
498        json_object = super()._parse_json_object()
499        array_kv_pair = seq_get(json_object.expressions, 0)
500
501        # Converts BQ's "signature 2" of JSON_OBJECT into SQLGlot's canonical representation
502        # https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#json_object_signature2
503        if (
504            array_kv_pair
505            and isinstance(array_kv_pair.this, exp.Array)
506            and isinstance(array_kv_pair.expression, exp.Array)
507        ):
508            keys = array_kv_pair.this.expressions
509            values = array_kv_pair.expression.expressions
510
511            json_object.set(
512                "expressions",
513                [exp.JSONKeyValue(this=k, expression=v) for k, v in zip(keys, values)],
514            )
515
516        return json_object
517
518    def _parse_bracket(self, this: exp.Expr | None = None) -> exp.Expr | None:
519        bracket = super()._parse_bracket(this)
520
521        if isinstance(bracket, exp.Array):
522            bracket.set("struct_name_inheritance", True)
523
524        if this is bracket:
525            return bracket
526
527        if isinstance(bracket, exp.Bracket):
528            for expression in bracket.expressions:
529                name = expression.name.upper()
530
531                expressions = expression.expressions
532
533                if name not in self.BRACKET_OFFSETS or not expressions:
534                    break
535
536                offset, safe = self.BRACKET_OFFSETS[name]
537                bracket.set("offset", offset)
538                bracket.set("safe", safe)
539                expression.replace(expressions[0])
540
541        return bracket
542
543    def _parse_unnest(self, with_alias: bool = True) -> exp.Unnest | None:
544        unnest = super()._parse_unnest(with_alias=with_alias)
545
546        if not unnest:
547            return None
548
549        unnest_expr = seq_get(unnest.expressions, 0)
550        if unnest_expr:
551            from sqlglot.optimizer.annotate_types import annotate_types
552
553            unnest_expr = annotate_types(unnest_expr, dialect=self.dialect)
554
555            # Unnesting a nested array (i.e array of structs) explodes the top-level struct fields,
556            # in contrast to other dialects such as DuckDB which flattens only the array by default
557            if unnest_expr.is_type(exp.DType.ARRAY) and any(
558                array_elem.is_type(exp.DType.STRUCT) for array_elem in unnest_expr._type.expressions
559            ):
560                unnest.set("explode_array", True)
561
562        return unnest
563
564    def _parse_make_interval(self) -> exp.MakeInterval:
565        expr = exp.MakeInterval()
566
567        for arg_key in MAKE_INTERVAL_KWARGS:
568            value = self._parse_lambda()
569
570            if not value:
571                break
572
573            # Non-named arguments are filled sequentially, (optionally) followed by named arguments
574            # that can appear in any order e.g MAKE_INTERVAL(1, minute => 5, day => 2)
575            if isinstance(value, exp.Kwarg):
576                arg_key = value.this.name
577
578            expr.set(arg_key, value)
579
580            self._match(TokenType.COMMA)
581
582        return expr
583
584    def _parse_ml(self, expr_type: type[E], **kwargs: t.Any) -> E:
585        self._match_text_seq("MODEL")
586        this = self._parse_table()
587
588        self._match(TokenType.COMMA)
589        self._match_text_seq("TABLE")
590
591        # Certain functions like ML.FORECAST require a STRUCT argument but not a TABLE/SELECT one
592        expression = (
593            self._parse_table() if not self._match(TokenType.STRUCT, advance=False) else None
594        )
595
596        self._match(TokenType.COMMA)
597
598        return self.expression(
599            expr_type(
600                this=this, expression=expression, params_struct=self._parse_bitwise(), **kwargs
601            )
602        )
603
604    def _parse_generate(self, expr_type: type[E], **kwargs: t.Any) -> E:
605        self._match_text_seq("MODEL")
606        this = self._parse_table()
607
608        self._match(TokenType.COMMA)
609
610        if self._match_text_seq("TABLE"):
611            expression = self._parse_table()
612        elif self._match(TokenType.L_PAREN, advance=False):
613            expression = self._parse_table()
614        else:
615            expression = self._parse_bitwise()
616
617        params_struct = self._match(TokenType.COMMA) and self._parse_bitwise()
618
619        return self.expression(
620            expr_type(this=this, expression=expression, params_struct=params_struct, **kwargs)
621        )
622
623    def _parse_translate(self) -> exp.Translate | exp.MLTranslate:
624        # Check if this is ML.TRANSLATE by looking at previous tokens
625        token = seq_get(self._tokens, self._index - 4)
626        if token and token.text.upper() == "ML":
627            return self._parse_ml(exp.MLTranslate)
628
629        return exp.Translate.from_arg_list(self._parse_function_args())
630
631    def _parse_forecast(self) -> exp.AIForecast | exp.MLForecast:
632        # Check if this is ML.FORECAST by looking at previous tokens.
633        token = seq_get(self._tokens, self._index - 4)
634        if token and token.text.upper() == "ML":
635            return self._parse_ml(exp.MLForecast)
636
637        # AI.FORECAST is a TVF, where the first argument is either TABLE <table>
638        # or a parenthesized query statement, followed by named arguments.
639        self._match(TokenType.TABLE)
640        this = self._parse_table()
641        if not this:
642            self.raise_error("Expected table or query statement")
643
644        expr = self.expression(exp.AIForecast(this=this))
645        while self._match(TokenType.COMMA):
646            arg = self._parse_lambda()
647            if isinstance(arg, exp.Kwarg):
648                expr.set(arg.this.name, arg)
649            else:
650                self.raise_error(f"Expected key => value syntax for AI.FORECAST, got {arg}")
651                break
652
653        return expr
654
655    def _parse_features_at_time(self) -> exp.FeaturesAtTime:
656        self._match(TokenType.TABLE)
657        this = self._parse_table()
658
659        expr = self.expression(exp.FeaturesAtTime(this=this))
660
661        while self._match(TokenType.COMMA):
662            arg = self._parse_lambda()
663
664            # Get the LHS of the Kwarg and set the arg to that value, e.g
665            # "num_rows => 1" sets the expr's `num_rows` arg
666            if arg:
667                expr.set(arg.this.name, arg)
668
669        return expr
670
671    def _parse_vector_search(self) -> exp.VectorSearch:
672        self._match(TokenType.TABLE)
673        base_table = self._parse_table()
674
675        self._match(TokenType.COMMA)
676
677        column_to_search = self._parse_bitwise()
678        self._match(TokenType.COMMA)
679
680        self._match(TokenType.TABLE)
681        query_table = self._parse_table()
682
683        expr = self.expression(
684            exp.VectorSearch(
685                this=base_table, column_to_search=column_to_search, query_table=query_table
686            )
687        )
688
689        while self._match(TokenType.COMMA):
690            # query_column_to_search can be named argument or positional
691            if self._match(TokenType.STRING, advance=False):
692                query_column = self._parse_string()
693                expr.set("query_column_to_search", query_column)
694            else:
695                arg = self._parse_lambda()
696                if arg:
697                    expr.set(arg.this.name, arg)
698
699        return expr
700
701    def _parse_export_data(self) -> exp.Export:
702        self._match_text_seq("DATA")
703
704        return self.expression(
705            exp.Export(
706                connection=self._match_text_seq("WITH", "CONNECTION") and self._parse_table_parts(),
707                options=self._parse_properties(),
708                this=self._match_text_seq("AS") and self._parse_select(nested=True),
709            )
710        )
711
712    def _parse_column_ops(self, this: exp.Expr | None) -> exp.Expr | None:
713        func_index = self._index + 1
714        this = super()._parse_column_ops(this)
715
716        if isinstance(this, exp.Dot) and isinstance(this.expression, exp.Func):
717            prefix = this.this.name.upper()
718
719            func: type[exp.Func] | None = None
720            if prefix == "NET":
721                func = exp.NetFunc
722            elif prefix == "SAFE":
723                func = exp.SafeFunc
724
725            if func:
726                # Retreat to try and parse a known function instead of an anonymous one,
727                # which is parsed by the base column ops parser due to anonymous_func=true
728                self._retreat(func_index)
729                this = func(this=self._parse_function(any_token=True))
730            elif prefix in ("AI", "ML"):
731                # AI.* and ML.* function calls can use custom BigQuery signatures that rely on
732                # function parsers, so re-parse the function in non-anonymous mode.
733                self._retreat(func_index)
734                parsed = self._parse_function(any_token=True)
735                if parsed:
736                    this = self.expression(exp.Dot(this=this.this, expression=parsed))
737
738        return this