sqlglot.typing.bigquery

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.typing import EXPRESSION_METADATA, TIMESTAMP_EXPRESSIONS, annotate_int_as_double
  7
  8if t.TYPE_CHECKING:
  9    from sqlglot.optimizer.annotate_types import TypeAnnotator
 10
 11
 12# DATE_ADD / DATE_SUB / *_TRUNC return the type of their first argument. BigQuery
 13# implicitly casts a string literal first arg to the function's own temporal type,
 14# so map each to that type (e.g. DATE_ADD('2020-01-01', ...) -> DATE,
 15# TIMESTAMP_TRUNC('...') -> TIMESTAMP).
 16_DateFunc = t.Union[
 17    exp.DateAdd,
 18    exp.DateSub,
 19    exp.DateTrunc,
 20    exp.DatetimeTrunc,
 21    exp.TimestampTrunc,
 22]
 23
 24_DATE_FUNC_LITERAL_TYPE: dict[type[_DateFunc], exp.DType] = {
 25    exp.DateAdd: exp.DType.DATE,
 26    exp.DateSub: exp.DType.DATE,
 27    exp.DateTrunc: exp.DType.DATE,
 28    exp.DatetimeTrunc: exp.DType.DATETIME,
 29    exp.TimestampTrunc: exp.DType.TIMESTAMPTZ,
 30}
 31
 32
 33def _annotate_date_func(self: TypeAnnotator, expression: _DateFunc) -> exp.Expr:
 34    """Annotate DATE_ADD / DATE_SUB / *_TRUNC, which return their first arg's type.
 35
 36    A typed first argument keeps its exact type (e.g. DATE_ADD(DATETIME, ...) ->
 37    DATETIME). For a string literal first argument, BigQuery implicitly casts it to
 38    the function's own temporal type, so the result is that type (e.g.
 39    DATE_ADD('2020-01-01', INTERVAL 1 DAY) -> DATE).
 40    """
 41    this = expression.this
 42
 43    # BigQuery rejects expressions like DATE_ADD(c, ...); it requires the first argument to be a literal
 44    if isinstance(this, exp.Literal) and this.is_string:
 45        return self._set_type(expression, _DATE_FUNC_LITERAL_TYPE[type(expression)])
 46
 47    return self._annotate_by_args(expression, "this")
 48
 49
 50def _annotate_safe_divide(self: TypeAnnotator, expression: exp.SafeDivide) -> exp.Expr:
 51    """
 52    +------------+------------+------------+-------------+---------+
 53    | INPUT      | INT64      | NUMERIC    | BIGNUMERIC  | FLOAT64 |
 54    +------------+------------+------------+-------------+---------+
 55    | INT64      | FLOAT64    | NUMERIC    | BIGNUMERIC  | FLOAT64 |
 56    | NUMERIC    | NUMERIC    | NUMERIC    | BIGNUMERIC  | FLOAT64 |
 57    | BIGNUMERIC | BIGNUMERIC | BIGNUMERIC | BIGNUMERIC  | FLOAT64 |
 58    | FLOAT64    | FLOAT64    | FLOAT64    | FLOAT64     | FLOAT64 |
 59    +------------+------------+------------+-------------+---------+
 60    """
 61    if expression.this.is_type(*exp.DataType.INTEGER_TYPES) and expression.expression.is_type(
 62        *exp.DataType.INTEGER_TYPES
 63    ):
 64        return self._set_type(expression, exp.DType.DOUBLE)
 65
 66    return _annotate_by_args_with_coerce(self, expression)
 67
 68
 69def _annotate_by_args_with_coerce(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
 70    """
 71    +------------+------------+------------+-------------+---------+
 72    | INPUT      | INT64      | NUMERIC    | BIGNUMERIC  | FLOAT64 |
 73    +------------+------------+------------+-------------+---------+
 74    | INT64      | INT64      | NUMERIC    | BIGNUMERIC  | FLOAT64 |
 75    | NUMERIC    | NUMERIC    | NUMERIC    | BIGNUMERIC  | FLOAT64 |
 76    | BIGNUMERIC | BIGNUMERIC | BIGNUMERIC | BIGNUMERIC  | FLOAT64 |
 77    | FLOAT64    | FLOAT64    | FLOAT64    | FLOAT64     | FLOAT64 |
 78    +------------+------------+------------+-------------+---------+
 79    """
 80    self._set_type(expression, self._maybe_coerce(expression.this.type, expression.expression.type))
 81    return expression
 82
 83
 84def _annotate_by_args_approx_top(self: TypeAnnotator, expression: exp.ApproxTopK) -> exp.ApproxTopK:
 85    if not expression.this.type or expression.this.is_type(exp.DType.UNKNOWN):
 86        self._set_type(expression, exp.DType.UNKNOWN)
 87        return expression
 88
 89    struct_type = exp.DataType(
 90        this=exp.DType.STRUCT,
 91        expressions=[expression.this.type, exp.DataType(this=exp.DType.BIGINT)],
 92        nested=True,
 93    )
 94    self._set_type(
 95        expression,
 96        exp.DataType(this=exp.DType.ARRAY, expressions=[struct_type], nested=True),
 97    )
 98
 99    return expression
100
101
102def _annotate_concat(self: TypeAnnotator, expression: exp.Concat) -> exp.Concat:
103    annotated = self._annotate_by_args(expression, "expressions")
104
105    # Args must be BYTES or types that can be cast to STRING, return type is either BYTES or STRING
106    # https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#concat
107    if not annotated.is_type(exp.DType.BINARY, exp.DType.UNKNOWN):
108        self._set_type(annotated, exp.DType.VARCHAR)
109
110    return annotated
111
112
113def _all_struct_kind(node: exp.Expr) -> bool:
114    # True only if every leaf SELECT of this (possibly nested) UNION is SELECT AS STRUCT.
115    # Iterative, not recursive, since a long UNION chain can exceed the recursion limit.
116    stack = [node]
117    while stack:
118        n = stack.pop()
119        if isinstance(n, exp.SetOperation):
120            stack.append(n.this.unnest())
121            stack.append(n.expression.unnest())
122        elif not (isinstance(n, exp.Select) and n.args.get("kind") == "STRUCT"):
123            return False
124
125    return True
126
127
128def _annotate_array(self: TypeAnnotator, expression: exp.Array) -> exp.Array:
129    array_args = expression.expressions
130
131    # BigQuery behaves as follows:
132    #
133    # SELECT t, TYPEOF(t) FROM (SELECT 'foo') AS t            -- foo, STRUCT<STRING>
134    # SELECT ARRAY(SELECT 'foo'), TYPEOF(ARRAY(SELECT 'foo')) -- foo, ARRAY<STRING>
135    # ARRAY(SELECT ... UNION ALL SELECT ...)                  -- ARRAY<type from coerced projections>
136    # ARRAY(SELECT AS STRUCT 1 AS a, 'b' AS b)                -- ARRAY<STRUCT<INT64, STRING>>
137    if len(array_args) == 1:
138        unnested = array_args[0].unnest()
139        projection_type: exp.DataType | exp.DType | None = None
140
141        # Handle ARRAY(SELECT ...) - single SELECT query
142        if isinstance(unnested, exp.Select):
143            query_type = unnested.meta_get("query_type")
144
145            if query_type and query_type.is_type(exp.DType.STRUCT):
146                query_exprs = query_type.expressions
147
148                col_defs = [
149                    e
150                    for e in query_exprs
151                    if isinstance(e, exp.ColumnDef)
152                    and not (e.kind and e.kind.is_type(exp.DType.UNKNOWN))
153                ]
154
155                if len(col_defs) == len(query_exprs):
156                    if unnested.args.get("kind") == "STRUCT":
157                        # ARRAY(SELECT AS STRUCT ...) -> ARRAY<STRUCT<col1, col2, ...>>
158                        projection_type = query_type
159                    elif len(col_defs) == 1 and (col_type := col_defs[0].kind):
160                        # ARRAY(SELECT col FROM ...) -> ARRAY<col_type>
161                        projection_type = col_type
162
163        # Handle ARRAY(SELECT ... UNION ALL SELECT ...) - set operations
164        elif isinstance(unnested, exp.SetOperation):
165            col_types = self._get_setop_column_types(unnested)
166
167            if col_types and _all_struct_kind(unnested):
168                # Each branch's row is itself the struct value, not N separate columns
169                field_types = [(name, exp.DataType.build(kind)) for name, kind in col_types]
170
171                if not any(kind.is_type(exp.DType.UNKNOWN) for _, kind in field_types):
172                    projection_type = exp.DataType(
173                        this=exp.DType.STRUCT,
174                        expressions=[
175                            exp.ColumnDef(this=exp.to_identifier(name), kind=kind)
176                            for name, kind in field_types
177                        ],
178                        nested=True,
179                    )
180            # For ARRAY constructor, there should only be one projection
181            # https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/array_functions#array
182            elif len(col_types) == 1:
183                projection_type = col_types[0][1]
184
185        # If we successfully determine a projection type and it's not UNKNOWN, wrap it in ARRAY
186        if projection_type and not (
187            (
188                isinstance(projection_type, exp.DataType)
189                and projection_type.is_type(exp.DType.UNKNOWN)
190            )
191            or projection_type == exp.DType.UNKNOWN
192        ):
193            element_type = (
194                projection_type.copy()
195                if isinstance(projection_type, exp.DataType)
196                else exp.DataType(this=projection_type)
197            )
198            array_type = exp.DataType(
199                this=exp.DType.ARRAY,
200                expressions=[element_type],
201                nested=True,
202            )
203            self._set_type(expression, array_type)
204            return expression
205
206    return self._annotate_by_args(expression, "expressions", array=True)
207
208
209EXPRESSION_METADATA = {
210    **EXPRESSION_METADATA,
211    **{
212        expr_type: {"annotator": annotate_int_as_double}
213        for expr_type in {
214            exp.Avg,
215            exp.Ceil,
216            exp.Exp,
217            exp.Floor,
218            exp.Ln,
219            exp.Log,
220            exp.Round,
221            exp.Sqrt,
222        }
223    },
224    **{
225        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
226        for expr_type in {
227            exp.ArgMax,
228            exp.ArgMin,
229            exp.GroupConcat,
230            exp.IgnoreNulls,
231            exp.JSONExtract,
232            exp.Left,
233            exp.Lower,
234            exp.NetFunc,
235            exp.Pad,
236            exp.PercentileDisc,
237            exp.RegexpExtract,
238            exp.RegexpReplace,
239            exp.Repeat,
240            exp.Replace,
241            exp.RespectNulls,
242            exp.Reverse,
243            exp.Right,
244            exp.SafeFunc,
245            exp.SafeNegate,
246            exp.Sign,
247            exp.Substring,
248            exp.Translate,
249            exp.Trim,
250            exp.Upper,
251        }
252    },
253    **{
254        expr_type: {"annotator": lambda self, e: _annotate_date_func(self, e)}
255        for expr_type in _DATE_FUNC_LITERAL_TYPE
256    },
257    **{
258        expr_type: {"returns": exp.DType.BIGINT}
259        for expr_type in {
260            exp.BitwiseAndAgg,
261            exp.BitwiseCount,
262            exp.BitwiseOrAgg,
263            exp.BitwiseXorAgg,
264            exp.ByteLength,
265            exp.FarmFingerprint,
266            exp.Grouping,
267            exp.LaxInt64,
268            exp.Length,
269            exp.RangeBucket,
270            exp.RegexpInstr,
271            exp.UnixDate,
272        }
273    },
274    **{
275        expr_type: {"returns": exp.DType.BINARY}
276        for expr_type in {
277            exp.ByteString,
278            exp.CodePointsToBytes,
279            exp.MD5Digest,
280            exp.SHA,
281            exp.SHA2,
282            exp.SHA1Digest,
283            exp.SHA2Digest,
284            exp.Unhex,
285        }
286    },
287    **{
288        expr_type: {"returns": exp.DType.BOOLEAN}
289        for expr_type in {
290            exp.JSONBool,
291            exp.LaxBool,
292        }
293    },
294    **{
295        expr_type: {"returns": exp.DType.DATETIME}
296        for expr_type in {
297            exp.ParseDatetime,
298            exp.TimestampFromParts,
299        }
300    },
301    **{
302        expr_type: {"returns": exp.DType.DOUBLE}
303        for expr_type in {
304            exp.Atan2,
305            exp.Corr,
306            exp.CosineDistance,
307            exp.Coth,
308            exp.Csc,
309            exp.Csch,
310            exp.EuclideanDistance,
311            exp.Float64,
312            exp.LaxFloat64,
313            exp.Sec,
314            exp.Sech,
315        }
316    },
317    **{
318        expr_type: {"returns": exp.DType.JSON}
319        for expr_type in {
320            exp.JSONArray,
321            exp.JSONArrayAppend,
322            exp.JSONArrayInsert,
323            exp.JSONObject,
324            exp.JSONRemove,
325            exp.JSONSet,
326            exp.JSONStripNulls,
327        }
328    },
329    **{
330        expr_type: {"returns": exp.DType.TIME}
331        for expr_type in {
332            exp.ParseTime,
333            exp.TimeFromParts,
334            exp.TimeTrunc,
335            exp.TsOrDsToTime,
336        }
337    },
338    **{
339        expr_type: {"returns": exp.DType.VARCHAR}
340        for expr_type in {
341            exp.CodePointsToString,
342            exp.Format,
343            exp.Host,
344            exp.JSONExtractScalar,
345            exp.JSONType,
346            exp.LaxString,
347            exp.LowerHex,
348            exp.Normalize,
349            exp.RegDomain,
350            exp.SafeConvertBytesToString,
351            exp.Soundex,
352            exp.Uuid,
353        }
354    },
355    **{
356        expr_type: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)}
357        for expr_type in {
358            exp.PercentileCont,
359            exp.SafeAdd,
360            exp.SafeDivide,
361            exp.SafeMultiply,
362            exp.SafeSubtract,
363        }
364    },
365    **{
366        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)}
367        for expr_type in {
368            exp.ApproxQuantiles,
369            exp.JSONExtractArray,
370            exp.RegexpExtractAll,
371            exp.Split,
372        }
373    },
374    **{expr_type: {"returns": exp.DType.TIMESTAMPTZ} for expr_type in TIMESTAMP_EXPRESSIONS},
375    exp.ApproxTopK: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
376    exp.ApproxTopSum: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
377    exp.Array: {"annotator": _annotate_array},
378    exp.Concat: {"annotator": _annotate_concat},
379    exp.DateFromUnixDate: {"returns": exp.DType.DATE},
380    exp.GenerateTimestampArray: {
381        "annotator": lambda self, e: self._set_type(
382            e, exp.DataType.from_str("ARRAY<TIMESTAMP>", dialect="bigquery")
383        )
384    },
385    exp.JSONFormat: {
386        "annotator": lambda self, e: self._set_type(
387            e, exp.DType.JSON if e.args.get("to_json") else exp.DType.VARCHAR
388        )
389    },
390    exp.JSONKeysAtDepth: {
391        "annotator": lambda self, e: self._set_type(
392            e, exp.DataType.from_str("ARRAY<VARCHAR>", dialect="bigquery")
393        )
394    },
395    exp.JSONValueArray: {
396        "annotator": lambda self, e: self._set_type(
397            e, exp.DataType.from_str("ARRAY<VARCHAR>", dialect="bigquery")
398        )
399    },
400    exp.ParseBignumeric: {"returns": exp.DType.BIGDECIMAL},
401    exp.ParseNumeric: {"returns": exp.DType.DECIMAL},
402    exp.SafeDivide: {"annotator": lambda self, e: _annotate_safe_divide(self, e)},
403    exp.ToCodePoints: {
404        "annotator": lambda self, e: self._set_type(
405            e, exp.DataType.from_str("ARRAY<BIGINT>", dialect="bigquery")
406        )
407    },
408}