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.ArgMax,
167            exp.ArgMin,
168            exp.DateAdd,
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.NetFunc,
179            exp.NthValue,
180            exp.Pad,
181            exp.PercentileDisc,
182            exp.RegexpExtract,
183            exp.RegexpReplace,
184            exp.Repeat,
185            exp.Replace,
186            exp.RespectNulls,
187            exp.Reverse,
188            exp.Right,
189            exp.SafeFunc,
190            exp.SafeNegate,
191            exp.Sign,
192            exp.Substring,
193            exp.TimestampTrunc,
194            exp.Translate,
195            exp.Trim,
196            exp.Upper,
197        }
198    },
199    **{
200        expr_type: {"returns": exp.DataType.Type.BIGINT}
201        for expr_type in {
202            exp.BitwiseAndAgg,
203            exp.BitwiseCount,
204            exp.BitwiseOrAgg,
205            exp.BitwiseXorAgg,
206            exp.ByteLength,
207            exp.DenseRank,
208            exp.FarmFingerprint,
209            exp.Grouping,
210            exp.LaxInt64,
211            exp.Length,
212            exp.Ntile,
213            exp.Rank,
214            exp.RangeBucket,
215            exp.RegexpInstr,
216            exp.RowNumber,
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.JSONBool,
236            exp.LaxBool,
237        }
238    },
239    **{
240        expr_type: {"returns": exp.DataType.Type.DATETIME}
241        for expr_type in {
242            exp.ParseDatetime,
243            exp.TimestampFromParts,
244        }
245    },
246    **{
247        expr_type: {"returns": exp.DataType.Type.DOUBLE}
248        for expr_type in {
249            exp.Atan2,
250            exp.Corr,
251            exp.CosineDistance,
252            exp.Coth,
253            exp.CovarPop,
254            exp.CovarSamp,
255            exp.Csc,
256            exp.Csch,
257            exp.CumeDist,
258            exp.EuclideanDistance,
259            exp.Float64,
260            exp.LaxFloat64,
261            exp.PercentRank,
262            exp.Rand,
263            exp.Sec,
264            exp.Sech,
265        }
266    },
267    **{
268        expr_type: {"returns": exp.DataType.Type.JSON}
269        for expr_type in {
270            exp.JSONArray,
271            exp.JSONArrayAppend,
272            exp.JSONArrayInsert,
273            exp.JSONObject,
274            exp.JSONRemove,
275            exp.JSONSet,
276            exp.JSONStripNulls,
277        }
278    },
279    **{
280        expr_type: {"returns": exp.DataType.Type.TIME}
281        for expr_type in {
282            exp.ParseTime,
283            exp.TimeFromParts,
284            exp.TimeTrunc,
285            exp.TsOrDsToTime,
286        }
287    },
288    **{
289        expr_type: {"returns": exp.DataType.Type.VARCHAR}
290        for expr_type in {
291            exp.CodePointsToString,
292            exp.Format,
293            exp.Host,
294            exp.JSONExtractScalar,
295            exp.JSONType,
296            exp.LaxString,
297            exp.LowerHex,
298            exp.Normalize,
299            exp.RegDomain,
300            exp.SafeConvertBytesToString,
301            exp.Soundex,
302            exp.Uuid,
303        }
304    },
305    **{
306        expr_type: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)}
307        for expr_type in {
308            exp.PercentileCont,
309            exp.SafeAdd,
310            exp.SafeDivide,
311            exp.SafeMultiply,
312            exp.SafeSubtract,
313        }
314    },
315    **{
316        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)}
317        for expr_type in {
318            exp.ApproxQuantiles,
319            exp.JSONExtractArray,
320            exp.RegexpExtractAll,
321            exp.Split,
322        }
323    },
324    **{
325        expr_type: {"returns": exp.DataType.Type.TIMESTAMPTZ} for expr_type in TIMESTAMP_EXPRESSIONS
326    },
327    exp.ApproxTopK: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
328    exp.ApproxTopSum: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
329    exp.Array: {"annotator": _annotate_array},
330    exp.Concat: {"annotator": _annotate_concat},
331    exp.DateFromUnixDate: {"returns": exp.DataType.Type.DATE},
332    exp.GenerateTimestampArray: {
333        "annotator": lambda self, e: self._set_type(
334            e, exp.DataType.build("ARRAY<TIMESTAMP>", dialect="bigquery")
335        )
336    },
337    exp.JSONFormat: {
338        "annotator": lambda self, e: self._set_type(
339            e, exp.DataType.Type.JSON if e.args.get("to_json") else exp.DataType.Type.VARCHAR
340        )
341    },
342    exp.JSONKeysAtDepth: {
343        "annotator": lambda self, e: self._set_type(
344            e, exp.DataType.build("ARRAY<VARCHAR>", dialect="bigquery")
345        )
346    },
347    exp.JSONValueArray: {
348        "annotator": lambda self, e: self._set_type(
349            e, exp.DataType.build("ARRAY<VARCHAR>", dialect="bigquery")
350        )
351    },
352    exp.Lag: {"annotator": lambda self, e: self._annotate_by_args(e, "this", "default")},
353    exp.ParseBignumeric: {"returns": exp.DataType.Type.BIGDECIMAL},
354    exp.ParseNumeric: {"returns": exp.DataType.Type.DECIMAL},
355    exp.SafeDivide: {"annotator": lambda self, e: _annotate_safe_divide(self, e)},
356    exp.ToCodePoints: {
357        "annotator": lambda self, e: self._set_type(
358            e, exp.DataType.build("ARRAY<BIGINT>", dialect="bigquery")
359        )
360    },
361}