sqlglot.typing.snowflake

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.helper import seq_get
  7from sqlglot.typing import EXPRESSION_METADATA
  8
  9if t.TYPE_CHECKING:
 10    from sqlglot.optimizer.annotate_types import TypeAnnotator
 11
 12DATE_PARTS = {"DAY", "WEEK", "MONTH", "QUARTER", "YEAR"}
 13
 14MAX_PRECISION = 38
 15
 16MAX_SCALE = 37
 17
 18
 19def _annotate_reverse(self: TypeAnnotator, expression: exp.Reverse) -> exp.Reverse:
 20    expression = self._annotate_by_args(expression, "this")
 21    if expression.is_type(exp.DataType.Type.NULL):
 22        # Snowflake treats REVERSE(NULL) as a VARCHAR
 23        self._set_type(expression, exp.DataType.Type.VARCHAR)
 24
 25    return expression
 26
 27
 28def _annotate_timestamp_from_parts(
 29    self: TypeAnnotator, expression: exp.TimestampFromParts
 30) -> exp.TimestampFromParts:
 31    """Annotate TimestampFromParts with correct type based on arguments.
 32    TIMESTAMP_FROM_PARTS with time_zone -> TIMESTAMPTZ
 33    TIMESTAMP_FROM_PARTS without time_zone -> TIMESTAMP (defaults to TIMESTAMP_NTZ)
 34    """
 35    if expression.args.get("zone"):
 36        self._set_type(expression, exp.DataType.Type.TIMESTAMPTZ)
 37    else:
 38        self._set_type(expression, exp.DataType.Type.TIMESTAMP)
 39
 40    return expression
 41
 42
 43def _annotate_date_or_time_add(self: TypeAnnotator, expression: exp.Expression) -> exp.Expression:
 44    if (
 45        expression.this.is_type(exp.DataType.Type.DATE)
 46        and expression.text("unit").upper() not in DATE_PARTS
 47    ):
 48        self._set_type(expression, exp.DataType.Type.TIMESTAMPNTZ)
 49    else:
 50        self._annotate_by_args(expression, "this")
 51    return expression
 52
 53
 54def _annotate_decode_case(self: TypeAnnotator, expression: exp.DecodeCase) -> exp.DecodeCase:
 55    """Annotate DecodeCase with the type inferred from return values only.
 56
 57    DECODE uses the format: DECODE(expr, val1, ret1, val2, ret2, ..., default)
 58    We only look at the return values (ret1, ret2, ..., default) to determine the type,
 59    not the comparison values (val1, val2, ...) or the expression being compared.
 60    """
 61    expressions = expression.expressions
 62
 63    # Return values are at indices 2, 4, 6, ... and the last element (if even length)
 64    # DECODE(expr, val1, ret1, val2, ret2, ..., default)
 65    return_types = [expressions[i].type for i in range(2, len(expressions), 2)]
 66
 67    # If the total number of expressions is even, the last one is the default
 68    # Example:
 69    #   DECODE(x, 1, 'a', 2, 'b')             -> len=5 (odd), no default
 70    #   DECODE(x, 1, 'a', 2, 'b', 'default')  -> len=6 (even), has default
 71    if len(expressions) % 2 == 0:
 72        return_types.append(expressions[-1].type)
 73
 74    # Determine the common type from all return values
 75    last_type = None
 76    for ret_type in return_types:
 77        last_type = self._maybe_coerce(last_type or ret_type, ret_type)
 78
 79    self._set_type(expression, last_type)
 80    return expression
 81
 82
 83def _annotate_arg_max_min(self, expression):
 84    self._set_type(
 85        expression,
 86        exp.DataType.Type.ARRAY if expression.args.get("count") else expression.this.type,
 87    )
 88    return expression
 89
 90
 91def _annotate_within_group(self: TypeAnnotator, expression: exp.WithinGroup) -> exp.WithinGroup:
 92    """Annotate WithinGroup with correct type based on the inner function.
 93
 94    1) Annotate args first
 95    2) Check if this is PercentileDisc/PercentileCont and if so, re-annotate its type to match the ordered expression's type
 96    """
 97
 98    if (
 99        isinstance(expression.this, (exp.PercentileDisc, exp.PercentileCont))
100        and isinstance(order_expr := expression.expression, exp.Order)
101        and len(order_expr.expressions) == 1
102        and isinstance(ordered_expr := order_expr.expressions[0], exp.Ordered)
103    ):
104        self._set_type(expression, ordered_expr.this.type)
105
106    return expression
107
108
109def _annotate_median(self: TypeAnnotator, expression: exp.Median) -> exp.Median:
110    """Annotate MEDIAN function with correct return type.
111
112    Based on Snowflake documentation:
113    - If the expr is FLOAT/DOUBLE -> annotate as DOUBLE (FLOAT is a synonym for DOUBLE)
114    - If the expr is NUMBER(p, s) -> annotate as NUMBER(min(p+3, 38), min(s+3, 37))
115    """
116    # First annotate the argument to get its type
117    expression = self._annotate_by_args(expression, "this")
118
119    # Get the input type
120    input_type = expression.this.type
121
122    if input_type.is_type(exp.DataType.Type.DOUBLE):
123        # If input is FLOAT/DOUBLE, return DOUBLE (FLOAT is normalized to DOUBLE in Snowflake)
124        self._set_type(expression, exp.DataType.Type.DOUBLE)
125    else:
126        # If input is NUMBER(p, s), return NUMBER(min(p+3, 38), min(s+3, 37))
127        exprs = input_type.expressions
128
129        precision_expr = seq_get(exprs, 0)
130        precision = precision_expr.this.to_py() if precision_expr else MAX_PRECISION
131
132        scale_expr = seq_get(exprs, 1)
133        scale = scale_expr.this.to_py() if scale_expr else 0
134
135        new_precision = min(precision + 3, MAX_PRECISION)
136        new_scale = min(scale + 3, MAX_SCALE)
137
138        # Build the new NUMBER type
139        new_type = exp.DataType.build(f"NUMBER({new_precision}, {new_scale})", dialect="snowflake")
140        self._set_type(expression, new_type)
141
142    return expression
143
144
145def _annotate_variance(self: TypeAnnotator, expression: exp.Expression) -> exp.Expression:
146    """Annotate variance functions (VAR_POP, VAR_SAMP, VARIANCE, VARIANCE_POP) with correct return type.
147
148    Based on Snowflake behavior:
149    - DECFLOAT -> DECFLOAT(38)
150    - FLOAT/DOUBLE -> FLOAT
151    - INT, NUMBER(p, 0) -> NUMBER(38, 6)
152    - NUMBER(p, s) -> NUMBER(38, max(12, s))
153    """
154    # First annotate the argument to get its type
155    expression = self._annotate_by_args(expression, "this")
156
157    # Get the input type
158    input_type = expression.this.type
159
160    # Special case: DECFLOAT -> DECFLOAT(38)
161    if input_type.is_type(exp.DataType.Type.DECFLOAT):
162        self._set_type(expression, exp.DataType.build("DECFLOAT", dialect="snowflake"))
163    # Special case: FLOAT/DOUBLE -> DOUBLE
164    elif input_type.is_type(exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE):
165        self._set_type(expression, exp.DataType.Type.DOUBLE)
166    # For NUMBER types: determine the scale
167    else:
168        exprs = input_type.expressions
169        scale_expr = seq_get(exprs, 1)
170        scale = scale_expr.this.to_py() if scale_expr else 0
171
172        # If scale is 0 (INT, BIGINT, NUMBER(p,0)): return NUMBER(38, 6)
173        # Otherwise, Snowflake appears to assign scale through the formula MAX(12, s)
174        new_scale = 6 if scale == 0 else max(12, scale)
175
176        # Build the new NUMBER type
177        new_type = exp.DataType.build(f"NUMBER({MAX_PRECISION}, {new_scale})", dialect="snowflake")
178        self._set_type(expression, new_type)
179
180    return expression
181
182
183def _annotate_kurtosis(self: TypeAnnotator, expression: exp.Kurtosis) -> exp.Kurtosis:
184    """Annotate KURTOSIS with correct return type.
185
186    Based on Snowflake behavior:
187    - DECFLOAT input -> DECFLOAT
188    - DOUBLE or FLOAT input -> DOUBLE
189    - Other numeric types (INT, NUMBER) -> NUMBER(38, 12)
190    """
191    expression = self._annotate_by_args(expression, "this")
192    input_type = expression.this.type
193
194    if input_type.is_type(exp.DataType.Type.DECFLOAT):
195        self._set_type(expression, exp.DataType.build("DECFLOAT", dialect="snowflake"))
196    elif input_type.is_type(exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE):
197        self._set_type(expression, exp.DataType.Type.DOUBLE)
198    else:
199        self._set_type(
200            expression, exp.DataType.build(f"NUMBER({MAX_PRECISION}, 12)", dialect="snowflake")
201        )
202
203    return expression
204
205
206def _annotate_math_with_float_decfloat(
207    self: TypeAnnotator, expression: exp.Expression
208) -> exp.Expression:
209    """Annotate math functions that preserve  DECFLOAT but return DOUBLE for others.
210
211    In Snowflake, trigonometric and exponential math functions:
212    - If input is DECFLOAT -> return DECFLOAT
213    - For integer types (INT, BIGINT, etc.) -> return DOUBLE
214    - For other numeric types (NUMBER, DECIMAL, DOUBLE) -> return DOUBLE
215    """
216    expression = self._annotate_by_args(expression, "this")
217
218    # If input is DECFLOAT, preserve
219    if expression.this.is_type(exp.DataType.Type.DECFLOAT):
220        self._set_type(expression, expression.this.type)
221    else:
222        # For all other types (integers, decimals, etc.), return DOUBLE
223        self._set_type(expression, exp.DataType.Type.DOUBLE)
224
225    return expression
226
227
228def _annotate_str_to_time(self: TypeAnnotator, expression: exp.StrToTime) -> exp.StrToTime:
229    # target_type is stored as a DataType instance
230    target_type_arg = expression.args.get("target_type")
231    target_type = (
232        target_type_arg.this
233        if isinstance(target_type_arg, exp.DataType)
234        else exp.DataType.Type.TIMESTAMP
235    )
236    self._set_type(expression, target_type)
237    return expression
238
239
240EXPRESSION_METADATA = {
241    **EXPRESSION_METADATA,
242    **{
243        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
244        for expr_type in {
245            exp.AddMonths,
246            exp.Ceil,
247            exp.DateTrunc,
248            exp.Floor,
249            exp.Left,
250            exp.Mode,
251            exp.Pad,
252            exp.Right,
253            exp.Round,
254            exp.Stuff,
255            exp.Substring,
256            exp.TimeSlice,
257            exp.TimestampTrunc,
258        }
259    },
260    **{
261        expr_type: {"returns": exp.DataType.Type.ARRAY}
262        for expr_type in (
263            exp.ApproxTopK,
264            exp.ApproxTopKEstimate,
265            exp.Array,
266            exp.ArrayAgg,
267            exp.ArrayAppend,
268            exp.ArrayConcat,
269            exp.ArrayConstructCompact,
270            exp.ArrayPrepend,
271            exp.ArrayRemove,
272            exp.ArrayUniqueAgg,
273            exp.ArrayUnionAgg,
274            exp.MapKeys,
275            exp.RegexpExtractAll,
276            exp.Split,
277            exp.StringToArray,
278        )
279    },
280    **{
281        expr_type: {"returns": exp.DataType.Type.BIGINT}
282        for expr_type in {
283            exp.BitmapBitPosition,
284            exp.BitmapBucketNumber,
285            exp.BitmapCount,
286            exp.Factorial,
287            exp.GroupingId,
288            exp.MD5NumberLower64,
289            exp.MD5NumberUpper64,
290            exp.Rand,
291            exp.Zipf,
292        }
293    },
294    **{
295        expr_type: {"returns": exp.DataType.Type.BINARY}
296        for expr_type in {
297            exp.Base64DecodeBinary,
298            exp.BitmapConstructAgg,
299            exp.BitmapOrAgg,
300            exp.Compress,
301            exp.DecompressBinary,
302            exp.Decrypt,
303            exp.DecryptRaw,
304            exp.Encrypt,
305            exp.EncryptRaw,
306            exp.HexString,
307            exp.MD5Digest,
308            exp.SHA1Digest,
309            exp.SHA2Digest,
310            exp.ToBinary,
311            exp.TryBase64DecodeBinary,
312            exp.TryHexDecodeBinary,
313            exp.Unhex,
314        }
315    },
316    **{
317        expr_type: {"returns": exp.DataType.Type.BOOLEAN}
318        for expr_type in {
319            exp.Booland,
320            exp.Boolnot,
321            exp.Boolor,
322            exp.BoolxorAgg,
323            exp.EqualNull,
324            exp.IsNullValue,
325            exp.MapContainsKey,
326            exp.Search,
327            exp.SearchIp,
328            exp.ToBoolean,
329        }
330    },
331    **{
332        expr_type: {"returns": exp.DataType.Type.DATE}
333        for expr_type in {
334            exp.NextDay,
335            exp.PreviousDay,
336        }
337    },
338    **{
339        expr_type: {
340            "annotator": lambda self, e: self._set_type(
341                e, exp.DataType.build("NUMBER", dialect="snowflake")
342            )
343        }
344        for expr_type in (
345            exp.BitwiseAndAgg,
346            exp.BitwiseOrAgg,
347            exp.BitwiseXorAgg,
348            exp.RegexpCount,
349            exp.RegexpInstr,
350            exp.ToNumber,
351        )
352    },
353    **{
354        expr_type: {"returns": exp.DataType.Type.DOUBLE}
355        for expr_type in {
356            exp.ApproxPercentileEstimate,
357            exp.ApproximateSimilarity,
358            exp.Asinh,
359            exp.Atanh,
360            exp.Cbrt,
361            exp.Cosh,
362            exp.CosineDistance,
363            exp.CovarPop,
364            exp.CovarSamp,
365            exp.DotProduct,
366            exp.EuclideanDistance,
367            exp.ManhattanDistance,
368            exp.MonthsBetween,
369            exp.Normal,
370            exp.Sinh,
371        }
372    },
373    exp.Kurtosis: {"annotator": _annotate_kurtosis},
374    **{
375        expr_type: {"returns": exp.DataType.Type.DECFLOAT}
376        for expr_type in {
377            exp.ToDecfloat,
378            exp.TryToDecfloat,
379        }
380    },
381    **{
382        expr_type: {"annotator": _annotate_math_with_float_decfloat}
383        for expr_type in {
384            exp.Acos,
385            exp.Asin,
386            exp.Atan,
387            exp.Atan2,
388            exp.Cos,
389            exp.Cot,
390            exp.Degrees,
391            exp.Exp,
392            exp.Ln,
393            exp.Log,
394            exp.Pow,
395            exp.Radians,
396            exp.RegrAvgx,
397            exp.RegrAvgy,
398            exp.RegrCount,
399            exp.RegrIntercept,
400            exp.RegrR2,
401            exp.RegrSlope,
402            exp.RegrSxx,
403            exp.RegrSxy,
404            exp.RegrSyy,
405            exp.RegrValx,
406            exp.RegrValy,
407            exp.Sin,
408            exp.Sqrt,
409            exp.Tan,
410            exp.Tanh,
411        }
412    },
413    **{
414        expr_type: {"returns": exp.DataType.Type.INT}
415        for expr_type in {
416            exp.Ascii,
417            exp.BitLength,
418            exp.ByteLength,
419            exp.Getbit,
420            exp.Grouping,
421            exp.Hour,
422            exp.JarowinklerSimilarity,
423            exp.Length,
424            exp.Levenshtein,
425            exp.MapSize,
426            exp.Minute,
427            exp.RtrimmedLength,
428            exp.Second,
429            exp.StrPosition,
430            exp.Unicode,
431            exp.WidthBucket,
432        }
433    },
434    **{
435        expr_type: {"returns": exp.DataType.Type.OBJECT}
436        for expr_type in {
437            exp.ApproxPercentileAccumulate,
438            exp.ApproxPercentileCombine,
439            exp.ApproxTopKAccumulate,
440            exp.ApproxTopKCombine,
441            exp.ObjectAgg,
442            exp.ParseIp,
443            exp.ParseUrl,
444            exp.XMLGet,
445        }
446    },
447    **{
448        expr_type: {"returns": exp.DataType.Type.MAP}
449        for expr_type in {
450            exp.MapCat,
451            exp.MapDelete,
452            exp.MapInsert,
453            exp.MapPick,
454        }
455    },
456    **{
457        expr_type: {"returns": exp.DataType.Type.FILE}
458        for expr_type in {
459            exp.ToFile,
460        }
461    },
462    **{
463        expr_type: {"returns": exp.DataType.Type.TIME}
464        for expr_type in {
465            exp.TimeFromParts,
466            exp.TsOrDsToTime,
467        }
468    },
469    **{
470        expr_type: {"returns": exp.DataType.Type.TIMESTAMPLTZ}
471        for expr_type in {
472            exp.CurrentTimestamp,
473            exp.Localtimestamp,
474        }
475    },
476    **{
477        expr_type: {"returns": exp.DataType.Type.VARCHAR}
478        for expr_type in {
479            exp.AIAgg,
480            exp.AIClassify,
481            exp.AISummarizeAgg,
482            exp.Base64DecodeString,
483            exp.Base64Encode,
484            exp.CheckJson,
485            exp.CheckXml,
486            exp.Chr,
487            exp.Collate,
488            exp.Collation,
489            exp.CurrentAccount,
490            exp.CurrentAccountName,
491            exp.CurrentAvailableRoles,
492            exp.CurrentClient,
493            exp.CurrentDatabase,
494            exp.CurrentIpAddress,
495            exp.CurrentSchemas,
496            exp.CurrentSecondaryRoles,
497            exp.CurrentSession,
498            exp.CurrentStatement,
499            exp.CurrentVersion,
500            exp.CurrentTransaction,
501            exp.CurrentWarehouse,
502            exp.CurrentOrganizationUser,
503            exp.CurrentRegion,
504            exp.CurrentRole,
505            exp.CurrentRoleType,
506            exp.CurrentOrganizationName,
507            exp.DecompressString,
508            exp.HexDecodeString,
509            exp.HexEncode,
510            exp.MD5,
511            exp.Monthname,
512            exp.Randstr,
513            exp.RegexpExtract,
514            exp.RegexpReplace,
515            exp.Repeat,
516            exp.Replace,
517            exp.SHA,
518            exp.SHA2,
519            exp.Soundex,
520            exp.SoundexP123,
521            exp.Space,
522            exp.SplitPart,
523            exp.Translate,
524            exp.TryBase64DecodeString,
525            exp.TryHexDecodeString,
526            exp.Uuid,
527        }
528    },
529    **{
530        expr_type: {"returns": exp.DataType.Type.VARIANT}
531        for expr_type in {
532            exp.Minhash,
533            exp.MinhashCombine,
534        }
535    },
536    **{
537        expr_type: {"annotator": _annotate_variance}
538        for expr_type in (
539            exp.Variance,
540            exp.VariancePop,
541        )
542    },
543    exp.ArgMax: {"annotator": _annotate_arg_max_min},
544    exp.ArgMin: {"annotator": _annotate_arg_max_min},
545    exp.ConcatWs: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
546    exp.ConvertTimezone: {
547        "annotator": lambda self, e: self._set_type(
548            e,
549            exp.DataType.Type.TIMESTAMPNTZ
550            if e.args.get("source_tz")
551            else exp.DataType.Type.TIMESTAMPTZ,
552        )
553    },
554    exp.DateAdd: {"annotator": _annotate_date_or_time_add},
555    exp.DecodeCase: {"annotator": _annotate_decode_case},
556    exp.HashAgg: {
557        "annotator": lambda self, e: self._set_type(
558            e, exp.DataType.build("NUMBER(19, 0)", dialect="snowflake")
559        )
560    },
561    exp.Median: {"annotator": _annotate_median},
562    exp.Reverse: {"annotator": _annotate_reverse},
563    exp.StrToTime: {"annotator": _annotate_str_to_time},
564    exp.TimeAdd: {"annotator": _annotate_date_or_time_add},
565    exp.TimestampFromParts: {"annotator": _annotate_timestamp_from_parts},
566    exp.WithinGroup: {"annotator": _annotate_within_group},
567}