sqlglot.optimizer.annotate_types API documentation

Edit on GitHub

# sqlglot.optimizer.annotate_types

View Source

      1from __future__ import annotations
      2
      3import functools
      4import typing as t
      5
      6from sqlglot import exp
      7from sqlglot.helper import (
      8    ensure_list,
      9    is_date_unit,
     10    is_iso_date,
     11    is_iso_datetime,
     12    seq_get,
     13    subclasses,
     14)
     15from sqlglot.optimizer.scope import Scope, traverse_scope
     16from sqlglot.schema import Schema, ensure_schema
     17
     18if t.TYPE_CHECKING:
     19    from sqlglot._typing import B, E
     20
     21    BinaryCoercionFunc = t.Callable[[exp.Expression, exp.Expression], exp.DataType.Type]
     22    BinaryCoercions = t.Dict[
     23        t.Tuple[exp.DataType.Type, exp.DataType.Type],
     24        BinaryCoercionFunc,
     25    ]
     26
     27
     28def annotate_types(
     29    expression: E,
     30    schema: t.Optional[t.Dict | Schema] = None,
     31    annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
     32    coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
     33) -> E:
     34    """
     35    Infers the types of an expression, annotating its AST accordingly.
     36
     37    Example:
     38        >>> import sqlglot
     39        >>> schema = {"y": {"cola": "SMALLINT"}}
     40        >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
     41        >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
     42        >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
     43        <Type.DOUBLE: 'DOUBLE'>
     44
     45    Args:
     46        expression: Expression to annotate.
     47        schema: Database schema.
     48        annotators: Maps expression type to corresponding annotation function.
     49        coerces_to: Maps expression type to set of types that it can be coerced into.
     50
     51    Returns:
     52        The expression annotated with types.
     53    """
     54
     55    schema = ensure_schema(schema)
     56
     57    return TypeAnnotator(schema, annotators, coerces_to).annotate(expression)
     58
     59
     60def _annotate_with_type_lambda(data_type: exp.DataType.Type) -> t.Callable[[TypeAnnotator, E], E]:
     61    return lambda self, e: self._annotate_with_type(e, data_type)
     62
     63
     64def _coerce_date_literal(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
     65    date_text = l.name
     66    is_iso_date_ = is_iso_date(date_text)
     67
     68    if is_iso_date_ and is_date_unit(unit):
     69        return exp.DataType.Type.DATE
     70
     71    # An ISO date is also an ISO datetime, but not vice versa
     72    if is_iso_date_ or is_iso_datetime(date_text):
     73        return exp.DataType.Type.DATETIME
     74
     75    return exp.DataType.Type.UNKNOWN
     76
     77
     78def _coerce_date(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
     79    if not is_date_unit(unit):
     80        return exp.DataType.Type.DATETIME
     81    return l.type.this if l.type else exp.DataType.Type.UNKNOWN
     82
     83
     84def swap_args(func: BinaryCoercionFunc) -> BinaryCoercionFunc:
     85    @functools.wraps(func)
     86    def _swapped(l: exp.Expression, r: exp.Expression) -> exp.DataType.Type:
     87        return func(r, l)
     88
     89    return _swapped
     90
     91
     92def swap_all(coercions: BinaryCoercions) -> BinaryCoercions:
     93    return {**coercions, **{(b, a): swap_args(func) for (a, b), func in coercions.items()}}
     94
     95
     96class _TypeAnnotator(type):
     97    def __new__(cls, clsname, bases, attrs):
     98        klass = super().__new__(cls, clsname, bases, attrs)
     99
    100        # Highest-to-lowest type precedence, as specified in Spark's docs (ANSI):
    101        # https://spark.apache.org/docs/3.2.0/sql-ref-ansi-compliance.html
    102        text_precedence = (
    103            exp.DataType.Type.TEXT,
    104            exp.DataType.Type.NVARCHAR,
    105            exp.DataType.Type.VARCHAR,
    106            exp.DataType.Type.NCHAR,
    107            exp.DataType.Type.CHAR,
    108        )
    109        numeric_precedence = (
    110            exp.DataType.Type.DOUBLE,
    111            exp.DataType.Type.FLOAT,
    112            exp.DataType.Type.DECIMAL,
    113            exp.DataType.Type.BIGINT,
    114            exp.DataType.Type.INT,
    115            exp.DataType.Type.SMALLINT,
    116            exp.DataType.Type.TINYINT,
    117        )
    118        timelike_precedence = (
    119            exp.DataType.Type.TIMESTAMPLTZ,
    120            exp.DataType.Type.TIMESTAMPTZ,
    121            exp.DataType.Type.TIMESTAMP,
    122            exp.DataType.Type.DATETIME,
    123            exp.DataType.Type.DATE,
    124        )
    125
    126        for type_precedence in (text_precedence, numeric_precedence, timelike_precedence):
    127            coerces_to = set()
    128            for data_type in type_precedence:
    129                klass.COERCES_TO[data_type] = coerces_to.copy()
    130                coerces_to |= {data_type}
    131
    132        return klass
    133
    134
    135class TypeAnnotator(metaclass=_TypeAnnotator):
    136    TYPE_TO_EXPRESSIONS: t.Dict[exp.DataType.Type, t.Set[t.Type[exp.Expression]]] = {
    137        exp.DataType.Type.BIGINT: {
    138            exp.ApproxDistinct,
    139            exp.ArraySize,
    140            exp.Count,
    141            exp.Length,
    142        },
    143        exp.DataType.Type.BOOLEAN: {
    144            exp.Between,
    145            exp.Boolean,
    146            exp.In,
    147            exp.RegexpLike,
    148        },
    149        exp.DataType.Type.DATE: {
    150            exp.CurrentDate,
    151            exp.Date,
    152            exp.DateFromParts,
    153            exp.DateStrToDate,
    154            exp.DiToDate,
    155            exp.StrToDate,
    156            exp.TimeStrToDate,
    157            exp.TsOrDsToDate,
    158        },
    159        exp.DataType.Type.DATETIME: {
    160            exp.CurrentDatetime,
    161            exp.DatetimeAdd,
    162            exp.DatetimeSub,
    163        },
    164        exp.DataType.Type.DOUBLE: {
    165            exp.ApproxQuantile,
    166            exp.Avg,
    167            exp.Div,
    168            exp.Exp,
    169            exp.Ln,
    170            exp.Log,
    171            exp.Log2,
    172            exp.Log10,
    173            exp.Pow,
    174            exp.Quantile,
    175            exp.Round,
    176            exp.SafeDivide,
    177            exp.Sqrt,
    178            exp.Stddev,
    179            exp.StddevPop,
    180            exp.StddevSamp,
    181            exp.Variance,
    182            exp.VariancePop,
    183        },
    184        exp.DataType.Type.INT: {
    185            exp.Ceil,
    186            exp.DatetimeDiff,
    187            exp.DateDiff,
    188            exp.Extract,
    189            exp.TimestampDiff,
    190            exp.TimeDiff,
    191            exp.DateToDi,
    192            exp.Floor,
    193            exp.Levenshtein,
    194            exp.Sign,
    195            exp.StrPosition,
    196            exp.TsOrDiToDi,
    197        },
    198        exp.DataType.Type.JSON: {
    199            exp.ParseJSON,
    200        },
    201        exp.DataType.Type.TIMESTAMP: {
    202            exp.CurrentTime,
    203            exp.CurrentTimestamp,
    204            exp.StrToTime,
    205            exp.TimeAdd,
    206            exp.TimeStrToTime,
    207            exp.TimeSub,
    208            exp.TimestampAdd,
    209            exp.TimestampSub,
    210            exp.UnixToTime,
    211        },
    212        exp.DataType.Type.TINYINT: {
    213            exp.Day,
    214            exp.Month,
    215            exp.Week,
    216            exp.Year,
    217        },
    218        exp.DataType.Type.VARCHAR: {
    219            exp.ArrayConcat,
    220            exp.Concat,
    221            exp.ConcatWs,
    222            exp.DateToDateStr,
    223            exp.GroupConcat,
    224            exp.Initcap,
    225            exp.Lower,
    226            exp.Substring,
    227            exp.TimeToStr,
    228            exp.TimeToTimeStr,
    229            exp.Trim,
    230            exp.TsOrDsToDateStr,
    231            exp.UnixToStr,
    232            exp.UnixToTimeStr,
    233            exp.Upper,
    234        },
    235    }
    236
    237    ANNOTATORS: t.Dict = {
    238        **{
    239            expr_type: lambda self, e: self._annotate_unary(e)
    240            for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
    241        },
    242        **{
    243            expr_type: lambda self, e: self._annotate_binary(e)
    244            for expr_type in subclasses(exp.__name__, exp.Binary)
    245        },
    246        **{
    247            expr_type: _annotate_with_type_lambda(data_type)
    248            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
    249            for expr_type in expressions
    250        },
    251        exp.Abs: lambda self, e: self._annotate_by_args(e, "this"),
    252        exp.Anonymous: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    253        exp.Array: lambda self, e: self._annotate_by_args(e, "expressions", array=True),
    254        exp.ArrayAgg: lambda self, e: self._annotate_by_args(e, "this", array=True),
    255        exp.ArrayConcat: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    256        exp.Bracket: lambda self, e: self._annotate_bracket(e),
    257        exp.Cast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    258        exp.Case: lambda self, e: self._annotate_by_args(e, "default", "ifs"),
    259        exp.Coalesce: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    260        exp.DataType: lambda self, e: self._annotate_with_type(e, e.copy()),
    261        exp.DateAdd: lambda self, e: self._annotate_timeunit(e),
    262        exp.DateSub: lambda self, e: self._annotate_timeunit(e),
    263        exp.DateTrunc: lambda self, e: self._annotate_timeunit(e),
    264        exp.Distinct: lambda self, e: self._annotate_by_args(e, "expressions"),
    265        exp.Div: lambda self, e: self._annotate_div(e),
    266        exp.Dot: lambda self, e: self._annotate_dot(e),
    267        exp.Explode: lambda self, e: self._annotate_explode(e),
    268        exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
    269        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
    270        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
    271        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
    272        exp.Literal: lambda self, e: self._annotate_literal(e),
    273        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    274        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    275        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    276        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
    277        exp.Nullif: lambda self, e: self._annotate_by_args(e, "this", "expression"),
    278        exp.PropertyEQ: lambda self, e: self._annotate_by_args(e, "expression"),
    279        exp.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    280        exp.Struct: lambda self, e: self._annotate_by_args(e, "expressions", struct=True),
    281        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
    282        exp.Timestamp: lambda self, e: self._annotate_with_type(
    283            e,
    284            exp.DataType.Type.TIMESTAMPTZ if e.args.get("with_tz") else exp.DataType.Type.TIMESTAMP,
    285        ),
    286        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    287        exp.Unnest: lambda self, e: self._annotate_unnest(e),
    288        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    289    }
    290
    291    NESTED_TYPES = {
    292        exp.DataType.Type.ARRAY,
    293    }
    294
    295    # Specifies what types a given type can be coerced into (autofilled)
    296    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
    297
    298    # Coercion functions for binary operations.
    299    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
    300    BINARY_COERCIONS: BinaryCoercions = {
    301        **swap_all(
    302            {
    303                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
    304                    l, r.args.get("unit")
    305                )
    306                for t in exp.DataType.TEXT_TYPES
    307            }
    308        ),
    309        **swap_all(
    310            {
    311                # text + numeric will yield the numeric type to match most dialects' semantics
    312                (text, numeric): lambda l, r: t.cast(
    313                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
    314                )
    315                for text in exp.DataType.TEXT_TYPES
    316                for numeric in exp.DataType.NUMERIC_TYPES
    317            }
    318        ),
    319        **swap_all(
    320            {
    321                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
    322                    l, r.args.get("unit")
    323                ),
    324            }
    325        ),
    326    }
    327
    328    def __init__(
    329        self,
    330        schema: Schema,
    331        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    332        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    333        binary_coercions: t.Optional[BinaryCoercions] = None,
    334    ) -> None:
    335        self.schema = schema
    336        self.annotators = annotators or self.ANNOTATORS
    337        self.coerces_to = coerces_to or self.COERCES_TO
    338        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    339
    340        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    341        self._visited: t.Set[int] = set()
    342
    343    def _set_type(
    344        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
    345    ) -> None:
    346        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
    347        self._visited.add(id(expression))
    348
    349    def annotate(self, expression: E) -> E:
    350        for scope in traverse_scope(expression):
    351            selects = {}
    352            for name, source in scope.sources.items():
    353                if not isinstance(source, Scope):
    354                    continue
    355                if isinstance(source.expression, exp.UDTF):
    356                    values = []
    357
    358                    if isinstance(source.expression, exp.Lateral):
    359                        if isinstance(source.expression.this, exp.Explode):
    360                            values = [source.expression.this.this]
    361                    else:
    362                        values = source.expression.expressions[0].expressions
    363
    364                    if not values:
    365                        continue
    366
    367                    selects[name] = {
    368                        alias: column
    369                        for alias, column in zip(
    370                            source.expression.alias_column_names,
    371                            values,
    372                        )
    373                    }
    374                else:
    375                    selects[name] = {
    376                        select.alias_or_name: select for select in source.expression.selects
    377                    }
    378
    379            # First annotate the current scope's column references
    380            for col in scope.columns:
    381                if not col.table:
    382                    continue
    383
    384                source = scope.sources.get(col.table)
    385                if isinstance(source, exp.Table):
    386                    self._set_type(col, self.schema.get_column_type(source, col))
    387                elif source:
    388                    if col.table in selects and col.name in selects[col.table]:
    389                        self._set_type(col, selects[col.table][col.name].type)
    390                    elif isinstance(source.expression, exp.Unnest):
    391                        self._set_type(col, source.expression.type)
    392
    393            # Then (possibly) annotate the remaining expressions in the scope
    394            self._maybe_annotate(scope.expression)
    395
    396        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
    397
    398    def _maybe_annotate(self, expression: E) -> E:
    399        if id(expression) in self._visited:
    400            return expression  # We've already inferred the expression's type
    401
    402        annotator = self.annotators.get(expression.__class__)
    403
    404        return (
    405            annotator(self, expression)
    406            if annotator
    407            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
    408        )
    409
    410    def _annotate_args(self, expression: E) -> E:
    411        for _, value in expression.iter_expressions():
    412            self._maybe_annotate(value)
    413
    414        return expression
    415
    416    def _maybe_coerce(
    417        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
    418    ) -> exp.DataType | exp.DataType.Type:
    419        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
    420        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
    421
    422        # We propagate the NULL / UNKNOWN types upwards if found
    423        if exp.DataType.Type.NULL in (type1_value, type2_value):
    424            return exp.DataType.Type.NULL
    425        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
    426            return exp.DataType.Type.UNKNOWN
    427
    428        if type1_value in self.NESTED_TYPES:
    429            return type1
    430        if type2_value in self.NESTED_TYPES:
    431            return type2
    432
    433        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
    434
    435    # Note: the following "no_type_check" decorators were added because mypy was yelling due
    436    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
    437    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
    438
    439    @t.no_type_check
    440    def _annotate_binary(self, expression: B) -> B:
    441        self._annotate_args(expression)
    442
    443        left, right = expression.left, expression.right
    444        left_type, right_type = left.type.this, right.type.this
    445
    446        if isinstance(expression, exp.Connector):
    447            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
    448                self._set_type(expression, exp.DataType.Type.NULL)
    449            elif exp.DataType.Type.NULL in (left_type, right_type):
    450                self._set_type(
    451                    expression,
    452                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
    453                )
    454            else:
    455                self._set_type(expression, exp.DataType.Type.BOOLEAN)
    456        elif isinstance(expression, exp.Predicate):
    457            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    458        elif (left_type, right_type) in self.binary_coercions:
    459            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
    460        else:
    461            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    462
    463        return expression
    464
    465    @t.no_type_check
    466    def _annotate_unary(self, expression: E) -> E:
    467        self._annotate_args(expression)
    468
    469        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
    470            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    471        else:
    472            self._set_type(expression, expression.this.type)
    473
    474        return expression
    475
    476    @t.no_type_check
    477    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
    478        if expression.is_string:
    479            self._set_type(expression, exp.DataType.Type.VARCHAR)
    480        elif expression.is_int:
    481            self._set_type(expression, exp.DataType.Type.INT)
    482        else:
    483            self._set_type(expression, exp.DataType.Type.DOUBLE)
    484
    485        return expression
    486
    487    @t.no_type_check
    488    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
    489        self._set_type(expression, target_type)
    490        return self._annotate_args(expression)
    491
    492    @t.no_type_check
    493    def _annotate_struct_value(
    494        self, expression: exp.Expression
    495    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
    496        alias = expression.args.get("alias")
    497        if alias:
    498            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
    499
    500        # Case: key = value or key := value
    501        if expression.expression:
    502            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
    503
    504        return expression.type
    505
    506    @t.no_type_check
    507    def _annotate_by_args(
    508        self,
    509        expression: E,
    510        *args: str,
    511        promote: bool = False,
    512        array: bool = False,
    513        struct: bool = False,
    514    ) -> E:
    515        self._annotate_args(expression)
    516
    517        expressions: t.List[exp.Expression] = []
    518        for arg in args:
    519            arg_expr = expression.args.get(arg)
    520            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
    521
    522        last_datatype = None
    523        for expr in expressions:
    524            expr_type = expr.type
    525
    526            # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
    527            if expr_type.args.get("nested"):
    528                last_datatype = expr_type
    529                break
    530
    531            last_datatype = self._maybe_coerce(last_datatype or expr_type, expr_type)
    532
    533        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
    534
    535        if promote:
    536            if expression.type.this in exp.DataType.INTEGER_TYPES:
    537                self._set_type(expression, exp.DataType.Type.BIGINT)
    538            elif expression.type.this in exp.DataType.FLOAT_TYPES:
    539                self._set_type(expression, exp.DataType.Type.DOUBLE)
    540
    541        if array:
    542            self._set_type(
    543                expression,
    544                exp.DataType(
    545                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
    546                ),
    547            )
    548
    549        if struct:
    550            self._set_type(
    551                expression,
    552                exp.DataType(
    553                    this=exp.DataType.Type.STRUCT,
    554                    expressions=[self._annotate_struct_value(expr) for expr in expressions],
    555                    nested=True,
    556                ),
    557            )
    558
    559        return expression
    560
    561    def _annotate_timeunit(
    562        self, expression: exp.TimeUnit | exp.DateTrunc
    563    ) -> exp.TimeUnit | exp.DateTrunc:
    564        self._annotate_args(expression)
    565
    566        if expression.this.type.this in exp.DataType.TEXT_TYPES:
    567            datatype = _coerce_date_literal(expression.this, expression.unit)
    568        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
    569            datatype = _coerce_date(expression.this, expression.unit)
    570        else:
    571            datatype = exp.DataType.Type.UNKNOWN
    572
    573        self._set_type(expression, datatype)
    574        return expression
    575
    576    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
    577        self._annotate_args(expression)
    578
    579        bracket_arg = expression.expressions[0]
    580        this = expression.this
    581
    582        if isinstance(bracket_arg, exp.Slice):
    583            self._set_type(expression, this.type)
    584        elif this.type.is_type(exp.DataType.Type.ARRAY):
    585            self._set_type(expression, seq_get(this.type.expressions, 0))
    586        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
    587            index = this.keys.index(bracket_arg)
    588            value = seq_get(this.values, index)
    589            self._set_type(expression, value.type if value else None)
    590        else:
    591            self._set_type(expression, exp.DataType.Type.UNKNOWN)
    592
    593        return expression
    594
    595    def _annotate_div(self, expression: exp.Div) -> exp.Div:
    596        self._annotate_args(expression)
    597
    598        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
    599
    600        if (
    601            expression.args.get("typed")
    602            and left_type in exp.DataType.INTEGER_TYPES
    603            and right_type in exp.DataType.INTEGER_TYPES
    604        ):
    605            self._set_type(expression, exp.DataType.Type.BIGINT)
    606        else:
    607            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    608
    609        return expression
    610
    611    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
    612        self._annotate_args(expression)
    613        self._set_type(expression, None)
    614        this_type = expression.this.type
    615
    616        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
    617            for e in this_type.expressions:
    618                if e.name == expression.expression.name:
    619                    self._set_type(expression, e.kind)
    620                    break
    621
    622        return expression
    623
    624    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
    625        self._annotate_args(expression)
    626        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
    627        return expression
    628
    629    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
    630        self._annotate_args(expression)
    631        child = seq_get(expression.expressions, 0)
    632        self._set_type(expression, child and seq_get(child.type.expressions, 0))
    633        return expression

def annotate_types( expression: ~E, schema: Union[Dict, sqlglot.schema.Schema, NoneType] = None, annotators: Optional[Dict[Type[~E], Callable[[TypeAnnotator, ~E], ~E]]] = None, coerces_to: Optional[Dict[sqlglot.expressions.DataType.Type, Set[sqlglot.expressions.DataType.Type]]] = None) -> ~E: View Source

    29def annotate_types(
    30    expression: E,
    31    schema: t.Optional[t.Dict | Schema] = None,
    32    annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    33    coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    34) -> E:
    35    """
    36    Infers the types of an expression, annotating its AST accordingly.
    37
    38    Example:
    39        >>> import sqlglot
    40        >>> schema = {"y": {"cola": "SMALLINT"}}
    41        >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
    42        >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
    43        >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
    44        <Type.DOUBLE: 'DOUBLE'>
    45
    46    Args:
    47        expression: Expression to annotate.
    48        schema: Database schema.
    49        annotators: Maps expression type to corresponding annotation function.
    50        coerces_to: Maps expression type to set of types that it can be coerced into.
    51
    52    Returns:
    53        The expression annotated with types.
    54    """
    55
    56    schema = ensure_schema(schema)
    57
    58    return TypeAnnotator(schema, annotators, coerces_to).annotate(expression)

Infers the types of an expression, annotating its AST accordingly.

###### Example:

>     >>> import sqlglot
>     >>> schema = {"y": {"cola": "SMALLINT"}}
>     >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
>     >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
>     >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
>     <Type.DOUBLE: 'DOUBLE'>

###### Arguments:

- expression: Expression to annotate.
- schema: Database schema.
- annotators: Maps expression type to corresponding annotation function.
- coerces_to: Maps expression type to set of types that it can be coerced into.

###### Returns:

> The expression annotated with types.

def swap_args( func: Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]) -> Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]: View Source

    85def swap_args(func: BinaryCoercionFunc) -> BinaryCoercionFunc:
    86    @functools.wraps(func)
    87    def _swapped(l: exp.Expression, r: exp.Expression) -> exp.DataType.Type:
    88        return func(r, l)
    89
    90    return _swapped

def swap_all( coercions: Dict[Tuple[sqlglot.expressions.DataType.Type, sqlglot.expressions.DataType.Type], Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]]) -> Dict[Tuple[sqlglot.expressions.DataType.Type, sqlglot.expressions.DataType.Type], Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]]: View Source

    93def swap_all(coercions: BinaryCoercions) -> BinaryCoercions:
    94    return {**coercions, **{(b, a): swap_args(func) for (a, b), func in coercions.items()}}

class TypeAnnotator: View Source

    136class TypeAnnotator(metaclass=_TypeAnnotator):
    137    TYPE_TO_EXPRESSIONS: t.Dict[exp.DataType.Type, t.Set[t.Type[exp.Expression]]] = {
    138        exp.DataType.Type.BIGINT: {
    139            exp.ApproxDistinct,
    140            exp.ArraySize,
    141            exp.Count,
    142            exp.Length,
    143        },
    144        exp.DataType.Type.BOOLEAN: {
    145            exp.Between,
    146            exp.Boolean,
    147            exp.In,
    148            exp.RegexpLike,
    149        },
    150        exp.DataType.Type.DATE: {
    151            exp.CurrentDate,
    152            exp.Date,
    153            exp.DateFromParts,
    154            exp.DateStrToDate,
    155            exp.DiToDate,
    156            exp.StrToDate,
    157            exp.TimeStrToDate,
    158            exp.TsOrDsToDate,
    159        },
    160        exp.DataType.Type.DATETIME: {
    161            exp.CurrentDatetime,
    162            exp.DatetimeAdd,
    163            exp.DatetimeSub,
    164        },
    165        exp.DataType.Type.DOUBLE: {
    166            exp.ApproxQuantile,
    167            exp.Avg,
    168            exp.Div,
    169            exp.Exp,
    170            exp.Ln,
    171            exp.Log,
    172            exp.Log2,
    173            exp.Log10,
    174            exp.Pow,
    175            exp.Quantile,
    176            exp.Round,
    177            exp.SafeDivide,
    178            exp.Sqrt,
    179            exp.Stddev,
    180            exp.StddevPop,
    181            exp.StddevSamp,
    182            exp.Variance,
    183            exp.VariancePop,
    184        },
    185        exp.DataType.Type.INT: {
    186            exp.Ceil,
    187            exp.DatetimeDiff,
    188            exp.DateDiff,
    189            exp.Extract,
    190            exp.TimestampDiff,
    191            exp.TimeDiff,
    192            exp.DateToDi,
    193            exp.Floor,
    194            exp.Levenshtein,
    195            exp.Sign,
    196            exp.StrPosition,
    197            exp.TsOrDiToDi,
    198        },
    199        exp.DataType.Type.JSON: {
    200            exp.ParseJSON,
    201        },
    202        exp.DataType.Type.TIMESTAMP: {
    203            exp.CurrentTime,
    204            exp.CurrentTimestamp,
    205            exp.StrToTime,
    206            exp.TimeAdd,
    207            exp.TimeStrToTime,
    208            exp.TimeSub,
    209            exp.TimestampAdd,
    210            exp.TimestampSub,
    211            exp.UnixToTime,
    212        },
    213        exp.DataType.Type.TINYINT: {
    214            exp.Day,
    215            exp.Month,
    216            exp.Week,
    217            exp.Year,
    218        },
    219        exp.DataType.Type.VARCHAR: {
    220            exp.ArrayConcat,
    221            exp.Concat,
    222            exp.ConcatWs,
    223            exp.DateToDateStr,
    224            exp.GroupConcat,
    225            exp.Initcap,
    226            exp.Lower,
    227            exp.Substring,
    228            exp.TimeToStr,
    229            exp.TimeToTimeStr,
    230            exp.Trim,
    231            exp.TsOrDsToDateStr,
    232            exp.UnixToStr,
    233            exp.UnixToTimeStr,
    234            exp.Upper,
    235        },
    236    }
    237
    238    ANNOTATORS: t.Dict = {
    239        **{
    240            expr_type: lambda self, e: self._annotate_unary(e)
    241            for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
    242        },
    243        **{
    244            expr_type: lambda self, e: self._annotate_binary(e)
    245            for expr_type in subclasses(exp.__name__, exp.Binary)
    246        },
    247        **{
    248            expr_type: _annotate_with_type_lambda(data_type)
    249            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
    250            for expr_type in expressions
    251        },
    252        exp.Abs: lambda self, e: self._annotate_by_args(e, "this"),
    253        exp.Anonymous: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    254        exp.Array: lambda self, e: self._annotate_by_args(e, "expressions", array=True),
    255        exp.ArrayAgg: lambda self, e: self._annotate_by_args(e, "this", array=True),
    256        exp.ArrayConcat: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    257        exp.Bracket: lambda self, e: self._annotate_bracket(e),
    258        exp.Cast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    259        exp.Case: lambda self, e: self._annotate_by_args(e, "default", "ifs"),
    260        exp.Coalesce: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    261        exp.DataType: lambda self, e: self._annotate_with_type(e, e.copy()),
    262        exp.DateAdd: lambda self, e: self._annotate_timeunit(e),
    263        exp.DateSub: lambda self, e: self._annotate_timeunit(e),
    264        exp.DateTrunc: lambda self, e: self._annotate_timeunit(e),
    265        exp.Distinct: lambda self, e: self._annotate_by_args(e, "expressions"),
    266        exp.Div: lambda self, e: self._annotate_div(e),
    267        exp.Dot: lambda self, e: self._annotate_dot(e),
    268        exp.Explode: lambda self, e: self._annotate_explode(e),
    269        exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
    270        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
    271        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
    272        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
    273        exp.Literal: lambda self, e: self._annotate_literal(e),
    274        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    275        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    276        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    277        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
    278        exp.Nullif: lambda self, e: self._annotate_by_args(e, "this", "expression"),
    279        exp.PropertyEQ: lambda self, e: self._annotate_by_args(e, "expression"),
    280        exp.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    281        exp.Struct: lambda self, e: self._annotate_by_args(e, "expressions", struct=True),
    282        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
    283        exp.Timestamp: lambda self, e: self._annotate_with_type(
    284            e,
    285            exp.DataType.Type.TIMESTAMPTZ if e.args.get("with_tz") else exp.DataType.Type.TIMESTAMP,
    286        ),
    287        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    288        exp.Unnest: lambda self, e: self._annotate_unnest(e),
    289        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    290    }
    291
    292    NESTED_TYPES = {
    293        exp.DataType.Type.ARRAY,
    294    }
    295
    296    # Specifies what types a given type can be coerced into (autofilled)
    297    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
    298
    299    # Coercion functions for binary operations.
    300    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
    301    BINARY_COERCIONS: BinaryCoercions = {
    302        **swap_all(
    303            {
    304                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
    305                    l, r.args.get("unit")
    306                )
    307                for t in exp.DataType.TEXT_TYPES
    308            }
    309        ),
    310        **swap_all(
    311            {
    312                # text + numeric will yield the numeric type to match most dialects' semantics
    313                (text, numeric): lambda l, r: t.cast(
    314                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
    315                )
    316                for text in exp.DataType.TEXT_TYPES
    317                for numeric in exp.DataType.NUMERIC_TYPES
    318            }
    319        ),
    320        **swap_all(
    321            {
    322                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
    323                    l, r.args.get("unit")
    324                ),
    325            }
    326        ),
    327    }
    328
    329    def __init__(
    330        self,
    331        schema: Schema,
    332        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    333        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    334        binary_coercions: t.Optional[BinaryCoercions] = None,
    335    ) -> None:
    336        self.schema = schema
    337        self.annotators = annotators or self.ANNOTATORS
    338        self.coerces_to = coerces_to or self.COERCES_TO
    339        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    340
    341        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    342        self._visited: t.Set[int] = set()
    343
    344    def _set_type(
    345        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
    346    ) -> None:
    347        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
    348        self._visited.add(id(expression))
    349
    350    def annotate(self, expression: E) -> E:
    351        for scope in traverse_scope(expression):
    352            selects = {}
    353            for name, source in scope.sources.items():
    354                if not isinstance(source, Scope):
    355                    continue
    356                if isinstance(source.expression, exp.UDTF):
    357                    values = []
    358
    359                    if isinstance(source.expression, exp.Lateral):
    360                        if isinstance(source.expression.this, exp.Explode):
    361                            values = [source.expression.this.this]
    362                    else:
    363                        values = source.expression.expressions[0].expressions
    364
    365                    if not values:
    366                        continue
    367
    368                    selects[name] = {
    369                        alias: column
    370                        for alias, column in zip(
    371                            source.expression.alias_column_names,
    372                            values,
    373                        )
    374                    }
    375                else:
    376                    selects[name] = {
    377                        select.alias_or_name: select for select in source.expression.selects
    378                    }
    379
    380            # First annotate the current scope's column references
    381            for col in scope.columns:
    382                if not col.table:
    383                    continue
    384
    385                source = scope.sources.get(col.table)
    386                if isinstance(source, exp.Table):
    387                    self._set_type(col, self.schema.get_column_type(source, col))
    388                elif source:
    389                    if col.table in selects and col.name in selects[col.table]:
    390                        self._set_type(col, selects[col.table][col.name].type)
    391                    elif isinstance(source.expression, exp.Unnest):
    392                        self._set_type(col, source.expression.type)
    393
    394            # Then (possibly) annotate the remaining expressions in the scope
    395            self._maybe_annotate(scope.expression)
    396
    397        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
    398
    399    def _maybe_annotate(self, expression: E) -> E:
    400        if id(expression) in self._visited:
    401            return expression  # We've already inferred the expression's type
    402
    403        annotator = self.annotators.get(expression.__class__)
    404
    405        return (
    406            annotator(self, expression)
    407            if annotator
    408            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
    409        )
    410
    411    def _annotate_args(self, expression: E) -> E:
    412        for _, value in expression.iter_expressions():
    413            self._maybe_annotate(value)
    414
    415        return expression
    416
    417    def _maybe_coerce(
    418        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
    419    ) -> exp.DataType | exp.DataType.Type:
    420        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
    421        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
    422
    423        # We propagate the NULL / UNKNOWN types upwards if found
    424        if exp.DataType.Type.NULL in (type1_value, type2_value):
    425            return exp.DataType.Type.NULL
    426        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
    427            return exp.DataType.Type.UNKNOWN
    428
    429        if type1_value in self.NESTED_TYPES:
    430            return type1
    431        if type2_value in self.NESTED_TYPES:
    432            return type2
    433
    434        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
    435
    436    # Note: the following "no_type_check" decorators were added because mypy was yelling due
    437    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
    438    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
    439
    440    @t.no_type_check
    441    def _annotate_binary(self, expression: B) -> B:
    442        self._annotate_args(expression)
    443
    444        left, right = expression.left, expression.right
    445        left_type, right_type = left.type.this, right.type.this
    446
    447        if isinstance(expression, exp.Connector):
    448            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
    449                self._set_type(expression, exp.DataType.Type.NULL)
    450            elif exp.DataType.Type.NULL in (left_type, right_type):
    451                self._set_type(
    452                    expression,
    453                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
    454                )
    455            else:
    456                self._set_type(expression, exp.DataType.Type.BOOLEAN)
    457        elif isinstance(expression, exp.Predicate):
    458            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    459        elif (left_type, right_type) in self.binary_coercions:
    460            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
    461        else:
    462            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    463
    464        return expression
    465
    466    @t.no_type_check
    467    def _annotate_unary(self, expression: E) -> E:
    468        self._annotate_args(expression)
    469
    470        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
    471            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    472        else:
    473            self._set_type(expression, expression.this.type)
    474
    475        return expression
    476
    477    @t.no_type_check
    478    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
    479        if expression.is_string:
    480            self._set_type(expression, exp.DataType.Type.VARCHAR)
    481        elif expression.is_int:
    482            self._set_type(expression, exp.DataType.Type.INT)
    483        else:
    484            self._set_type(expression, exp.DataType.Type.DOUBLE)
    485
    486        return expression
    487
    488    @t.no_type_check
    489    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
    490        self._set_type(expression, target_type)
    491        return self._annotate_args(expression)
    492
    493    @t.no_type_check
    494    def _annotate_struct_value(
    495        self, expression: exp.Expression
    496    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
    497        alias = expression.args.get("alias")
    498        if alias:
    499            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
    500
    501        # Case: key = value or key := value
    502        if expression.expression:
    503            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
    504
    505        return expression.type
    506
    507    @t.no_type_check
    508    def _annotate_by_args(
    509        self,
    510        expression: E,
    511        *args: str,
    512        promote: bool = False,
    513        array: bool = False,
    514        struct: bool = False,
    515    ) -> E:
    516        self._annotate_args(expression)
    517
    518        expressions: t.List[exp.Expression] = []
    519        for arg in args:
    520            arg_expr = expression.args.get(arg)
    521            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
    522
    523        last_datatype = None
    524        for expr in expressions:
    525            expr_type = expr.type
    526
    527            # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
    528            if expr_type.args.get("nested"):
    529                last_datatype = expr_type
    530                break
    531
    532            last_datatype = self._maybe_coerce(last_datatype or expr_type, expr_type)
    533
    534        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
    535
    536        if promote:
    537            if expression.type.this in exp.DataType.INTEGER_TYPES:
    538                self._set_type(expression, exp.DataType.Type.BIGINT)
    539            elif expression.type.this in exp.DataType.FLOAT_TYPES:
    540                self._set_type(expression, exp.DataType.Type.DOUBLE)
    541
    542        if array:
    543            self._set_type(
    544                expression,
    545                exp.DataType(
    546                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
    547                ),
    548            )
    549
    550        if struct:
    551            self._set_type(
    552                expression,
    553                exp.DataType(
    554                    this=exp.DataType.Type.STRUCT,
    555                    expressions=[self._annotate_struct_value(expr) for expr in expressions],
    556                    nested=True,
    557                ),
    558            )
    559
    560        return expression
    561
    562    def _annotate_timeunit(
    563        self, expression: exp.TimeUnit | exp.DateTrunc
    564    ) -> exp.TimeUnit | exp.DateTrunc:
    565        self._annotate_args(expression)
    566
    567        if expression.this.type.this in exp.DataType.TEXT_TYPES:
    568            datatype = _coerce_date_literal(expression.this, expression.unit)
    569        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
    570            datatype = _coerce_date(expression.this, expression.unit)
    571        else:
    572            datatype = exp.DataType.Type.UNKNOWN
    573
    574        self._set_type(expression, datatype)
    575        return expression
    576
    577    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
    578        self._annotate_args(expression)
    579
    580        bracket_arg = expression.expressions[0]
    581        this = expression.this
    582
    583        if isinstance(bracket_arg, exp.Slice):
    584            self._set_type(expression, this.type)
    585        elif this.type.is_type(exp.DataType.Type.ARRAY):
    586            self._set_type(expression, seq_get(this.type.expressions, 0))
    587        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
    588            index = this.keys.index(bracket_arg)
    589            value = seq_get(this.values, index)
    590            self._set_type(expression, value.type if value else None)
    591        else:
    592            self._set_type(expression, exp.DataType.Type.UNKNOWN)
    593
    594        return expression
    595
    596    def _annotate_div(self, expression: exp.Div) -> exp.Div:
    597        self._annotate_args(expression)
    598
    599        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
    600
    601        if (
    602            expression.args.get("typed")
    603            and left_type in exp.DataType.INTEGER_TYPES
    604            and right_type in exp.DataType.INTEGER_TYPES
    605        ):
    606            self._set_type(expression, exp.DataType.Type.BIGINT)
    607        else:
    608            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    609
    610        return expression
    611
    612    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
    613        self._annotate_args(expression)
    614        self._set_type(expression, None)
    615        this_type = expression.this.type
    616
    617        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
    618            for e in this_type.expressions:
    619                if e.name == expression.expression.name:
    620                    self._set_type(expression, e.kind)
    621                    break
    622
    623        return expression
    624
    625    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
    626        self._annotate_args(expression)
    627        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
    628        return expression
    629
    630    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
    631        self._annotate_args(expression)
    632        child = seq_get(expression.expressions, 0)
    633        self._set_type(expression, child and seq_get(child.type.expressions, 0))
    634        return expression

TypeAnnotator( schema: sqlglot.schema.Schema, annotators: Optional[Dict[Type[~E], Callable[[TypeAnnotator, ~E], ~E]]] = None, coerces_to: Optional[Dict[sqlglot.expressions.DataType.Type, Set[sqlglot.expressions.DataType.Type]]] = None, binary_coercions: Optional[Dict[Tuple[sqlglot.expressions.DataType.Type, sqlglot.expressions.DataType.Type], Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]]] = None) View Source

    329    def __init__(
    330        self,
    331        schema: Schema,
    332        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    333        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    334        binary_coercions: t.Optional[BinaryCoercions] = None,
    335    ) -> None:
    336        self.schema = schema
    337        self.annotators = annotators or self.ANNOTATORS
    338        self.coerces_to = coerces_to or self.COERCES_TO
    339        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    340
    341        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    342        self._visited: t.Set[int] = set()

TYPE_TO_EXPRESSIONS: Dict[sqlglot.expressions.DataType.Type, Set[Type[sqlglot.expressions.Expression]]] = {<Type.BIGINT: 'BIGINT'>: {<class 'sqlglot.expressions.ArraySize'>, <class 'sqlglot.expressions.Count'>, <class 'sqlglot.expressions.Length'>, <class 'sqlglot.expressions.ApproxDistinct'>}, <Type.BOOLEAN: 'BOOLEAN'>: {<class 'sqlglot.expressions.In'>, <class 'sqlglot.expressions.Boolean'>, <class 'sqlglot.expressions.Between'>, <class 'sqlglot.expressions.RegexpLike'>}, <Type.DATE: 'DATE'>: {<class 'sqlglot.expressions.DateStrToDate'>, <class 'sqlglot.expressions.DateFromParts'>, <class 'sqlglot.expressions.DiToDate'>, <class 'sqlglot.expressions.TsOrDsToDate'>, <class 'sqlglot.expressions.CurrentDate'>, <class 'sqlglot.expressions.Date'>, <class 'sqlglot.expressions.TimeStrToDate'>, <class 'sqlglot.expressions.StrToDate'>}, <Type.DATETIME: 'DATETIME'>: {<class 'sqlglot.expressions.DatetimeAdd'>, <class 'sqlglot.expressions.CurrentDatetime'>, <class 'sqlglot.expressions.DatetimeSub'>}, <Type.DOUBLE: 'DOUBLE'>: {<class 'sqlglot.expressions.Variance'>, <class 'sqlglot.expressions.ApproxQuantile'>, <class 'sqlglot.expressions.Ln'>, <class 'sqlglot.expressions.Exp'>, <class 'sqlglot.expressions.Quantile'>, <class 'sqlglot.expressions.SafeDivide'>, <class 'sqlglot.expressions.StddevSamp'>, <class 'sqlglot.expressions.Round'>, <class 'sqlglot.expressions.Avg'>, <class 'sqlglot.expressions.Div'>, <class 'sqlglot.expressions.Pow'>, <class 'sqlglot.expressions.Stddev'>, <class 'sqlglot.expressions.StddevPop'>, <class 'sqlglot.expressions.Log10'>, <class 'sqlglot.expressions.Sqrt'>, <class 'sqlglot.expressions.Log2'>, <class 'sqlglot.expressions.VariancePop'>, <class 'sqlglot.expressions.Log'>}, <Type.INT: 'INT'>: {<class 'sqlglot.expressions.StrPosition'>, <class 'sqlglot.expressions.TimestampDiff'>, <class 'sqlglot.expressions.Levenshtein'>, <class 'sqlglot.expressions.DatetimeDiff'>, <class 'sqlglot.expressions.Floor'>, <class 'sqlglot.expressions.DateDiff'>, <class 'sqlglot.expressions.Extract'>, <class 'sqlglot.expressions.TimeDiff'>, <class 'sqlglot.expressions.TsOrDiToDi'>, <class 'sqlglot.expressions.Ceil'>, <class 'sqlglot.expressions.Sign'>, <class 'sqlglot.expressions.DateToDi'>}, <Type.JSON: 'JSON'>: {<class 'sqlglot.expressions.ParseJSON'>}, <Type.TIMESTAMP: 'TIMESTAMP'>: {<class 'sqlglot.expressions.TimestampSub'>, <class 'sqlglot.expressions.CurrentTimestamp'>, <class 'sqlglot.expressions.TimestampAdd'>, <class 'sqlglot.expressions.CurrentTime'>, <class 'sqlglot.expressions.UnixToTime'>, <class 'sqlglot.expressions.TimeStrToTime'>, <class 'sqlglot.expressions.TimeSub'>, <class 'sqlglot.expressions.StrToTime'>, <class 'sqlglot.expressions.TimeAdd'>}, <Type.TINYINT: 'TINYINT'>: {<class 'sqlglot.expressions.Week'>, <class 'sqlglot.expressions.Month'>, <class 'sqlglot.expressions.Day'>, <class 'sqlglot.expressions.Year'>}, <Type.VARCHAR: 'VARCHAR'>: {<class 'sqlglot.expressions.DateToDateStr'>, <class 'sqlglot.expressions.GroupConcat'>, <class 'sqlglot.expressions.TimeToTimeStr'>, <class 'sqlglot.expressions.TsOrDsToDateStr'>, <class 'sqlglot.expressions.Initcap'>, <class 'sqlglot.expressions.Upper'>, <class 'sqlglot.expressions.TimeToStr'>, <class 'sqlglot.expressions.ConcatWs'>, <class 'sqlglot.expressions.Lower'>, <class 'sqlglot.expressions.UnixToStr'>, <class 'sqlglot.expressions.Concat'>, <class 'sqlglot.expressions.UnixToTimeStr'>, <class 'sqlglot.expressions.Substring'>, <class 'sqlglot.expressions.Trim'>, <class 'sqlglot.expressions.ArrayConcat'>}}
ANNOTATORS: Dict = {<class 'sqlglot.expressions.Alias'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.BitwiseNot'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Neg'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Not'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Paren'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.PivotAlias'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Unary'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Add'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.And'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.ArrayContained'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.ArrayContains'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.ArrayOverlaps'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Binary'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.BitwiseAnd'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.BitwiseLeftShift'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.BitwiseOr'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.BitwiseRightShift'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.BitwiseXor'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Collate'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Connector'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.DPipe'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Distance'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Div'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Dot'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.EQ'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Escape'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.GT'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.GTE'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Glob'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.ILike'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.ILikeAny'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.IntDiv'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Is'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.JSONArrayContains'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.JSONBContains'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.JSONBExtract'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.JSONBExtractScalar'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.JSONExtract'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.JSONExtractScalar'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Kwarg'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.LT'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.LTE'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Like'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.LikeAny'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Mod'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Mul'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.NEQ'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.NullSafeEQ'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.NullSafeNEQ'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Operator'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Or'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Overlaps'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Pow'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.PropertyEQ'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.RegexpILike'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.RegexpLike'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.SimilarTo'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Slice'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Sub'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.Xor'>: <function TypeAnnotator.<dictcomp>.<lambda>>, <class 'sqlglot.expressions.ArraySize'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Count'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Length'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ApproxDistinct'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.In'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Boolean'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Between'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateStrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateFromParts'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DiToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDsToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Date'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeStrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeAdd'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentDatetime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeSub'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Variance'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ApproxQuantile'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Ln'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Exp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Quantile'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.SafeDivide'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StddevSamp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Round'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Avg'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Stddev'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StddevPop'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log10'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Sqrt'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log2'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.VariancePop'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StrPosition'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimestampDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Levenshtein'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Floor'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Extract'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDiToDi'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Ceil'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Sign'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateToDi'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ParseJSON'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimestampSub'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentTimestamp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimestampAdd'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentTime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToTime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeStrToTime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeSub'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StrToTime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeAdd'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Week'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Month'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Day'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Year'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateToDateStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.GroupConcat'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeToTimeStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDsToDateStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Initcap'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Upper'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeToStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ConcatWs'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Lower'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Concat'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToTimeStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Substring'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Trim'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ArrayConcat'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Abs'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Anonymous'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Array'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.ArrayAgg'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Bracket'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Cast'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Case'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Coalesce'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.DataType'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.DateAdd'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.DateSub'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.DateTrunc'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Distinct'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Explode'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Filter'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.If'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Interval'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Least'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Literal'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Map'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Max'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Min'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Null'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Nullif'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Struct'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Sum'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Timestamp'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.TryCast'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Unnest'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.VarMap'>: <function TypeAnnotator.<lambda>>}
NESTED_TYPES = {<Type.ARRAY: 'ARRAY'>}
COERCES_TO: Dict[sqlglot.expressions.DataType.Type, Set[sqlglot.expressions.DataType.Type]] = {<Type.TEXT: 'TEXT'>: set(), <Type.NVARCHAR: 'NVARCHAR'>: {<Type.TEXT: 'TEXT'>}, <Type.VARCHAR: 'VARCHAR'>: {<Type.TEXT: 'TEXT'>, <Type.NVARCHAR: 'NVARCHAR'>}, <Type.NCHAR: 'NCHAR'>: {<Type.TEXT: 'TEXT'>, <Type.VARCHAR: 'VARCHAR'>, <Type.NVARCHAR: 'NVARCHAR'>}, <Type.CHAR: 'CHAR'>: {<Type.TEXT: 'TEXT'>, <Type.NCHAR: 'NCHAR'>, <Type.VARCHAR: 'VARCHAR'>, <Type.NVARCHAR: 'NVARCHAR'>}, <Type.DOUBLE: 'DOUBLE'>: set(), <Type.FLOAT: 'FLOAT'>: {<Type.DOUBLE: 'DOUBLE'>}, <Type.DECIMAL: 'DECIMAL'>: {<Type.FLOAT: 'FLOAT'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.BIGINT: 'BIGINT'>: {<Type.FLOAT: 'FLOAT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.INT: 'INT'>: {<Type.BIGINT: 'BIGINT'>, <Type.FLOAT: 'FLOAT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.SMALLINT: 'SMALLINT'>: {<Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.INT: 'INT'>, <Type.FLOAT: 'FLOAT'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.TINYINT: 'TINYINT'>: {<Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.INT: 'INT'>, <Type.SMALLINT: 'SMALLINT'>, <Type.FLOAT: 'FLOAT'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>: set(), <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>: {<Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>}, <Type.TIMESTAMP: 'TIMESTAMP'>: {<Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>}, <Type.DATETIME: 'DATETIME'>: {<Type.TIMESTAMP: 'TIMESTAMP'>, <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>}, <Type.DATE: 'DATE'>: {<Type.DATETIME: 'DATETIME'>, <Type.TIMESTAMP: 'TIMESTAMP'>, <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>}}
BINARY_COERCIONS: Dict[Tuple[sqlglot.expressions.DataType.Type, sqlglot.expressions.DataType.Type], Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]] = {(<Type.VARCHAR: 'VARCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DATE: 'DATE'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.DATE: 'DATE'>): <function TypeAnnotator.<lambda>>}
schema
annotators
coerces_to
binary_coercions
def annotate(self, expression: ~E) -> ~E: View Source

    350    def annotate(self, expression: E) -> E:
    351        for scope in traverse_scope(expression):
    352            selects = {}
    353            for name, source in scope.sources.items():
    354                if not isinstance(source, Scope):
    355                    continue
    356                if isinstance(source.expression, exp.UDTF):
    357                    values = []
    358
    359                    if isinstance(source.expression, exp.Lateral):
    360                        if isinstance(source.expression.this, exp.Explode):
    361                            values = [source.expression.this.this]
    362                    else:
    363                        values = source.expression.expressions[0].expressions
    364
    365                    if not values:
    366                        continue
    367
    368                    selects[name] = {
    369                        alias: column
    370                        for alias, column in zip(
    371                            source.expression.alias_column_names,
    372                            values,
    373                        )
    374                    }
    375                else:
    376                    selects[name] = {
    377                        select.alias_or_name: select for select in source.expression.selects
    378                    }
    379
    380            # First annotate the current scope's column references
    381            for col in scope.columns:
    382                if not col.table:
    383                    continue
    384
    385                source = scope.sources.get(col.table)
    386                if isinstance(source, exp.Table):
    387                    self._set_type(col, self.schema.get_column_type(source, col))
    388                elif source:
    389                    if col.table in selects and col.name in selects[col.table]:
    390                        self._set_type(col, selects[col.table][col.name].type)
    391                    elif isinstance(source.expression, exp.Unnest):
    392                        self._set_type(col, source.expression.type)
    393
    394            # Then (possibly) annotate the remaining expressions in the scope
    395            self._maybe_annotate(scope.expression)
    396
    397        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions