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        return expression
252
253    def annotate(self, expression: E, annotate_scope: bool = True) -> E:
254        # This flag is used to avoid costly scope traversals when we only care about annotating
255        # non-column expressions (partial type inference), e.g., when simplifying in the optimizer
256        if annotate_scope:
257            for scope in traverse_scope(expression):
258                self.annotate_scope(scope)
259
260        # This takes care of non-traversable expressions
261        self._annotate_expression(expression)
262
263        # Replace NULL type with the default type of the targeted dialect, since the former is not an actual type;
264        # it is mostly used to aid type coercion, e.g. in query set operations.
265        for expr in self._null_expressions.values():
266            expr.type = self.dialect.DEFAULT_NULL_TYPE
267
268        return expression
269
270    def _get_scope_selects(self, scope: Scope) -> t.Dict[str, t.Dict[str, t.Any]]:
271        if scope not in self._scope_selects:
272            selects = {}
273            for name, source in scope.sources.items():
274                if not isinstance(source, Scope):
275                    continue
276
277                expression = source.expression
278                if isinstance(expression, exp.UDTF):
279                    values = []
280
281                    if isinstance(expression, exp.Lateral):
282                        if isinstance(expression.this, exp.Explode):
283                            values = [expression.this.this]
284                    elif isinstance(expression, exp.Unnest):
285                        values = [expression]
286                    elif not isinstance(expression, exp.TableFromRows):
287                        values = expression.expressions[0].expressions
288
289                    if not values:
290                        continue
291
292                    alias_column_names = expression.alias_column_names
293
294                    if (
295                        isinstance(expression, exp.Unnest)
296                        and expression.type
297                        and expression.type.is_type(exp.DataType.Type.STRUCT)
298                    ):
299                        selects[name] = {
300                            col_def.name: t.cast(
301                                t.Union[exp.DataType, exp.DataType.Type], col_def.kind
302                            )
303                            for col_def in expression.type.expressions
304                            if isinstance(col_def, exp.ColumnDef) and col_def.kind
305                        }
306                    else:
307                        selects[name] = {
308                            alias: column.type for alias, column in zip(alias_column_names, values)
309                        }
310                elif isinstance(expression, exp.SetOperation) and len(
311                    expression.left.selects
312                ) == len(expression.right.selects):
313                    selects[name] = self._get_setop_column_types(expression)
314                else:
315                    selects[name] = {s.alias_or_name: s.type for s in expression.selects}
316
317            self._scope_selects[scope] = selects
318
319        return self._scope_selects[scope]
320
321    def annotate_scope(self, scope: Scope) -> None:
322        if isinstance(self.schema, MappingSchema):
323            for table_column in scope.table_columns:
324                source = scope.sources.get(table_column.name)
325
326                if isinstance(source, exp.Table):
327                    schema = self.schema.find(
328                        source, raise_on_missing=False, ensure_data_types=True
329                    )
330                    if not isinstance(schema, dict):
331                        continue
332
333                    struct_type = exp.DataType(
334                        this=exp.DataType.Type.STRUCT,
335                        expressions=[
336                            exp.ColumnDef(this=exp.to_identifier(c), kind=kind)
337                            for c, kind in schema.items()
338                        ],
339                        nested=True,
340                    )
341                    self._set_type(table_column, struct_type)
342                elif (
343                    isinstance(source, Scope)
344                    and isinstance(source.expression, exp.Query)
345                    and (
346                        source.expression.meta.get("query_type") or exp.DataType.build("UNKNOWN")
347                    ).is_type(exp.DataType.Type.STRUCT)
348                ):
349                    self._set_type(table_column, source.expression.meta["query_type"])
350
351        # Iterate through all the expressions of the current scope in post-order, and annotate
352        self._annotate_expression(scope.expression, scope)
353
354        if self.dialect.QUERY_RESULTS_ARE_STRUCTS and isinstance(scope.expression, exp.Query):
355            struct_type = exp.DataType(
356                this=exp.DataType.Type.STRUCT,
357                expressions=[
358                    exp.ColumnDef(
359                        this=exp.to_identifier(select.output_name),
360                        kind=select.type.copy() if select.type else None,
361                    )
362                    for select in scope.expression.selects
363                ],
364                nested=True,
365            )
366
367            if not any(
368                cd.kind.is_type(exp.DataType.Type.UNKNOWN)
369                for cd in struct_type.expressions
370                if cd.kind
371            ):
372                # We don't use `_set_type` on purpose here. If we annotated the query directly, then
373                # using it in other contexts (e.g., ARRAY(<query>)) could result in incorrect type
374                # annotations, i.e., it shouldn't be interpreted as a STRUCT value.
375                scope.expression.meta["query_type"] = struct_type
376
377    def _annotate_expression(
378        self,
379        expression: exp.Expression,
380        scope: t.Optional[Scope] = None,
381    ) -> None:
382        stack = [(expression, False)]
383
384        while stack:
385            expr, children_annotated = stack.pop()
386
387            if id(expr) in self._visited or (
388                not self._overwrite_types
389                and expr.type
390                and not expr.is_type(exp.DataType.Type.UNKNOWN)
391            ):
392                continue  # We've already inferred the expression's type
393
394            if not children_annotated:
395                stack.append((expr, True))
396                for child_expr in expr.iter_expressions():
397                    stack.append((child_expr, False))
398                continue
399
400            if scope and isinstance(expr, exp.Column) and expr.table:
401                source = None
402                source_scope = scope
403                while source_scope and not source:
404                    source = source_scope.sources.get(expr.table)
405                    if not source:
406                        source_scope = source_scope.parent
407
408                if isinstance(source, exp.Table):
409                    self._set_type(expr, self.schema.get_column_type(source, expr))
410                elif source:
411                    col_type = (
412                        self._get_scope_selects(source_scope).get(expr.table, {}).get(expr.name)
413                    )
414                    if col_type:
415                        self._set_type(expr, col_type)
416                    elif isinstance(source.expression, exp.Unnest):
417                        self._set_type(expr, source.expression.type)
418                    else:
419                        self._set_type(expr, exp.DataType.Type.UNKNOWN)
420                else:
421                    self._set_type(expr, exp.DataType.Type.UNKNOWN)
422
423                if expr.is_type(exp.DataType.Type.JSON) and (
424                    dot_parts := expr.meta.get("dot_parts")
425                ):
426                    # JSON dot access is case sensitive across all dialects, so we need to undo the normalization.
427                    i = iter(dot_parts)
428                    parent = expr.parent
429                    while isinstance(parent, exp.Dot):
430                        parent.expression.replace(exp.to_identifier(next(i), quoted=True))
431                        parent = parent.parent
432
433                    expr.meta.pop("dot_parts", None)
434
435                if expr.type and expr.type.args.get("nullable") is False:
436                    expr.meta["nonnull"] = True
437                continue
438
439            spec = self.expression_metadata.get(expr.__class__)
440
441            if spec and (annotator := spec.get("annotator")):
442                annotator(self, expr)
443            elif spec and (returns := spec.get("returns")):
444                self._set_type(expr, t.cast(exp.DataType.Type, returns))
445            else:
446                self._set_type(expr, exp.DataType.Type.UNKNOWN)
447
448    def _maybe_coerce(
449        self,
450        type1: exp.DataType | exp.DataType.Type,
451        type2: exp.DataType | exp.DataType.Type,
452    ) -> exp.DataType | exp.DataType.Type:
453        """
454        Returns type2 if type1 can be coerced into it, otherwise type1.
455
456        If either type is parameterized (e.g. DECIMAL(18, 2) contains two parameters),
457        we assume type1 does not coerce into type2, so we also return it in this case.
458        """
459        if isinstance(type1, exp.DataType):
460            if type1.expressions:
461                return type1
462            type1_value = type1.this
463        else:
464            type1_value = type1
465
466        if isinstance(type2, exp.DataType):
467            if type2.expressions:
468                return type2
469            type2_value = type2.this
470        else:
471            type2_value = type2
472
473        # We propagate the UNKNOWN type upwards if found
474        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
475            return exp.DataType.Type.UNKNOWN
476
477        if type1_value == exp.DataType.Type.NULL:
478            return type2_value
479        if type2_value == exp.DataType.Type.NULL:
480            return type1_value
481
482        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value
483
484    def _get_setop_column_types(
485        self, setop: exp.SetOperation
486    ) -> t.Dict[str, exp.DataType | exp.DataType.Type]:
487        """
488        Computes and returns the coerced column types for a SetOperation.
489
490        This handles UNION, INTERSECT, EXCEPT, etc., coercing types across
491        left and right operands for all projections/columns.
492
493        Args:
494            setop: The SetOperation expression to analyze
495
496        Returns:
497            Dictionary mapping column names to their coerced types
498        """
499        setop_id = id(setop)
500        if setop_id in self._setop_column_types:
501            return self._setop_column_types[setop_id]
502
503        col_types: t.Dict[str, exp.DataType | exp.DataType.Type] = {}
504
505        # Validate that left and right have same number of projections
506        if not (
507            isinstance(setop, exp.SetOperation)
508            and setop.left.selects
509            and setop.right.selects
510            and len(setop.left.selects) == len(setop.right.selects)
511        ):
512            return col_types
513
514        # Process a chain / sub-tree of set operations
515        for set_op in setop.walk(
516            prune=lambda n: not isinstance(n, (exp.SetOperation, exp.Subquery))
517        ):
518            if not isinstance(set_op, exp.SetOperation):
519                continue
520
521            if set_op.args.get("by_name"):
522                r_type_by_select = {s.alias_or_name: s.type for s in set_op.right.selects}
523                setop_cols = {
524                    s.alias_or_name: self._maybe_coerce(
525                        t.cast(exp.DataType, s.type),
526                        r_type_by_select.get(s.alias_or_name) or exp.DataType.Type.UNKNOWN,
527                    )
528                    for s in set_op.left.selects
529                }
530            else:
531                setop_cols = {
532                    ls.alias_or_name: self._maybe_coerce(
533                        t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
534                    )
535                    for ls, rs in zip(set_op.left.selects, set_op.right.selects)
536                }
537
538            # Coerce intermediate results with the previously registered types, if they exist
539            for col_name, col_type in setop_cols.items():
540                col_types[col_name] = self._maybe_coerce(
541                    col_type, col_types.get(col_name, exp.DataType.Type.NULL)
542                )
543
544        self._setop_column_types[setop_id] = col_types
545        return col_types
546
547    def _annotate_binary(self, expression: B) -> B:
548        left, right = expression.left, expression.right
549        if not left or not right:
550            expression_sql = expression.sql(self.dialect)
551            logger.warning(f"Failed to annotate badly formed binary expression: {expression_sql}")
552            self._set_type(expression, None)
553            return expression
554
555        left_type, right_type = left.type.this, right.type.this  # type: ignore
556
557        if isinstance(expression, (exp.Connector, exp.Predicate)):
558            self._set_type(expression, exp.DataType.Type.BOOLEAN)
559        elif (left_type, right_type) in self.binary_coercions:
560            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
561        else:
562            self._annotate_by_args(expression, left, right)
563
564        if isinstance(expression, exp.Is) or (
565            left.meta.get("nonnull") is True and right.meta.get("nonnull") is True
566        ):
567            expression.meta["nonnull"] = True
568
569        return expression
570
571    def _annotate_unary(self, expression: E) -> E:
572        if isinstance(expression, exp.Not):
573            self._set_type(expression, exp.DataType.Type.BOOLEAN)
574        else:
575            self._set_type(expression, expression.this.type)
576
577        if expression.this.meta.get("nonnull") is True:
578            expression.meta["nonnull"] = True
579
580        return expression
581
582    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
583        if expression.is_string:
584            self._set_type(expression, exp.DataType.Type.VARCHAR)
585        elif expression.is_int:
586            self._set_type(expression, exp.DataType.Type.INT)
587        else:
588            self._set_type(expression, exp.DataType.Type.DOUBLE)
589
590        expression.meta["nonnull"] = True
591
592        return expression
593
594    @t.no_type_check
595    def _annotate_by_args(
596        self,
597        expression: E,
598        *args: str | exp.Expression,
599        promote: bool = False,
600        array: bool = False,
601    ) -> E:
602        literal_type = None
603        non_literal_type = None
604        nested_type = None
605
606        for arg in args:
607            if isinstance(arg, str):
608                expressions = expression.args.get(arg)
609            else:
610                expressions = arg
611
612            for expr in ensure_list(expressions):
613                expr_type = expr.type
614
615                # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
616                if expr_type.args.get("nested"):
617                    nested_type = expr_type
618                    break
619
620                if isinstance(expr, exp.Literal):
621                    literal_type = self._maybe_coerce(literal_type or expr_type, expr_type)
622                else:
623                    non_literal_type = self._maybe_coerce(non_literal_type or expr_type, expr_type)
624
625            if nested_type:
626                break
627
628        result_type = None
629
630        if nested_type:
631            result_type = nested_type
632        elif literal_type and non_literal_type:
633            if self.dialect.PRIORITIZE_NON_LITERAL_TYPES:
634                literal_this_type = (
635                    literal_type.this if isinstance(literal_type, exp.DataType) else literal_type
636                )
637                non_literal_this_type = (
638                    non_literal_type.this
639                    if isinstance(non_literal_type, exp.DataType)
640                    else non_literal_type
641                )
642                if (
643                    literal_this_type in exp.DataType.INTEGER_TYPES
644                    and non_literal_this_type in exp.DataType.INTEGER_TYPES
645                ) or (
646                    literal_this_type in exp.DataType.REAL_TYPES
647                    and non_literal_this_type in exp.DataType.REAL_TYPES
648                ):
649                    result_type = non_literal_type
650        else:
651            result_type = literal_type or non_literal_type or exp.DataType.Type.UNKNOWN
652
653        self._set_type(
654            expression, result_type or self._maybe_coerce(non_literal_type, literal_type)
655        )
656
657        if promote:
658            if expression.type.this in exp.DataType.INTEGER_TYPES:
659                self._set_type(expression, exp.DataType.Type.BIGINT)
660            elif expression.type.this in exp.DataType.FLOAT_TYPES:
661                self._set_type(expression, exp.DataType.Type.DOUBLE)
662
663        if array:
664            self._set_type(
665                expression,
666                exp.DataType(
667                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
668                ),
669            )
670
671        return expression
672
673    def _annotate_timeunit(
674        self, expression: exp.TimeUnit | exp.DateTrunc
675    ) -> exp.TimeUnit | exp.DateTrunc:
676        if expression.this.type.this in exp.DataType.TEXT_TYPES:
677            datatype = _coerce_date_literal(expression.this, expression.unit)
678        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
679            datatype = _coerce_date(expression.this, expression.unit)
680        else:
681            datatype = exp.DataType.Type.UNKNOWN
682
683        self._set_type(expression, datatype)
684        return expression
685
686    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
687        bracket_arg = expression.expressions[0]
688        this = expression.this
689
690        if isinstance(bracket_arg, exp.Slice):
691            self._set_type(expression, this.type)
692        elif this.type.is_type(exp.DataType.Type.ARRAY):
693            self._set_type(expression, seq_get(this.type.expressions, 0))
694        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
695            index = this.keys.index(bracket_arg)
696            value = seq_get(this.values, index)
697            self._set_type(expression, value.type if value else None)
698        else:
699            self._set_type(expression, exp.DataType.Type.UNKNOWN)
700
701        return expression
702
703    def _annotate_div(self, expression: exp.Div) -> exp.Div:
704        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
705
706        if (
707            expression.args.get("typed")
708            and left_type in exp.DataType.INTEGER_TYPES
709            and right_type in exp.DataType.INTEGER_TYPES
710        ):
711            self._set_type(expression, exp.DataType.Type.BIGINT)
712        else:
713            self._set_type(expression, self._maybe_coerce(left_type, right_type))
714            if expression.type and expression.type.this not in exp.DataType.REAL_TYPES:
715                self._set_type(
716                    expression, self._maybe_coerce(expression.type, exp.DataType.Type.DOUBLE)
717                )
718
719        return expression
720
721    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
722        self._set_type(expression, None)
723
724        # Propagate type from qualified UDF calls (e.g., db.my_udf(...))
725        if isinstance(expression.expression, exp.Anonymous):
726            self._set_type(expression, expression.expression.type)
727            return expression
728
729        this_type = expression.this.type
730
731        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
732            for e in this_type.expressions:
733                if e.name == expression.expression.name:
734                    self._set_type(expression, e.kind)
735                    break
736
737        return expression
738
739    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
740        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
741        return expression
742
743    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
744        child = seq_get(expression.expressions, 0)
745
746        if child and child.is_type(exp.DataType.Type.ARRAY):
747            expr_type = seq_get(child.type.expressions, 0)
748        else:
749            expr_type = None
750
751        self._set_type(expression, expr_type)
752        return expression
753
754    def _annotate_subquery(self, expression: exp.Subquery) -> exp.Subquery:
755        # For scalar subqueries (subqueries with a single projection), infer the type
756        # from that single projection. This allows type propagation in cases like:
757        # SELECT (SELECT 1 AS c) AS c
758        query = expression.unnest()
759
760        if isinstance(query, exp.Query):
761            selects = query.selects
762            if len(selects) == 1:
763                self._set_type(expression, selects[0].type)
764                return expression
765
766        self._set_type(expression, exp.DataType.Type.UNKNOWN)
767        return expression
768
769    def _annotate_struct_value(
770        self, expression: exp.Expression
771    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
772        # Case: STRUCT(key AS value)
773        this: t.Optional[exp.Expression] = None
774        kind = expression.type
775
776        if alias := expression.args.get("alias"):
777            this = alias.copy()
778        elif expression.expression:
779            # Case: STRUCT(key = value) or STRUCT(key := value)
780            this = expression.this.copy()
781            kind = expression.expression.type
782        elif isinstance(expression, exp.Column):
783            # Case: STRUCT(c)
784            this = expression.this.copy()
785
786        if kind and kind.is_type(exp.DataType.Type.UNKNOWN):
787            return None
788
789        if this:
790            return exp.ColumnDef(this=this, kind=kind)
791
792        return kind
793
794    def _annotate_struct(self, expression: exp.Struct) -> exp.Struct:
795        expressions = []
796        for expr in expression.expressions:
797            struct_field_type = self._annotate_struct_value(expr)
798            if struct_field_type is None:
799                self._set_type(expression, None)
800                return expression
801
802            expressions.append(struct_field_type)
803
804        self._set_type(
805            expression,
806            exp.DataType(this=exp.DataType.Type.STRUCT, expressions=expressions, nested=True),
807        )
808        return expression
809
810    @t.overload
811    def _annotate_map(self, expression: exp.Map) -> exp.Map: ...
812
813    @t.overload
814    def _annotate_map(self, expression: exp.VarMap) -> exp.VarMap: ...
815
816    def _annotate_map(self, expression):
817        keys = expression.args.get("keys")
818        values = expression.args.get("values")
819
820        map_type = exp.DataType(this=exp.DataType.Type.MAP)
821        if isinstance(keys, exp.Array) and isinstance(values, exp.Array):
822            key_type = seq_get(keys.type.expressions, 0) or exp.DataType.Type.UNKNOWN
823            value_type = seq_get(values.type.expressions, 0) or exp.DataType.Type.UNKNOWN
824
825            if key_type != exp.DataType.Type.UNKNOWN and value_type != exp.DataType.Type.UNKNOWN:
826                map_type.set("expressions", [key_type, value_type])
827                map_type.set("nested", True)
828
829        self._set_type(expression, map_type)
830        return expression
831
832    def _annotate_to_map(self, expression: exp.ToMap) -> exp.ToMap:
833        map_type = exp.DataType(this=exp.DataType.Type.MAP)
834        arg = expression.this
835        if arg.is_type(exp.DataType.Type.STRUCT):
836            for coldef in arg.type.expressions:
837                kind = coldef.kind
838                if kind != exp.DataType.Type.UNKNOWN:
839                    map_type.set("expressions", [exp.DataType.build("varchar"), kind])
840                    map_type.set("nested", True)
841                    break
842
843        self._set_type(expression, map_type)
844        return expression
845
846    def _annotate_extract(self, expression: exp.Extract) -> exp.Extract:
847        part = expression.name
848        if part == "TIME":
849            self._set_type(expression, exp.DataType.Type.TIME)
850        elif part == "DATE":
851            self._set_type(expression, exp.DataType.Type.DATE)
852        elif part in BIGINT_EXTRACT_DATE_PARTS:
853            self._set_type(expression, exp.DataType.Type.BIGINT)
854        else:
855            self._set_type(expression, exp.DataType.Type.INT)
856        return expression
857
858    def _annotate_by_array_element(self, expression: exp.Expression) -> exp.Expression:
859        array_arg = expression.this
860        if array_arg.type.is_type(exp.DataType.Type.ARRAY):
861            element_type = seq_get(array_arg.type.expressions, 0) or exp.DataType.Type.UNKNOWN
862            self._set_type(expression, element_type)
863        else:
864            self._set_type(expression, exp.DataType.Type.UNKNOWN)
865
866        return expression