sqlglot.optimizer.annotate_types

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot._typing import E
  7from sqlglot.helper import ensure_list, subclasses
  8from sqlglot.optimizer.scope import Scope, traverse_scope
  9from sqlglot.schema import Schema, ensure_schema
 10
 11if t.TYPE_CHECKING:
 12    B = t.TypeVar("B", bound=exp.Binary)
 13
 14
 15def annotate_types(
 16    expression: E,
 17    schema: t.Optional[t.Dict | Schema] = None,
 18    annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
 19    coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
 20) -> E:
 21    """
 22    Infers the types of an expression, annotating its AST accordingly.
 23
 24    Example:
 25        >>> import sqlglot
 26        >>> schema = {"y": {"cola": "SMALLINT"}}
 27        >>> sql = "SELECT x.cola + 2.5 AS cola FROM (SELECT y.cola AS cola FROM y AS y) AS x"
 28        >>> annotated_expr = annotate_types(sqlglot.parse_one(sql), schema=schema)
 29        >>> annotated_expr.expressions[0].type.this  # Get the type of "x.cola + 2.5 AS cola"
 30        <Type.DOUBLE: 'DOUBLE'>
 31
 32    Args:
 33        expression: Expression to annotate.
 34        schema: Database schema.
 35        annotators: Maps expression type to corresponding annotation function.
 36        coerces_to: Maps expression type to set of types that it can be coerced into.
 37
 38    Returns:
 39        The expression annotated with types.
 40    """
 41
 42    schema = ensure_schema(schema)
 43
 44    return TypeAnnotator(schema, annotators, coerces_to).annotate(expression)
 45
 46
 47def _annotate_with_type_lambda(data_type: exp.DataType.Type) -> t.Callable[[TypeAnnotator, E], E]:
 48    return lambda self, e: self._annotate_with_type(e, data_type)
 49
 50
 51class _TypeAnnotator(type):
 52    def __new__(cls, clsname, bases, attrs):
 53        klass = super().__new__(cls, clsname, bases, attrs)
 54
 55        # Highest-to-lowest type precedence, as specified in Spark's docs (ANSI):
 56        # https://spark.apache.org/docs/3.2.0/sql-ref-ansi-compliance.html
 57        text_precedence = (
 58            exp.DataType.Type.TEXT,
 59            exp.DataType.Type.NVARCHAR,
 60            exp.DataType.Type.VARCHAR,
 61            exp.DataType.Type.NCHAR,
 62            exp.DataType.Type.CHAR,
 63        )
 64        numeric_precedence = (
 65            exp.DataType.Type.DOUBLE,
 66            exp.DataType.Type.FLOAT,
 67            exp.DataType.Type.DECIMAL,
 68            exp.DataType.Type.BIGINT,
 69            exp.DataType.Type.INT,
 70            exp.DataType.Type.SMALLINT,
 71            exp.DataType.Type.TINYINT,
 72        )
 73        timelike_precedence = (
 74            exp.DataType.Type.TIMESTAMPLTZ,
 75            exp.DataType.Type.TIMESTAMPTZ,
 76            exp.DataType.Type.TIMESTAMP,
 77            exp.DataType.Type.DATETIME,
 78            exp.DataType.Type.DATE,
 79        )
 80
 81        for type_precedence in (text_precedence, numeric_precedence, timelike_precedence):
 82            coerces_to = set()
 83            for data_type in type_precedence:
 84                klass.COERCES_TO[data_type] = coerces_to.copy()
 85                coerces_to |= {data_type}
 86
 87        return klass
 88
 89
 90class TypeAnnotator(metaclass=_TypeAnnotator):
 91    TYPE_TO_EXPRESSIONS: t.Dict[exp.DataType.Type, t.Set[t.Type[exp.Expression]]] = {
 92        exp.DataType.Type.BIGINT: {
 93            exp.ApproxDistinct,
 94            exp.ArraySize,
 95            exp.Count,
 96            exp.Length,
 97        },
 98        exp.DataType.Type.BOOLEAN: {
 99            exp.Between,
100            exp.Boolean,
101            exp.In,
102            exp.RegexpLike,
103        },
104        exp.DataType.Type.DATE: {
105            exp.CurrentDate,
106            exp.Date,
107            exp.DateAdd,
108            exp.DateStrToDate,
109            exp.DateSub,
110            exp.DateTrunc,
111            exp.DiToDate,
112            exp.StrToDate,
113            exp.TimeStrToDate,
114            exp.TsOrDsToDate,
115        },
116        exp.DataType.Type.DATETIME: {
117            exp.CurrentDatetime,
118            exp.DatetimeAdd,
119            exp.DatetimeSub,
120        },
121        exp.DataType.Type.DOUBLE: {
122            exp.ApproxQuantile,
123            exp.Avg,
124            exp.Exp,
125            exp.Ln,
126            exp.Log,
127            exp.Log2,
128            exp.Log10,
129            exp.Pow,
130            exp.Quantile,
131            exp.Round,
132            exp.SafeDivide,
133            exp.Sqrt,
134            exp.Stddev,
135            exp.StddevPop,
136            exp.StddevSamp,
137            exp.Variance,
138            exp.VariancePop,
139        },
140        exp.DataType.Type.INT: {
141            exp.Ceil,
142            exp.DateDiff,
143            exp.DatetimeDiff,
144            exp.Extract,
145            exp.TimestampDiff,
146            exp.TimeDiff,
147            exp.DateToDi,
148            exp.Floor,
149            exp.Levenshtein,
150            exp.StrPosition,
151            exp.TsOrDiToDi,
152        },
153        exp.DataType.Type.TIMESTAMP: {
154            exp.CurrentTime,
155            exp.CurrentTimestamp,
156            exp.StrToTime,
157            exp.TimeAdd,
158            exp.TimeStrToTime,
159            exp.TimeSub,
160            exp.TimestampAdd,
161            exp.TimestampSub,
162            exp.UnixToTime,
163        },
164        exp.DataType.Type.TINYINT: {
165            exp.Day,
166            exp.Month,
167            exp.Week,
168            exp.Year,
169        },
170        exp.DataType.Type.VARCHAR: {
171            exp.ArrayConcat,
172            exp.Concat,
173            exp.ConcatWs,
174            exp.DateToDateStr,
175            exp.GroupConcat,
176            exp.Initcap,
177            exp.Lower,
178            exp.SafeConcat,
179            exp.Substring,
180            exp.TimeToStr,
181            exp.TimeToTimeStr,
182            exp.Trim,
183            exp.TsOrDsToDateStr,
184            exp.UnixToStr,
185            exp.UnixToTimeStr,
186            exp.Upper,
187        },
188    }
189
190    ANNOTATORS: t.Dict = {
191        **{
192            expr_type: lambda self, e: self._annotate_unary(e)
193            for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
194        },
195        **{
196            expr_type: lambda self, e: self._annotate_binary(e)
197            for expr_type in subclasses(exp.__name__, exp.Binary)
198        },
199        **{
200            expr_type: _annotate_with_type_lambda(data_type)
201            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
202            for expr_type in expressions
203        },
204        exp.Anonymous: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
205        exp.Cast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
206        exp.Case: lambda self, e: self._annotate_by_args(e, "default", "ifs"),
207        exp.Coalesce: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
208        exp.DataType: lambda self, e: self._annotate_with_type(e, e.copy()),
209        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
210        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
211        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
212        exp.Literal: lambda self, e: self._annotate_literal(e),
213        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
214        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
215        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
216        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
217        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
218        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
219        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
220    }
221
222    # Specifies what types a given type can be coerced into (autofilled)
223    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
224
225    def __init__(
226        self,
227        schema: Schema,
228        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
229        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
230    ) -> None:
231        self.schema = schema
232        self.annotators = annotators or self.ANNOTATORS
233        self.coerces_to = coerces_to or self.COERCES_TO
234
235    def annotate(self, expression: E) -> E:
236        for scope in traverse_scope(expression):
237            selects = {}
238            for name, source in scope.sources.items():
239                if not isinstance(source, Scope):
240                    continue
241                if isinstance(source.expression, exp.UDTF):
242                    values = []
243
244                    if isinstance(source.expression, exp.Lateral):
245                        if isinstance(source.expression.this, exp.Explode):
246                            values = [source.expression.this.this]
247                    else:
248                        values = source.expression.expressions[0].expressions
249
250                    if not values:
251                        continue
252
253                    selects[name] = {
254                        alias: column
255                        for alias, column in zip(
256                            source.expression.alias_column_names,
257                            values,
258                        )
259                    }
260                else:
261                    selects[name] = {
262                        select.alias_or_name: select for select in source.expression.selects
263                    }
264
265            # First annotate the current scope's column references
266            for col in scope.columns:
267                if not col.table:
268                    continue
269
270                source = scope.sources.get(col.table)
271                if isinstance(source, exp.Table):
272                    col.type = self.schema.get_column_type(source, col)
273                elif source and col.table in selects and col.name in selects[col.table]:
274                    col.type = selects[col.table][col.name].type
275
276            # Then (possibly) annotate the remaining expressions in the scope
277            self._maybe_annotate(scope.expression)
278
279        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
280
281    def _maybe_annotate(self, expression: E) -> E:
282        if expression.type:
283            return expression  # We've already inferred the expression's type
284
285        annotator = self.annotators.get(expression.__class__)
286
287        return (
288            annotator(self, expression)
289            if annotator
290            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
291        )
292
293    def _annotate_args(self, expression: E) -> E:
294        for _, value in expression.iter_expressions():
295            self._maybe_annotate(value)
296
297        return expression
298
299    def _maybe_coerce(
300        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
301    ) -> exp.DataType.Type:
302        # We propagate the NULL / UNKNOWN types upwards if found
303        if isinstance(type1, exp.DataType):
304            type1 = type1.this
305        if isinstance(type2, exp.DataType):
306            type2 = type2.this
307
308        if exp.DataType.Type.NULL in (type1, type2):
309            return exp.DataType.Type.NULL
310        if exp.DataType.Type.UNKNOWN in (type1, type2):
311            return exp.DataType.Type.UNKNOWN
312
313        return type2 if type2 in self.coerces_to.get(type1, {}) else type1  # type: ignore
314
315    # Note: the following "no_type_check" decorators were added because mypy was yelling due
316    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
317    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
318
319    @t.no_type_check
320    def _annotate_binary(self, expression: B) -> B:
321        self._annotate_args(expression)
322
323        left_type = expression.left.type.this
324        right_type = expression.right.type.this
325
326        if isinstance(expression, exp.Connector):
327            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
328                expression.type = exp.DataType.Type.NULL
329            elif exp.DataType.Type.NULL in (left_type, right_type):
330                expression.type = exp.DataType.build(
331                    "NULLABLE", expressions=exp.DataType.build("BOOLEAN")
332                )
333            else:
334                expression.type = exp.DataType.Type.BOOLEAN
335        elif isinstance(expression, exp.Predicate):
336            expression.type = exp.DataType.Type.BOOLEAN
337        else:
338            expression.type = self._maybe_coerce(left_type, right_type)
339
340        return expression
341
342    @t.no_type_check
343    def _annotate_unary(self, expression: E) -> E:
344        self._annotate_args(expression)
345
346        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
347            expression.type = exp.DataType.Type.BOOLEAN
348        else:
349            expression.type = expression.this.type
350
351        return expression
352
353    @t.no_type_check
354    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
355        if expression.is_string:
356            expression.type = exp.DataType.Type.VARCHAR
357        elif expression.is_int:
358            expression.type = exp.DataType.Type.INT
359        else:
360            expression.type = exp.DataType.Type.DOUBLE
361
362        return expression
363
364    @t.no_type_check
365    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
366        expression.type = target_type
367        return self._annotate_args(expression)
368
369    @t.no_type_check
370    def _annotate_by_args(self, expression: E, *args: str, promote: bool = False) -> E:
371        self._annotate_args(expression)
372
373        expressions: t.List[exp.Expression] = []
374        for arg in args:
375            arg_expr = expression.args.get(arg)
376            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
377
378        last_datatype = None
379        for expr in expressions:
380            last_datatype = self._maybe_coerce(last_datatype or expr.type, expr.type)
381
382        expression.type = last_datatype or exp.DataType.Type.UNKNOWN
383
384        if promote:
385            if expression.type.this in exp.DataType.INTEGER_TYPES:
386                expression.type = exp.DataType.Type.BIGINT
387            elif expression.type.this in exp.DataType.FLOAT_TYPES:
388                expression.type = exp.DataType.Type.DOUBLE
389
390        return expression