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)
 14from sqlglot.optimizer.scope import Scope, traverse_scope
 15from sqlglot.schema import MappingSchema, Schema, ensure_schema
 16from sqlglot.dialects.dialect import Dialect
 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    from sqlglot.dialects.dialect import DialectType, AnnotatorsType
 28
 29
 30def annotate_types(
 31    expression: E,
 32    schema: t.Optional[t.Dict | Schema] = None,
 33    annotators: t.Optional[AnnotatorsType] = None,
 34    coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
 35    dialect: DialectType = None,
 36) -> E:
 37    """
 38    Infers the types of an expression, annotating its AST accordingly.
 39
 40    Example:
 41        >>> import sqlglot
 42        >>> schema = {"y": {"cola": "SMALLINT"}}
 43        >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
 44        >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
 45        >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
 46        <Type.DOUBLE: 'DOUBLE'>
 47
 48    Args:
 49        expression: Expression to annotate.
 50        schema: Database schema.
 51        annotators: Maps expression type to corresponding annotation function.
 52        coerces_to: Maps expression type to set of types that it can be coerced into.
 53
 54    Returns:
 55        The expression annotated with types.
 56    """
 57
 58    schema = ensure_schema(schema, dialect=dialect)
 59
 60    return TypeAnnotator(schema, annotators, coerces_to).annotate(expression)
 61
 62
 63def _coerce_date_literal(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
 64    date_text = l.name
 65    is_iso_date_ = is_iso_date(date_text)
 66
 67    if is_iso_date_ and is_date_unit(unit):
 68        return exp.DataType.Type.DATE
 69
 70    # An ISO date is also an ISO datetime, but not vice versa
 71    if is_iso_date_ or is_iso_datetime(date_text):
 72        return exp.DataType.Type.DATETIME
 73
 74    return exp.DataType.Type.UNKNOWN
 75
 76
 77def _coerce_date(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
 78    if not is_date_unit(unit):
 79        return exp.DataType.Type.DATETIME
 80    return l.type.this if l.type else exp.DataType.Type.UNKNOWN
 81
 82
 83def swap_args(func: BinaryCoercionFunc) -> BinaryCoercionFunc:
 84    @functools.wraps(func)
 85    def _swapped(l: exp.Expression, r: exp.Expression) -> exp.DataType.Type:
 86        return func(r, l)
 87
 88    return _swapped
 89
 90
 91def swap_all(coercions: BinaryCoercions) -> BinaryCoercions:
 92    return {**coercions, **{(b, a): swap_args(func) for (a, b), func in coercions.items()}}
 93
 94
 95class _TypeAnnotator(type):
 96    def __new__(cls, clsname, bases, attrs):
 97        klass = super().__new__(cls, clsname, bases, attrs)
 98
 99        # Highest-to-lowest type precedence, as specified in Spark's docs (ANSI):
100        # https://spark.apache.org/docs/3.2.0/sql-ref-ansi-compliance.html
101        text_precedence = (
102            exp.DataType.Type.TEXT,
103            exp.DataType.Type.NVARCHAR,
104            exp.DataType.Type.VARCHAR,
105            exp.DataType.Type.NCHAR,
106            exp.DataType.Type.CHAR,
107        )
108        numeric_precedence = (
109            exp.DataType.Type.DOUBLE,
110            exp.DataType.Type.FLOAT,
111            exp.DataType.Type.DECIMAL,
112            exp.DataType.Type.BIGINT,
113            exp.DataType.Type.INT,
114            exp.DataType.Type.SMALLINT,
115            exp.DataType.Type.TINYINT,
116        )
117        timelike_precedence = (
118            exp.DataType.Type.TIMESTAMPLTZ,
119            exp.DataType.Type.TIMESTAMPTZ,
120            exp.DataType.Type.TIMESTAMP,
121            exp.DataType.Type.DATETIME,
122            exp.DataType.Type.DATE,
123        )
124
125        for type_precedence in (text_precedence, numeric_precedence, timelike_precedence):
126            coerces_to = set()
127            for data_type in type_precedence:
128                klass.COERCES_TO[data_type] = coerces_to.copy()
129                coerces_to |= {data_type}
130
131        # NULL can be coerced to any type, so e.g. NULL + 1 will have type INT
132        klass.COERCES_TO[exp.DataType.Type.NULL] = {
133            *text_precedence,
134            *numeric_precedence,
135            *timelike_precedence,
136        }
137
138        return klass
139
140
141class TypeAnnotator(metaclass=_TypeAnnotator):
142    NESTED_TYPES = {
143        exp.DataType.Type.ARRAY,
144    }
145
146    # Specifies what types a given type can be coerced into (autofilled)
147    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
148
149    # Coercion functions for binary operations.
150    # Map of type pairs to a callable that takes both sides of the binary operation and returns the resulting type.
151    BINARY_COERCIONS: BinaryCoercions = {
152        **swap_all(
153            {
154                (t, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date_literal(
155                    l, r.args.get("unit")
156                )
157                for t in exp.DataType.TEXT_TYPES
158            }
159        ),
160        **swap_all(
161            {
162                # text + numeric will yield the numeric type to match most dialects' semantics
163                (text, numeric): lambda l, r: t.cast(
164                    exp.DataType.Type, l.type if l.type in exp.DataType.NUMERIC_TYPES else r.type
165                )
166                for text in exp.DataType.TEXT_TYPES
167                for numeric in exp.DataType.NUMERIC_TYPES
168            }
169        ),
170        **swap_all(
171            {
172                (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL): lambda l, r: _coerce_date(
173                    l, r.args.get("unit")
174                ),
175            }
176        ),
177    }
178
179    def __init__(
180        self,
181        schema: Schema,
182        annotators: t.Optional[AnnotatorsType] = None,
183        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
184        binary_coercions: t.Optional[BinaryCoercions] = None,
185    ) -> None:
186        self.schema = schema
187        self.annotators = annotators or Dialect.get_or_raise(schema.dialect).ANNOTATORS
188        self.coerces_to = (
189            coerces_to or Dialect.get_or_raise(schema.dialect).COERCES_TO or self.COERCES_TO
190        )
191        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
192
193        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
194        self._visited: t.Set[int] = set()
195
196        # Maps an exp.SetOperation's id (e.g. UNION) to its projection types. This is computed if the
197        # exp.SetOperation is the expression of a scope source, as selecting from it multiple times
198        # would reprocess the entire subtree to coerce the types of its operands' projections
199        self._setop_column_types: t.Dict[int, t.Dict[str, exp.DataType | exp.DataType.Type]] = {}
200
201    def _set_type(
202        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
203    ) -> None:
204        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
205        self._visited.add(id(expression))
206
207    def annotate(self, expression: E) -> E:
208        for scope in traverse_scope(expression):
209            self.annotate_scope(scope)
210        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
211
212    def annotate_scope(self, scope: Scope) -> None:
213        selects = {}
214        for name, source in scope.sources.items():
215            if not isinstance(source, Scope):
216                continue
217
218            expression = source.expression
219            if isinstance(expression, exp.UDTF):
220                values = []
221
222                if isinstance(expression, exp.Lateral):
223                    if isinstance(expression.this, exp.Explode):
224                        values = [expression.this.this]
225                elif isinstance(expression, exp.Unnest):
226                    values = [expression]
227                elif not isinstance(expression, exp.TableFromRows):
228                    values = expression.expressions[0].expressions
229
230                if not values:
231                    continue
232
233                selects[name] = {
234                    alias: column.type
235                    for alias, column in zip(expression.alias_column_names, values)
236                }
237            elif isinstance(expression, exp.SetOperation) and len(expression.left.selects) == len(
238                expression.right.selects
239            ):
240                selects[name] = col_types = self._setop_column_types.setdefault(id(expression), {})
241
242                if not col_types:
243                    # Process a chain / sub-tree of set operations
244                    for set_op in expression.walk(
245                        prune=lambda n: not isinstance(n, (exp.SetOperation, exp.Subquery))
246                    ):
247                        if not isinstance(set_op, exp.SetOperation):
248                            continue
249
250                        if set_op.args.get("by_name"):
251                            r_type_by_select = {
252                                s.alias_or_name: s.type for s in set_op.right.selects
253                            }
254                            setop_cols = {
255                                s.alias_or_name: self._maybe_coerce(
256                                    t.cast(exp.DataType, s.type),
257                                    r_type_by_select.get(s.alias_or_name)
258                                    or exp.DataType.Type.UNKNOWN,
259                                )
260                                for s in set_op.left.selects
261                            }
262                        else:
263                            setop_cols = {
264                                ls.alias_or_name: self._maybe_coerce(
265                                    t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
266                                )
267                                for ls, rs in zip(set_op.left.selects, set_op.right.selects)
268                            }
269
270                        # Coerce intermediate results with the previously registered types, if they exist
271                        for col_name, col_type in setop_cols.items():
272                            col_types[col_name] = self._maybe_coerce(
273                                col_type, col_types.get(col_name, exp.DataType.Type.NULL)
274                            )
275
276            else:
277                selects[name] = {s.alias_or_name: s.type for s in expression.selects}
278
279        # First annotate the current scope's column references
280        for col in scope.columns:
281            if not col.table:
282                continue
283
284            source = scope.sources.get(col.table)
285            if isinstance(source, exp.Table):
286                self._set_type(col, self.schema.get_column_type(source, col))
287            elif source:
288                if col.table in selects and col.name in selects[col.table]:
289                    self._set_type(col, selects[col.table][col.name])
290                elif isinstance(source.expression, exp.Unnest):
291                    self._set_type(col, source.expression.type)
292
293        if isinstance(self.schema, MappingSchema):
294            for table_column in scope.table_columns:
295                source = scope.sources.get(table_column.name)
296
297                if isinstance(source, exp.Table):
298                    schema = self.schema.find(
299                        source, raise_on_missing=False, ensure_data_types=True
300                    )
301                    if not isinstance(schema, dict):
302                        continue
303
304                    struct_type = exp.DataType(
305                        this=exp.DataType.Type.STRUCT,
306                        expressions=[
307                            exp.ColumnDef(this=exp.to_identifier(c), kind=kind)
308                            for c, kind in schema.items()
309                        ],
310                        nested=True,
311                    )
312                    self._set_type(table_column, struct_type)
313                elif (
314                    isinstance(source, Scope)
315                    and isinstance(source.expression, exp.Query)
316                    and (
317                        source.expression.meta.get("query_type") or exp.DataType.build("UNKNOWN")
318                    ).is_type(exp.DataType.Type.STRUCT)
319                ):
320                    self._set_type(table_column, source.expression.meta["query_type"])
321
322        # Then (possibly) annotate the remaining expressions in the scope
323        self._maybe_annotate(scope.expression)
324
325        if self.schema.dialect == "bigquery" and isinstance(scope.expression, exp.Query):
326            struct_type = exp.DataType(
327                this=exp.DataType.Type.STRUCT,
328                expressions=[
329                    exp.ColumnDef(this=exp.to_identifier(select.output_name), kind=select.type)
330                    for select in scope.expression.selects
331                ],
332                nested=True,
333            )
334
335            if not any(
336                cd.kind.is_type(exp.DataType.Type.UNKNOWN)
337                for cd in struct_type.expressions
338                if cd.kind
339            ):
340                # We don't use `_set_type` on purpose here. If we annotated the query directly, then
341                # using it in other contexts (e.g., ARRAY(<query>)) could result in incorrect type
342                # annotations, i.e., it shouldn't be interpreted as a STRUCT value.
343                scope.expression.meta["query_type"] = struct_type
344
345    def _maybe_annotate(self, expression: E) -> E:
346        if id(expression) in self._visited:
347            return expression  # We've already inferred the expression's type
348
349        annotator = self.annotators.get(expression.__class__)
350
351        return (
352            annotator(self, expression)
353            if annotator
354            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
355        )
356
357    def _annotate_args(self, expression: E) -> E:
358        for value in expression.iter_expressions():
359            self._maybe_annotate(value)
360
361        return expression
362
363    def _maybe_coerce(
364        self,
365        type1: exp.DataType | exp.DataType.Type,
366        type2: exp.DataType | exp.DataType.Type,
367    ) -> exp.DataType | exp.DataType.Type:
368        """
369        Returns type2 if type1 can be coerced into it, otherwise type1.
370
371        If either type is parameterized (e.g. DECIMAL(18, 2) contains two parameters),
372        we assume type1 does not coerce into type2, so we also return it in this case.
373        """
374        if isinstance(type1, exp.DataType):
375            if type1.expressions:
376                return type1
377            type1_value = type1.this
378        else:
379            type1_value = type1
380
381        if isinstance(type2, exp.DataType):
382            if type2.expressions:
383                return type2
384            type2_value = type2.this
385        else:
386            type2_value = type2
387
388        # We propagate the UNKNOWN type upwards if found
389        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
390            return exp.DataType.Type.UNKNOWN
391
392        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value
393
394    def _annotate_binary(self, expression: B) -> B:
395        self._annotate_args(expression)
396
397        left, right = expression.left, expression.right
398        left_type, right_type = left.type.this, right.type.this  # type: ignore
399
400        if isinstance(expression, (exp.Connector, exp.Predicate)):
401            self._set_type(expression, exp.DataType.Type.BOOLEAN)
402        elif (left_type, right_type) in self.binary_coercions:
403            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
404        else:
405            self._set_type(expression, self._maybe_coerce(left_type, right_type))
406
407        return expression
408
409    def _annotate_unary(self, expression: E) -> E:
410        self._annotate_args(expression)
411
412        if isinstance(expression, exp.Not):
413            self._set_type(expression, exp.DataType.Type.BOOLEAN)
414        else:
415            self._set_type(expression, expression.this.type)
416
417        return expression
418
419    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
420        if expression.is_string:
421            self._set_type(expression, exp.DataType.Type.VARCHAR)
422        elif expression.is_int:
423            self._set_type(expression, exp.DataType.Type.INT)
424        else:
425            self._set_type(expression, exp.DataType.Type.DOUBLE)
426
427        return expression
428
429    def _annotate_with_type(
430        self, expression: E, target_type: exp.DataType | exp.DataType.Type
431    ) -> E:
432        self._set_type(expression, target_type)
433        return self._annotate_args(expression)
434
435    @t.no_type_check
436    def _annotate_by_args(
437        self,
438        expression: E,
439        *args: str,
440        promote: bool = False,
441        array: bool = False,
442    ) -> E:
443        self._annotate_args(expression)
444
445        expressions: t.List[exp.Expression] = []
446        for arg in args:
447            arg_expr = expression.args.get(arg)
448            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
449
450        last_datatype = None
451        for expr in expressions:
452            expr_type = expr.type
453
454            # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
455            if expr_type.args.get("nested"):
456                last_datatype = expr_type
457                break
458
459            if not expr_type.is_type(exp.DataType.Type.UNKNOWN):
460                last_datatype = self._maybe_coerce(last_datatype or expr_type, expr_type)
461
462        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
463
464        if promote:
465            if expression.type.this in exp.DataType.INTEGER_TYPES:
466                self._set_type(expression, exp.DataType.Type.BIGINT)
467            elif expression.type.this in exp.DataType.FLOAT_TYPES:
468                self._set_type(expression, exp.DataType.Type.DOUBLE)
469
470        if array:
471            self._set_type(
472                expression,
473                exp.DataType(
474                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
475                ),
476            )
477
478        return expression
479
480    def _annotate_timeunit(
481        self, expression: exp.TimeUnit | exp.DateTrunc
482    ) -> exp.TimeUnit | exp.DateTrunc:
483        self._annotate_args(expression)
484
485        if expression.this.type.this in exp.DataType.TEXT_TYPES:
486            datatype = _coerce_date_literal(expression.this, expression.unit)
487        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
488            datatype = _coerce_date(expression.this, expression.unit)
489        else:
490            datatype = exp.DataType.Type.UNKNOWN
491
492        self._set_type(expression, datatype)
493        return expression
494
495    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
496        self._annotate_args(expression)
497
498        bracket_arg = expression.expressions[0]
499        this = expression.this
500
501        if isinstance(bracket_arg, exp.Slice):
502            self._set_type(expression, this.type)
503        elif this.type.is_type(exp.DataType.Type.ARRAY):
504            self._set_type(expression, seq_get(this.type.expressions, 0))
505        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
506            index = this.keys.index(bracket_arg)
507            value = seq_get(this.values, index)
508            self._set_type(expression, value.type if value else None)
509        else:
510            self._set_type(expression, exp.DataType.Type.UNKNOWN)
511
512        return expression
513
514    def _annotate_div(self, expression: exp.Div) -> exp.Div:
515        self._annotate_args(expression)
516
517        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
518
519        if (
520            expression.args.get("typed")
521            and left_type in exp.DataType.INTEGER_TYPES
522            and right_type in exp.DataType.INTEGER_TYPES
523        ):
524            self._set_type(expression, exp.DataType.Type.BIGINT)
525        else:
526            self._set_type(expression, self._maybe_coerce(left_type, right_type))
527            if expression.type and expression.type.this not in exp.DataType.REAL_TYPES:
528                self._set_type(
529                    expression, self._maybe_coerce(expression.type, exp.DataType.Type.DOUBLE)
530                )
531
532        return expression
533
534    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
535        self._annotate_args(expression)
536        self._set_type(expression, None)
537        this_type = expression.this.type
538
539        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
540            for e in this_type.expressions:
541                if e.name == expression.expression.name:
542                    self._set_type(expression, e.kind)
543                    break
544
545        return expression
546
547    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
548        self._annotate_args(expression)
549        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
550        return expression
551
552    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
553        self._annotate_args(expression)
554        child = seq_get(expression.expressions, 0)
555
556        if child and child.is_type(exp.DataType.Type.ARRAY):
557            expr_type = seq_get(child.type.expressions, 0)
558        else:
559            expr_type = None
560
561        self._set_type(expression, expr_type)
562        return expression
563
564    def _annotate_struct_value(
565        self, expression: exp.Expression
566    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
567        alias = expression.args.get("alias")
568        if alias:
569            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
570
571        # Case: key = value or key := value
572        if expression.expression:
573            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
574
575        return expression.type
576
577    def _annotate_struct(self, expression: exp.Struct) -> exp.Struct:
578        self._annotate_args(expression)
579        self._set_type(
580            expression,
581            exp.DataType(
582                this=exp.DataType.Type.STRUCT,
583                expressions=[self._annotate_struct_value(expr) for expr in expression.expressions],
584                nested=True,
585            ),
586        )
587        return expression
588
589    @t.overload
590    def _annotate_map(self, expression: exp.Map) -> exp.Map: ...
591
592    @t.overload
593    def _annotate_map(self, expression: exp.VarMap) -> exp.VarMap: ...
594
595    def _annotate_map(self, expression):
596        self._annotate_args(expression)
597
598        keys = expression.args.get("keys")
599        values = expression.args.get("values")
600
601        map_type = exp.DataType(this=exp.DataType.Type.MAP)
602        if isinstance(keys, exp.Array) and isinstance(values, exp.Array):
603            key_type = seq_get(keys.type.expressions, 0) or exp.DataType.Type.UNKNOWN
604            value_type = seq_get(values.type.expressions, 0) or exp.DataType.Type.UNKNOWN
605
606            if key_type != exp.DataType.Type.UNKNOWN and value_type != exp.DataType.Type.UNKNOWN:
607                map_type.set("expressions", [key_type, value_type])
608                map_type.set("nested", True)
609
610        self._set_type(expression, map_type)
611        return expression
612
613    def _annotate_to_map(self, expression: exp.ToMap) -> exp.ToMap:
614        self._annotate_args(expression)
615
616        map_type = exp.DataType(this=exp.DataType.Type.MAP)
617        arg = expression.this
618        if arg.is_type(exp.DataType.Type.STRUCT):
619            for coldef in arg.type.expressions:
620                kind = coldef.kind
621                if kind != exp.DataType.Type.UNKNOWN:
622                    map_type.set("expressions", [exp.DataType.build("varchar"), kind])
623                    map_type.set("nested", True)
624                    break
625
626        self._set_type(expression, map_type)
627        return expression
628
629    def _annotate_extract(self, expression: exp.Extract) -> exp.Extract:
630        self._annotate_args(expression)
631        part = expression.name
632        if part == "TIME":
633            self._set_type(expression, exp.DataType.Type.TIME)
634        elif part == "DATE":
635            self._set_type(expression, exp.DataType.Type.DATE)
636        else:
637            self._set_type(expression, exp.DataType.Type.INT)
638        return expression
639
640    def _annotate_by_array_element(self, expression: exp.Expression) -> exp.Expression:
641        self._annotate_args(expression)
642
643        array_arg = expression.this
644        if array_arg.type.is_type(exp.DataType.Type.ARRAY):
645            element_type = seq_get(array_arg.type.expressions, 0) or exp.DataType.Type.UNKNOWN
646            self._set_type(expression, element_type)
647        else:
648            self._set_type(expression, exp.DataType.Type.UNKNOWN)
649
650        return expression