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