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