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