sqlglot.optimizer.annotate_types

  1from __future__ import annotations
  2
  3import functools
  4import logging
  5import typing as t
  6
  7from sqlglot import exp
  8from sqlglot.dialects.dialect import Dialect
  9from sqlglot.helper import (
 10    ensure_list,
 11    is_date_unit,
 12    is_iso_date,
 13    is_iso_datetime,
 14    seq_get,
 15)
 16from sqlglot.optimizer.scope import Scope, traverse_scope
 17from sqlglot.schema import MappingSchema, Schema, ensure_schema
 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
 29    from sqlglot.typing import ExpressionMetadataType
 30
 31logger = logging.getLogger("sqlglot")
 32
 33# EXTRACT/DATE_PART specifiers that return BIGINT instead of INT
 34BIGINT_EXTRACT_DATE_PARTS = {
 35    "EPOCH_SECOND",
 36    "EPOCH_MILLISECOND",
 37    "EPOCH_MICROSECOND",
 38    "EPOCH_NANOSECOND",
 39    "NANOSECOND",
 40}
 41
 42
 43def annotate_types(
 44    expression: E,
 45    schema: t.Optional[t.Dict | Schema] = None,
 46    expression_metadata: t.Optional[ExpressionMetadataType] = None,
 47    coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
 48    dialect: DialectType = None,
 49    overwrite_types: bool = True,
 50) -> E:
 51    """
 52    Infers the types of an expression, annotating its AST accordingly.
 53
 54    Example:
 55        >>> import sqlglot
 56        >>> schema = {"y": {"cola": "SMALLINT"}}
 57        >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
 58        >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
 59        >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
 60        <Type.DOUBLE: 'DOUBLE'>
 61
 62    Args:
 63        expression: Expression to annotate.
 64        schema: Database schema.
 65        expression_metadata: Maps expression type to corresponding annotation function.
 66        coerces_to: Maps expression type to set of types that it can be coerced into.
 67        overwrite_types: Re-annotate the existing AST types.
 68
 69    Returns:
 70        The expression annotated with types.
 71    """
 72
 73    schema = ensure_schema(schema, dialect=dialect)
 74
 75    return TypeAnnotator(
 76        schema=schema,
 77        expression_metadata=expression_metadata,
 78        coerces_to=coerces_to,
 79        overwrite_types=overwrite_types,
 80    ).annotate(expression)
 81
 82
 83def _coerce_date_literal(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
 84    date_text = l.name
 85    is_iso_date_ = is_iso_date(date_text)
 86
 87    if is_iso_date_ and is_date_unit(unit):
 88        return exp.DataType.Type.DATE
 89
 90    # An ISO date is also an ISO datetime, but not vice versa
 91    if is_iso_date_ or is_iso_datetime(date_text):
 92        return exp.DataType.Type.DATETIME
 93
 94    return exp.DataType.Type.UNKNOWN
 95
 96
 97def _coerce_date(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
 98    if not is_date_unit(unit):
 99        return exp.DataType.Type.DATETIME
100    return l.type.this if l.type else exp.DataType.Type.UNKNOWN
101
102
103def swap_args(func: BinaryCoercionFunc) -> BinaryCoercionFunc:
104    @functools.wraps(func)
105    def _swapped(l: exp.Expression, r: exp.Expression) -> exp.DataType.Type:
106        return func(r, l)
107
108    return _swapped
109
110
111def swap_all(coercions: BinaryCoercions) -> BinaryCoercions:
112    return {**coercions, **{(b, a): swap_args(func) for (a, b), func in coercions.items()}}
113
114
115class _TypeAnnotator(type):
116    def __new__(cls, clsname, bases, attrs):
117        klass = super().__new__(cls, clsname, bases, attrs)
118
119        # Highest-to-lowest type precedence, as specified in Spark's docs (ANSI):
120        # https://spark.apache.org/docs/3.2.0/sql-ref-ansi-compliance.html
121        text_precedence = (
122            exp.DataType.Type.TEXT,
123            exp.DataType.Type.NVARCHAR,
124            exp.DataType.Type.VARCHAR,
125            exp.DataType.Type.NCHAR,
126            exp.DataType.Type.CHAR,
127        )
128        numeric_precedence = (
129            exp.DataType.Type.DECFLOAT,
130            exp.DataType.Type.DOUBLE,
131            exp.DataType.Type.FLOAT,
132            exp.DataType.Type.BIGDECIMAL,
133            exp.DataType.Type.DECIMAL,
134            exp.DataType.Type.BIGINT,
135            exp.DataType.Type.INT,
136            exp.DataType.Type.SMALLINT,
137            exp.DataType.Type.TINYINT,
138        )
139        timelike_precedence = (
140            exp.DataType.Type.TIMESTAMPLTZ,
141            exp.DataType.Type.TIMESTAMPTZ,
142            exp.DataType.Type.TIMESTAMP,
143            exp.DataType.Type.DATETIME,
144            exp.DataType.Type.DATE,
145        )
146
147        for type_precedence in (text_precedence, numeric_precedence, timelike_precedence):
148            coerces_to = set()
149            for data_type in type_precedence:
150                klass.COERCES_TO[data_type] = coerces_to.copy()
151                coerces_to |= {data_type}
152        return klass
153
154
155class TypeAnnotator(metaclass=_TypeAnnotator):
156    NESTED_TYPES = {
157        exp.DataType.Type.ARRAY,
158    }
159
160    # Specifies what types a given type can be coerced into (autofilled)
161    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
162
163    # Coercion functions for binary operations.
164    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
165    BINARY_COERCIONS: BinaryCoercions = {
166        **swap_all(
167            {
168                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
169                    l, r.args.get("unit")
170                )
171                for t in exp.DataType.TEXT_TYPES
172            }
173        ),
174        **swap_all(
175            {
176                # text + numeric will yield the numeric type to match most dialects' semantics
177                (text, numeric): lambda l, r: t.cast(
178                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
179                )
180                for text in exp.DataType.TEXT_TYPES
181                for numeric in exp.DataType.NUMERIC_TYPES
182            }
183        ),
184        **swap_all(
185            {
186                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
187                    l, r.args.get("unit")
188                ),
189            }
190        ),
191    }
192
193    def __init__(
194        self,
195        schema: Schema,
196        expression_metadata: t.Optional[ExpressionMetadataType] = None,
197        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
198        binary_coercions: t.Optional[BinaryCoercions] = None,
199        overwrite_types: bool = True,
200    ) -> None:
201        self.schema = schema
202        dialect = schema.dialect or Dialect()
203        self.dialect = dialect
204        self.expression_metadata = expression_metadata or dialect.EXPRESSION_METADATA
205        self.coerces_to = coerces_to or dialect.COERCES_TO or self.COERCES_TO
206        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
207
208        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
209        self._visited: t.Set[int] = set()
210
211        # Caches NULL-annotated expressions to set them to UNKNOWN after type inference is completed
212        self._null_expressions: t.Dict[int, exp.Expression] = {}
213
214        # Databricks and Spark ≥v3 actually support NULL (i.e., VOID) as a type
215        self._supports_null_type = dialect.SUPPORTS_NULL_TYPE
216
217        # Maps an exp.SetOperation's id (e.g. UNION) to its projection types. This is computed if the
218        # exp.SetOperation is the expression of a scope source, as selecting from it multiple times
219        # would reprocess the entire subtree to coerce the types of its operands' projections
220        self._setop_column_types: t.Dict[int, t.Dict[str, exp.DataType | exp.DataType.Type]] = {}
221
222        # When set to False, this enables partial annotation by skipping already-annotated nodes
223        self._overwrite_types = overwrite_types
224
225        # Maps Scope to its corresponding selected sources
226        self._scope_selects: t.Dict[Scope, t.Dict[str, t.Dict[str, t.Any]]] = {}
227
228    def clear(self) -> None:
229        self._visited.clear()
230        self._null_expressions.clear()
231        self._setop_column_types.clear()
232        self._scope_selects.clear()
233
234    def _set_type(
235        self, expression: E, target_type: t.Optional[exp.DataType | exp.DataType.Type]
236    ) -> E:
237        prev_type = expression.type
238        expression_id = id(expression)
239
240        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
241        self._visited.add(expression_id)
242
243        if (
244            not self._supports_null_type
245            and t.cast(exp.DataType, expression.type).this == exp.DataType.Type.NULL
246        ):
247            self._null_expressions[expression_id] = expression
248        elif prev_type and t.cast(exp.DataType, prev_type).this == exp.DataType.Type.NULL:
249            self._null_expressions.pop(expression_id, None)
250
251        if (
252            isinstance(expression, exp.Column)
253            and expression.is_type(exp.DataType.Type.JSON)
254            and (dot_parts := expression.meta.get("dot_parts"))
255        ):
256            # JSON dot access is case sensitive across all dialects, so we need to undo the normalization.
257            i = iter(dot_parts)
258            parent = expression.parent
259            while isinstance(parent, exp.Dot):
260                parent.expression.replace(exp.to_identifier(next(i), quoted=True))
261                parent = parent.parent
262
263            expression.meta.pop("dot_parts", None)
264
265        return expression
266
267    def annotate(self, expression: E, annotate_scope: bool = True) -> E:
268        # This flag is used to avoid costly scope traversals when we only care about annotating
269        # non-column expressions (partial type inference), e.g., when simplifying in the optimizer
270        if annotate_scope:
271            for scope in traverse_scope(expression):
272                self.annotate_scope(scope)
273
274        # This takes care of non-traversable expressions
275        self._annotate_expression(expression)
276
277        # Replace NULL type with the default type of the targeted dialect, since the former is not an actual type;
278        # it is mostly used to aid type coercion, e.g. in query set operations.
279        for expr in self._null_expressions.values():
280            expr.type = self.dialect.DEFAULT_NULL_TYPE
281
282        return expression
283
284    def _get_scope_selects(self, scope: Scope) -> t.Dict[str, t.Dict[str, t.Any]]:
285        if scope not in self._scope_selects:
286            selects = {}
287            for name, source in scope.sources.items():
288                if not isinstance(source, Scope):
289                    continue
290
291                expression = source.expression
292                if isinstance(expression, exp.UDTF):
293                    values = []
294
295                    if isinstance(expression, exp.Lateral):
296                        if isinstance(expression.this, exp.Explode):
297                            values = [expression.this.this]
298                    elif isinstance(expression, exp.Unnest):
299                        values = [expression]
300                    elif not isinstance(expression, exp.TableFromRows):
301                        values = expression.expressions[0].expressions
302
303                    if not values:
304                        continue
305
306                    alias_column_names = expression.alias_column_names
307
308                    if (
309                        isinstance(expression, exp.Unnest)
310                        and expression.type
311                        and expression.type.is_type(exp.DataType.Type.STRUCT)
312                    ):
313                        selects[name] = {
314                            col_def.name: t.cast(
315                                t.Union[exp.DataType, exp.DataType.Type], col_def.kind
316                            )
317                            for col_def in expression.type.expressions
318                            if isinstance(col_def, exp.ColumnDef) and col_def.kind
319                        }
320                    else:
321                        selects[name] = {
322                            alias: column.type for alias, column in zip(alias_column_names, values)
323                        }
324                elif isinstance(expression, exp.SetOperation) and len(
325                    expression.left.selects
326                ) == len(expression.right.selects):
327                    selects[name] = self._get_setop_column_types(expression)
328                else:
329                    selects[name] = {s.alias_or_name: s.type for s in expression.selects}
330
331            self._scope_selects[scope] = selects
332
333        return self._scope_selects[scope]
334
335    def annotate_scope(self, scope: Scope) -> None:
336        if isinstance(self.schema, MappingSchema):
337            for table_column in scope.table_columns:
338                source = scope.sources.get(table_column.name)
339
340                if isinstance(source, exp.Table):
341                    schema = self.schema.find(
342                        source, raise_on_missing=False, ensure_data_types=True
343                    )
344                    if not isinstance(schema, dict):
345                        continue
346
347                    struct_type = exp.DataType(
348                        this=exp.DataType.Type.STRUCT,
349                        expressions=[
350                            exp.ColumnDef(this=exp.to_identifier(c), kind=kind)
351                            for c, kind in schema.items()
352                        ],
353                        nested=True,
354                    )
355                    self._set_type(table_column, struct_type)
356                elif (
357                    isinstance(source, Scope)
358                    and isinstance(source.expression, exp.Query)
359                    and (
360                        source.expression.meta.get("query_type") or exp.DataType.build("UNKNOWN")
361                    ).is_type(exp.DataType.Type.STRUCT)
362                ):
363                    self._set_type(table_column, source.expression.meta["query_type"])
364
365        # Iterate through all the expressions of the current scope in post-order, and annotate
366        self._annotate_expression(scope.expression, scope)
367
368        if self.dialect.QUERY_RESULTS_ARE_STRUCTS and isinstance(scope.expression, exp.Query):
369            struct_type = exp.DataType(
370                this=exp.DataType.Type.STRUCT,
371                expressions=[
372                    exp.ColumnDef(
373                        this=exp.to_identifier(select.output_name),
374                        kind=select.type.copy() if select.type else None,
375                    )
376                    for select in scope.expression.selects
377                ],
378                nested=True,
379            )
380
381            if not any(
382                cd.kind.is_type(exp.DataType.Type.UNKNOWN)
383                for cd in struct_type.expressions
384                if cd.kind
385            ):
386                # We don't use `_set_type` on purpose here. If we annotated the query directly, then
387                # using it in other contexts (e.g., ARRAY(<query>)) could result in incorrect type
388                # annotations, i.e., it shouldn't be interpreted as a STRUCT value.
389                scope.expression.meta["query_type"] = struct_type
390
391    def _annotate_expression(
392        self,
393        expression: exp.Expression,
394        scope: t.Optional[Scope] = None,
395    ) -> None:
396        stack = [(expression, False)]
397
398        while stack:
399            expr, children_annotated = stack.pop()
400
401            if id(expr) in self._visited or (
402                not self._overwrite_types
403                and expr.type
404                and not expr.is_type(exp.DataType.Type.UNKNOWN)
405            ):
406                continue  # We've already inferred the expression's type
407
408            if not children_annotated:
409                stack.append((expr, True))
410                for child_expr in expr.iter_expressions():
411                    stack.append((child_expr, False))
412                continue
413
414            if scope and isinstance(expr, exp.Column) and expr.table:
415                source = None
416                source_scope = scope
417                while source_scope and not source:
418                    source = source_scope.sources.get(expr.table)
419                    if not source:
420                        source_scope = source_scope.parent
421
422                if isinstance(source, exp.Table):
423                    self._set_type(expr, self.schema.get_column_type(source, expr))
424                elif source:
425                    col_type = (
426                        self._get_scope_selects(source_scope).get(expr.table, {}).get(expr.name)
427                    )
428                    if col_type:
429                        self._set_type(expr, col_type)
430                    elif isinstance(source.expression, exp.Unnest):
431                        self._set_type(expr, source.expression.type)
432                    else:
433                        self._set_type(expr, exp.DataType.Type.UNKNOWN)
434                else:
435                    self._set_type(expr, exp.DataType.Type.UNKNOWN)
436
437                if expr.type and expr.type.args.get("nullable") is False:
438                    expr.meta["nonnull"] = True
439                continue
440
441            spec = self.expression_metadata.get(expr.__class__)
442
443            if spec and (annotator := spec.get("annotator")):
444                annotator(self, expr)
445            elif spec and (returns := spec.get("returns")):
446                self._set_type(expr, t.cast(exp.DataType.Type, returns))
447            else:
448                self._set_type(expr, exp.DataType.Type.UNKNOWN)
449
450    def _maybe_coerce(
451        self,
452        type1: exp.DataType | exp.DataType.Type,
453        type2: exp.DataType | exp.DataType.Type,
454    ) -> exp.DataType | exp.DataType.Type:
455        """
456        Returns type2 if type1 can be coerced into it, otherwise type1.
457
458        If either type is parameterized (e.g. DECIMAL(18, 2) contains two parameters),
459        we assume type1 does not coerce into type2, so we also return it in this case.
460        """
461        if isinstance(type1, exp.DataType):
462            if type1.expressions:
463                return type1
464            type1_value = type1.this
465        else:
466            type1_value = type1
467
468        if isinstance(type2, exp.DataType):
469            if type2.expressions:
470                return type2
471            type2_value = type2.this
472        else:
473            type2_value = type2
474
475        # We propagate the UNKNOWN type upwards if found
476        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
477            return exp.DataType.Type.UNKNOWN
478
479        if type1_value == exp.DataType.Type.NULL:
480            return type2_value
481        if type2_value == exp.DataType.Type.NULL:
482            return type1_value
483
484        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value
485
486    def _get_setop_column_types(
487        self, setop: exp.SetOperation
488    ) -> t.Dict[str, exp.DataType | exp.DataType.Type]:
489        """
490        Computes and returns the coerced column types for a SetOperation.
491
492        This handles UNION, INTERSECT, EXCEPT, etc., coercing types across
493        left and right operands for all projections/columns.
494
495        Args:
496            setop: The SetOperation expression to analyze
497
498        Returns:
499            Dictionary mapping column names to their coerced types
500        """
501        setop_id = id(setop)
502        if setop_id in self._setop_column_types:
503            return self._setop_column_types[setop_id]
504
505        col_types: t.Dict[str, exp.DataType | exp.DataType.Type] = {}
506
507        # Validate that left and right have same number of projections
508        if not (
509            isinstance(setop, exp.SetOperation)
510            and setop.left.selects
511            and setop.right.selects
512            and len(setop.left.selects) == len(setop.right.selects)
513        ):
514            return col_types
515
516        # Process a chain / sub-tree of set operations
517        for set_op in setop.walk(
518            prune=lambda n: not isinstance(n, (exp.SetOperation, exp.Subquery))
519        ):
520            if not isinstance(set_op, exp.SetOperation):
521                continue
522
523            if set_op.args.get("by_name"):
524                r_type_by_select = {s.alias_or_name: s.type for s in set_op.right.selects}
525                setop_cols = {
526                    s.alias_or_name: self._maybe_coerce(
527                        t.cast(exp.DataType, s.type),
528                        r_type_by_select.get(s.alias_or_name) or exp.DataType.Type.UNKNOWN,
529                    )
530                    for s in set_op.left.selects
531                }
532            else:
533                setop_cols = {
534                    ls.alias_or_name: self._maybe_coerce(
535                        t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
536                    )
537                    for ls, rs in zip(set_op.left.selects, set_op.right.selects)
538                }
539
540            # Coerce intermediate results with the previously registered types, if they exist
541            for col_name, col_type in setop_cols.items():
542                col_types[col_name] = self._maybe_coerce(
543                    col_type, col_types.get(col_name, exp.DataType.Type.NULL)
544                )
545
546        self._setop_column_types[setop_id] = col_types
547        return col_types
548
549    def _annotate_binary(self, expression: B) -> B:
550        left, right = expression.left, expression.right
551        if not left or not right:
552            expression_sql = expression.sql(self.dialect)
553            logger.warning(f"Failed to annotate badly formed binary expression: {expression_sql}")
554            self._set_type(expression, None)
555            return expression
556
557        left_type, right_type = left.type.this, right.type.this  # type: ignore
558
559        if isinstance(expression, (exp.Connector, exp.Predicate)):
560            self._set_type(expression, exp.DataType.Type.BOOLEAN)
561        elif (left_type, right_type) in self.binary_coercions:
562            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
563        else:
564            self._annotate_by_args(expression, left, right)
565
566        if isinstance(expression, exp.Is) or (
567            left.meta.get("nonnull") is True and right.meta.get("nonnull") is True
568        ):
569            expression.meta["nonnull"] = True
570
571        return expression
572
573    def _annotate_unary(self, expression: E) -> E:
574        if isinstance(expression, exp.Not):
575            self._set_type(expression, exp.DataType.Type.BOOLEAN)
576        else:
577            self._set_type(expression, expression.this.type)
578
579        if expression.this.meta.get("nonnull") is True:
580            expression.meta["nonnull"] = True
581
582        return expression
583
584    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
585        if expression.is_string:
586            self._set_type(expression, exp.DataType.Type.VARCHAR)
587        elif expression.is_int:
588            self._set_type(expression, exp.DataType.Type.INT)
589        else:
590            self._set_type(expression, exp.DataType.Type.DOUBLE)
591
592        expression.meta["nonnull"] = True
593
594        return expression
595
596    @t.no_type_check
597    def _annotate_by_args(
598        self,
599        expression: E,
600        *args: str | exp.Expression,
601        promote: bool = False,
602        array: bool = False,
603    ) -> E:
604        literal_type = None
605        non_literal_type = None
606        nested_type = None
607
608        for arg in args:
609            if isinstance(arg, str):
610                expressions = expression.args.get(arg)
611            else:
612                expressions = arg
613
614            for expr in ensure_list(expressions):
615                expr_type = expr.type
616
617                # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
618                if expr_type.args.get("nested"):
619                    nested_type = expr_type
620                    break
621
622                if isinstance(expr, exp.Literal):
623                    literal_type = self._maybe_coerce(literal_type or expr_type, expr_type)
624                else:
625                    non_literal_type = self._maybe_coerce(non_literal_type or expr_type, expr_type)
626
627            if nested_type:
628                break
629
630        result_type = None
631
632        if nested_type:
633            result_type = nested_type
634        elif literal_type and non_literal_type:
635            if self.dialect.PRIORITIZE_NON_LITERAL_TYPES:
636                literal_this_type = (
637                    literal_type.this if isinstance(literal_type, exp.DataType) else literal_type
638                )
639                non_literal_this_type = (
640                    non_literal_type.this
641                    if isinstance(non_literal_type, exp.DataType)
642                    else non_literal_type
643                )
644                if (
645                    literal_this_type in exp.DataType.INTEGER_TYPES
646                    and non_literal_this_type in exp.DataType.INTEGER_TYPES
647                ) or (
648                    literal_this_type in exp.DataType.REAL_TYPES
649                    and non_literal_this_type in exp.DataType.REAL_TYPES
650                ):
651                    result_type = non_literal_type
652        else:
653            result_type = literal_type or non_literal_type or exp.DataType.Type.UNKNOWN
654
655        self._set_type(
656            expression, result_type or self._maybe_coerce(non_literal_type, literal_type)
657        )
658
659        if promote:
660            if expression.type.this in exp.DataType.INTEGER_TYPES:
661                self._set_type(expression, exp.DataType.Type.BIGINT)
662            elif expression.type.this in exp.DataType.FLOAT_TYPES:
663                self._set_type(expression, exp.DataType.Type.DOUBLE)
664
665        if array:
666            self._set_type(
667                expression,
668                exp.DataType(
669                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
670                ),
671            )
672
673        return expression
674
675    def _annotate_timeunit(
676        self, expression: exp.TimeUnit | exp.DateTrunc
677    ) -> exp.TimeUnit | exp.DateTrunc:
678        if expression.this.type.this in exp.DataType.TEXT_TYPES:
679            datatype = _coerce_date_literal(expression.this, expression.unit)
680        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
681            datatype = _coerce_date(expression.this, expression.unit)
682        else:
683            datatype = exp.DataType.Type.UNKNOWN
684
685        self._set_type(expression, datatype)
686        return expression
687
688    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
689        bracket_arg = expression.expressions[0]
690        this = expression.this
691
692        if isinstance(bracket_arg, exp.Slice):
693            self._set_type(expression, this.type)
694        elif this.type.is_type(exp.DataType.Type.ARRAY):
695            self._set_type(expression, seq_get(this.type.expressions, 0))
696        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
697            index = this.keys.index(bracket_arg)
698            value = seq_get(this.values, index)
699            self._set_type(expression, value.type if value else None)
700        else:
701            self._set_type(expression, exp.DataType.Type.UNKNOWN)
702
703        return expression
704
705    def _annotate_div(self, expression: exp.Div) -> exp.Div:
706        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
707
708        if (
709            expression.args.get("typed")
710            and left_type in exp.DataType.INTEGER_TYPES
711            and right_type in exp.DataType.INTEGER_TYPES
712        ):
713            self._set_type(expression, exp.DataType.Type.BIGINT)
714        else:
715            self._set_type(expression, self._maybe_coerce(left_type, right_type))
716            if expression.type and expression.type.this not in exp.DataType.REAL_TYPES:
717                self._set_type(
718                    expression, self._maybe_coerce(expression.type, exp.DataType.Type.DOUBLE)
719                )
720
721        return expression
722
723    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
724        self._set_type(expression, None)
725
726        # Propagate type from qualified UDF calls (e.g., db.my_udf(...))
727        if isinstance(expression.expression, exp.Anonymous):
728            self._set_type(expression, expression.expression.type)
729            return expression
730
731        this_type = expression.this.type
732
733        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
734            for e in this_type.expressions:
735                if e.name == expression.expression.name:
736                    self._set_type(expression, e.kind)
737                    break
738
739        return expression
740
741    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
742        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
743        return expression
744
745    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
746        child = seq_get(expression.expressions, 0)
747
748        if child and child.is_type(exp.DataType.Type.ARRAY):
749            expr_type = seq_get(child.type.expressions, 0)
750        else:
751            expr_type = None
752
753        self._set_type(expression, expr_type)
754        return expression
755
756    def _annotate_subquery(self, expression: exp.Subquery) -> exp.Subquery:
757        # For scalar subqueries (subqueries with a single projection), infer the type
758        # from that single projection. This allows type propagation in cases like:
759        # SELECT (SELECT 1 AS c) AS c
760        query = expression.unnest()
761
762        if isinstance(query, exp.Query):
763            selects = query.selects
764            if len(selects) == 1:
765                self._set_type(expression, selects[0].type)
766                return expression
767
768        self._set_type(expression, exp.DataType.Type.UNKNOWN)
769        return expression
770
771    def _annotate_struct_value(
772        self, expression: exp.Expression
773    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
774        # Case: STRUCT(key AS value)
775        this: t.Optional[exp.Expression] = None
776        kind = expression.type
777
778        if alias := expression.args.get("alias"):
779            this = alias.copy()
780        elif expression.expression:
781            # Case: STRUCT(key = value) or STRUCT(key := value)
782            this = expression.this.copy()
783            kind = expression.expression.type
784        elif isinstance(expression, exp.Column):
785            # Case: STRUCT(c)
786            this = expression.this.copy()
787
788        if kind and kind.is_type(exp.DataType.Type.UNKNOWN):
789            return None
790
791        if this:
792            return exp.ColumnDef(this=this, kind=kind)
793
794        return kind
795
796    def _annotate_struct(self, expression: exp.Struct) -> exp.Struct:
797        expressions = []
798        for expr in expression.expressions:
799            struct_field_type = self._annotate_struct_value(expr)
800            if struct_field_type is None:
801                self._set_type(expression, None)
802                return expression
803
804            expressions.append(struct_field_type)
805
806        self._set_type(
807            expression,
808            exp.DataType(this=exp.DataType.Type.STRUCT, expressions=expressions, nested=True),
809        )
810        return expression
811
812    @t.overload
813    def _annotate_map(self, expression: exp.Map) -> exp.Map: ...
814
815    @t.overload
816    def _annotate_map(self, expression: exp.VarMap) -> exp.VarMap: ...
817
818    def _annotate_map(self, expression):
819        keys = expression.args.get("keys")
820        values = expression.args.get("values")
821
822        map_type = exp.DataType(this=exp.DataType.Type.MAP)
823        if isinstance(keys, exp.Array) and isinstance(values, exp.Array):
824            key_type = seq_get(keys.type.expressions, 0) or exp.DataType.Type.UNKNOWN
825            value_type = seq_get(values.type.expressions, 0) or exp.DataType.Type.UNKNOWN
826
827            if key_type != exp.DataType.Type.UNKNOWN and value_type != exp.DataType.Type.UNKNOWN:
828                map_type.set("expressions", [key_type, value_type])
829                map_type.set("nested", True)
830
831        self._set_type(expression, map_type)
832        return expression
833
834    def _annotate_to_map(self, expression: exp.ToMap) -> exp.ToMap:
835        map_type = exp.DataType(this=exp.DataType.Type.MAP)
836        arg = expression.this
837        if arg.is_type(exp.DataType.Type.STRUCT):
838            for coldef in arg.type.expressions:
839                kind = coldef.kind
840                if kind != exp.DataType.Type.UNKNOWN:
841                    map_type.set("expressions", [exp.DataType.build("varchar"), kind])
842                    map_type.set("nested", True)
843                    break
844
845        self._set_type(expression, map_type)
846        return expression
847
848    def _annotate_extract(self, expression: exp.Extract) -> exp.Extract:
849        part = expression.name
850        if part == "TIME":
851            self._set_type(expression, exp.DataType.Type.TIME)
852        elif part == "DATE":
853            self._set_type(expression, exp.DataType.Type.DATE)
854        elif part in BIGINT_EXTRACT_DATE_PARTS:
855            self._set_type(expression, exp.DataType.Type.BIGINT)
856        else:
857            self._set_type(expression, exp.DataType.Type.INT)
858        return expression
859
860    def _annotate_by_array_element(self, expression: exp.Expression) -> exp.Expression:
861        array_arg = expression.this
862        if array_arg.type.is_type(exp.DataType.Type.ARRAY):
863            element_type = seq_get(array_arg.type.expressions, 0) or exp.DataType.Type.UNKNOWN
864            self._set_type(expression, element_type)
865        else:
866            self._set_type(expression, exp.DataType.Type.UNKNOWN)
867
868        return expression