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