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.ArrayCompact,
269            exp.ArrayConcat,
270            exp.ArrayConstructCompact,
271            exp.ArrayPrepend,
272            exp.ArrayRemove,
273            exp.ArraysZip,
274            exp.ArrayUniqueAgg,
275            exp.ArrayUnionAgg,
276            exp.MapKeys,
277            exp.RegexpExtractAll,
278            exp.Split,
279            exp.StringToArray,
280        )
281    },
282    **{
283        expr_type: {"returns": exp.DataType.Type.BIGINT}
284        for expr_type in {
285            exp.BitmapBitPosition,
286            exp.BitmapBucketNumber,
287            exp.BitmapCount,
288            exp.Factorial,
289            exp.GroupingId,
290            exp.MD5NumberLower64,
291            exp.MD5NumberUpper64,
292            exp.Rand,
293            exp.Seq8,
294            exp.Zipf,
295        }
296    },
297    **{
298        expr_type: {"returns": exp.DataType.Type.BINARY}
299        for expr_type in {
300            exp.Base64DecodeBinary,
301            exp.BitmapConstructAgg,
302            exp.BitmapOrAgg,
303            exp.Compress,
304            exp.DecompressBinary,
305            exp.Decrypt,
306            exp.DecryptRaw,
307            exp.Encrypt,
308            exp.EncryptRaw,
309            exp.HexString,
310            exp.MD5Digest,
311            exp.SHA1Digest,
312            exp.SHA2Digest,
313            exp.ToBinary,
314            exp.TryBase64DecodeBinary,
315            exp.TryHexDecodeBinary,
316            exp.Unhex,
317        }
318    },
319    **{
320        expr_type: {"returns": exp.DataType.Type.BOOLEAN}
321        for expr_type in {
322            exp.Booland,
323            exp.Boolnot,
324            exp.Boolor,
325            exp.BoolxorAgg,
326            exp.EqualNull,
327            exp.IsNullValue,
328            exp.MapContainsKey,
329            exp.Search,
330            exp.SearchIp,
331            exp.ToBoolean,
332        }
333    },
334    **{
335        expr_type: {"returns": exp.DataType.Type.DATE}
336        for expr_type in {
337            exp.NextDay,
338            exp.PreviousDay,
339        }
340    },
341    **{
342        expr_type: {
343            "annotator": lambda self, e: self._set_type(
344                e, exp.DataType.build("NUMBER", dialect="snowflake")
345            )
346        }
347        for expr_type in (
348            exp.BitwiseAndAgg,
349            exp.BitwiseOrAgg,
350            exp.BitwiseXorAgg,
351            exp.RegexpCount,
352            exp.RegexpInstr,
353            exp.ToNumber,
354        )
355    },
356    **{
357        expr_type: {"returns": exp.DataType.Type.DOUBLE}
358        for expr_type in {
359            exp.ApproxPercentileEstimate,
360            exp.ApproximateSimilarity,
361            exp.CosineDistance,
362            exp.CovarPop,
363            exp.CovarSamp,
364            exp.DotProduct,
365            exp.EuclideanDistance,
366            exp.ManhattanDistance,
367            exp.MonthsBetween,
368            exp.Normal,
369        }
370    },
371    exp.Kurtosis: {"annotator": _annotate_kurtosis},
372    **{
373        expr_type: {"returns": exp.DataType.Type.DECFLOAT}
374        for expr_type in {
375            exp.ToDecfloat,
376            exp.TryToDecfloat,
377        }
378    },
379    **{
380        expr_type: {"annotator": _annotate_math_with_float_decfloat}
381        for expr_type in {
382            exp.Acos,
383            exp.Asin,
384            exp.Atan,
385            exp.Atan2,
386            exp.Cbrt,
387            exp.Cos,
388            exp.Cot,
389            exp.Degrees,
390            exp.Exp,
391            exp.Ln,
392            exp.Log,
393            exp.Pow,
394            exp.Radians,
395            exp.RegrAvgx,
396            exp.RegrAvgy,
397            exp.RegrCount,
398            exp.RegrIntercept,
399            exp.RegrR2,
400            exp.RegrSlope,
401            exp.RegrSxx,
402            exp.RegrSxy,
403            exp.RegrSyy,
404            exp.RegrValx,
405            exp.RegrValy,
406            exp.Sin,
407            exp.Sqrt,
408            exp.Tan,
409            exp.Tanh,
410        }
411    },
412    **{
413        expr_type: {"returns": exp.DataType.Type.INT}
414        for expr_type in {
415            exp.ByteLength,
416            exp.Grouping,
417            exp.JarowinklerSimilarity,
418            exp.MapSize,
419            exp.Minute,
420            exp.RtrimmedLength,
421            exp.Second,
422            exp.Seq1,
423            exp.Seq2,
424            exp.Seq4,
425            exp.WidthBucket,
426        }
427    },
428    **{
429        expr_type: {"returns": exp.DataType.Type.OBJECT}
430        for expr_type in {
431            exp.ApproxPercentileAccumulate,
432            exp.ApproxPercentileCombine,
433            exp.ApproxTopKAccumulate,
434            exp.ApproxTopKCombine,
435            exp.ObjectAgg,
436            exp.ParseIp,
437            exp.ParseUrl,
438            exp.XMLGet,
439        }
440    },
441    **{
442        expr_type: {"returns": exp.DataType.Type.MAP}
443        for expr_type in {
444            exp.MapCat,
445            exp.MapDelete,
446            exp.MapInsert,
447            exp.MapPick,
448        }
449    },
450    **{
451        expr_type: {"returns": exp.DataType.Type.FILE}
452        for expr_type in {
453            exp.ToFile,
454        }
455    },
456    **{
457        expr_type: {"returns": exp.DataType.Type.TIME}
458        for expr_type in {
459            exp.TimeFromParts,
460            exp.TsOrDsToTime,
461        }
462    },
463    **{
464        expr_type: {"returns": exp.DataType.Type.TIMESTAMPLTZ}
465        for expr_type in {
466            exp.CurrentTimestamp,
467            exp.Localtimestamp,
468        }
469    },
470    **{
471        expr_type: {"returns": exp.DataType.Type.TINYINT}
472        for expr_type in {
473            exp.Quarter,
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.Collate,
487            exp.Collation,
488            exp.CurrentAccount,
489            exp.CurrentAccountName,
490            exp.CurrentAvailableRoles,
491            exp.CurrentClient,
492            exp.CurrentDatabase,
493            exp.CurrentIpAddress,
494            exp.CurrentSchemas,
495            exp.CurrentSecondaryRoles,
496            exp.CurrentSession,
497            exp.CurrentStatement,
498            exp.CurrentVersion,
499            exp.CurrentTransaction,
500            exp.CurrentWarehouse,
501            exp.CurrentOrganizationUser,
502            exp.CurrentRegion,
503            exp.CurrentRole,
504            exp.CurrentRoleType,
505            exp.CurrentOrganizationName,
506            exp.DecompressString,
507            exp.HexDecodeString,
508            exp.HexEncode,
509            exp.Monthname,
510            exp.Randstr,
511            exp.RegexpExtract,
512            exp.RegexpReplace,
513            exp.Repeat,
514            exp.Replace,
515            exp.Soundex,
516            exp.SoundexP123,
517            exp.SplitPart,
518            exp.Translate,
519            exp.TryBase64DecodeString,
520            exp.TryHexDecodeString,
521            exp.Uuid,
522        }
523    },
524    **{
525        expr_type: {"returns": exp.DataType.Type.VARIANT}
526        for expr_type in {
527            exp.Minhash,
528            exp.MinhashCombine,
529        }
530    },
531    **{
532        expr_type: {"annotator": _annotate_variance}
533        for expr_type in (
534            exp.Variance,
535            exp.VariancePop,
536        )
537    },
538    exp.ArgMax: {"annotator": _annotate_arg_max_min},
539    exp.ArgMin: {"annotator": _annotate_arg_max_min},
540    exp.ConcatWs: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
541    exp.ConvertTimezone: {
542        "annotator": lambda self, e: self._set_type(
543            e,
544            exp.DataType.Type.TIMESTAMPNTZ
545            if e.args.get("source_tz")
546            else exp.DataType.Type.TIMESTAMPTZ,
547        )
548    },
549    exp.DateAdd: {"annotator": _annotate_date_or_time_add},
550    exp.DecodeCase: {"annotator": _annotate_decode_case},
551    exp.HashAgg: {
552        "annotator": lambda self, e: self._set_type(
553            e, exp.DataType.build("NUMBER(19, 0)", dialect="snowflake")
554        )
555    },
556    exp.Median: {"annotator": _annotate_median},
557    exp.Reverse: {"annotator": _annotate_reverse},
558    exp.StrToTime: {"annotator": _annotate_str_to_time},
559    exp.TimeAdd: {"annotator": _annotate_date_or_time_add},
560    exp.TimestampFromParts: {"annotator": _annotate_timestamp_from_parts},
561    exp.WithinGroup: {"annotator": _annotate_within_group},
562}