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