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.StrPosition,
    195            exp.TsOrDiToDi,
    196        },
    197        exp.DataType.Type.JSON: {
    198            exp.ParseJSON,
    199        },
    200        exp.DataType.Type.TIMESTAMP: {
    201            exp.CurrentTime,
    202            exp.CurrentTimestamp,
    203            exp.StrToTime,
    204            exp.TimeAdd,
    205            exp.TimeStrToTime,
    206            exp.TimeSub,
    207            exp.Timestamp,
    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.Explode: lambda self, e: self._annotate_explode(e),
    267        exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
    268        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
    269        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
    270        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
    271        exp.Literal: lambda self, e: self._annotate_literal(e),
    272        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    273        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    274        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    275        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
    276        exp.Nullif: lambda self, e: self._annotate_by_args(e, "this", "expression"),
    277        exp.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    278        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
    279        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    280        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    281        exp.Struct: lambda self, e: self._annotate_by_args(e, "expressions", struct=True),
    282    }
    283
    284    NESTED_TYPES = {
    285        exp.DataType.Type.ARRAY,
    286    }
    287
    288    # Specifies what types a given type can be coerced into (autofilled)
    289    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
    290
    291    # Coercion functions for binary operations.
    292    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
    293    BINARY_COERCIONS: BinaryCoercions = {
    294        **swap_all(
    295            {
    296                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
    297                    l, r.args.get("unit")
    298                )
    299                for t in exp.DataType.TEXT_TYPES
    300            }
    301        ),
    302        **swap_all(
    303            {
    304                # text + numeric will yield the numeric type to match most dialects' semantics
    305                (text, numeric): lambda l, r: t.cast(
    306                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
    307                )
    308                for text in exp.DataType.TEXT_TYPES
    309                for numeric in exp.DataType.NUMERIC_TYPES
    310            }
    311        ),
    312        **swap_all(
    313            {
    314                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
    315                    l, r.args.get("unit")
    316                ),
    317            }
    318        ),
    319    }
    320
    321    def __init__(
    322        self,
    323        schema: Schema,
    324        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    325        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    326        binary_coercions: t.Optional[BinaryCoercions] = None,
    327    ) -> None:
    328        self.schema = schema
    329        self.annotators = annotators or self.ANNOTATORS
    330        self.coerces_to = coerces_to or self.COERCES_TO
    331        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    332
    333        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    334        self._visited: t.Set[int] = set()
    335
    336    def _set_type(
    337        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
    338    ) -> None:
    339        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
    340        self._visited.add(id(expression))
    341
    342    def annotate(self, expression: E) -> E:
    343        for scope in traverse_scope(expression):
    344            selects = {}
    345            for name, source in scope.sources.items():
    346                if not isinstance(source, Scope):
    347                    continue
    348                if isinstance(source.expression, exp.UDTF):
    349                    values = []
    350
    351                    if isinstance(source.expression, exp.Lateral):
    352                        if isinstance(source.expression.this, exp.Explode):
    353                            values = [source.expression.this.this]
    354                    else:
    355                        values = source.expression.expressions[0].expressions
    356
    357                    if not values:
    358                        continue
    359
    360                    selects[name] = {
    361                        alias: column
    362                        for alias, column in zip(
    363                            source.expression.alias_column_names,
    364                            values,
    365                        )
    366                    }
    367                else:
    368                    selects[name] = {
    369                        select.alias_or_name: select for select in source.expression.selects
    370                    }
    371
    372            # First annotate the current scope's column references
    373            for col in scope.columns:
    374                if not col.table:
    375                    continue
    376
    377                source = scope.sources.get(col.table)
    378                if isinstance(source, exp.Table):
    379                    self._set_type(col, self.schema.get_column_type(source, col))
    380                elif source and col.table in selects and col.name in selects[col.table]:
    381                    self._set_type(col, selects[col.table][col.name].type)
    382
    383            # Then (possibly) annotate the remaining expressions in the scope
    384            self._maybe_annotate(scope.expression)
    385
    386        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
    387
    388    def _maybe_annotate(self, expression: E) -> E:
    389        if id(expression) in self._visited:
    390            return expression  # We've already inferred the expression's type
    391
    392        annotator = self.annotators.get(expression.__class__)
    393
    394        return (
    395            annotator(self, expression)
    396            if annotator
    397            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
    398        )
    399
    400    def _annotate_args(self, expression: E) -> E:
    401        for _, value in expression.iter_expressions():
    402            self._maybe_annotate(value)
    403
    404        return expression
    405
    406    def _maybe_coerce(
    407        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
    408    ) -> exp.DataType | exp.DataType.Type:
    409        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
    410        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
    411
    412        # We propagate the NULL / UNKNOWN types upwards if found
    413        if exp.DataType.Type.NULL in (type1_value, type2_value):
    414            return exp.DataType.Type.NULL
    415        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
    416            return exp.DataType.Type.UNKNOWN
    417
    418        if type1_value in self.NESTED_TYPES:
    419            return type1
    420        if type2_value in self.NESTED_TYPES:
    421            return type2
    422
    423        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
    424
    425    # Note: the following "no_type_check" decorators were added because mypy was yelling due
    426    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
    427    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
    428
    429    @t.no_type_check
    430    def _annotate_binary(self, expression: B) -> B:
    431        self._annotate_args(expression)
    432
    433        left, right = expression.left, expression.right
    434        left_type, right_type = left.type.this, right.type.this
    435
    436        if isinstance(expression, exp.Connector):
    437            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
    438                self._set_type(expression, exp.DataType.Type.NULL)
    439            elif exp.DataType.Type.NULL in (left_type, right_type):
    440                self._set_type(
    441                    expression,
    442                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
    443                )
    444            else:
    445                self._set_type(expression, exp.DataType.Type.BOOLEAN)
    446        elif isinstance(expression, exp.Predicate):
    447            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    448        elif (left_type, right_type) in self.binary_coercions:
    449            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
    450        else:
    451            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    452
    453        return expression
    454
    455    @t.no_type_check
    456    def _annotate_unary(self, expression: E) -> E:
    457        self._annotate_args(expression)
    458
    459        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
    460            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    461        else:
    462            self._set_type(expression, expression.this.type)
    463
    464        return expression
    465
    466    @t.no_type_check
    467    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
    468        if expression.is_string:
    469            self._set_type(expression, exp.DataType.Type.VARCHAR)
    470        elif expression.is_int:
    471            self._set_type(expression, exp.DataType.Type.INT)
    472        else:
    473            self._set_type(expression, exp.DataType.Type.DOUBLE)
    474
    475        return expression
    476
    477    @t.no_type_check
    478    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
    479        self._set_type(expression, target_type)
    480        return self._annotate_args(expression)
    481
    482    @t.no_type_check
    483    def _annotate_struct_value(
    484        self, expression: exp.Expression
    485    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
    486        alias = expression.args.get("alias")
    487        if alias:
    488            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
    489
    490        # Case: key = value or key := value
    491        if expression.expression:
    492            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
    493
    494        return expression.type
    495
    496    @t.no_type_check
    497    def _annotate_by_args(
    498        self,
    499        expression: E,
    500        *args: str,
    501        promote: bool = False,
    502        array: bool = False,
    503        struct: bool = False,
    504    ) -> E:
    505        self._annotate_args(expression)
    506
    507        expressions: t.List[exp.Expression] = []
    508        for arg in args:
    509            arg_expr = expression.args.get(arg)
    510            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
    511
    512        last_datatype = None
    513        for expr in expressions:
    514            last_datatype = self._maybe_coerce(last_datatype or expr.type, expr.type)
    515
    516        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
    517
    518        if promote:
    519            if expression.type.this in exp.DataType.INTEGER_TYPES:
    520                self._set_type(expression, exp.DataType.Type.BIGINT)
    521            elif expression.type.this in exp.DataType.FLOAT_TYPES:
    522                self._set_type(expression, exp.DataType.Type.DOUBLE)
    523
    524        if array:
    525            self._set_type(
    526                expression,
    527                exp.DataType(
    528                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
    529                ),
    530            )
    531
    532        if struct:
    533            self._set_type(
    534                expression,
    535                exp.DataType(
    536                    this=exp.DataType.Type.STRUCT,
    537                    expressions=[self._annotate_struct_value(expr) for expr in expressions],
    538                    nested=True,
    539                ),
    540            )
    541
    542        return expression
    543
    544    def _annotate_timeunit(
    545        self, expression: exp.TimeUnit | exp.DateTrunc
    546    ) -> exp.TimeUnit | exp.DateTrunc:
    547        self._annotate_args(expression)
    548
    549        if expression.this.type.this in exp.DataType.TEXT_TYPES:
    550            datatype = _coerce_date_literal(expression.this, expression.unit)
    551        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
    552            datatype = _coerce_date(expression.this, expression.unit)
    553        else:
    554            datatype = exp.DataType.Type.UNKNOWN
    555
    556        self._set_type(expression, datatype)
    557        return expression
    558
    559    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
    560        self._annotate_args(expression)
    561
    562        bracket_arg = expression.expressions[0]
    563        this = expression.this
    564
    565        if isinstance(bracket_arg, exp.Slice):
    566            self._set_type(expression, this.type)
    567        elif this.type.is_type(exp.DataType.Type.ARRAY):
    568            self._set_type(expression, seq_get(this.type.expressions, 0))
    569        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
    570            index = this.keys.index(bracket_arg)
    571            value = seq_get(this.values, index)
    572            self._set_type(expression, value.type if value else None)
    573        else:
    574            self._set_type(expression, exp.DataType.Type.UNKNOWN)
    575
    576        return expression
    577
    578    def _annotate_div(self, expression: exp.Div) -> exp.Div:
    579        self._annotate_args(expression)
    580
    581        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
    582
    583        if (
    584            expression.args.get("typed")
    585            and left_type in exp.DataType.INTEGER_TYPES
    586            and right_type in exp.DataType.INTEGER_TYPES
    587        ):
    588            self._set_type(expression, exp.DataType.Type.BIGINT)
    589        else:
    590            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    591
    592        return expression
    593
    594    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
    595        self._annotate_args(expression)
    596        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
    597        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.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.Timestamp,
    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.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.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    279        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
    280        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    281        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    282        exp.Struct: lambda self, e: self._annotate_by_args(e, "expressions", struct=True),
    283    }
    284
    285    NESTED_TYPES = {
    286        exp.DataType.Type.ARRAY,
    287    }
    288
    289    # Specifies what types a given type can be coerced into (autofilled)
    290    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
    291
    292    # Coercion functions for binary operations.
    293    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
    294    BINARY_COERCIONS: BinaryCoercions = {
    295        **swap_all(
    296            {
    297                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
    298                    l, r.args.get("unit")
    299                )
    300                for t in exp.DataType.TEXT_TYPES
    301            }
    302        ),
    303        **swap_all(
    304            {
    305                # text + numeric will yield the numeric type to match most dialects' semantics
    306                (text, numeric): lambda l, r: t.cast(
    307                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
    308                )
    309                for text in exp.DataType.TEXT_TYPES
    310                for numeric in exp.DataType.NUMERIC_TYPES
    311            }
    312        ),
    313        **swap_all(
    314            {
    315                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
    316                    l, r.args.get("unit")
    317                ),
    318            }
    319        ),
    320    }
    321
    322    def __init__(
    323        self,
    324        schema: Schema,
    325        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    326        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    327        binary_coercions: t.Optional[BinaryCoercions] = None,
    328    ) -> None:
    329        self.schema = schema
    330        self.annotators = annotators or self.ANNOTATORS
    331        self.coerces_to = coerces_to or self.COERCES_TO
    332        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    333
    334        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    335        self._visited: t.Set[int] = set()
    336
    337    def _set_type(
    338        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
    339    ) -> None:
    340        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
    341        self._visited.add(id(expression))
    342
    343    def annotate(self, expression: E) -> E:
    344        for scope in traverse_scope(expression):
    345            selects = {}
    346            for name, source in scope.sources.items():
    347                if not isinstance(source, Scope):
    348                    continue
    349                if isinstance(source.expression, exp.UDTF):
    350                    values = []
    351
    352                    if isinstance(source.expression, exp.Lateral):
    353                        if isinstance(source.expression.this, exp.Explode):
    354                            values = [source.expression.this.this]
    355                    else:
    356                        values = source.expression.expressions[0].expressions
    357
    358                    if not values:
    359                        continue
    360
    361                    selects[name] = {
    362                        alias: column
    363                        for alias, column in zip(
    364                            source.expression.alias_column_names,
    365                            values,
    366                        )
    367                    }
    368                else:
    369                    selects[name] = {
    370                        select.alias_or_name: select for select in source.expression.selects
    371                    }
    372
    373            # First annotate the current scope's column references
    374            for col in scope.columns:
    375                if not col.table:
    376                    continue
    377
    378                source = scope.sources.get(col.table)
    379                if isinstance(source, exp.Table):
    380                    self._set_type(col, self.schema.get_column_type(source, col))
    381                elif source and col.table in selects and col.name in selects[col.table]:
    382                    self._set_type(col, selects[col.table][col.name].type)
    383
    384            # Then (possibly) annotate the remaining expressions in the scope
    385            self._maybe_annotate(scope.expression)
    386
    387        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
    388
    389    def _maybe_annotate(self, expression: E) -> E:
    390        if id(expression) in self._visited:
    391            return expression  # We've already inferred the expression's type
    392
    393        annotator = self.annotators.get(expression.__class__)
    394
    395        return (
    396            annotator(self, expression)
    397            if annotator
    398            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
    399        )
    400
    401    def _annotate_args(self, expression: E) -> E:
    402        for _, value in expression.iter_expressions():
    403            self._maybe_annotate(value)
    404
    405        return expression
    406
    407    def _maybe_coerce(
    408        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
    409    ) -> exp.DataType | exp.DataType.Type:
    410        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
    411        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
    412
    413        # We propagate the NULL / UNKNOWN types upwards if found
    414        if exp.DataType.Type.NULL in (type1_value, type2_value):
    415            return exp.DataType.Type.NULL
    416        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
    417            return exp.DataType.Type.UNKNOWN
    418
    419        if type1_value in self.NESTED_TYPES:
    420            return type1
    421        if type2_value in self.NESTED_TYPES:
    422            return type2
    423
    424        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
    425
    426    # Note: the following "no_type_check" decorators were added because mypy was yelling due
    427    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
    428    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
    429
    430    @t.no_type_check
    431    def _annotate_binary(self, expression: B) -> B:
    432        self._annotate_args(expression)
    433
    434        left, right = expression.left, expression.right
    435        left_type, right_type = left.type.this, right.type.this
    436
    437        if isinstance(expression, exp.Connector):
    438            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
    439                self._set_type(expression, exp.DataType.Type.NULL)
    440            elif exp.DataType.Type.NULL in (left_type, right_type):
    441                self._set_type(
    442                    expression,
    443                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
    444                )
    445            else:
    446                self._set_type(expression, exp.DataType.Type.BOOLEAN)
    447        elif isinstance(expression, exp.Predicate):
    448            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    449        elif (left_type, right_type) in self.binary_coercions:
    450            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
    451        else:
    452            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    453
    454        return expression
    455
    456    @t.no_type_check
    457    def _annotate_unary(self, expression: E) -> E:
    458        self._annotate_args(expression)
    459
    460        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
    461            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    462        else:
    463            self._set_type(expression, expression.this.type)
    464
    465        return expression
    466
    467    @t.no_type_check
    468    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
    469        if expression.is_string:
    470            self._set_type(expression, exp.DataType.Type.VARCHAR)
    471        elif expression.is_int:
    472            self._set_type(expression, exp.DataType.Type.INT)
    473        else:
    474            self._set_type(expression, exp.DataType.Type.DOUBLE)
    475
    476        return expression
    477
    478    @t.no_type_check
    479    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
    480        self._set_type(expression, target_type)
    481        return self._annotate_args(expression)
    482
    483    @t.no_type_check
    484    def _annotate_struct_value(
    485        self, expression: exp.Expression
    486    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
    487        alias = expression.args.get("alias")
    488        if alias:
    489            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
    490
    491        # Case: key = value or key := value
    492        if expression.expression:
    493            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
    494
    495        return expression.type
    496
    497    @t.no_type_check
    498    def _annotate_by_args(
    499        self,
    500        expression: E,
    501        *args: str,
    502        promote: bool = False,
    503        array: bool = False,
    504        struct: bool = False,
    505    ) -> E:
    506        self._annotate_args(expression)
    507
    508        expressions: t.List[exp.Expression] = []
    509        for arg in args:
    510            arg_expr = expression.args.get(arg)
    511            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
    512
    513        last_datatype = None
    514        for expr in expressions:
    515            last_datatype = self._maybe_coerce(last_datatype or expr.type, expr.type)
    516
    517        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
    518
    519        if promote:
    520            if expression.type.this in exp.DataType.INTEGER_TYPES:
    521                self._set_type(expression, exp.DataType.Type.BIGINT)
    522            elif expression.type.this in exp.DataType.FLOAT_TYPES:
    523                self._set_type(expression, exp.DataType.Type.DOUBLE)
    524
    525        if array:
    526            self._set_type(
    527                expression,
    528                exp.DataType(
    529                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
    530                ),
    531            )
    532
    533        if struct:
    534            self._set_type(
    535                expression,
    536                exp.DataType(
    537                    this=exp.DataType.Type.STRUCT,
    538                    expressions=[self._annotate_struct_value(expr) for expr in expressions],
    539                    nested=True,
    540                ),
    541            )
    542
    543        return expression
    544
    545    def _annotate_timeunit(
    546        self, expression: exp.TimeUnit | exp.DateTrunc
    547    ) -> exp.TimeUnit | exp.DateTrunc:
    548        self._annotate_args(expression)
    549
    550        if expression.this.type.this in exp.DataType.TEXT_TYPES:
    551            datatype = _coerce_date_literal(expression.this, expression.unit)
    552        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
    553            datatype = _coerce_date(expression.this, expression.unit)
    554        else:
    555            datatype = exp.DataType.Type.UNKNOWN
    556
    557        self._set_type(expression, datatype)
    558        return expression
    559
    560    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
    561        self._annotate_args(expression)
    562
    563        bracket_arg = expression.expressions[0]
    564        this = expression.this
    565
    566        if isinstance(bracket_arg, exp.Slice):
    567            self._set_type(expression, this.type)
    568        elif this.type.is_type(exp.DataType.Type.ARRAY):
    569            self._set_type(expression, seq_get(this.type.expressions, 0))
    570        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
    571            index = this.keys.index(bracket_arg)
    572            value = seq_get(this.values, index)
    573            self._set_type(expression, value.type if value else None)
    574        else:
    575            self._set_type(expression, exp.DataType.Type.UNKNOWN)
    576
    577        return expression
    578
    579    def _annotate_div(self, expression: exp.Div) -> exp.Div:
    580        self._annotate_args(expression)
    581
    582        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
    583
    584        if (
    585            expression.args.get("typed")
    586            and left_type in exp.DataType.INTEGER_TYPES
    587            and right_type in exp.DataType.INTEGER_TYPES
    588        ):
    589            self._set_type(expression, exp.DataType.Type.BIGINT)
    590        else:
    591            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    592
    593        return expression
    594
    595    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
    596        self._annotate_args(expression)
    597        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
    598        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

    322    def __init__(
    323        self,
    324        schema: Schema,
    325        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    326        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    327        binary_coercions: t.Optional[BinaryCoercions] = None,
    328    ) -> None:
    329        self.schema = schema
    330        self.annotators = annotators or self.ANNOTATORS
    331        self.coerces_to = coerces_to or self.COERCES_TO
    332        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    333
    334        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    335        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.Length'>, <class 'sqlglot.expressions.ApproxDistinct'>, <class 'sqlglot.expressions.Count'>, <class 'sqlglot.expressions.ArraySize'>}, <Type.BOOLEAN: 'BOOLEAN'>: {<class 'sqlglot.expressions.In'>, <class 'sqlglot.expressions.RegexpLike'>, <class 'sqlglot.expressions.Between'>, <class 'sqlglot.expressions.Boolean'>}, <Type.DATE: 'DATE'>: {<class 'sqlglot.expressions.TsOrDsToDate'>, <class 'sqlglot.expressions.DiToDate'>, <class 'sqlglot.expressions.DateFromParts'>, <class 'sqlglot.expressions.Date'>, <class 'sqlglot.expressions.TimeStrToDate'>, <class 'sqlglot.expressions.StrToDate'>, <class 'sqlglot.expressions.CurrentDate'>, <class 'sqlglot.expressions.DateStrToDate'>}, <Type.DATETIME: 'DATETIME'>: {<class 'sqlglot.expressions.DatetimeAdd'>, <class 'sqlglot.expressions.DatetimeSub'>, <class 'sqlglot.expressions.CurrentDatetime'>}, <Type.DOUBLE: 'DOUBLE'>: {<class 'sqlglot.expressions.Log2'>, <class 'sqlglot.expressions.Avg'>, <class 'sqlglot.expressions.Log'>, <class 'sqlglot.expressions.Pow'>, <class 'sqlglot.expressions.Round'>, <class 'sqlglot.expressions.StddevSamp'>, <class 'sqlglot.expressions.Ln'>, <class 'sqlglot.expressions.StddevPop'>, <class 'sqlglot.expressions.Stddev'>, <class 'sqlglot.expressions.Sqrt'>, <class 'sqlglot.expressions.ApproxQuantile'>, <class 'sqlglot.expressions.SafeDivide'>, <class 'sqlglot.expressions.VariancePop'>, <class 'sqlglot.expressions.Quantile'>, <class 'sqlglot.expressions.Variance'>, <class 'sqlglot.expressions.Log10'>, <class 'sqlglot.expressions.Exp'>, <class 'sqlglot.expressions.Div'>}, <Type.INT: 'INT'>: {<class 'sqlglot.expressions.Floor'>, <class 'sqlglot.expressions.Levenshtein'>, <class 'sqlglot.expressions.Extract'>, <class 'sqlglot.expressions.DatetimeDiff'>, <class 'sqlglot.expressions.TimeDiff'>, <class 'sqlglot.expressions.TsOrDiToDi'>, <class 'sqlglot.expressions.DateToDi'>, <class 'sqlglot.expressions.Ceil'>, <class 'sqlglot.expressions.StrPosition'>, <class 'sqlglot.expressions.TimestampDiff'>, <class 'sqlglot.expressions.DateDiff'>}, <Type.JSON: 'JSON'>: {<class 'sqlglot.expressions.ParseJSON'>}, <Type.TIMESTAMP: 'TIMESTAMP'>: {<class 'sqlglot.expressions.TimestampSub'>, <class 'sqlglot.expressions.TimestampAdd'>, <class 'sqlglot.expressions.UnixToTime'>, <class 'sqlglot.expressions.Timestamp'>, <class 'sqlglot.expressions.CurrentTimestamp'>, <class 'sqlglot.expressions.TimeStrToTime'>, <class 'sqlglot.expressions.TimeSub'>, <class 'sqlglot.expressions.CurrentTime'>, <class 'sqlglot.expressions.StrToTime'>, <class 'sqlglot.expressions.TimeAdd'>}, <Type.TINYINT: 'TINYINT'>: {<class 'sqlglot.expressions.Month'>, <class 'sqlglot.expressions.Week'>, <class 'sqlglot.expressions.Day'>, <class 'sqlglot.expressions.Year'>}, <Type.VARCHAR: 'VARCHAR'>: {<class 'sqlglot.expressions.TimeToStr'>, <class 'sqlglot.expressions.UnixToStr'>, <class 'sqlglot.expressions.ConcatWs'>, <class 'sqlglot.expressions.UnixToTimeStr'>, <class 'sqlglot.expressions.Substring'>, <class 'sqlglot.expressions.Trim'>, <class 'sqlglot.expressions.Concat'>, <class 'sqlglot.expressions.ArrayConcat'>, <class 'sqlglot.expressions.GroupConcat'>, <class 'sqlglot.expressions.Initcap'>, <class 'sqlglot.expressions.Lower'>, <class 'sqlglot.expressions.TsOrDsToDateStr'>, <class 'sqlglot.expressions.DateToDateStr'>, <class 'sqlglot.expressions.TimeToTimeStr'>, <class 'sqlglot.expressions.Upper'>}}
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.<dictcomp>.<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.<dictcomp>.<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.Length'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ApproxDistinct'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Count'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ArraySize'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.In'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Between'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Boolean'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDsToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DiToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateFromParts'>: <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.CurrentDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateStrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeAdd'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeSub'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentDatetime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log2'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Avg'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Round'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StddevSamp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Ln'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StddevPop'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Stddev'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Sqrt'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ApproxQuantile'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.SafeDivide'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.VariancePop'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Quantile'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Variance'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log10'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Exp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Floor'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Levenshtein'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Extract'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeDiff'>: <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.DateToDi'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Ceil'>: <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.DateDiff'>: <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.TimestampAdd'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToTime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Timestamp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentTimestamp'>: <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.CurrentTime'>: <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.Month'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Week'>: <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.TimeToStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ConcatWs'>: <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.Concat'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ArrayConcat'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.GroupConcat'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Initcap'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Lower'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDsToDateStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateToDateStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeToTimeStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Upper'>: <function _annotate_with_type_lambda.<locals>.<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.Sum'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.TryCast'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.VarMap'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.Struct'>: <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.VARCHAR: 'VARCHAR'>, <Type.NCHAR: 'NCHAR'>, <Type.NVARCHAR: 'NVARCHAR'>}, <Type.DOUBLE: 'DOUBLE'>: set(), <Type.FLOAT: 'FLOAT'>: {<Type.DOUBLE: 'DOUBLE'>}, <Type.DECIMAL: 'DECIMAL'>: {<Type.DOUBLE: 'DOUBLE'>, <Type.FLOAT: 'FLOAT'>}, <Type.BIGINT: 'BIGINT'>: {<Type.DOUBLE: 'DOUBLE'>, <Type.FLOAT: 'FLOAT'>, <Type.DECIMAL: 'DECIMAL'>}, <Type.INT: 'INT'>: {<Type.DOUBLE: 'DOUBLE'>, <Type.FLOAT: 'FLOAT'>, <Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>}, <Type.SMALLINT: 'SMALLINT'>: {<Type.DOUBLE: 'DOUBLE'>, <Type.INT: 'INT'>, <Type.FLOAT: 'FLOAT'>, <Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>}, <Type.TINYINT: 'TINYINT'>: {<Type.DOUBLE: 'DOUBLE'>, <Type.INT: 'INT'>, <Type.FLOAT: 'FLOAT'>, <Type.SMALLINT: 'SMALLINT'>, <Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>}, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>: set(), <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>: {<Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>}, <Type.TIMESTAMP: 'TIMESTAMP'>: {<Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>, <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>}, <Type.DATETIME: 'DATETIME'>: {<Type.TIMESTAMP: 'TIMESTAMP'>, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>, <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>}, <Type.DATE: 'DATE'>: {<Type.TIMESTAMP: 'TIMESTAMP'>, <Type.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>, <Type.TIMESTAMPTZ: 'TIMESTAMPTZ'>, <Type.DATETIME: 'DATETIME'>}}
BINARY_COERCIONS: Dict[Tuple[sqlglot.expressions.DataType.Type, sqlglot.expressions.DataType.Type], Callable[[sqlglot.expressions.Expression, sqlglot.expressions.Expression], sqlglot.expressions.DataType.Type]] = {(<Type.TEXT: 'TEXT'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <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.INTERVAL: 'INTERVAL'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT: 'INT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.FLOAT: 'FLOAT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.TINYINT: 'TINYINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT: 'INT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.FLOAT: 'FLOAT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.SMALLINT: 'SMALLINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <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

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