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    struct_type = exp.DataType(
 86        this=exp.DType.STRUCT,
 87        expressions=[expression.this.type, exp.DataType(this=exp.DType.BIGINT)],
 88        nested=True,
 89    )
 90    self._set_type(
 91        expression,
 92        exp.DataType(this=exp.DType.ARRAY, expressions=[struct_type], nested=True),
 93    )
 94
 95    return expression
 96
 97
 98def _annotate_concat(self: TypeAnnotator, expression: exp.Concat) -> exp.Concat:
 99    annotated = self._annotate_by_args(expression, "expressions")
100
101    # Args must be BYTES or types that can be cast to STRING, return type is either BYTES or STRING
102    # https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#concat
103    if not annotated.is_type(exp.DType.BINARY, exp.DType.UNKNOWN):
104        self._set_type(annotated, exp.DType.VARCHAR)
105
106    return annotated
107
108
109def _annotate_array(self: TypeAnnotator, expression: exp.Array) -> exp.Array:
110    array_args = expression.expressions
111
112    # BigQuery behaves as follows:
113    #
114    # SELECT t, TYPEOF(t) FROM (SELECT 'foo') AS t            -- foo, STRUCT<STRING>
115    # SELECT ARRAY(SELECT 'foo'), TYPEOF(ARRAY(SELECT 'foo')) -- foo, ARRAY<STRING>
116    # ARRAY(SELECT ... UNION ALL SELECT ...)                  -- ARRAY<type from coerced projections>
117    # ARRAY(SELECT AS STRUCT 1 AS a, 'b' AS b)                -- ARRAY<STRUCT<INT64, STRING>>
118    if len(array_args) == 1:
119        unnested = array_args[0].unnest()
120        projection_type: exp.DataType | exp.DType | None = None
121
122        # Handle ARRAY(SELECT ...) - single SELECT query
123        if isinstance(unnested, exp.Select):
124            query_type = unnested.meta_get("query_type")
125
126            if query_type and query_type.is_type(exp.DType.STRUCT):
127                query_exprs = query_type.expressions
128
129                col_defs = [
130                    e
131                    for e in query_exprs
132                    if isinstance(e, exp.ColumnDef)
133                    and not (e.kind and e.kind.is_type(exp.DType.UNKNOWN))
134                ]
135
136                if len(col_defs) == len(query_exprs):
137                    if unnested.args.get("kind") == "STRUCT":
138                        # ARRAY(SELECT AS STRUCT ...) -> ARRAY<STRUCT<col1, col2, ...>>
139                        projection_type = query_type
140                    elif len(col_defs) == 1 and (col_type := col_defs[0].kind):
141                        # ARRAY(SELECT col FROM ...) -> ARRAY<col_type>
142                        projection_type = col_type
143
144        # Handle ARRAY(SELECT ... UNION ALL SELECT ...) - set operations
145        elif isinstance(unnested, exp.SetOperation):
146            # Get all column types for the SetOperation
147            col_types = self._get_setop_column_types(unnested)
148            # For ARRAY constructor, there should only be one projection
149            # https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/array_functions#array
150            if col_types and unnested.left.selects:
151                first_col_name = unnested.left.selects[0].alias_or_name
152                projection_type = col_types.get(first_col_name)
153
154        # If we successfully determine a projection type and it's not UNKNOWN, wrap it in ARRAY
155        if projection_type and not (
156            (
157                isinstance(projection_type, exp.DataType)
158                and projection_type.is_type(exp.DType.UNKNOWN)
159            )
160            or projection_type == exp.DType.UNKNOWN
161        ):
162            element_type = (
163                projection_type.copy()
164                if isinstance(projection_type, exp.DataType)
165                else exp.DataType(this=projection_type)
166            )
167            array_type = exp.DataType(
168                this=exp.DType.ARRAY,
169                expressions=[element_type],
170                nested=True,
171            )
172            self._set_type(expression, array_type)
173            return expression
174
175    return self._annotate_by_args(expression, "expressions", array=True)
176
177
178EXPRESSION_METADATA = {
179    **EXPRESSION_METADATA,
180    **{
181        expr_type: {"annotator": annotate_int_as_double}
182        for expr_type in {
183            exp.Avg,
184            exp.Ceil,
185            exp.Exp,
186            exp.Floor,
187            exp.Ln,
188            exp.Log,
189            exp.Round,
190            exp.Sqrt,
191        }
192    },
193    **{
194        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
195        for expr_type in {
196            exp.ArgMax,
197            exp.ArgMin,
198            exp.GroupConcat,
199            exp.IgnoreNulls,
200            exp.JSONExtract,
201            exp.Left,
202            exp.Lower,
203            exp.NetFunc,
204            exp.Pad,
205            exp.PercentileDisc,
206            exp.RegexpExtract,
207            exp.RegexpReplace,
208            exp.Repeat,
209            exp.Replace,
210            exp.RespectNulls,
211            exp.Reverse,
212            exp.Right,
213            exp.SafeFunc,
214            exp.SafeNegate,
215            exp.Sign,
216            exp.Substring,
217            exp.Translate,
218            exp.Trim,
219            exp.Upper,
220        }
221    },
222    **{
223        expr_type: {"annotator": lambda self, e: _annotate_date_func(self, e)}
224        for expr_type in _DATE_FUNC_LITERAL_TYPE
225    },
226    **{
227        expr_type: {"returns": exp.DType.BIGINT}
228        for expr_type in {
229            exp.BitwiseAndAgg,
230            exp.BitwiseCount,
231            exp.BitwiseOrAgg,
232            exp.BitwiseXorAgg,
233            exp.ByteLength,
234            exp.FarmFingerprint,
235            exp.Grouping,
236            exp.LaxInt64,
237            exp.Length,
238            exp.RangeBucket,
239            exp.RegexpInstr,
240            exp.UnixDate,
241        }
242    },
243    **{
244        expr_type: {"returns": exp.DType.BINARY}
245        for expr_type in {
246            exp.ByteString,
247            exp.CodePointsToBytes,
248            exp.MD5Digest,
249            exp.SHA,
250            exp.SHA2,
251            exp.SHA1Digest,
252            exp.SHA2Digest,
253            exp.Unhex,
254        }
255    },
256    **{
257        expr_type: {"returns": exp.DType.BOOLEAN}
258        for expr_type in {
259            exp.JSONBool,
260            exp.LaxBool,
261        }
262    },
263    **{
264        expr_type: {"returns": exp.DType.DATETIME}
265        for expr_type in {
266            exp.ParseDatetime,
267            exp.TimestampFromParts,
268        }
269    },
270    **{
271        expr_type: {"returns": exp.DType.DOUBLE}
272        for expr_type in {
273            exp.Atan2,
274            exp.Corr,
275            exp.CosineDistance,
276            exp.Coth,
277            exp.Csc,
278            exp.Csch,
279            exp.EuclideanDistance,
280            exp.Float64,
281            exp.LaxFloat64,
282            exp.Sec,
283            exp.Sech,
284        }
285    },
286    **{
287        expr_type: {"returns": exp.DType.JSON}
288        for expr_type in {
289            exp.JSONArray,
290            exp.JSONArrayAppend,
291            exp.JSONArrayInsert,
292            exp.JSONObject,
293            exp.JSONRemove,
294            exp.JSONSet,
295            exp.JSONStripNulls,
296        }
297    },
298    **{
299        expr_type: {"returns": exp.DType.TIME}
300        for expr_type in {
301            exp.ParseTime,
302            exp.TimeFromParts,
303            exp.TimeTrunc,
304            exp.TsOrDsToTime,
305        }
306    },
307    **{
308        expr_type: {"returns": exp.DType.VARCHAR}
309        for expr_type in {
310            exp.CodePointsToString,
311            exp.Format,
312            exp.Host,
313            exp.JSONExtractScalar,
314            exp.JSONType,
315            exp.LaxString,
316            exp.LowerHex,
317            exp.Normalize,
318            exp.RegDomain,
319            exp.SafeConvertBytesToString,
320            exp.Soundex,
321            exp.Uuid,
322        }
323    },
324    **{
325        expr_type: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)}
326        for expr_type in {
327            exp.PercentileCont,
328            exp.SafeAdd,
329            exp.SafeDivide,
330            exp.SafeMultiply,
331            exp.SafeSubtract,
332        }
333    },
334    **{
335        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)}
336        for expr_type in {
337            exp.ApproxQuantiles,
338            exp.JSONExtractArray,
339            exp.RegexpExtractAll,
340            exp.Split,
341        }
342    },
343    **{expr_type: {"returns": exp.DType.TIMESTAMPTZ} for expr_type in TIMESTAMP_EXPRESSIONS},
344    exp.ApproxTopK: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
345    exp.ApproxTopSum: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
346    exp.Array: {"annotator": _annotate_array},
347    exp.Concat: {"annotator": _annotate_concat},
348    exp.DateFromUnixDate: {"returns": exp.DType.DATE},
349    exp.GenerateTimestampArray: {
350        "annotator": lambda self, e: self._set_type(
351            e, exp.DataType.from_str("ARRAY<TIMESTAMP>", dialect="bigquery")
352        )
353    },
354    exp.JSONFormat: {
355        "annotator": lambda self, e: self._set_type(
356            e, exp.DType.JSON if e.args.get("to_json") else exp.DType.VARCHAR
357        )
358    },
359    exp.JSONKeysAtDepth: {
360        "annotator": lambda self, e: self._set_type(
361            e, exp.DataType.from_str("ARRAY<VARCHAR>", dialect="bigquery")
362        )
363    },
364    exp.JSONValueArray: {
365        "annotator": lambda self, e: self._set_type(
366            e, exp.DataType.from_str("ARRAY<VARCHAR>", dialect="bigquery")
367        )
368    },
369    exp.ParseBignumeric: {"returns": exp.DType.BIGDECIMAL},
370    exp.ParseNumeric: {"returns": exp.DType.DECIMAL},
371    exp.SafeDivide: {"annotator": lambda self, e: _annotate_safe_divide(self, e)},
372    exp.ToCodePoints: {
373        "annotator": lambda self, e: self._set_type(
374            e, exp.DataType.from_str("ARRAY<BIGINT>", dialect="bigquery")
375        )
376    },
377}