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