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.DateFromParts,
109            exp.DateStrToDate,
110            exp.DateSub,
111            exp.DateTrunc,
112            exp.DiToDate,
113            exp.StrToDate,
114            exp.TimeStrToDate,
115            exp.TsOrDsToDate,
116        },
117        exp.DataType.Type.DATETIME: {
118            exp.CurrentDatetime,
119            exp.DatetimeAdd,
120            exp.DatetimeSub,
121        },
122        exp.DataType.Type.DOUBLE: {
123            exp.ApproxQuantile,
124            exp.Avg,
125            exp.Exp,
126            exp.Ln,
127            exp.Log,
128            exp.Log2,
129            exp.Log10,
130            exp.Pow,
131            exp.Quantile,
132            exp.Round,
133            exp.SafeDivide,
134            exp.Sqrt,
135            exp.Stddev,
136            exp.StddevPop,
137            exp.StddevSamp,
138            exp.Variance,
139            exp.VariancePop,
140        },
141        exp.DataType.Type.INT: {
142            exp.Ceil,
143            exp.DateDiff,
144            exp.DatetimeDiff,
145            exp.Extract,
146            exp.TimestampDiff,
147            exp.TimeDiff,
148            exp.DateToDi,
149            exp.Floor,
150            exp.Levenshtein,
151            exp.StrPosition,
152            exp.TsOrDiToDi,
153        },
154        exp.DataType.Type.TIMESTAMP: {
155            exp.CurrentTime,
156            exp.CurrentTimestamp,
157            exp.StrToTime,
158            exp.TimeAdd,
159            exp.TimeStrToTime,
160            exp.TimeSub,
161            exp.TimestampAdd,
162            exp.TimestampSub,
163            exp.UnixToTime,
164        },
165        exp.DataType.Type.TINYINT: {
166            exp.Day,
167            exp.Month,
168            exp.Week,
169            exp.Year,
170        },
171        exp.DataType.Type.VARCHAR: {
172            exp.ArrayConcat,
173            exp.Concat,
174            exp.ConcatWs,
175            exp.DateToDateStr,
176            exp.GroupConcat,
177            exp.Initcap,
178            exp.Lower,
179            exp.SafeConcat,
180            exp.Substring,
181            exp.TimeToStr,
182            exp.TimeToTimeStr,
183            exp.Trim,
184            exp.TsOrDsToDateStr,
185            exp.UnixToStr,
186            exp.UnixToTimeStr,
187            exp.Upper,
188        },
189    }
190
191    ANNOTATORS: t.Dict = {
192        **{
193            expr_type: lambda self, e: self._annotate_unary(e)
194            for expr_type in subclasses(exp.__name__, (exp.Unary, exp.Alias))
195        },
196        **{
197            expr_type: lambda self, e: self._annotate_binary(e)
198            for expr_type in subclasses(exp.__name__, exp.Binary)
199        },
200        **{
201            expr_type: _annotate_with_type_lambda(data_type)
202            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
203            for expr_type in expressions
204        },
205        exp.Anonymous: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.UNKNOWN),
206        exp.Array: lambda self, e: self._annotate_by_args(e, "expressions", array=True),
207        exp.ArrayAgg: lambda self, e: self._annotate_by_args(e, "this", array=True),
208        exp.ArrayConcat: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
209        exp.Cast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
210        exp.Case: lambda self, e: self._annotate_by_args(e, "default", "ifs"),
211        exp.Coalesce: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
212        exp.DataType: lambda self, e: self._annotate_with_type(e, e.copy()),
213        exp.Distinct: lambda self, e: self._annotate_by_args(e, "expressions"),
214        exp.Filter: lambda self, e: self._annotate_by_args(e, "this"),
215        exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
216        exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
217        exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
218        exp.Literal: lambda self, e: self._annotate_literal(e),
219        exp.Map: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
220        exp.Max: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
221        exp.Min: lambda self, e: self._annotate_by_args(e, "this", "expressions"),
222        exp.Null: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.NULL),
223        exp.Sum: lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True),
224        exp.TryCast: lambda self, e: self._annotate_with_type(e, e.args["to"]),
225        exp.VarMap: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.MAP),
226    }
227
228    NESTED_TYPES = {
229        exp.DataType.Type.ARRAY,
230    }
231
232    # Specifies what types a given type can be coerced into (autofilled)
233    COERCES_TO: t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]] = {}
234
235    def __init__(
236        self,
237        schema: Schema,
238        annotators: t.Optional[t.Dict[t.Type[E], t.Callable[[TypeAnnotator, E], E]]] = None,
239        coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
240    ) -> None:
241        self.schema = schema
242        self.annotators = annotators or self.ANNOTATORS
243        self.coerces_to = coerces_to or self.COERCES_TO
244
245    def annotate(self, expression: E) -> E:
246        for scope in traverse_scope(expression):
247            selects = {}
248            for name, source in scope.sources.items():
249                if not isinstance(source, Scope):
250                    continue
251                if isinstance(source.expression, exp.UDTF):
252                    values = []
253
254                    if isinstance(source.expression, exp.Lateral):
255                        if isinstance(source.expression.this, exp.Explode):
256                            values = [source.expression.this.this]
257                    else:
258                        values = source.expression.expressions[0].expressions
259
260                    if not values:
261                        continue
262
263                    selects[name] = {
264                        alias: column
265                        for alias, column in zip(
266                            source.expression.alias_column_names,
267                            values,
268                        )
269                    }
270                else:
271                    selects[name] = {
272                        select.alias_or_name: select for select in source.expression.selects
273                    }
274
275            # First annotate the current scope's column references
276            for col in scope.columns:
277                if not col.table:
278                    continue
279
280                source = scope.sources.get(col.table)
281                if isinstance(source, exp.Table):
282                    col.type = self.schema.get_column_type(source, col)
283                elif source and col.table in selects and col.name in selects[col.table]:
284                    col.type = selects[col.table][col.name].type
285
286            # Then (possibly) annotate the remaining expressions in the scope
287            self._maybe_annotate(scope.expression)
288
289        return self._maybe_annotate(expression)  # This takes care of non-traversable expressions
290
291    def _maybe_annotate(self, expression: E) -> E:
292        if expression.type:
293            return expression  # We've already inferred the expression's type
294
295        annotator = self.annotators.get(expression.__class__)
296
297        return (
298            annotator(self, expression)
299            if annotator
300            else self._annotate_with_type(expression, exp.DataType.Type.UNKNOWN)
301        )
302
303    def _annotate_args(self, expression: E) -> E:
304        for _, value in expression.iter_expressions():
305            self._maybe_annotate(value)
306
307        return expression
308
309    def _maybe_coerce(
310        self, type1: exp.DataType | exp.DataType.Type, type2: exp.DataType | exp.DataType.Type
311    ) -> exp.DataType | exp.DataType.Type:
312        type1_value = type1.this if isinstance(type1, exp.DataType) else type1
313        type2_value = type2.this if isinstance(type2, exp.DataType) else type2
314
315        # We propagate the NULL / UNKNOWN types upwards if found
316        if exp.DataType.Type.NULL in (type1_value, type2_value):
317            return exp.DataType.Type.NULL
318        if exp.DataType.Type.UNKNOWN in (type1_value, type2_value):
319            return exp.DataType.Type.UNKNOWN
320
321        if type1_value in self.NESTED_TYPES:
322            return type1
323        if type2_value in self.NESTED_TYPES:
324            return type2
325
326        return type2_value if type2_value in self.coerces_to.get(type1_value, {}) else type1_value  # type: ignore
327
328    # Note: the following "no_type_check" decorators were added because mypy was yelling due
329    # to assigning Type values to expression.type (since its getter returns Optional[DataType]).
330    # This is a known mypy issue: https://github.com/python/mypy/issues/3004
331
332    @t.no_type_check
333    def _annotate_binary(self, expression: B) -> B:
334        self._annotate_args(expression)
335
336        left_type = expression.left.type.this
337        right_type = expression.right.type.this
338
339        if isinstance(expression, exp.Connector):
340            if left_type == exp.DataType.Type.NULL and right_type == exp.DataType.Type.NULL:
341                expression.type = exp.DataType.Type.NULL
342            elif exp.DataType.Type.NULL in (left_type, right_type):
343                expression.type = exp.DataType.build(
344                    "NULLABLE", expressions=exp.DataType.build("BOOLEAN")
345                )
346            else:
347                expression.type = exp.DataType.Type.BOOLEAN
348        elif isinstance(expression, exp.Predicate):
349            expression.type = exp.DataType.Type.BOOLEAN
350        else:
351            expression.type = self._maybe_coerce(left_type, right_type)
352
353        return expression
354
355    @t.no_type_check
356    def _annotate_unary(self, expression: E) -> E:
357        self._annotate_args(expression)
358
359        if isinstance(expression, exp.Condition) and not isinstance(expression, exp.Paren):
360            expression.type = exp.DataType.Type.BOOLEAN
361        else:
362            expression.type = expression.this.type
363
364        return expression
365
366    @t.no_type_check
367    def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
368        if expression.is_string:
369            expression.type = exp.DataType.Type.VARCHAR
370        elif expression.is_int:
371            expression.type = exp.DataType.Type.INT
372        else:
373            expression.type = exp.DataType.Type.DOUBLE
374
375        return expression
376
377    @t.no_type_check
378    def _annotate_with_type(self, expression: E, target_type: exp.DataType.Type) -> E:
379        expression.type = target_type
380        return self._annotate_args(expression)
381
382    @t.no_type_check
383    def _annotate_by_args(
384        self, expression: E, *args: str, promote: bool = False, array: bool = False
385    ) -> E:
386        self._annotate_args(expression)
387
388        expressions: t.List[exp.Expression] = []
389        for arg in args:
390            arg_expr = expression.args.get(arg)
391            expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
392
393        last_datatype = None
394        for expr in expressions:
395            last_datatype = self._maybe_coerce(last_datatype or expr.type, expr.type)
396
397        expression.type = last_datatype or exp.DataType.Type.UNKNOWN
398
399        if promote:
400            if expression.type.this in exp.DataType.INTEGER_TYPES:
401                expression.type = exp.DataType.Type.BIGINT
402            elif expression.type.this in exp.DataType.FLOAT_TYPES:
403                expression.type = exp.DataType.Type.DOUBLE
404
405        if array:
406            expression.type = exp.DataType(
407                this=exp.DataType.Type.ARRAY, expressions=[expression.type], nested=True
408            )
409
410        return expression