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.DType.NULL):
 22        # Snowflake treats REVERSE(NULL) as a VARCHAR
 23        self._set_type(expression, exp.DType.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.DType.TIMESTAMPTZ)
 37    else:
 38        self._set_type(expression, exp.DType.TIMESTAMP)
 39
 40    return expression
 41
 42
 43def _annotate_date_or_time_add(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
 44    if (
 45        expression.this.is_type(exp.DType.DATE)
 46        and expression.text("unit").upper() not in DATE_PARTS
 47    ):
 48        self._set_type(expression, exp.DType.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.DType.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    else:
106        self._set_type(expression, expression.this.type)
107
108    return expression
109
110
111def _annotate_median(self: TypeAnnotator, expression: exp.Median) -> exp.Median:
112    """Annotate MEDIAN function with correct return type.
113
114    Based on Snowflake documentation:
115    - If the expr is FLOAT/DOUBLE -> annotate as DOUBLE (FLOAT is a synonym for DOUBLE)
116    - If the expr is NUMBER(p, s) -> annotate as NUMBER(min(p+3, 38), min(s+3, 37))
117    """
118    # First annotate the argument to get its type
119    expression = self._annotate_by_args(expression, "this")
120
121    # Get the input type
122    input_type = expression.this.type
123
124    if input_type.is_type(exp.DType.DOUBLE):
125        # If input is FLOAT/DOUBLE, return DOUBLE (FLOAT is normalized to DOUBLE in Snowflake)
126        self._set_type(expression, exp.DType.DOUBLE)
127    else:
128        # If input is NUMBER(p, s), return NUMBER(min(p+3, 38), min(s+3, 37))
129        exprs = input_type.expressions
130
131        precision_expr = seq_get(exprs, 0)
132        precision = precision_expr.this.to_py() if precision_expr else MAX_PRECISION
133
134        scale_expr = seq_get(exprs, 1)
135        scale = scale_expr.this.to_py() if scale_expr else 0
136
137        new_precision = min(precision + 3, MAX_PRECISION)
138        new_scale = min(scale + 3, MAX_SCALE)
139
140        # Build the new NUMBER type
141        new_type = exp.DataType.from_str(
142            f"NUMBER({new_precision}, {new_scale})", dialect="snowflake"
143        )
144        self._set_type(expression, new_type)
145
146    return expression
147
148
149def _annotate_variance(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
150    """Annotate variance functions (VAR_POP, VAR_SAMP, VARIANCE, VARIANCE_POP) with correct return type.
151
152    Based on Snowflake behavior:
153    - DECFLOAT -> DECFLOAT(38)
154    - FLOAT/DOUBLE -> FLOAT
155    - INT, NUMBER(p, 0) -> NUMBER(38, 6)
156    - NUMBER(p, s) -> NUMBER(38, max(12, s))
157    """
158    # First annotate the argument to get its type
159    expression = self._annotate_by_args(expression, "this")
160
161    # Get the input type
162    input_type = expression.this.type
163
164    # Special case: DECFLOAT -> DECFLOAT(38)
165    if input_type.is_type(exp.DType.DECFLOAT):
166        self._set_type(expression, exp.DataType.from_str("DECFLOAT", dialect="snowflake"))
167    # Special case: FLOAT/DOUBLE -> DOUBLE
168    elif input_type.is_type(exp.DType.FLOAT, exp.DType.DOUBLE):
169        self._set_type(expression, exp.DType.DOUBLE)
170    # For NUMBER types: determine the scale
171    else:
172        exprs = input_type.expressions
173        scale_expr = seq_get(exprs, 1)
174        scale = scale_expr.this.to_py() if scale_expr else 0
175
176        # If scale is 0 (INT, BIGINT, NUMBER(p,0)): return NUMBER(38, 6)
177        # Otherwise, Snowflake appears to assign scale through the formula MAX(12, s)
178        new_scale = 6 if scale == 0 else max(12, scale)
179
180        # Build the new NUMBER type
181        new_type = exp.DataType.from_str(
182            f"NUMBER({MAX_PRECISION}, {new_scale})", dialect="snowflake"
183        )
184        self._set_type(expression, new_type)
185
186    return expression
187
188
189def _annotate_kurtosis(self: TypeAnnotator, expression: exp.Kurtosis) -> exp.Kurtosis:
190    """Annotate KURTOSIS with correct return type.
191
192    Based on Snowflake behavior:
193    - DECFLOAT input -> DECFLOAT
194    - DOUBLE or FLOAT input -> DOUBLE
195    - Other numeric types (INT, NUMBER) -> NUMBER(38, 12)
196    """
197    expression = self._annotate_by_args(expression, "this")
198    input_type = expression.this.type
199
200    if input_type.is_type(exp.DType.DECFLOAT):
201        self._set_type(expression, exp.DataType.from_str("DECFLOAT", dialect="snowflake"))
202    elif input_type.is_type(exp.DType.FLOAT, exp.DType.DOUBLE):
203        self._set_type(expression, exp.DType.DOUBLE)
204    else:
205        self._set_type(
206            expression, exp.DataType.from_str(f"NUMBER({MAX_PRECISION}, 12)", dialect="snowflake")
207        )
208
209    return expression
210
211
212def _annotate_math_with_float_decfloat(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
213    """Annotate math functions that preserve  DECFLOAT but return DOUBLE for others.
214
215    In Snowflake, trigonometric and exponential math functions:
216    - If input is DECFLOAT -> return DECFLOAT
217    - For integer types (INT, BIGINT, etc.) -> return DOUBLE
218    - For other numeric types (NUMBER, DECIMAL, DOUBLE) -> return DOUBLE
219    """
220    expression = self._annotate_by_args(expression, "this")
221
222    # If input is DECFLOAT, preserve
223    if expression.this.is_type(exp.DType.DECFLOAT):
224        self._set_type(expression, expression.this.type)
225    else:
226        # For all other types (integers, decimals, etc.), return DOUBLE
227        self._set_type(expression, exp.DType.DOUBLE)
228
229    return expression
230
231
232def _annotate_str_to_time(self: TypeAnnotator, expression: exp.StrToTime) -> exp.StrToTime:
233    # target_type is stored as a DataType instance
234    target_type_arg = expression.args.get("target_type")
235    target_type = (
236        target_type_arg.this if isinstance(target_type_arg, exp.DataType) else exp.DType.TIMESTAMP
237    )
238    self._set_type(expression, target_type)
239    return expression
240
241
242EXPRESSION_METADATA = {
243    **EXPRESSION_METADATA,
244    **{
245        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
246        for expr_type in {
247            exp.AddMonths,
248            exp.Ceil,
249            exp.DateTrunc,
250            exp.Floor,
251            exp.Left,
252            exp.Mode,
253            exp.Pad,
254            exp.Right,
255            exp.Round,
256            exp.Stuff,
257            exp.Substring,
258            exp.TimeSlice,
259            exp.TimestampTrunc,
260        }
261    },
262    **{
263        expr_type: {"returns": exp.DType.ARRAY}
264        for expr_type in (
265            exp.ApproxTopK,
266            exp.ApproxTopKEstimate,
267            exp.Array,
268            exp.ArrayAgg,
269            exp.ArrayAppend,
270            exp.ArrayCompact,
271            exp.ArrayConcat,
272            exp.ArrayConstructCompact,
273            exp.ArrayPrepend,
274            exp.ArrayRemove,
275            exp.ArraysZip,
276            exp.ArrayUniqueAgg,
277            exp.ArrayUnionAgg,
278            exp.MapKeys,
279            exp.RegexpExtractAll,
280            exp.Split,
281            exp.StringToArray,
282            exp.StrtokToArray,
283        )
284    },
285    **{
286        expr_type: {"returns": exp.DType.BIGINT}
287        for expr_type in {
288            exp.BitmapBitPosition,
289            exp.BitmapBucketNumber,
290            exp.BitmapCount,
291            exp.Factorial,
292            exp.GroupingId,
293            exp.MD5NumberLower64,
294            exp.MD5NumberUpper64,
295            exp.Rand,
296            exp.Seq8,
297            exp.Zipf,
298        }
299    },
300    **{
301        expr_type: {"returns": exp.DType.BINARY}
302        for expr_type in {
303            exp.Base64DecodeBinary,
304            exp.BitmapConstructAgg,
305            exp.BitmapOrAgg,
306            exp.Compress,
307            exp.DecompressBinary,
308            exp.Decrypt,
309            exp.DecryptRaw,
310            exp.Encrypt,
311            exp.EncryptRaw,
312            exp.HexString,
313            exp.MD5Digest,
314            exp.SHA1Digest,
315            exp.SHA2Digest,
316            exp.ToBinary,
317            exp.TryBase64DecodeBinary,
318            exp.TryHexDecodeBinary,
319            exp.Unhex,
320        }
321    },
322    **{
323        expr_type: {"returns": exp.DType.BOOLEAN}
324        for expr_type in {
325            exp.Booland,
326            exp.Boolnot,
327            exp.Boolor,
328            exp.BoolxorAgg,
329            exp.EqualNull,
330            exp.IsNullValue,
331            exp.MapContainsKey,
332            exp.Search,
333            exp.SearchIp,
334            exp.ToBoolean,
335        }
336    },
337    **{
338        expr_type: {"returns": exp.DType.DATE}
339        for expr_type in {
340            exp.NextDay,
341            exp.PreviousDay,
342        }
343    },
344    **{
345        expr_type: {
346            "annotator": lambda self, e: self._set_type(
347                e, exp.DataType.from_str("NUMBER", dialect="snowflake")
348            )
349        }
350        for expr_type in (
351            exp.BitwiseAndAgg,
352            exp.BitwiseOrAgg,
353            exp.BitwiseXorAgg,
354            exp.RegexpCount,
355            exp.RegexpInstr,
356            exp.ToNumber,
357        )
358    },
359    **{
360        expr_type: {"returns": exp.DType.DOUBLE}
361        for expr_type in {
362            exp.ApproxPercentileEstimate,
363            exp.ApproximateSimilarity,
364            exp.CosineDistance,
365            exp.CovarPop,
366            exp.CovarSamp,
367            exp.DotProduct,
368            exp.EuclideanDistance,
369            exp.ManhattanDistance,
370            exp.MonthsBetween,
371            exp.Normal,
372        }
373    },
374    exp.Kurtosis: {"annotator": _annotate_kurtosis},
375    **{
376        expr_type: {"returns": exp.DType.DECFLOAT}
377        for expr_type in {
378            exp.ToDecfloat,
379            exp.TryToDecfloat,
380        }
381    },
382    **{
383        expr_type: {"annotator": _annotate_math_with_float_decfloat}
384        for expr_type in {
385            exp.Acos,
386            exp.Asin,
387            exp.Atan,
388            exp.Atan2,
389            exp.Cbrt,
390            exp.Cos,
391            exp.Cot,
392            exp.Degrees,
393            exp.Exp,
394            exp.Ln,
395            exp.Log,
396            exp.Pow,
397            exp.Radians,
398            exp.RegrAvgx,
399            exp.RegrAvgy,
400            exp.RegrCount,
401            exp.RegrIntercept,
402            exp.RegrR2,
403            exp.RegrSlope,
404            exp.RegrSxx,
405            exp.RegrSxy,
406            exp.RegrSyy,
407            exp.RegrValx,
408            exp.RegrValy,
409            exp.Sin,
410            exp.Sqrt,
411            exp.Tan,
412            exp.Tanh,
413        }
414    },
415    **{
416        expr_type: {"returns": exp.DType.INT}
417        for expr_type in {
418            exp.ByteLength,
419            exp.DenseRank,
420            exp.Grouping,
421            exp.JarowinklerSimilarity,
422            exp.MapSize,
423            exp.Minute,
424            exp.Ntile,
425            exp.Rank,
426            exp.RowNumber,
427            exp.RtrimmedLength,
428            exp.Second,
429            exp.Seq1,
430            exp.Seq2,
431            exp.Seq4,
432            exp.WidthBucket,
433        }
434    },
435    **{
436        expr_type: {"returns": exp.DType.OBJECT}
437        for expr_type in {
438            exp.ApproxPercentileAccumulate,
439            exp.ApproxPercentileCombine,
440            exp.ApproxTopKAccumulate,
441            exp.ApproxTopKCombine,
442            exp.ObjectAgg,
443            exp.ParseIp,
444            exp.ParseUrl,
445            exp.XMLGet,
446        }
447    },
448    **{
449        expr_type: {"returns": exp.DType.MAP}
450        for expr_type in {
451            exp.MapCat,
452            exp.MapDelete,
453            exp.MapInsert,
454            exp.MapPick,
455        }
456    },
457    **{
458        expr_type: {"returns": exp.DType.FILE}
459        for expr_type in {
460            exp.ToFile,
461        }
462    },
463    **{
464        expr_type: {"returns": exp.DType.TIME}
465        for expr_type in {
466            exp.TimeFromParts,
467            exp.TsOrDsToTime,
468        }
469    },
470    **{
471        expr_type: {"returns": exp.DType.TIMESTAMPLTZ}
472        for expr_type in {
473            exp.CurrentTimestamp,
474            exp.Localtimestamp,
475        }
476    },
477    **{
478        expr_type: {"returns": exp.DType.TINYINT}
479        for expr_type in {
480            exp.DayOfMonth,
481            exp.DayOfWeek,
482            exp.DayOfYear,
483            exp.Quarter,
484        }
485    },
486    **{
487        expr_type: {"returns": exp.DType.VARCHAR}
488        for expr_type in {
489            exp.AIAgg,
490            exp.AIClassify,
491            exp.AISummarizeAgg,
492            exp.Base64DecodeString,
493            exp.Base64Encode,
494            exp.CheckJson,
495            exp.CheckXml,
496            exp.Collate,
497            exp.Collation,
498            exp.CurrentAccount,
499            exp.CurrentAccountName,
500            exp.CurrentAvailableRoles,
501            exp.CurrentClient,
502            exp.CurrentDatabase,
503            exp.CurrentIpAddress,
504            exp.CurrentSchemas,
505            exp.CurrentSecondaryRoles,
506            exp.CurrentSession,
507            exp.CurrentStatement,
508            exp.CurrentTransaction,
509            exp.CurrentWarehouse,
510            exp.CurrentOrganizationUser,
511            exp.CurrentRegion,
512            exp.CurrentRole,
513            exp.CurrentRoleType,
514            exp.CurrentOrganizationName,
515            exp.DecompressString,
516            exp.HexDecodeString,
517            exp.Hex,
518            exp.Randstr,
519            exp.RegexpExtract,
520            exp.RegexpReplace,
521            exp.Repeat,
522            exp.Replace,
523            exp.Soundex,
524            exp.SoundexP123,
525            exp.SplitPart,
526            exp.Strtok,
527            exp.TryBase64DecodeString,
528            exp.TryHexDecodeString,
529            exp.Uuid,
530        }
531    },
532    **{
533        expr_type: {"returns": exp.DType.VARIANT}
534        for expr_type in {
535            exp.Minhash,
536            exp.MinhashCombine,
537        }
538    },
539    **{
540        expr_type: {"annotator": _annotate_variance}
541        for expr_type in (
542            exp.Variance,
543            exp.VariancePop,
544        )
545    },
546    exp.ArgMax: {"annotator": _annotate_arg_max_min},
547    exp.ArgMin: {"annotator": _annotate_arg_max_min},
548    exp.ConcatWs: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
549    exp.ConvertTimezone: {
550        "annotator": lambda self, e: self._set_type(
551            e,
552            exp.DType.TIMESTAMPNTZ if e.args.get("source_tz") else exp.DType.TIMESTAMPTZ,
553        )
554    },
555    exp.DateAdd: {"annotator": _annotate_date_or_time_add},
556    exp.DecodeCase: {"annotator": _annotate_decode_case},
557    exp.HashAgg: {
558        "annotator": lambda self, e: self._set_type(
559            e, exp.DataType.from_str("NUMBER(19, 0)", dialect="snowflake")
560        )
561    },
562    exp.Median: {"annotator": _annotate_median},
563    exp.Reverse: {"annotator": _annotate_reverse},
564    exp.StrToTime: {"annotator": _annotate_str_to_time},
565    exp.TimeAdd: {"annotator": _annotate_date_or_time_add},
566    exp.TimestampFromParts: {"annotator": _annotate_timestamp_from_parts},
567    exp.WithinGroup: {"annotator": _annotate_within_group},
568}