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