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.build(f"NUMBER({new_precision}, {new_scale})", dialect="snowflake")
142        self._set_type(expression, new_type)
143
144    return expression
145
146
147def _annotate_variance(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
148    """Annotate variance functions (VAR_POP, VAR_SAMP, VARIANCE, VARIANCE_POP) with correct return type.
149
150    Based on Snowflake behavior:
151    - DECFLOAT -> DECFLOAT(38)
152    - FLOAT/DOUBLE -> FLOAT
153    - INT, NUMBER(p, 0) -> NUMBER(38, 6)
154    - NUMBER(p, s) -> NUMBER(38, max(12, s))
155    """
156    # First annotate the argument to get its type
157    expression = self._annotate_by_args(expression, "this")
158
159    # Get the input type
160    input_type = expression.this.type
161
162    # Special case: DECFLOAT -> DECFLOAT(38)
163    if input_type.is_type(exp.DType.DECFLOAT):
164        self._set_type(expression, exp.DataType.build("DECFLOAT", dialect="snowflake"))
165    # Special case: FLOAT/DOUBLE -> DOUBLE
166    elif input_type.is_type(exp.DType.FLOAT, exp.DType.DOUBLE):
167        self._set_type(expression, exp.DType.DOUBLE)
168    # For NUMBER types: determine the scale
169    else:
170        exprs = input_type.expressions
171        scale_expr = seq_get(exprs, 1)
172        scale = scale_expr.this.to_py() if scale_expr else 0
173
174        # If scale is 0 (INT, BIGINT, NUMBER(p,0)): return NUMBER(38, 6)
175        # Otherwise, Snowflake appears to assign scale through the formula MAX(12, s)
176        new_scale = 6 if scale == 0 else max(12, scale)
177
178        # Build the new NUMBER type
179        new_type = exp.DataType.build(f"NUMBER({MAX_PRECISION}, {new_scale})", dialect="snowflake")
180        self._set_type(expression, new_type)
181
182    return expression
183
184
185def _annotate_kurtosis(self: TypeAnnotator, expression: exp.Kurtosis) -> exp.Kurtosis:
186    """Annotate KURTOSIS with correct return type.
187
188    Based on Snowflake behavior:
189    - DECFLOAT input -> DECFLOAT
190    - DOUBLE or FLOAT input -> DOUBLE
191    - Other numeric types (INT, NUMBER) -> NUMBER(38, 12)
192    """
193    expression = self._annotate_by_args(expression, "this")
194    input_type = expression.this.type
195
196    if input_type.is_type(exp.DType.DECFLOAT):
197        self._set_type(expression, exp.DataType.build("DECFLOAT", dialect="snowflake"))
198    elif input_type.is_type(exp.DType.FLOAT, exp.DType.DOUBLE):
199        self._set_type(expression, exp.DType.DOUBLE)
200    else:
201        self._set_type(
202            expression, exp.DataType.build(f"NUMBER({MAX_PRECISION}, 12)", dialect="snowflake")
203        )
204
205    return expression
206
207
208def _annotate_math_with_float_decfloat(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
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.DType.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.DType.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 if isinstance(target_type_arg, exp.DataType) else exp.DType.TIMESTAMP
233    )
234    self._set_type(expression, target_type)
235    return expression
236
237
238EXPRESSION_METADATA = {
239    **EXPRESSION_METADATA,
240    **{
241        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
242        for expr_type in {
243            exp.AddMonths,
244            exp.Ceil,
245            exp.DateTrunc,
246            exp.Floor,
247            exp.Left,
248            exp.Mode,
249            exp.Pad,
250            exp.Right,
251            exp.Round,
252            exp.Stuff,
253            exp.Substring,
254            exp.TimeSlice,
255            exp.TimestampTrunc,
256        }
257    },
258    **{
259        expr_type: {"returns": exp.DType.ARRAY}
260        for expr_type in (
261            exp.ApproxTopK,
262            exp.ApproxTopKEstimate,
263            exp.Array,
264            exp.ArrayAgg,
265            exp.ArrayAppend,
266            exp.ArrayCompact,
267            exp.ArrayConcat,
268            exp.ArrayConstructCompact,
269            exp.ArrayPrepend,
270            exp.ArrayRemove,
271            exp.ArraysZip,
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.DType.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.Seq8,
292            exp.Zipf,
293        }
294    },
295    **{
296        expr_type: {"returns": exp.DType.BINARY}
297        for expr_type in {
298            exp.Base64DecodeBinary,
299            exp.BitmapConstructAgg,
300            exp.BitmapOrAgg,
301            exp.Compress,
302            exp.DecompressBinary,
303            exp.Decrypt,
304            exp.DecryptRaw,
305            exp.Encrypt,
306            exp.EncryptRaw,
307            exp.HexString,
308            exp.MD5Digest,
309            exp.SHA1Digest,
310            exp.SHA2Digest,
311            exp.ToBinary,
312            exp.TryBase64DecodeBinary,
313            exp.TryHexDecodeBinary,
314            exp.Unhex,
315        }
316    },
317    **{
318        expr_type: {"returns": exp.DType.BOOLEAN}
319        for expr_type in {
320            exp.Booland,
321            exp.Boolnot,
322            exp.Boolor,
323            exp.BoolxorAgg,
324            exp.EqualNull,
325            exp.IsNullValue,
326            exp.MapContainsKey,
327            exp.Search,
328            exp.SearchIp,
329            exp.ToBoolean,
330        }
331    },
332    **{
333        expr_type: {"returns": exp.DType.DATE}
334        for expr_type in {
335            exp.NextDay,
336            exp.PreviousDay,
337        }
338    },
339    **{
340        expr_type: {
341            "annotator": lambda self, e: self._set_type(
342                e, exp.DataType.build("NUMBER", dialect="snowflake")
343            )
344        }
345        for expr_type in (
346            exp.BitwiseAndAgg,
347            exp.BitwiseOrAgg,
348            exp.BitwiseXorAgg,
349            exp.RegexpCount,
350            exp.RegexpInstr,
351            exp.ToNumber,
352        )
353    },
354    **{
355        expr_type: {"returns": exp.DType.DOUBLE}
356        for expr_type in {
357            exp.ApproxPercentileEstimate,
358            exp.ApproximateSimilarity,
359            exp.CosineDistance,
360            exp.CovarPop,
361            exp.CovarSamp,
362            exp.DotProduct,
363            exp.EuclideanDistance,
364            exp.ManhattanDistance,
365            exp.MonthsBetween,
366            exp.Normal,
367        }
368    },
369    exp.Kurtosis: {"annotator": _annotate_kurtosis},
370    **{
371        expr_type: {"returns": exp.DType.DECFLOAT}
372        for expr_type in {
373            exp.ToDecfloat,
374            exp.TryToDecfloat,
375        }
376    },
377    **{
378        expr_type: {"annotator": _annotate_math_with_float_decfloat}
379        for expr_type in {
380            exp.Acos,
381            exp.Asin,
382            exp.Atan,
383            exp.Atan2,
384            exp.Cbrt,
385            exp.Cos,
386            exp.Cot,
387            exp.Degrees,
388            exp.Exp,
389            exp.Ln,
390            exp.Log,
391            exp.Pow,
392            exp.Radians,
393            exp.RegrAvgx,
394            exp.RegrAvgy,
395            exp.RegrCount,
396            exp.RegrIntercept,
397            exp.RegrR2,
398            exp.RegrSlope,
399            exp.RegrSxx,
400            exp.RegrSxy,
401            exp.RegrSyy,
402            exp.RegrValx,
403            exp.RegrValy,
404            exp.Sin,
405            exp.Sqrt,
406            exp.Tan,
407            exp.Tanh,
408        }
409    },
410    **{
411        expr_type: {"returns": exp.DType.INT}
412        for expr_type in {
413            exp.ByteLength,
414            exp.Grouping,
415            exp.JarowinklerSimilarity,
416            exp.MapSize,
417            exp.Minute,
418            exp.RtrimmedLength,
419            exp.Second,
420            exp.Seq1,
421            exp.Seq2,
422            exp.Seq4,
423            exp.WidthBucket,
424        }
425    },
426    **{
427        expr_type: {"returns": exp.DType.OBJECT}
428        for expr_type in {
429            exp.ApproxPercentileAccumulate,
430            exp.ApproxPercentileCombine,
431            exp.ApproxTopKAccumulate,
432            exp.ApproxTopKCombine,
433            exp.ObjectAgg,
434            exp.ParseIp,
435            exp.ParseUrl,
436            exp.XMLGet,
437        }
438    },
439    **{
440        expr_type: {"returns": exp.DType.MAP}
441        for expr_type in {
442            exp.MapCat,
443            exp.MapDelete,
444            exp.MapInsert,
445            exp.MapPick,
446        }
447    },
448    **{
449        expr_type: {"returns": exp.DType.FILE}
450        for expr_type in {
451            exp.ToFile,
452        }
453    },
454    **{
455        expr_type: {"returns": exp.DType.TIME}
456        for expr_type in {
457            exp.TimeFromParts,
458            exp.TsOrDsToTime,
459        }
460    },
461    **{
462        expr_type: {"returns": exp.DType.TIMESTAMPLTZ}
463        for expr_type in {
464            exp.CurrentTimestamp,
465            exp.Localtimestamp,
466        }
467    },
468    **{
469        expr_type: {"returns": exp.DType.TINYINT}
470        for expr_type in {
471            exp.DayOfMonth,
472            exp.DayOfWeek,
473            exp.DayOfYear,
474            exp.Quarter,
475        }
476    },
477    **{
478        expr_type: {"returns": exp.DType.VARCHAR}
479        for expr_type in {
480            exp.AIAgg,
481            exp.AIClassify,
482            exp.AISummarizeAgg,
483            exp.Base64DecodeString,
484            exp.Base64Encode,
485            exp.CheckJson,
486            exp.CheckXml,
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.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.Randstr,
510            exp.RegexpExtract,
511            exp.RegexpReplace,
512            exp.Repeat,
513            exp.Replace,
514            exp.Soundex,
515            exp.SoundexP123,
516            exp.SplitPart,
517            exp.Strtok,
518            exp.TryBase64DecodeString,
519            exp.TryHexDecodeString,
520            exp.Uuid,
521        }
522    },
523    **{
524        expr_type: {"returns": exp.DType.VARIANT}
525        for expr_type in {
526            exp.Minhash,
527            exp.MinhashCombine,
528        }
529    },
530    **{
531        expr_type: {"annotator": _annotate_variance}
532        for expr_type in (
533            exp.Variance,
534            exp.VariancePop,
535        )
536    },
537    exp.ArgMax: {"annotator": _annotate_arg_max_min},
538    exp.ArgMin: {"annotator": _annotate_arg_max_min},
539    exp.ConcatWs: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
540    exp.ConvertTimezone: {
541        "annotator": lambda self, e: self._set_type(
542            e,
543            exp.DType.TIMESTAMPNTZ if e.args.get("source_tz") else exp.DType.TIMESTAMPTZ,
544        )
545    },
546    exp.DateAdd: {"annotator": _annotate_date_or_time_add},
547    exp.DecodeCase: {"annotator": _annotate_decode_case},
548    exp.HashAgg: {
549        "annotator": lambda self, e: self._set_type(
550            e, exp.DataType.build("NUMBER(19, 0)", dialect="snowflake")
551        )
552    },
553    exp.Median: {"annotator": _annotate_median},
554    exp.Reverse: {"annotator": _annotate_reverse},
555    exp.StrToTime: {"annotator": _annotate_str_to_time},
556    exp.TimeAdd: {"annotator": _annotate_date_or_time_add},
557    exp.TimestampFromParts: {"annotator": _annotate_timestamp_from_parts},
558    exp.WithinGroup: {"annotator": _annotate_within_group},
559}