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, annotate_by_numeric_arg
  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
242def _annotate_floor_ceil(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
243    if expression.this.is_type(*exp.DataType.TEXT_TYPES):
244        # A scale argument changes Snowflake's implicit string coercion from FLOAT to NUMBER.
245        return self._set_type(
246            expression,
247            exp.DType.DOUBLE if expression.args.get("decimals") is None else exp.DType.UNKNOWN,
248        )
249    if expression.this.is_type(exp.DType.VARIANT):
250        return self._set_type(expression, exp.DType.DOUBLE)
251    return annotate_by_numeric_arg(self, expression)
252
253
254EXPRESSION_METADATA = {
255    **EXPRESSION_METADATA,
256    **{expr_type: {"annotator": _annotate_floor_ceil} for expr_type in {exp.Ceil, exp.Floor}},
257    exp.Sign: {
258        "annotator": lambda self, e: self._set_type(
259            e, exp.DType.DOUBLE if e.this.is_type(*exp.DataType.FLOAT_TYPES) else exp.DType.INT
260        )
261    },
262    **{
263        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
264        for expr_type in {
265            exp.AddMonths,
266            exp.DateTrunc,
267            exp.Left,
268            exp.Mode,
269            exp.Pad,
270            exp.Right,
271            exp.Round,
272            exp.Stuff,
273            exp.Substring,
274            exp.TimeSlice,
275            exp.TimestampTrunc,
276        }
277    },
278    **{
279        expr_type: {"returns": exp.DType.ARRAY}
280        for expr_type in (
281            exp.ApproxTopK,
282            exp.ApproxTopKEstimate,
283            exp.Array,
284            exp.ArrayAgg,
285            exp.ArrayAppend,
286            exp.ArrayCompact,
287            exp.ArrayConcat,
288            exp.ArrayConstructCompact,
289            exp.ArrayPrepend,
290            exp.ArrayRemove,
291            exp.ArraysZip,
292            exp.ArrayUniqueAgg,
293            exp.ArrayUnionAgg,
294            exp.MapKeys,
295            exp.RegexpExtractAll,
296            exp.Split,
297            exp.StringToArray,
298            exp.StrtokToArray,
299        )
300    },
301    **{
302        expr_type: {"returns": exp.DType.BIGINT}
303        for expr_type in {
304            exp.BitmapBitPosition,
305            exp.BitmapBucketNumber,
306            exp.BitmapCount,
307            exp.Factorial,
308            exp.GroupingId,
309            exp.MD5NumberLower64,
310            exp.MD5NumberUpper64,
311            exp.Rand,
312            exp.Seq8,
313            exp.Zipf,
314        }
315    },
316    **{
317        expr_type: {"returns": exp.DType.BINARY}
318        for expr_type in {
319            exp.Base64DecodeBinary,
320            exp.BitmapConstructAgg,
321            exp.BitmapOrAgg,
322            exp.Compress,
323            exp.DecompressBinary,
324            exp.Decrypt,
325            exp.DecryptRaw,
326            exp.Encrypt,
327            exp.EncryptRaw,
328            exp.HexString,
329            exp.MD5Digest,
330            exp.SHA1Digest,
331            exp.SHA2Digest,
332            exp.ToBinary,
333            exp.TryBase64DecodeBinary,
334            exp.TryHexDecodeBinary,
335            exp.Unhex,
336        }
337    },
338    **{
339        expr_type: {"returns": exp.DType.BOOLEAN}
340        for expr_type in {
341            exp.Booland,
342            exp.Boolnot,
343            exp.Boolor,
344            exp.BoolxorAgg,
345            exp.EqualNull,
346            exp.IsNullValue,
347            exp.MapContainsKey,
348            exp.Search,
349            exp.SearchIp,
350            exp.ToBoolean,
351        }
352    },
353    **{
354        expr_type: {"returns": exp.DType.DATE}
355        for expr_type in {
356            exp.NextDay,
357            exp.PreviousDay,
358        }
359    },
360    **{
361        expr_type: {
362            "annotator": lambda self, e: self._set_type(
363                e, exp.DataType.from_str("NUMBER", dialect="snowflake")
364            )
365        }
366        for expr_type in (
367            exp.BitwiseAndAgg,
368            exp.BitwiseOrAgg,
369            exp.BitwiseXorAgg,
370            exp.RegexpCount,
371            exp.RegexpInstr,
372            exp.ToNumber,
373        )
374    },
375    **{
376        expr_type: {"returns": exp.DType.DOUBLE}
377        for expr_type in {
378            exp.ApproxPercentileEstimate,
379            exp.ApproximateSimilarity,
380            exp.CosineDistance,
381            exp.DotProduct,
382            exp.EuclideanDistance,
383            exp.ManhattanDistance,
384            exp.MonthsBetween,
385            exp.Normal,
386        }
387    },
388    exp.Kurtosis: {"annotator": _annotate_kurtosis},
389    **{
390        expr_type: {"returns": exp.DType.DECFLOAT}
391        for expr_type in {
392            exp.ToDecfloat,
393            exp.TryToDecfloat,
394        }
395    },
396    **{
397        expr_type: {"annotator": _annotate_math_with_float_decfloat}
398        for expr_type in {
399            exp.Acos,
400            exp.Asin,
401            exp.Atan,
402            exp.Atan2,
403            exp.Cbrt,
404            exp.Cos,
405            exp.Cot,
406            exp.Degrees,
407            exp.Exp,
408            exp.Ln,
409            exp.Log,
410            exp.Pow,
411            exp.Radians,
412            exp.RegrAvgx,
413            exp.RegrAvgy,
414            exp.RegrCount,
415            exp.RegrIntercept,
416            exp.RegrR2,
417            exp.RegrSlope,
418            exp.RegrSxx,
419            exp.RegrSxy,
420            exp.RegrSyy,
421            exp.RegrValx,
422            exp.RegrValy,
423            exp.Sin,
424            exp.Sqrt,
425            exp.Tan,
426            exp.Tanh,
427        }
428    },
429    **{
430        expr_type: {"returns": exp.DType.INT}
431        for expr_type in {
432            exp.ByteLength,
433            exp.DenseRank,
434            exp.Grouping,
435            exp.JarowinklerSimilarity,
436            exp.MapSize,
437            exp.Minute,
438            exp.Ntile,
439            exp.Rank,
440            exp.RowNumber,
441            exp.RtrimmedLength,
442            exp.Second,
443            exp.Seq1,
444            exp.Seq2,
445            exp.Seq4,
446            exp.WidthBucket,
447        }
448    },
449    **{
450        expr_type: {"returns": exp.DType.OBJECT}
451        for expr_type in {
452            exp.ApproxPercentileAccumulate,
453            exp.ApproxPercentileCombine,
454            exp.ApproxTopKAccumulate,
455            exp.ApproxTopKCombine,
456            exp.ObjectAgg,
457            exp.ParseIp,
458            exp.ParseUrl,
459            exp.XMLGet,
460        }
461    },
462    **{
463        expr_type: {"returns": exp.DType.MAP}
464        for expr_type in {
465            exp.MapCat,
466            exp.MapDelete,
467            exp.MapInsert,
468            exp.MapPick,
469        }
470    },
471    **{
472        expr_type: {"returns": exp.DType.FILE}
473        for expr_type in {
474            exp.ToFile,
475        }
476    },
477    **{
478        expr_type: {"returns": exp.DType.TIME}
479        for expr_type in {
480            exp.TimeFromParts,
481            exp.TsOrDsToTime,
482        }
483    },
484    **{
485        expr_type: {"returns": exp.DType.TIMESTAMPLTZ}
486        for expr_type in {
487            exp.CurrentTimestamp,
488            exp.Localtimestamp,
489        }
490    },
491    **{
492        expr_type: {"returns": exp.DType.TINYINT}
493        for expr_type in {
494            exp.DayOfMonth,
495            exp.DayOfWeek,
496            exp.DayOfYear,
497            exp.Quarter,
498        }
499    },
500    **{
501        expr_type: {"returns": exp.DType.VARCHAR}
502        for expr_type in {
503            exp.AIAgg,
504            exp.AIClassify,
505            exp.AISummarizeAgg,
506            exp.Base64DecodeString,
507            exp.Base64Encode,
508            exp.CheckJson,
509            exp.CheckXml,
510            exp.Collate,
511            exp.Collation,
512            exp.CurrentAccount,
513            exp.CurrentAccountName,
514            exp.CurrentAvailableRoles,
515            exp.CurrentClient,
516            exp.CurrentDatabase,
517            exp.CurrentIpAddress,
518            exp.CurrentSchemas,
519            exp.CurrentSecondaryRoles,
520            exp.CurrentSession,
521            exp.CurrentStatement,
522            exp.CurrentTransaction,
523            exp.CurrentWarehouse,
524            exp.CurrentOrganizationUser,
525            exp.CurrentRegion,
526            exp.CurrentRoleType,
527            exp.CurrentOrganizationName,
528            exp.DecompressString,
529            exp.HexDecodeString,
530            exp.Hex,
531            exp.Randstr,
532            exp.RegexpExtract,
533            exp.RegexpReplace,
534            exp.Replace,
535            exp.Soundex,
536            exp.SoundexP123,
537            exp.SplitPart,
538            exp.Strtok,
539            exp.TryBase64DecodeString,
540            exp.TryHexDecodeString,
541            exp.Uuid,
542        }
543    },
544    **{
545        expr_type: {"returns": exp.DType.VARIANT}
546        for expr_type in {
547            exp.Minhash,
548            exp.MinhashCombine,
549        }
550    },
551    **{
552        expr_type: {"annotator": _annotate_variance}
553        for expr_type in (
554            exp.Variance,
555            exp.VariancePop,
556        )
557    },
558    exp.ArgMax: {"annotator": _annotate_arg_max_min},
559    exp.ArgMin: {"annotator": _annotate_arg_max_min},
560    exp.ConcatWs: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
561    exp.ConvertTimezone: {
562        "annotator": lambda self, e: self._set_type(
563            e,
564            exp.DType.TIMESTAMPNTZ if e.args.get("source_tz") else exp.DType.TIMESTAMPTZ,
565        )
566    },
567    exp.DateAdd: {"annotator": _annotate_date_or_time_add},
568    exp.DecodeCase: {"annotator": _annotate_decode_case},
569    exp.HashAgg: {
570        "annotator": lambda self, e: self._set_type(
571            e, exp.DataType.from_str("NUMBER(19, 0)", dialect="snowflake")
572        )
573    },
574    exp.Median: {"annotator": _annotate_median},
575    exp.Reverse: {"annotator": _annotate_reverse},
576    exp.StrToTime: {"annotator": _annotate_str_to_time},
577    exp.TimeAdd: {"annotator": _annotate_date_or_time_add},
578    exp.TimestampFromParts: {"annotator": _annotate_timestamp_from_parts},
579    exp.WithinGroup: {"annotator": _annotate_within_group},
580}