sqlglot.optimizer.annotate_types

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