Edit on GitHub

sqlglot.optimizer.annotate_types

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

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