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

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

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