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 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: t.Optional[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)
 59
 60    return TypeAnnotator(schema, annotators, coerces_to, dialect=dialect).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        dialect: t.Optional[DialectType] = None,
186    ) -> None:
187        self.schema = schema
188        self.annotators = annotators or Dialect.get_or_raise(dialect).ANNOTATORS
189        self.coerces_to = coerces_to or self.COERCES_TO
190        self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
191
192        # Caches the ids of annotated sub-Expressions, to ensure we only visit them once
193        self._visited: t.Set[int] = set()
194
195        # Maps an exp.SetOperation's id (e.g. UNION) to its projection types. This is computed if the
196        # exp.SetOperation is the expression of a scope source, as selecting from it multiple times
197        # would reprocess the entire subtree to coerce the types of its operands' projections
198        self._setop_column_types: t.Dict[int, t.Dict[str, exp.DataType.Type]] = {}
199
200    def _set_type(
201        self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
202    ) -> None:
203        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
204        self._visited.add(id(expression))
205
206    def annotate(self, expression: E) -> E:
207        for scope in traverse_scope(expression):
208            self.annotate_scope(scope)
209        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
210
211    def annotate_scope(self, scope: Scope) -> None:
212        selects = {}
213        for name, source in scope.sources.items():
214            if not isinstance(source, Scope):
215                continue
216
217            expression = source.expression
218            if isinstance(expression, exp.UDTF):
219                values = []
220
221                if isinstance(expression, exp.Lateral):
222                    if isinstance(expression.this, exp.Explode):
223                        values = [expression.this.this]
224                elif isinstance(expression, exp.Unnest):
225                    values = [expression]
226                else:
227                    values = expression.expressions[0].expressions
228
229                if not values:
230                    continue
231
232                selects[name] = {
233                    alias: column.type
234                    for alias, column in zip(expression.alias_column_names, values)
235                }
236            elif isinstance(expression, exp.SetOperation) and len(expression.left.selects) == len(
237                expression.right.selects
238            ):
239                selects[name] = col_types = self._setop_column_types.setdefault(id(expression), {})
240
241                if not col_types:
242                    # Process a chain / sub-tree of set operations
243                    for set_op in expression.walk(
244                        prune=lambda n: not isinstance(n, (exp.SetOperation, exp.Subquery))
245                    ):
246                        if not isinstance(set_op, exp.SetOperation):
247                            continue
248
249                        if set_op.args.get("by_name"):
250                            r_type_by_select = {
251                                s.alias_or_name: s.type for s in set_op.right.selects
252                            }
253                            setop_cols = {
254                                s.alias_or_name: self._maybe_coerce(
255                                    t.cast(exp.DataType, s.type),
256                                    r_type_by_select.get(s.alias_or_name)
257                                    or exp.DataType.Type.UNKNOWN,
258                                )
259                                for s in set_op.left.selects
260                            }
261                        else:
262                            setop_cols = {
263                                ls.alias_or_name: self._maybe_coerce(
264                                    t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
265                                )
266                                for ls, rs in zip(set_op.left.selects, set_op.right.selects)
267                            }
268
269                        # Coerce intermediate results with the previously registered types, if they exist
270                        for col_name, col_type in setop_cols.items():
271                            col_types[col_name] = self._maybe_coerce(
272                                col_type, col_types.get(col_name, exp.DataType.Type.NULL)
273                            )
274
275            else:
276                selects[name] = {s.alias_or_name: s.type for s in expression.selects}
277
278        # First annotate the current scope's column references
279        for col in scope.columns:
280            if not col.table:
281                continue
282
283            source = scope.sources.get(col.table)
284            if isinstance(source, exp.Table):
285                self._set_type(col, self.schema.get_column_type(source, col))
286            elif source:
287                if col.table in selects and col.name in selects[col.table]:
288                    self._set_type(col, selects[col.table][col.name])
289                elif isinstance(source.expression, exp.Unnest):
290                    self._set_type(col, source.expression.type)
291
292        # Then (possibly) annotate the remaining expressions in the scope
293        self._maybe_annotate(scope.expression)
294
295    def _maybe_annotate(self, expression: E) -> E:
296        if id(expression) in self._visited:
297            return expression  # We've already inferred the expression's type
298
299        annotator = self.annotators.get(expression.__class__)
300
301        return (
302            annotator(self, expression)
303            if annotator
304            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
305        )
306
307    def _annotate_args(self, expression: E) -> E:
308        for value in expression.iter_expressions():
309            self._maybe_annotate(value)
310
311        return expression
312
313    def _maybe_coerce(
314        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
315    ) -> exp.DataType.Type:
316        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
317        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
318
319        # We propagate the UNKNOWN type upwards if found
320        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
321            return exp.DataType.Type.UNKNOWN
322
323        return t.cast(
324            exp.DataType.Type,
325            type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value,
326        )
327
328    def _annotate_binary(self, expression: B) -> B:
329        self._annotate_args(expression)
330
331        left, right = expression.left, expression.right
332        left_type, right_type = left.type.this, right.type.this  # type: ignore
333
334        if isinstance(expression, (exp.Connector, exp.Predicate)):
335            self._set_type(expression, exp.DataType.Type.BOOLEAN)
336        elif (left_type, right_type) in self.binary_coercions:
337            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
338        else:
339            self._set_type(expression, self._maybe_coerce(left_type, right_type))
340
341        return expression
342
343    def _annotate_unary(self, expression: E) -> E:
344        self._annotate_args(expression)
345
346        if isinstance(expression, exp.Not):
347            self._set_type(expression, exp.DataType.Type.BOOLEAN)
348        else:
349            self._set_type(expression, expression.this.type)
350
351        return expression
352
353    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
354        if expression.is_string:
355            self._set_type(expression, exp.DataType.Type.VARCHAR)
356        elif expression.is_int:
357            self._set_type(expression, exp.DataType.Type.INT)
358        else:
359            self._set_type(expression, exp.DataType.Type.DOUBLE)
360
361        return expression
362
363    def _annotate_with_type(
364        self, expression: E, target_type: exp.DataType | exp.DataType.Type
365    ) -> E:
366        self._set_type(expression, target_type)
367        return self._annotate_args(expression)
368
369    @t.no_type_check
370    def _annotate_by_args(
371        self,
372        expression: E,
373        *args: str,
374        promote: bool = False,
375        array: bool = False,
376    ) -> E:
377        self._annotate_args(expression)
378
379        expressions: t.List[exp.Expression] = []
380        for arg in args:
381            arg_expr = expression.args.get(arg)
382            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
383
384        last_datatype = None
385        for expr in expressions:
386            expr_type = expr.type
387
388            # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
389            if expr_type.args.get("nested"):
390                last_datatype = expr_type
391                break
392
393            if not expr_type.is_type(exp.DataType.Type.UNKNOWN):
394                last_datatype = self._maybe_coerce(last_datatype or expr_type, expr_type)
395
396        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
397
398        if promote:
399            if expression.type.this in exp.DataType.INTEGER_TYPES:
400                self._set_type(expression, exp.DataType.Type.BIGINT)
401            elif expression.type.this in exp.DataType.FLOAT_TYPES:
402                self._set_type(expression, exp.DataType.Type.DOUBLE)
403
404        if array:
405            self._set_type(
406                expression,
407                exp.DataType(
408                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
409                ),
410            )
411
412        return expression
413
414    def _annotate_timeunit(
415        self, expression: exp.TimeUnit | exp.DateTrunc
416    ) -> exp.TimeUnit | exp.DateTrunc:
417        self._annotate_args(expression)
418
419        if expression.this.type.this in exp.DataType.TEXT_TYPES:
420            datatype = _coerce_date_literal(expression.this, expression.unit)
421        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
422            datatype = _coerce_date(expression.this, expression.unit)
423        else:
424            datatype = exp.DataType.Type.UNKNOWN
425
426        self._set_type(expression, datatype)
427        return expression
428
429    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
430        self._annotate_args(expression)
431
432        bracket_arg = expression.expressions[0]
433        this = expression.this
434
435        if isinstance(bracket_arg, exp.Slice):
436            self._set_type(expression, this.type)
437        elif this.type.is_type(exp.DataType.Type.ARRAY):
438            self._set_type(expression, seq_get(this.type.expressions, 0))
439        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
440            index = this.keys.index(bracket_arg)
441            value = seq_get(this.values, index)
442            self._set_type(expression, value.type if value else None)
443        else:
444            self._set_type(expression, exp.DataType.Type.UNKNOWN)
445
446        return expression
447
448    def _annotate_div(self, expression: exp.Div) -> exp.Div:
449        self._annotate_args(expression)
450
451        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
452
453        if (
454            expression.args.get("typed")
455            and left_type in exp.DataType.INTEGER_TYPES
456            and right_type in exp.DataType.INTEGER_TYPES
457        ):
458            self._set_type(expression, exp.DataType.Type.BIGINT)
459        else:
460            self._set_type(expression, self._maybe_coerce(left_type, right_type))
461            if expression.type and expression.type.this not in exp.DataType.REAL_TYPES:
462                self._set_type(
463                    expression, self._maybe_coerce(expression.type, exp.DataType.Type.DOUBLE)
464                )
465
466        return expression
467
468    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
469        self._annotate_args(expression)
470        self._set_type(expression, None)
471        this_type = expression.this.type
472
473        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
474            for e in this_type.expressions:
475                if e.name == expression.expression.name:
476                    self._set_type(expression, e.kind)
477                    break
478
479        return expression
480
481    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
482        self._annotate_args(expression)
483        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
484        return expression
485
486    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
487        self._annotate_args(expression)
488        child = seq_get(expression.expressions, 0)
489
490        if child and child.is_type(exp.DataType.Type.ARRAY):
491            expr_type = seq_get(child.type.expressions, 0)
492        else:
493            expr_type = None
494
495        self._set_type(expression, expr_type)
496        return expression
497
498    def _annotate_struct_value(
499        self, expression: exp.Expression
500    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
501        alias = expression.args.get("alias")
502        if alias:
503            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
504
505        # Case: key = value or key := value
506        if expression.expression:
507            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
508
509        return expression.type
510
511    def _annotate_struct(self, expression: exp.Struct) -> exp.Struct:
512        self._annotate_args(expression)
513        self._set_type(
514            expression,
515            exp.DataType(
516                this=exp.DataType.Type.STRUCT,
517                expressions=[self._annotate_struct_value(expr) for expr in expression.expressions],
518                nested=True,
519            ),
520        )
521        return expression
522
523    @t.overload
524    def _annotate_map(self, expression: exp.Map) -> exp.Map: ...
525
526    @t.overload
527    def _annotate_map(self, expression: exp.VarMap) -> exp.VarMap: ...
528
529    def _annotate_map(self, expression):
530        self._annotate_args(expression)
531
532        keys = expression.args.get("keys")
533        values = expression.args.get("values")
534
535        map_type = exp.DataType(this=exp.DataType.Type.MAP)
536        if isinstance(keys, exp.Array) and isinstance(values, exp.Array):
537            key_type = seq_get(keys.type.expressions, 0) or exp.DataType.Type.UNKNOWN
538            value_type = seq_get(values.type.expressions, 0) or exp.DataType.Type.UNKNOWN
539
540            if key_type != exp.DataType.Type.UNKNOWN and value_type != exp.DataType.Type.UNKNOWN:
541                map_type.set("expressions", [key_type, value_type])
542                map_type.set("nested", True)
543
544        self._set_type(expression, map_type)
545        return expression
546
547    def _annotate_to_map(self, expression: exp.ToMap) -> exp.ToMap:
548        self._annotate_args(expression)
549
550        map_type = exp.DataType(this=exp.DataType.Type.MAP)
551        arg = expression.this
552        if arg.is_type(exp.DataType.Type.STRUCT):
553            for coldef in arg.type.expressions:
554                kind = coldef.kind
555                if kind != exp.DataType.Type.UNKNOWN:
556                    map_type.set("expressions", [exp.DataType.build("varchar"), kind])
557                    map_type.set("nested", True)
558                    break
559
560        self._set_type(expression, map_type)
561        return expression
562
563    def _annotate_extract(self, expression: exp.Extract) -> exp.Extract:
564        self._annotate_args(expression)
565        part = expression.name
566        if part == "TIME":
567            self._set_type(expression, exp.DataType.Type.TIME)
568        elif part == "DATE":
569            self._set_type(expression, exp.DataType.Type.DATE)
570        else:
571            self._set_type(expression, exp.DataType.Type.INT)
572        return expression