sqlglot.optimizer.annotate_types

  1from __future__ import annotations
  2
  3import functools
  4import typing as t
  5
  6from sqlglot import exp
  7from sqlglot.helper import (
  8    ensure_list,
  9    is_date_unit,
 10    is_iso_date,
 11    is_iso_datetime,
 12    seq_get,
 13    subclasses,
 14)
 15from sqlglot.optimizer.scope import Scope, traverse_scope
 16from sqlglot.schema import Schema, ensure_schema
 17
 18if t.TYPE_CHECKING:
 19    from sqlglot._typing import B, E
 20
 21    BinaryCoercionFunc = t.Callable[[exp.Expression, exp.Expression], exp.DataType.Type]
 22    BinaryCoercions = t.Dict[
 23        t.Tuple[exp.DataType.Type, exp.DataType.Type],
 24        BinaryCoercionFunc,
 25    ]
 26
 27
 28def annotate_types(
 29    expression: E,
 30    schema: t.Optional[t.Dict | Schema] = None,
 31    annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
 32    coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
 33) -> E:
 34    """
 35    Infers the types of an expression, annotating its AST accordingly.
 36
 37    Example:
 38        >>> import sqlglot
 39        >>> schema = {"y": {"cola": "SMALLINT"}}
 40        >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
 41        >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
 42        >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
 43        <Type.DOUBLE: 'DOUBLE'>
 44
 45    Args:
 46        expression: Expression to annotate.
 47        schema: Database schema.
 48        annotators: Maps expression type to corresponding annotation function.
 49        coerces_to: Maps expression type to set of types that it can be coerced into.
 50
 51    Returns:
 52        The expression annotated with types.
 53    """
 54
 55    schema = ensure_schema(schema)
 56
 57    return TypeAnnotator(schema, annotators, coerces_to).annotate(expression)
 58
 59
 60def _annotate_with_type_lambda(data_type: exp.DataType.Type) -> t.Callable[[TypeAnnotator, E], E]:
 61    return lambda self, e: self._annotate_with_type(e, data_type)
 62
 63
 64def _coerce_date_literal(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
 65    date_text = l.name
 66    is_iso_date_ = is_iso_date(date_text)
 67
 68    if is_iso_date_ and is_date_unit(unit):
 69        return exp.DataType.Type.DATE
 70
 71    # An ISO date is also an ISO datetime, but not vice versa
 72    if is_iso_date_ or is_iso_datetime(date_text):
 73        return exp.DataType.Type.DATETIME
 74
 75    return exp.DataType.Type.UNKNOWN
 76
 77
 78def _coerce_date(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
 79    if not is_date_unit(unit):
 80        return exp.DataType.Type.DATETIME
 81    return l.type.this if l.type else exp.DataType.Type.UNKNOWN
 82
 83
 84def swap_args(func: BinaryCoercionFunc) -> BinaryCoercionFunc:
 85    @functools.wraps(func)
 86    def _swapped(l: exp.Expression, r: exp.Expression) -> exp.DataType.Type:
 87        return func(r, l)
 88
 89    return _swapped
 90
 91
 92def swap_all(coercions: BinaryCoercions) -> BinaryCoercions:
 93    return {**coercions, **{(b, a): swap_args(func) for (a, b), func in coercions.items()}}
 94
 95
 96class _TypeAnnotator(type):
 97    def __new__(cls, clsname, bases, attrs):
 98        klass = super().__new__(cls, clsname, bases, attrs)
 99
100        # Highest-to-lowest type precedence, as specified in Spark's docs (ANSI):
101        # https://spark.apache.org/docs/3.2.0/sql-ref-ansi-compliance.html
102        text_precedence = (
103            exp.DataType.Type.TEXT,
104            exp.DataType.Type.NVARCHAR,
105            exp.DataType.Type.VARCHAR,
106            exp.DataType.Type.NCHAR,
107            exp.DataType.Type.CHAR,
108        )
109        numeric_precedence = (
110            exp.DataType.Type.DOUBLE,
111            exp.DataType.Type.FLOAT,
112            exp.DataType.Type.DECIMAL,
113            exp.DataType.Type.BIGINT,
114            exp.DataType.Type.INT,
115            exp.DataType.Type.SMALLINT,
116            exp.DataType.Type.TINYINT,
117        )
118        timelike_precedence = (
119            exp.DataType.Type.TIMESTAMPLTZ,
120            exp.DataType.Type.TIMESTAMPTZ,
121            exp.DataType.Type.TIMESTAMP,
122            exp.DataType.Type.DATETIME,
123            exp.DataType.Type.DATE,
124        )
125
126        for type_precedence in (text_precedence, numeric_precedence, timelike_precedence):
127            coerces_to = set()
128            for data_type in type_precedence:
129                klass.COERCES_TO[data_type] = coerces_to.copy()
130                coerces_to |= {data_type}
131
132        return klass
133
134
135class TypeAnnotator(metaclass=_TypeAnnotator):
136    TYPE_TO_EXPRESSIONS: t.Dict[exp.DataType.Type, t.Set[t.Type[exp.Expression]]] = {
137        exp.DataType.Type.BIGINT: {
138            exp.ApproxDistinct,
139            exp.ArraySize,
140            exp.Count,
141            exp.Length,
142        },
143        exp.DataType.Type.BOOLEAN: {
144            exp.Between,
145            exp.Boolean,
146            exp.In,
147            exp.RegexpLike,
148        },
149        exp.DataType.Type.DATE: {
150            exp.CurrentDate,
151            exp.Date,
152            exp.DateFromParts,
153            exp.DateStrToDate,
154            exp.DiToDate,
155            exp.StrToDate,
156            exp.TimeStrToDate,
157            exp.TsOrDsToDate,
158        },
159        exp.DataType.Type.DATETIME: {
160            exp.CurrentDatetime,
161            exp.DatetimeAdd,
162            exp.DatetimeSub,
163        },
164        exp.DataType.Type.DOUBLE: {
165            exp.ApproxQuantile,
166            exp.Avg,
167            exp.Div,
168            exp.Exp,
169            exp.Ln,
170            exp.Log,
171            exp.Pow,
172            exp.Quantile,
173            exp.Round,
174            exp.SafeDivide,
175            exp.Sqrt,
176            exp.Stddev,
177            exp.StddevPop,
178            exp.StddevSamp,
179            exp.Variance,
180            exp.VariancePop,
181        },
182        exp.DataType.Type.INT: {
183            exp.Ceil,
184            exp.DatetimeDiff,
185            exp.DateDiff,
186            exp.Extract,
187            exp.TimestampDiff,
188            exp.TimeDiff,
189            exp.DateToDi,
190            exp.Floor,
191            exp.Levenshtein,
192            exp.Sign,
193            exp.StrPosition,
194            exp.TsOrDiToDi,
195        },
196        exp.DataType.Type.JSON: {
197            exp.ParseJSON,
198        },
199        exp.DataType.Type.TIMESTAMP: {
200            exp.CurrentTime,
201            exp.CurrentTimestamp,
202            exp.StrToTime,
203            exp.TimeAdd,
204            exp.TimeStrToTime,
205            exp.TimeSub,
206            exp.TimestampAdd,
207            exp.TimestampSub,
208            exp.UnixToTime,
209        },
210        exp.DataType.Type.TINYINT: {
211            exp.Day,
212            exp.Month,
213            exp.Week,
214            exp.Year,
215            exp.Quarter,
216        },
217        exp.DataType.Type.VARCHAR: {
218            exp.ArrayConcat,
219            exp.Concat,
220            exp.ConcatWs,
221            exp.DateToDateStr,
222            exp.GroupConcat,
223            exp.Initcap,
224            exp.Lower,
225            exp.Substring,
226            exp.TimeToStr,
227            exp.TimeToTimeStr,
228            exp.Trim,
229            exp.TsOrDsToDateStr,
230            exp.UnixToStr,
231            exp.UnixToTimeStr,
232            exp.Upper,
233        },
234    }
235
236    ANNOTATORS: t.Dict = {
237        **{
238            expr_type: lambda self, e: self._annotate_unary(e)
239            for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
240        },
241        **{
242            expr_type: lambda self, e: self._annotate_binary(e)
243            for expr_type in subclasses(exp.__name__, exp.Binary)
244        },
245        **{
246            expr_type: _annotate_with_type_lambda(data_type)
247            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
248            for expr_type in expressions
249        },
250        exp.Abs: lambda self, e: self._annotate_by_args(e, "this"),
251        exp.Anonymous: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
252        exp.Array: lambda self, e: self._annotate_by_args(e, "expressions", array=True),
253        exp.ArrayAgg: lambda self, e: self._annotate_by_args(e, "this", array=True),
254        exp.ArrayConcat: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
255        exp.Bracket: lambda self, e: self._annotate_bracket(e),
256        exp.Cast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
257        exp.Case: lambda self, e: self._annotate_by_args(e, "default", "ifs"),
258        exp.Coalesce: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
259        exp.DataType: lambda self, e: self._annotate_with_type(e, e.copy()),
260        exp.DateAdd: lambda self, e: self._annotate_timeunit(e),
261        exp.DateSub: lambda self, e: self._annotate_timeunit(e),
262        exp.DateTrunc: lambda self, e: self._annotate_timeunit(e),
263        exp.Distinct: lambda self, e: self._annotate_by_args(e, "expressions"),
264        exp.Div: lambda self, e: self._annotate_div(e),
265        exp.Dot: lambda self, e: self._annotate_dot(e),
266        exp.Explode: lambda self, e: self._annotate_explode(e),
267        exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
268        exp.GenerateDateArray: lambda self, e: self._annotate_with_type(
269            e, exp.DataType.build("ARRAY<DATE>")
270        ),
271        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
272        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
273        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
274        exp.Literal: lambda self, e: self._annotate_literal(e),
275        exp.Map: lambda self, e: self._annotate_map(e),
276        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
277        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
278        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
279        exp.Nullif: lambda self, e: self._annotate_by_args(e, "this", "expression"),
280        exp.PropertyEQ: lambda self, e: self._annotate_by_args(e, "expression"),
281        exp.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
282        exp.Struct: lambda self, e: self._annotate_struct(e),
283        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
284        exp.Timestamp: lambda self, e: self._annotate_with_type(
285            e,
286            exp.DataType.Type.TIMESTAMPTZ if e.args.get("with_tz") else exp.DataType.Type.TIMESTAMP,
287        ),
288        exp.ToMap: lambda self, e: self._annotate_to_map(e),
289        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
290        exp.Unnest: lambda self, e: self._annotate_unnest(e),
291        exp.VarMap: lambda self, e: self._annotate_map(e),
292    }
293
294    NESTED_TYPES = {
295        exp.DataType.Type.ARRAY,
296    }
297
298    # Specifies what types a given type can be coerced into (autofilled)
299    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
300
301    # Coercion functions for binary operations.
302    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
303    BINARY_COERCIONS: BinaryCoercions = {
304        **swap_all(
305            {
306                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
307                    l, r.args.get("unit")
308                )
309                for t in exp.DataType.TEXT_TYPES
310            }
311        ),
312        **swap_all(
313            {
314                # text + numeric will yield the numeric type to match most dialects' semantics
315                (text, numeric): lambda l, r: t.cast(
316                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
317                )
318                for text in exp.DataType.TEXT_TYPES
319                for numeric in exp.DataType.NUMERIC_TYPES
320            }
321        ),
322        **swap_all(
323            {
324                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
325                    l, r.args.get("unit")
326                ),
327            }
328        ),
329    }
330
331    def __init__(
332        self,
333        schema: Schema,
334        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
335        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
336        binary_coercions: t.Optional[BinaryCoercions] = None,
337    ) -> None:
338        self.schema = schema
339        self.annotators = annotators or self.ANNOTATORS
340        self.coerces_to = coerces_to or self.COERCES_TO
341        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
342
343        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
344        self._visited: t.Set[int] = set()
345
346    def _set_type(
347        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
348    ) -> None:
349        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
350        self._visited.add(id(expression))
351
352    def annotate(self, expression: E) -> E:
353        for scope in traverse_scope(expression):
354            self.annotate_scope(scope)
355        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
356
357    def annotate_scope(self, scope: Scope) -> None:
358        selects = {}
359        for name, source in scope.sources.items():
360            if not isinstance(source, Scope):
361                continue
362            if isinstance(source.expression, exp.UDTF):
363                values = []
364
365                if isinstance(source.expression, exp.Lateral):
366                    if isinstance(source.expression.this, exp.Explode):
367                        values = [source.expression.this.this]
368                elif isinstance(source.expression, exp.Unnest):
369                    values = [source.expression]
370                else:
371                    values = source.expression.expressions[0].expressions
372
373                if not values:
374                    continue
375
376                selects[name] = {
377                    alias: column
378                    for alias, column in zip(
379                        source.expression.alias_column_names,
380                        values,
381                    )
382                }
383            else:
384                selects[name] = {
385                    select.alias_or_name: select for select in source.expression.selects
386                }
387
388        # First annotate the current scope's column references
389        for col in scope.columns:
390            if not col.table:
391                continue
392
393            source = scope.sources.get(col.table)
394            if isinstance(source, exp.Table):
395                self._set_type(col, self.schema.get_column_type(source, col))
396            elif source:
397                if col.table in selects and col.name in selects[col.table]:
398                    self._set_type(col, selects[col.table][col.name].type)
399                elif isinstance(source.expression, exp.Unnest):
400                    self._set_type(col, source.expression.type)
401
402        # Then (possibly) annotate the remaining expressions in the scope
403        self._maybe_annotate(scope.expression)
404
405    def _maybe_annotate(self, expression: E) -> E:
406        if id(expression) in self._visited:
407            return expression  # We've already inferred the expression's type
408
409        annotator = self.annotators.get(expression.__class__)
410
411        return (
412            annotator(self, expression)
413            if annotator
414            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
415        )
416
417    def _annotate_args(self, expression: E) -> E:
418        for value in expression.iter_expressions():
419            self._maybe_annotate(value)
420
421        return expression
422
423    def _maybe_coerce(
424        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
425    ) -> exp.DataType | exp.DataType.Type:
426        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
427        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
428
429        # We propagate the NULL / UNKNOWN types upwards if found
430        if exp.DataType.Type.NULL in (type1_value, type2_value):
431            return exp.DataType.Type.NULL
432        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
433            return exp.DataType.Type.UNKNOWN
434
435        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value
436
437    def _annotate_binary(self, expression: B) -> B:
438        self._annotate_args(expression)
439
440        left, right = expression.left, expression.right
441        left_type, right_type = left.type.this, right.type.this  # type: ignore
442
443        if isinstance(expression, exp.Connector):
444            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
445                self._set_type(expression, exp.DataType.Type.NULL)
446            elif exp.DataType.Type.NULL in (left_type, right_type):
447                self._set_type(
448                    expression,
449                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
450                )
451            else:
452                self._set_type(expression, exp.DataType.Type.BOOLEAN)
453        elif isinstance(expression, exp.Predicate):
454            self._set_type(expression, exp.DataType.Type.BOOLEAN)
455        elif (left_type, right_type) in self.binary_coercions:
456            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
457        else:
458            self._set_type(expression, self._maybe_coerce(left_type, right_type))
459
460        return expression
461
462    def _annotate_unary(self, expression: E) -> E:
463        self._annotate_args(expression)
464
465        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
466            self._set_type(expression, exp.DataType.Type.BOOLEAN)
467        else:
468            self._set_type(expression, expression.this.type)
469
470        return expression
471
472    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
473        if expression.is_string:
474            self._set_type(expression, exp.DataType.Type.VARCHAR)
475        elif expression.is_int:
476            self._set_type(expression, exp.DataType.Type.INT)
477        else:
478            self._set_type(expression, exp.DataType.Type.DOUBLE)
479
480        return expression
481
482    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
483        self._set_type(expression, target_type)
484        return self._annotate_args(expression)
485
486    @t.no_type_check
487    def _annotate_by_args(
488        self,
489        expression: E,
490        *args: str,
491        promote: bool = False,
492        array: bool = False,
493    ) -> E:
494        self._annotate_args(expression)
495
496        expressions: t.List[exp.Expression] = []
497        for arg in args:
498            arg_expr = expression.args.get(arg)
499            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
500
501        last_datatype = None
502        for expr in expressions:
503            expr_type = expr.type
504
505            # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
506            if expr_type.args.get("nested"):
507                last_datatype = expr_type
508                break
509
510            if not expr_type.is_type(exp.DataType.Type.NULL, exp.DataType.Type.UNKNOWN):
511                last_datatype = self._maybe_coerce(last_datatype or expr_type, expr_type)
512
513        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
514
515        if promote:
516            if expression.type.this in exp.DataType.INTEGER_TYPES:
517                self._set_type(expression, exp.DataType.Type.BIGINT)
518            elif expression.type.this in exp.DataType.FLOAT_TYPES:
519                self._set_type(expression, exp.DataType.Type.DOUBLE)
520
521        if array:
522            self._set_type(
523                expression,
524                exp.DataType(
525                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
526                ),
527            )
528
529        return expression
530
531    def _annotate_timeunit(
532        self, expression: exp.TimeUnit | exp.DateTrunc
533    ) -> exp.TimeUnit | exp.DateTrunc:
534        self._annotate_args(expression)
535
536        if expression.this.type.this in exp.DataType.TEXT_TYPES:
537            datatype = _coerce_date_literal(expression.this, expression.unit)
538        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
539            datatype = _coerce_date(expression.this, expression.unit)
540        else:
541            datatype = exp.DataType.Type.UNKNOWN
542
543        self._set_type(expression, datatype)
544        return expression
545
546    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
547        self._annotate_args(expression)
548
549        bracket_arg = expression.expressions[0]
550        this = expression.this
551
552        if isinstance(bracket_arg, exp.Slice):
553            self._set_type(expression, this.type)
554        elif this.type.is_type(exp.DataType.Type.ARRAY):
555            self._set_type(expression, seq_get(this.type.expressions, 0))
556        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
557            index = this.keys.index(bracket_arg)
558            value = seq_get(this.values, index)
559            self._set_type(expression, value.type if value else None)
560        else:
561            self._set_type(expression, exp.DataType.Type.UNKNOWN)
562
563        return expression
564
565    def _annotate_div(self, expression: exp.Div) -> exp.Div:
566        self._annotate_args(expression)
567
568        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
569
570        if (
571            expression.args.get("typed")
572            and left_type in exp.DataType.INTEGER_TYPES
573            and right_type in exp.DataType.INTEGER_TYPES
574        ):
575            self._set_type(expression, exp.DataType.Type.BIGINT)
576        else:
577            self._set_type(expression, self._maybe_coerce(left_type, right_type))
578            if expression.type and expression.type.this not in exp.DataType.REAL_TYPES:
579                self._set_type(
580                    expression, self._maybe_coerce(expression.type, exp.DataType.Type.DOUBLE)
581                )
582
583        return expression
584
585    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
586        self._annotate_args(expression)
587        self._set_type(expression, None)
588        this_type = expression.this.type
589
590        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
591            for e in this_type.expressions:
592                if e.name == expression.expression.name:
593                    self._set_type(expression, e.kind)
594                    break
595
596        return expression
597
598    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
599        self._annotate_args(expression)
600        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
601        return expression
602
603    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
604        self._annotate_args(expression)
605        child = seq_get(expression.expressions, 0)
606
607        if child and child.is_type(exp.DataType.Type.ARRAY):
608            expr_type = seq_get(child.type.expressions, 0)
609        else:
610            expr_type = None
611
612        self._set_type(expression, expr_type)
613        return expression
614
615    def _annotate_struct_value(
616        self, expression: exp.Expression
617    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
618        alias = expression.args.get("alias")
619        if alias:
620            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
621
622        # Case: key = value or key := value
623        if expression.expression:
624            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
625
626        return expression.type
627
628    def _annotate_struct(self, expression: exp.Struct) -> exp.Struct:
629        self._annotate_args(expression)
630        self._set_type(
631            expression,
632            exp.DataType(
633                this=exp.DataType.Type.STRUCT,
634                expressions=[self._annotate_struct_value(expr) for expr in expression.expressions],
635                nested=True,
636            ),
637        )
638        return expression
639
640    @t.overload
641    def _annotate_map(self, expression: exp.Map) -> exp.Map: ...
642
643    @t.overload
644    def _annotate_map(self, expression: exp.VarMap) -> exp.VarMap: ...
645
646    def _annotate_map(self, expression):
647        self._annotate_args(expression)
648
649        keys = expression.args.get("keys")
650        values = expression.args.get("values")
651
652        map_type = exp.DataType(this=exp.DataType.Type.MAP)
653        if isinstance(keys, exp.Array) and isinstance(values, exp.Array):
654            key_type = seq_get(keys.type.expressions, 0) or exp.DataType.Type.UNKNOWN
655            value_type = seq_get(values.type.expressions, 0) or exp.DataType.Type.UNKNOWN
656
657            if key_type != exp.DataType.Type.UNKNOWN and value_type != exp.DataType.Type.UNKNOWN:
658                map_type.set("expressions", [key_type, value_type])
659                map_type.set("nested", True)
660
661        self._set_type(expression, map_type)
662        return expression
663
664    def _annotate_to_map(self, expression: exp.ToMap) -> exp.ToMap:
665        self._annotate_args(expression)
666
667        map_type = exp.DataType(this=exp.DataType.Type.MAP)
668        arg = expression.this
669        if arg.is_type(exp.DataType.Type.STRUCT):
670            for coldef in arg.type.expressions:
671                kind = coldef.kind
672                if kind != exp.DataType.Type.UNKNOWN:
673                    map_type.set("expressions", [exp.DataType.build("varchar"), kind])
674                    map_type.set("nested", True)
675                    break
676
677        self._set_type(expression, map_type)
678        return expression