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._typing import E
      8from sqlglot.helper import (
      9    ensure_list,
     10    is_date_unit,
     11    is_iso_date,
     12    is_iso_datetime,
     13    seq_get,
     14    subclasses,
     15)
     16from sqlglot.optimizer.scope import Scope, traverse_scope
     17from sqlglot.schema import Schema, ensure_schema
     18
     19if t.TYPE_CHECKING:
     20    B = t.TypeVar("B", bound=exp.Binary)
     21
     22    BinaryCoercionFunc = t.Callable[[exp.Expression, exp.Expression], exp.DataType.Type]
     23    BinaryCoercions = t.Dict[
     24        t.Tuple[exp.DataType.Type, exp.DataType.Type],
     25        BinaryCoercionFunc,
     26    ]
     27
     28
     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)
     59
     60
     61def _annotate_with_type_lambda(data_type: exp.DataType.Type) -> t.Callable[[TypeAnnotator, E], E]:
     62    return lambda self, e: self._annotate_with_type(e, data_type)
     63
     64
     65def _coerce_date_literal(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
     66    date_text = l.name
     67    is_iso_date_ = is_iso_date(date_text)
     68
     69    if is_iso_date_ and is_date_unit(unit):
     70        return exp.DataType.Type.DATE
     71
     72    # An ISO date is also an ISO datetime, but not vice versa
     73    if is_iso_date_ or is_iso_datetime(date_text):
     74        return exp.DataType.Type.DATETIME
     75
     76    return exp.DataType.Type.UNKNOWN
     77
     78
     79def _coerce_date(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
     80    if not is_date_unit(unit):
     81        return exp.DataType.Type.DATETIME
     82    return l.type.this if l.type else exp.DataType.Type.UNKNOWN
     83
     84
     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
     91
     92
     93def swap_all(coercions: BinaryCoercions) -> BinaryCoercions:
     94    return {**coercions, **{(b, a): swap_args(func) for (a, b), func in coercions.items()}}
     95
     96
     97class _TypeAnnotator(type):
     98    def __new__(cls, clsname, bases, attrs):
     99        klass = super().__new__(cls, clsname, bases, attrs)
    100
    101        # Highest-to-lowest type precedence, as specified in Spark's docs (ANSI):
    102        # https://spark.apache.org/docs/3.2.0/sql-ref-ansi-compliance.html
    103        text_precedence = (
    104            exp.DataType.Type.TEXT,
    105            exp.DataType.Type.NVARCHAR,
    106            exp.DataType.Type.VARCHAR,
    107            exp.DataType.Type.NCHAR,
    108            exp.DataType.Type.CHAR,
    109        )
    110        numeric_precedence = (
    111            exp.DataType.Type.DOUBLE,
    112            exp.DataType.Type.FLOAT,
    113            exp.DataType.Type.DECIMAL,
    114            exp.DataType.Type.BIGINT,
    115            exp.DataType.Type.INT,
    116            exp.DataType.Type.SMALLINT,
    117            exp.DataType.Type.TINYINT,
    118        )
    119        timelike_precedence = (
    120            exp.DataType.Type.TIMESTAMPLTZ,
    121            exp.DataType.Type.TIMESTAMPTZ,
    122            exp.DataType.Type.TIMESTAMP,
    123            exp.DataType.Type.DATETIME,
    124            exp.DataType.Type.DATE,
    125        )
    126
    127        for type_precedence in (text_precedence, numeric_precedence, timelike_precedence):
    128            coerces_to = set()
    129            for data_type in type_precedence:
    130                klass.COERCES_TO[data_type] = coerces_to.copy()
    131                coerces_to |= {data_type}
    132
    133        return klass
    134
    135
    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.TIMESTAMP: {
    199            exp.CurrentTime,
    200            exp.CurrentTimestamp,
    201            exp.StrToTime,
    202            exp.TimeAdd,
    203            exp.TimeStrToTime,
    204            exp.TimeSub,
    205            exp.Timestamp,
    206            exp.TimestampAdd,
    207            exp.TimestampSub,
    208            exp.UnixToTime,
    209        },
    210        exp.DataType.Type.TINYINT: {
    211            exp.Day,
    212            exp.Month,
    213            exp.Week,
    214            exp.Year,
    215        },
    216        exp.DataType.Type.VARCHAR: {
    217            exp.ArrayConcat,
    218            exp.Concat,
    219            exp.ConcatWs,
    220            exp.DateToDateStr,
    221            exp.GroupConcat,
    222            exp.Initcap,
    223            exp.Lower,
    224            exp.Substring,
    225            exp.TimeToStr,
    226            exp.TimeToTimeStr,
    227            exp.Trim,
    228            exp.TsOrDsToDateStr,
    229            exp.UnixToStr,
    230            exp.UnixToTimeStr,
    231            exp.Upper,
    232        },
    233    }
    234
    235    ANNOTATORS: t.Dict = {
    236        **{
    237            expr_type: lambda self, e: self._annotate_unary(e)
    238            for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
    239        },
    240        **{
    241            expr_type: lambda self, e: self._annotate_binary(e)
    242            for expr_type in subclasses(exp.__name__, exp.Binary)
    243        },
    244        **{
    245            expr_type: _annotate_with_type_lambda(data_type)
    246            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
    247            for expr_type in expressions
    248        },
    249        exp.Abs: lambda self, e: self._annotate_by_args(e, "this"),
    250        exp.Anonymous: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    251        exp.Array: lambda self, e: self._annotate_by_args(e, "expressions", array=True),
    252        exp.ArrayAgg: lambda self, e: self._annotate_by_args(e, "this", array=True),
    253        exp.ArrayConcat: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    254        exp.Bracket: lambda self, e: self._annotate_bracket(e),
    255        exp.Cast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    256        exp.Case: lambda self, e: self._annotate_by_args(e, "default", "ifs"),
    257        exp.Coalesce: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    258        exp.DataType: lambda self, e: self._annotate_with_type(e, e.copy()),
    259        exp.DateAdd: lambda self, e: self._annotate_timeunit(e),
    260        exp.DateSub: lambda self, e: self._annotate_timeunit(e),
    261        exp.DateTrunc: lambda self, e: self._annotate_timeunit(e),
    262        exp.Distinct: lambda self, e: self._annotate_by_args(e, "expressions"),
    263        exp.Div: lambda self, e: self._annotate_div(e),
    264        exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
    265        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
    266        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
    267        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
    268        exp.Literal: lambda self, e: self._annotate_literal(e),
    269        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    270        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    271        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    272        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
    273        exp.Nullif: lambda self, e: self._annotate_by_args(e, "this", "expression"),
    274        exp.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    275        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
    276        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    277        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    278    }
    279
    280    NESTED_TYPES = {
    281        exp.DataType.Type.ARRAY,
    282    }
    283
    284    # Specifies what types a given type can be coerced into (autofilled)
    285    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
    286
    287    # Coercion functions for binary operations.
    288    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
    289    BINARY_COERCIONS: BinaryCoercions = {
    290        **swap_all(
    291            {
    292                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
    293                    l, r.args.get("unit")
    294                )
    295                for t in exp.DataType.TEXT_TYPES
    296            }
    297        ),
    298        **swap_all(
    299            {
    300                # text + numeric will yield the numeric type to match most dialects' semantics
    301                (text, numeric): lambda l, r: t.cast(
    302                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
    303                )
    304                for text in exp.DataType.TEXT_TYPES
    305                for numeric in exp.DataType.NUMERIC_TYPES
    306            }
    307        ),
    308        **swap_all(
    309            {
    310                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
    311                    l, r.args.get("unit")
    312                ),
    313            }
    314        ),
    315    }
    316
    317    def __init__(
    318        self,
    319        schema: Schema,
    320        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    321        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    322        binary_coercions: t.Optional[BinaryCoercions] = None,
    323    ) -> None:
    324        self.schema = schema
    325        self.annotators = annotators or self.ANNOTATORS
    326        self.coerces_to = coerces_to or self.COERCES_TO
    327        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    328
    329        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    330        self._visited: t.Set[int] = set()
    331
    332    def _set_type(
    333        self, expression: exp.Expression, target_type: exp.DataType | exp.DataType.Type
    334    ) -> None:
    335        expression.type = target_type  # type: ignore
    336        self._visited.add(id(expression))
    337
    338    def annotate(self, expression: E) -> E:
    339        for scope in traverse_scope(expression):
    340            selects = {}
    341            for name, source in scope.sources.items():
    342                if not isinstance(source, Scope):
    343                    continue
    344                if isinstance(source.expression, exp.UDTF):
    345                    values = []
    346
    347                    if isinstance(source.expression, exp.Lateral):
    348                        if isinstance(source.expression.this, exp.Explode):
    349                            values = [source.expression.this.this]
    350                    else:
    351                        values = source.expression.expressions[0].expressions
    352
    353                    if not values:
    354                        continue
    355
    356                    selects[name] = {
    357                        alias: column
    358                        for alias, column in zip(
    359                            source.expression.alias_column_names,
    360                            values,
    361                        )
    362                    }
    363                else:
    364                    selects[name] = {
    365                        select.alias_or_name: select for select in source.expression.selects
    366                    }
    367
    368            # First annotate the current scope's column references
    369            for col in scope.columns:
    370                if not col.table:
    371                    continue
    372
    373                source = scope.sources.get(col.table)
    374                if isinstance(source, exp.Table):
    375                    self._set_type(col, self.schema.get_column_type(source, col))
    376                elif source and col.table in selects and col.name in selects[col.table]:
    377                    self._set_type(col, selects[col.table][col.name].type)
    378
    379            # Then (possibly) annotate the remaining expressions in the scope
    380            self._maybe_annotate(scope.expression)
    381
    382        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
    383
    384    def _maybe_annotate(self, expression: E) -> E:
    385        if id(expression) in self._visited:
    386            return expression  # We've already inferred the expression's type
    387
    388        annotator = self.annotators.get(expression.__class__)
    389
    390        return (
    391            annotator(self, expression)
    392            if annotator
    393            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
    394        )
    395
    396    def _annotate_args(self, expression: E) -> E:
    397        for _, value in expression.iter_expressions():
    398            self._maybe_annotate(value)
    399
    400        return expression
    401
    402    def _maybe_coerce(
    403        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
    404    ) -> exp.DataType | exp.DataType.Type:
    405        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
    406        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
    407
    408        # We propagate the NULL / UNKNOWN types upwards if found
    409        if exp.DataType.Type.NULL in (type1_value, type2_value):
    410            return exp.DataType.Type.NULL
    411        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
    412            return exp.DataType.Type.UNKNOWN
    413
    414        if type1_value in self.NESTED_TYPES:
    415            return type1
    416        if type2_value in self.NESTED_TYPES:
    417            return type2
    418
    419        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
    420
    421    # Note: the following "no_type_check" decorators were added because mypy was yelling due
    422    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
    423    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
    424
    425    @t.no_type_check
    426    def _annotate_binary(self, expression: B) -> B:
    427        self._annotate_args(expression)
    428
    429        left, right = expression.left, expression.right
    430        left_type, right_type = left.type.this, right.type.this
    431
    432        if isinstance(expression, exp.Connector):
    433            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
    434                self._set_type(expression, exp.DataType.Type.NULL)
    435            elif exp.DataType.Type.NULL in (left_type, right_type):
    436                self._set_type(
    437                    expression,
    438                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
    439                )
    440            else:
    441                self._set_type(expression, exp.DataType.Type.BOOLEAN)
    442        elif isinstance(expression, exp.Predicate):
    443            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    444        elif (left_type, right_type) in self.binary_coercions:
    445            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
    446        else:
    447            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    448
    449        return expression
    450
    451    @t.no_type_check
    452    def _annotate_unary(self, expression: E) -> E:
    453        self._annotate_args(expression)
    454
    455        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
    456            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    457        else:
    458            self._set_type(expression, expression.this.type)
    459
    460        return expression
    461
    462    @t.no_type_check
    463    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
    464        if expression.is_string:
    465            self._set_type(expression, exp.DataType.Type.VARCHAR)
    466        elif expression.is_int:
    467            self._set_type(expression, exp.DataType.Type.INT)
    468        else:
    469            self._set_type(expression, exp.DataType.Type.DOUBLE)
    470
    471        return expression
    472
    473    @t.no_type_check
    474    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
    475        self._set_type(expression, target_type)
    476        return self._annotate_args(expression)
    477
    478    @t.no_type_check
    479    def _annotate_by_args(
    480        self, expression: E, *args: str, promote: bool = False, array: bool = False
    481    ) -> E:
    482        self._annotate_args(expression)
    483
    484        expressions: t.List[exp.Expression] = []
    485        for arg in args:
    486            arg_expr = expression.args.get(arg)
    487            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
    488
    489        last_datatype = None
    490        for expr in expressions:
    491            last_datatype = self._maybe_coerce(last_datatype or expr.type, expr.type)
    492
    493        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
    494
    495        if promote:
    496            if expression.type.this in exp.DataType.INTEGER_TYPES:
    497                self._set_type(expression, exp.DataType.Type.BIGINT)
    498            elif expression.type.this in exp.DataType.FLOAT_TYPES:
    499                self._set_type(expression, exp.DataType.Type.DOUBLE)
    500
    501        if array:
    502            self._set_type(
    503                expression,
    504                exp.DataType(
    505                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
    506                ),
    507            )
    508
    509        return expression
    510
    511    def _annotate_timeunit(
    512        self, expression: exp.TimeUnit | exp.DateTrunc
    513    ) -> exp.TimeUnit | exp.DateTrunc:
    514        self._annotate_args(expression)
    515
    516        if expression.this.type.this in exp.DataType.TEXT_TYPES:
    517            datatype = _coerce_date_literal(expression.this, expression.unit)
    518        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
    519            datatype = _coerce_date(expression.this, expression.unit)
    520        else:
    521            datatype = exp.DataType.Type.UNKNOWN
    522
    523        self._set_type(expression, datatype)
    524        return expression
    525
    526    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
    527        self._annotate_args(expression)
    528
    529        bracket_arg = expression.expressions[0]
    530        this = expression.this
    531
    532        if isinstance(bracket_arg, exp.Slice):
    533            self._set_type(expression, this.type)
    534        elif this.type.is_type(exp.DataType.Type.ARRAY):
    535            contained_type = seq_get(this.type.expressions, 0) or exp.DataType.Type.UNKNOWN
    536            self._set_type(expression, contained_type)
    537        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
    538            index = this.keys.index(bracket_arg)
    539            value = seq_get(this.values, index)
    540            value_type = value.type if value else exp.DataType.Type.UNKNOWN
    541            self._set_type(expression, value_type or exp.DataType.Type.UNKNOWN)
    542        else:
    543            self._set_type(expression, exp.DataType.Type.UNKNOWN)
    544
    545        return expression
    546
    547    def _annotate_div(self, expression: exp.Div) -> exp.Div:
    548        self._annotate_args(expression)
    549
    550        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
    551
    552        if (
    553            expression.args.get("typed")
    554            and left_type in exp.DataType.INTEGER_TYPES
    555            and right_type in exp.DataType.INTEGER_TYPES
    556        ):
    557            self._set_type(expression, exp.DataType.Type.BIGINT)
    558        else:
    559            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    560
    561        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

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

    86def swap_args(func: BinaryCoercionFunc) -> BinaryCoercionFunc:
    87    @functools.wraps(func)
    88    def _swapped(l: exp.Expression, r: exp.Expression) -> exp.DataType.Type:
    89        return func(r, l)
    90
    91    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

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

class TypeAnnotator: View Source

    137class TypeAnnotator(metaclass=_TypeAnnotator):
    138    TYPE_TO_EXPRESSIONS: t.Dict[exp.DataType.Type, t.Set[t.Type[exp.Expression]]] = {
    139        exp.DataType.Type.BIGINT: {
    140            exp.ApproxDistinct,
    141            exp.ArraySize,
    142            exp.Count,
    143            exp.Length,
    144        },
    145        exp.DataType.Type.BOOLEAN: {
    146            exp.Between,
    147            exp.Boolean,
    148            exp.In,
    149            exp.RegexpLike,
    150        },
    151        exp.DataType.Type.DATE: {
    152            exp.CurrentDate,
    153            exp.Date,
    154            exp.DateFromParts,
    155            exp.DateStrToDate,
    156            exp.DiToDate,
    157            exp.StrToDate,
    158            exp.TimeStrToDate,
    159            exp.TsOrDsToDate,
    160        },
    161        exp.DataType.Type.DATETIME: {
    162            exp.CurrentDatetime,
    163            exp.DatetimeAdd,
    164            exp.DatetimeSub,
    165        },
    166        exp.DataType.Type.DOUBLE: {
    167            exp.ApproxQuantile,
    168            exp.Avg,
    169            exp.Div,
    170            exp.Exp,
    171            exp.Ln,
    172            exp.Log,
    173            exp.Log2,
    174            exp.Log10,
    175            exp.Pow,
    176            exp.Quantile,
    177            exp.Round,
    178            exp.SafeDivide,
    179            exp.Sqrt,
    180            exp.Stddev,
    181            exp.StddevPop,
    182            exp.StddevSamp,
    183            exp.Variance,
    184            exp.VariancePop,
    185        },
    186        exp.DataType.Type.INT: {
    187            exp.Ceil,
    188            exp.DatetimeDiff,
    189            exp.DateDiff,
    190            exp.Extract,
    191            exp.TimestampDiff,
    192            exp.TimeDiff,
    193            exp.DateToDi,
    194            exp.Floor,
    195            exp.Levenshtein,
    196            exp.StrPosition,
    197            exp.TsOrDiToDi,
    198        },
    199        exp.DataType.Type.TIMESTAMP: {
    200            exp.CurrentTime,
    201            exp.CurrentTimestamp,
    202            exp.StrToTime,
    203            exp.TimeAdd,
    204            exp.TimeStrToTime,
    205            exp.TimeSub,
    206            exp.Timestamp,
    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.Filter: lambda self, e: self._annotate_by_args(e, "this"),
    266        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
    267        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
    268        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
    269        exp.Literal: lambda self, e: self._annotate_literal(e),
    270        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    271        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    272        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
    273        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
    274        exp.Nullif: lambda self, e: self._annotate_by_args(e, "this", "expression"),
    275        exp.Slice: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
    276        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
    277        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
    278        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
    279    }
    280
    281    NESTED_TYPES = {
    282        exp.DataType.Type.ARRAY,
    283    }
    284
    285    # Specifies what types a given type can be coerced into (autofilled)
    286    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
    287
    288    # Coercion functions for binary operations.
    289    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
    290    BINARY_COERCIONS: BinaryCoercions = {
    291        **swap_all(
    292            {
    293                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
    294                    l, r.args.get("unit")
    295                )
    296                for t in exp.DataType.TEXT_TYPES
    297            }
    298        ),
    299        **swap_all(
    300            {
    301                # text + numeric will yield the numeric type to match most dialects' semantics
    302                (text, numeric): lambda l, r: t.cast(
    303                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
    304                )
    305                for text in exp.DataType.TEXT_TYPES
    306                for numeric in exp.DataType.NUMERIC_TYPES
    307            }
    308        ),
    309        **swap_all(
    310            {
    311                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
    312                    l, r.args.get("unit")
    313                ),
    314            }
    315        ),
    316    }
    317
    318    def __init__(
    319        self,
    320        schema: Schema,
    321        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    322        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    323        binary_coercions: t.Optional[BinaryCoercions] = None,
    324    ) -> None:
    325        self.schema = schema
    326        self.annotators = annotators or self.ANNOTATORS
    327        self.coerces_to = coerces_to or self.COERCES_TO
    328        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    329
    330        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    331        self._visited: t.Set[int] = set()
    332
    333    def _set_type(
    334        self, expression: exp.Expression, target_type: exp.DataType | exp.DataType.Type
    335    ) -> None:
    336        expression.type = target_type  # type: ignore
    337        self._visited.add(id(expression))
    338
    339    def annotate(self, expression: E) -> E:
    340        for scope in traverse_scope(expression):
    341            selects = {}
    342            for name, source in scope.sources.items():
    343                if not isinstance(source, Scope):
    344                    continue
    345                if isinstance(source.expression, exp.UDTF):
    346                    values = []
    347
    348                    if isinstance(source.expression, exp.Lateral):
    349                        if isinstance(source.expression.this, exp.Explode):
    350                            values = [source.expression.this.this]
    351                    else:
    352                        values = source.expression.expressions[0].expressions
    353
    354                    if not values:
    355                        continue
    356
    357                    selects[name] = {
    358                        alias: column
    359                        for alias, column in zip(
    360                            source.expression.alias_column_names,
    361                            values,
    362                        )
    363                    }
    364                else:
    365                    selects[name] = {
    366                        select.alias_or_name: select for select in source.expression.selects
    367                    }
    368
    369            # First annotate the current scope's column references
    370            for col in scope.columns:
    371                if not col.table:
    372                    continue
    373
    374                source = scope.sources.get(col.table)
    375                if isinstance(source, exp.Table):
    376                    self._set_type(col, self.schema.get_column_type(source, col))
    377                elif source and col.table in selects and col.name in selects[col.table]:
    378                    self._set_type(col, selects[col.table][col.name].type)
    379
    380            # Then (possibly) annotate the remaining expressions in the scope
    381            self._maybe_annotate(scope.expression)
    382
    383        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
    384
    385    def _maybe_annotate(self, expression: E) -> E:
    386        if id(expression) in self._visited:
    387            return expression  # We've already inferred the expression's type
    388
    389        annotator = self.annotators.get(expression.__class__)
    390
    391        return (
    392            annotator(self, expression)
    393            if annotator
    394            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
    395        )
    396
    397    def _annotate_args(self, expression: E) -> E:
    398        for _, value in expression.iter_expressions():
    399            self._maybe_annotate(value)
    400
    401        return expression
    402
    403    def _maybe_coerce(
    404        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
    405    ) -> exp.DataType | exp.DataType.Type:
    406        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
    407        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
    408
    409        # We propagate the NULL / UNKNOWN types upwards if found
    410        if exp.DataType.Type.NULL in (type1_value, type2_value):
    411            return exp.DataType.Type.NULL
    412        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
    413            return exp.DataType.Type.UNKNOWN
    414
    415        if type1_value in self.NESTED_TYPES:
    416            return type1
    417        if type2_value in self.NESTED_TYPES:
    418            return type2
    419
    420        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
    421
    422    # Note: the following "no_type_check" decorators were added because mypy was yelling due
    423    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
    424    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
    425
    426    @t.no_type_check
    427    def _annotate_binary(self, expression: B) -> B:
    428        self._annotate_args(expression)
    429
    430        left, right = expression.left, expression.right
    431        left_type, right_type = left.type.this, right.type.this
    432
    433        if isinstance(expression, exp.Connector):
    434            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
    435                self._set_type(expression, exp.DataType.Type.NULL)
    436            elif exp.DataType.Type.NULL in (left_type, right_type):
    437                self._set_type(
    438                    expression,
    439                    exp.DataType.build("NULLABLE", expressions=exp.DataType.build("BOOLEAN")),
    440                )
    441            else:
    442                self._set_type(expression, exp.DataType.Type.BOOLEAN)
    443        elif isinstance(expression, exp.Predicate):
    444            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    445        elif (left_type, right_type) in self.binary_coercions:
    446            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
    447        else:
    448            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    449
    450        return expression
    451
    452    @t.no_type_check
    453    def _annotate_unary(self, expression: E) -> E:
    454        self._annotate_args(expression)
    455
    456        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
    457            self._set_type(expression, exp.DataType.Type.BOOLEAN)
    458        else:
    459            self._set_type(expression, expression.this.type)
    460
    461        return expression
    462
    463    @t.no_type_check
    464    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
    465        if expression.is_string:
    466            self._set_type(expression, exp.DataType.Type.VARCHAR)
    467        elif expression.is_int:
    468            self._set_type(expression, exp.DataType.Type.INT)
    469        else:
    470            self._set_type(expression, exp.DataType.Type.DOUBLE)
    471
    472        return expression
    473
    474    @t.no_type_check
    475    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
    476        self._set_type(expression, target_type)
    477        return self._annotate_args(expression)
    478
    479    @t.no_type_check
    480    def _annotate_by_args(
    481        self, expression: E, *args: str, promote: bool = False, array: bool = False
    482    ) -> E:
    483        self._annotate_args(expression)
    484
    485        expressions: t.List[exp.Expression] = []
    486        for arg in args:
    487            arg_expr = expression.args.get(arg)
    488            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
    489
    490        last_datatype = None
    491        for expr in expressions:
    492            last_datatype = self._maybe_coerce(last_datatype or expr.type, expr.type)
    493
    494        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
    495
    496        if promote:
    497            if expression.type.this in exp.DataType.INTEGER_TYPES:
    498                self._set_type(expression, exp.DataType.Type.BIGINT)
    499            elif expression.type.this in exp.DataType.FLOAT_TYPES:
    500                self._set_type(expression, exp.DataType.Type.DOUBLE)
    501
    502        if array:
    503            self._set_type(
    504                expression,
    505                exp.DataType(
    506                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
    507                ),
    508            )
    509
    510        return expression
    511
    512    def _annotate_timeunit(
    513        self, expression: exp.TimeUnit | exp.DateTrunc
    514    ) -> exp.TimeUnit | exp.DateTrunc:
    515        self._annotate_args(expression)
    516
    517        if expression.this.type.this in exp.DataType.TEXT_TYPES:
    518            datatype = _coerce_date_literal(expression.this, expression.unit)
    519        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
    520            datatype = _coerce_date(expression.this, expression.unit)
    521        else:
    522            datatype = exp.DataType.Type.UNKNOWN
    523
    524        self._set_type(expression, datatype)
    525        return expression
    526
    527    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
    528        self._annotate_args(expression)
    529
    530        bracket_arg = expression.expressions[0]
    531        this = expression.this
    532
    533        if isinstance(bracket_arg, exp.Slice):
    534            self._set_type(expression, this.type)
    535        elif this.type.is_type(exp.DataType.Type.ARRAY):
    536            contained_type = seq_get(this.type.expressions, 0) or exp.DataType.Type.UNKNOWN
    537            self._set_type(expression, contained_type)
    538        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
    539            index = this.keys.index(bracket_arg)
    540            value = seq_get(this.values, index)
    541            value_type = value.type if value else exp.DataType.Type.UNKNOWN
    542            self._set_type(expression, value_type or exp.DataType.Type.UNKNOWN)
    543        else:
    544            self._set_type(expression, exp.DataType.Type.UNKNOWN)
    545
    546        return expression
    547
    548    def _annotate_div(self, expression: exp.Div) -> exp.Div:
    549        self._annotate_args(expression)
    550
    551        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
    552
    553        if (
    554            expression.args.get("typed")
    555            and left_type in exp.DataType.INTEGER_TYPES
    556            and right_type in exp.DataType.INTEGER_TYPES
    557        ):
    558            self._set_type(expression, exp.DataType.Type.BIGINT)
    559        else:
    560            self._set_type(expression, self._maybe_coerce(left_type, right_type))
    561
    562        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

    318    def __init__(
    319        self,
    320        schema: Schema,
    321        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
    322        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
    323        binary_coercions: t.Optional[BinaryCoercions] = None,
    324    ) -> None:
    325        self.schema = schema
    326        self.annotators = annotators or self.ANNOTATORS
    327        self.coerces_to = coerces_to or self.COERCES_TO
    328        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
    329
    330        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
    331        self._visited: t.Set[int] = set()

TYPE_TO_EXPRESSIONS: Dict[sqlglot.expressions.DataType.Type, Set[Type[sqlglot.expressions.Expression]]] = {<Type.BIGINT: 'BIGINT'>: {<class 'sqlglot.expressions.ArraySize'>, <class 'sqlglot.expressions.Length'>, <class 'sqlglot.expressions.Count'>, <class 'sqlglot.expressions.ApproxDistinct'>}, <Type.BOOLEAN: 'BOOLEAN'>: {<class 'sqlglot.expressions.RegexpLike'>, <class 'sqlglot.expressions.Between'>, <class 'sqlglot.expressions.In'>, <class 'sqlglot.expressions.Boolean'>}, <Type.DATE: 'DATE'>: {<class 'sqlglot.expressions.DateStrToDate'>, <class 'sqlglot.expressions.DiToDate'>, <class 'sqlglot.expressions.StrToDate'>, <class 'sqlglot.expressions.DateFromParts'>, <class 'sqlglot.expressions.TimeStrToDate'>, <class 'sqlglot.expressions.TsOrDsToDate'>, <class 'sqlglot.expressions.CurrentDate'>, <class 'sqlglot.expressions.Date'>}, <Type.DATETIME: 'DATETIME'>: {<class 'sqlglot.expressions.CurrentDatetime'>, <class 'sqlglot.expressions.DatetimeAdd'>, <class 'sqlglot.expressions.DatetimeSub'>}, <Type.DOUBLE: 'DOUBLE'>: {<class 'sqlglot.expressions.Log'>, <class 'sqlglot.expressions.Round'>, <class 'sqlglot.expressions.Ln'>, <class 'sqlglot.expressions.StddevSamp'>, <class 'sqlglot.expressions.ApproxQuantile'>, <class 'sqlglot.expressions.StddevPop'>, <class 'sqlglot.expressions.Exp'>, <class 'sqlglot.expressions.Quantile'>, <class 'sqlglot.expressions.Stddev'>, <class 'sqlglot.expressions.Sqrt'>, <class 'sqlglot.expressions.Avg'>, <class 'sqlglot.expressions.SafeDivide'>, <class 'sqlglot.expressions.VariancePop'>, <class 'sqlglot.expressions.Pow'>, <class 'sqlglot.expressions.Log10'>, <class 'sqlglot.expressions.Variance'>, <class 'sqlglot.expressions.Div'>, <class 'sqlglot.expressions.Log2'>}, <Type.INT: 'INT'>: {<class 'sqlglot.expressions.DateDiff'>, <class 'sqlglot.expressions.DateToDi'>, <class 'sqlglot.expressions.Floor'>, <class 'sqlglot.expressions.Levenshtein'>, <class 'sqlglot.expressions.TimestampDiff'>, <class 'sqlglot.expressions.DatetimeDiff'>, <class 'sqlglot.expressions.TsOrDiToDi'>, <class 'sqlglot.expressions.TimeDiff'>, <class 'sqlglot.expressions.StrPosition'>, <class 'sqlglot.expressions.Ceil'>, <class 'sqlglot.expressions.Extract'>}, <Type.TIMESTAMP: 'TIMESTAMP'>: {<class 'sqlglot.expressions.TimeSub'>, <class 'sqlglot.expressions.TimeAdd'>, <class 'sqlglot.expressions.UnixToTime'>, <class 'sqlglot.expressions.StrToTime'>, <class 'sqlglot.expressions.TimestampSub'>, <class 'sqlglot.expressions.CurrentTimestamp'>, <class 'sqlglot.expressions.TimeStrToTime'>, <class 'sqlglot.expressions.TimestampAdd'>, <class 'sqlglot.expressions.CurrentTime'>, <class 'sqlglot.expressions.Timestamp'>}, <Type.TINYINT: 'TINYINT'>: {<class 'sqlglot.expressions.Week'>, <class 'sqlglot.expressions.Day'>, <class 'sqlglot.expressions.Month'>, <class 'sqlglot.expressions.Year'>}, <Type.VARCHAR: 'VARCHAR'>: {<class 'sqlglot.expressions.Upper'>, <class 'sqlglot.expressions.TimeToStr'>, <class 'sqlglot.expressions.Substring'>, <class 'sqlglot.expressions.UnixToTimeStr'>, <class 'sqlglot.expressions.GroupConcat'>, <class 'sqlglot.expressions.DateToDateStr'>, <class 'sqlglot.expressions.ConcatWs'>, <class 'sqlglot.expressions.Trim'>, <class 'sqlglot.expressions.Lower'>, <class 'sqlglot.expressions.Concat'>, <class 'sqlglot.expressions.UnixToStr'>, <class 'sqlglot.expressions.TsOrDsToDateStr'>, <class 'sqlglot.expressions.Initcap'>, <class 'sqlglot.expressions.ArrayConcat'>, <class 'sqlglot.expressions.TimeToTimeStr'>}}
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.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.ArraySize'>: <function _annotate_with_type_lambda.<locals>.<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.ApproxDistinct'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Between'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.In'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Boolean'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateStrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DiToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateFromParts'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeStrToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDsToDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.CurrentDate'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Date'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.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.Log'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Round'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Ln'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StddevSamp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ApproxQuantile'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StddevPop'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Exp'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Quantile'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Stddev'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Sqrt'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Avg'>: <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.Log10'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Variance'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Log2'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateToDi'>: <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.TimestampDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DatetimeDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TsOrDiToDi'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeDiff'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StrPosition'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Ceil'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Extract'>: <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.UnixToTime'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.StrToTime'>: <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.TimeStrToTime'>: <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.Timestamp'>: <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.Month'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Year'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Upper'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.TimeToStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Substring'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToTimeStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.GroupConcat'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.DateToDateStr'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.ConcatWs'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Trim'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Lower'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.Concat'>: <function _annotate_with_type_lambda.<locals>.<lambda>>, <class 'sqlglot.expressions.UnixToStr'>: <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.ArrayConcat'>: <function TypeAnnotator.<lambda>>, <class 'sqlglot.expressions.TimeToTimeStr'>: <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.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>>}
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.NVARCHAR: 'NVARCHAR'>, <Type.VARCHAR: 'VARCHAR'>}, <Type.CHAR: 'CHAR'>: {<Type.TEXT: 'TEXT'>, <Type.NVARCHAR: 'NVARCHAR'>, <Type.NCHAR: 'NCHAR'>, <Type.VARCHAR: 'VARCHAR'>}, <Type.DOUBLE: 'DOUBLE'>: set(), <Type.FLOAT: 'FLOAT'>: {<Type.DOUBLE: 'DOUBLE'>}, <Type.DECIMAL: 'DECIMAL'>: {<Type.FLOAT: 'FLOAT'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.BIGINT: 'BIGINT'>: {<Type.FLOAT: 'FLOAT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.INT: 'INT'>: {<Type.BIGINT: 'BIGINT'>, <Type.FLOAT: 'FLOAT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.SMALLINT: 'SMALLINT'>: {<Type.INT: 'INT'>, <Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.FLOAT: 'FLOAT'>, <Type.DOUBLE: 'DOUBLE'>}, <Type.TINYINT: 'TINYINT'>: {<Type.SMALLINT: 'SMALLINT'>, <Type.INT: 'INT'>, <Type.BIGINT: 'BIGINT'>, <Type.DECIMAL: 'DECIMAL'>, <Type.FLOAT: 'FLOAT'>, <Type.DOUBLE: 'DOUBLE'>}, <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.TIMESTAMPLTZ: 'TIMESTAMPLTZ'>, <Type.TIMESTAMP: 'TIMESTAMP'>, <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.CHAR: 'CHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INTERVAL: 'INTERVAL'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INTERVAL: 'INTERVAL'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.CHAR: 'CHAR'>, <Type.TINYINT: 'TINYINT'>): <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.TEXT: 'TEXT'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TEXT: 'TEXT'>, <Type.TINYINT: 'TINYINT'>): <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.VARCHAR: 'VARCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.VARCHAR: 'VARCHAR'>, <Type.TINYINT: 'TINYINT'>): <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.NCHAR: 'NCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NCHAR: 'NCHAR'>, <Type.TINYINT: 'TINYINT'>): <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.NVARCHAR: 'NVARCHAR'>, <Type.SMALLINT: 'SMALLINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.DOUBLE: 'DOUBLE'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT128: 'INT128'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.INT256: 'INT256'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.BIGINT: 'BIGINT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.BIT: 'BIT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.NVARCHAR: 'NVARCHAR'>, <Type.TINYINT: 'TINYINT'>): <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.SMALLINT: 'SMALLINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.CHAR: 'CHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <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.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.TEXT: 'TEXT'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <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.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.VARCHAR: 'VARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <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.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.NCHAR: 'NCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <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.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.DOUBLE: 'DOUBLE'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT128: 'INT128'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.INT256: 'INT256'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIGINT: 'BIGINT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.BIT: 'BIT'>, <Type.NVARCHAR: 'NVARCHAR'>): <function TypeAnnotator.<dictcomp>.<lambda>>, (<Type.TINYINT: 'TINYINT'>, <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.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

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