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 | 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 | exp.DataType.Type:
316        """
317        Returns type2 if type1 can be coerced into it, otherwise type1.
318
319        If either type is parameterized (e.g. DECIMAL(18, 2) contains two parameters),
320        we assume type1 does not coerce into type2, so we also return it in this case.
321        """
322        if isinstance(type1, exp.DataType):
323            if type1.expressions:
324                return type1
325            type1_value = type1.this
326        else:
327            type1_value = type1
328
329        if isinstance(type2, exp.DataType):
330            if type2.expressions:
331                return type1
332            type2_value = type2.this
333        else:
334            type2_value = type2
335
336        # We propagate the UNKNOWN type upwards if found
337        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
338            return exp.DataType.Type.UNKNOWN
339
340        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value
341
342    def _annotate_binary(self, expression: B) -> B:
343        self._annotate_args(expression)
344
345        left, right = expression.left, expression.right
346        left_type, right_type = left.type.this, right.type.this  # type: ignore
347
348        if isinstance(expression, (exp.Connector, exp.Predicate)):
349            self._set_type(expression, exp.DataType.Type.BOOLEAN)
350        elif (left_type, right_type) in self.binary_coercions:
351            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
352        else:
353            self._set_type(expression, self._maybe_coerce(left_type, right_type))
354
355        return expression
356
357    def _annotate_unary(self, expression: E) -> E:
358        self._annotate_args(expression)
359
360        if isinstance(expression, exp.Not):
361            self._set_type(expression, exp.DataType.Type.BOOLEAN)
362        else:
363            self._set_type(expression, expression.this.type)
364
365        return expression
366
367    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
368        if expression.is_string:
369            self._set_type(expression, exp.DataType.Type.VARCHAR)
370        elif expression.is_int:
371            self._set_type(expression, exp.DataType.Type.INT)
372        else:
373            self._set_type(expression, exp.DataType.Type.DOUBLE)
374
375        return expression
376
377    def _annotate_with_type(
378        self, expression: E, target_type: exp.DataType | exp.DataType.Type
379    ) -> E:
380        self._set_type(expression, target_type)
381        return self._annotate_args(expression)
382
383    @t.no_type_check
384    def _annotate_by_args(
385        self,
386        expression: E,
387        *args: str,
388        promote: bool = False,
389        array: bool = False,
390    ) -> E:
391        self._annotate_args(expression)
392
393        expressions: t.List[exp.Expression] = []
394        for arg in args:
395            arg_expr = expression.args.get(arg)
396            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
397
398        last_datatype = None
399        for expr in expressions:
400            expr_type = expr.type
401
402            # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
403            if expr_type.args.get("nested"):
404                last_datatype = expr_type
405                break
406
407            if not expr_type.is_type(exp.DataType.Type.UNKNOWN):
408                last_datatype = self._maybe_coerce(last_datatype or expr_type, expr_type)
409
410        self._set_type(expression, last_datatype or exp.DataType.Type.UNKNOWN)
411
412        if promote:
413            if expression.type.this in exp.DataType.INTEGER_TYPES:
414                self._set_type(expression, exp.DataType.Type.BIGINT)
415            elif expression.type.this in exp.DataType.FLOAT_TYPES:
416                self._set_type(expression, exp.DataType.Type.DOUBLE)
417
418        if array:
419            self._set_type(
420                expression,
421                exp.DataType(
422                    this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
423                ),
424            )
425
426        return expression
427
428    def _annotate_timeunit(
429        self, expression: exp.TimeUnit | exp.DateTrunc
430    ) -> exp.TimeUnit | exp.DateTrunc:
431        self._annotate_args(expression)
432
433        if expression.this.type.this in exp.DataType.TEXT_TYPES:
434            datatype = _coerce_date_literal(expression.this, expression.unit)
435        elif expression.this.type.this in exp.DataType.TEMPORAL_TYPES:
436            datatype = _coerce_date(expression.this, expression.unit)
437        else:
438            datatype = exp.DataType.Type.UNKNOWN
439
440        self._set_type(expression, datatype)
441        return expression
442
443    def _annotate_bracket(self, expression: exp.Bracket) -> exp.Bracket:
444        self._annotate_args(expression)
445
446        bracket_arg = expression.expressions[0]
447        this = expression.this
448
449        if isinstance(bracket_arg, exp.Slice):
450            self._set_type(expression, this.type)
451        elif this.type.is_type(exp.DataType.Type.ARRAY):
452            self._set_type(expression, seq_get(this.type.expressions, 0))
453        elif isinstance(this, (exp.Map, exp.VarMap)) and bracket_arg in this.keys:
454            index = this.keys.index(bracket_arg)
455            value = seq_get(this.values, index)
456            self._set_type(expression, value.type if value else None)
457        else:
458            self._set_type(expression, exp.DataType.Type.UNKNOWN)
459
460        return expression
461
462    def _annotate_div(self, expression: exp.Div) -> exp.Div:
463        self._annotate_args(expression)
464
465        left_type, right_type = expression.left.type.this, expression.right.type.this  # type: ignore
466
467        if (
468            expression.args.get("typed")
469            and left_type in exp.DataType.INTEGER_TYPES
470            and right_type in exp.DataType.INTEGER_TYPES
471        ):
472            self._set_type(expression, exp.DataType.Type.BIGINT)
473        else:
474            self._set_type(expression, self._maybe_coerce(left_type, right_type))
475            if expression.type and expression.type.this not in exp.DataType.REAL_TYPES:
476                self._set_type(
477                    expression, self._maybe_coerce(expression.type, exp.DataType.Type.DOUBLE)
478                )
479
480        return expression
481
482    def _annotate_dot(self, expression: exp.Dot) -> exp.Dot:
483        self._annotate_args(expression)
484        self._set_type(expression, None)
485        this_type = expression.this.type
486
487        if this_type and this_type.is_type(exp.DataType.Type.STRUCT):
488            for e in this_type.expressions:
489                if e.name == expression.expression.name:
490                    self._set_type(expression, e.kind)
491                    break
492
493        return expression
494
495    def _annotate_explode(self, expression: exp.Explode) -> exp.Explode:
496        self._annotate_args(expression)
497        self._set_type(expression, seq_get(expression.this.type.expressions, 0))
498        return expression
499
500    def _annotate_unnest(self, expression: exp.Unnest) -> exp.Unnest:
501        self._annotate_args(expression)
502        child = seq_get(expression.expressions, 0)
503
504        if child and child.is_type(exp.DataType.Type.ARRAY):
505            expr_type = seq_get(child.type.expressions, 0)
506        else:
507            expr_type = None
508
509        self._set_type(expression, expr_type)
510        return expression
511
512    def _annotate_struct_value(
513        self, expression: exp.Expression
514    ) -> t.Optional[exp.DataType] | exp.ColumnDef:
515        alias = expression.args.get("alias")
516        if alias:
517            return exp.ColumnDef(this=alias.copy(), kind=expression.type)
518
519        # Case: key = value or key := value
520        if expression.expression:
521            return exp.ColumnDef(this=expression.this.copy(), kind=expression.expression.type)
522
523        return expression.type
524
525    def _annotate_struct(self, expression: exp.Struct) -> exp.Struct:
526        self._annotate_args(expression)
527        self._set_type(
528            expression,
529            exp.DataType(
530                this=exp.DataType.Type.STRUCT,
531                expressions=[self._annotate_struct_value(expr) for expr in expression.expressions],
532                nested=True,
533            ),
534        )
535        return expression
536
537    @t.overload
538    def _annotate_map(self, expression: exp.Map) -> exp.Map: ...
539
540    @t.overload
541    def _annotate_map(self, expression: exp.VarMap) -> exp.VarMap: ...
542
543    def _annotate_map(self, expression):
544        self._annotate_args(expression)
545
546        keys = expression.args.get("keys")
547        values = expression.args.get("values")
548
549        map_type = exp.DataType(this=exp.DataType.Type.MAP)
550        if isinstance(keys, exp.Array) and isinstance(values, exp.Array):
551            key_type = seq_get(keys.type.expressions, 0) or exp.DataType.Type.UNKNOWN
552            value_type = seq_get(values.type.expressions, 0) or exp.DataType.Type.UNKNOWN
553
554            if key_type != exp.DataType.Type.UNKNOWN and value_type != exp.DataType.Type.UNKNOWN:
555                map_type.set("expressions", [key_type, value_type])
556                map_type.set("nested", True)
557
558        self._set_type(expression, map_type)
559        return expression
560
561    def _annotate_to_map(self, expression: exp.ToMap) -> exp.ToMap:
562        self._annotate_args(expression)
563
564        map_type = exp.DataType(this=exp.DataType.Type.MAP)
565        arg = expression.this
566        if arg.is_type(exp.DataType.Type.STRUCT):
567            for coldef in arg.type.expressions:
568                kind = coldef.kind
569                if kind != exp.DataType.Type.UNKNOWN:
570                    map_type.set("expressions", [exp.DataType.build("varchar"), kind])
571                    map_type.set("nested", True)
572                    break
573
574        self._set_type(expression, map_type)
575        return expression
576
577    def _annotate_extract(self, expression: exp.Extract) -> exp.Extract:
578        self._annotate_args(expression)
579        part = expression.name
580        if part == "TIME":
581            self._set_type(expression, exp.DataType.Type.TIME)
582        elif part == "DATE":
583            self._set_type(expression, exp.DataType.Type.DATE)
584        else:
585            self._set_type(expression, exp.DataType.Type.INT)
586        return expression