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