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.DotProduct,
366            exp.EuclideanDistance,
367            exp.ManhattanDistance,
368            exp.MonthsBetween,
369            exp.Normal,
370        }
371    },
372    exp.Kurtosis: {"annotator": _annotate_kurtosis},
373    **{
374        expr_type: {"returns": exp.DType.DECFLOAT}
375        for expr_type in {
376            exp.ToDecfloat,
377            exp.TryToDecfloat,
378        }
379    },
380    **{
381        expr_type: {"annotator": _annotate_math_with_float_decfloat}
382        for expr_type in {
383            exp.Acos,
384            exp.Asin,
385            exp.Atan,
386            exp.Atan2,
387            exp.Cbrt,
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.DType.INT}
415        for expr_type in {
416            exp.ByteLength,
417            exp.DenseRank,
418            exp.Grouping,
419            exp.JarowinklerSimilarity,
420            exp.MapSize,
421            exp.Minute,
422            exp.Ntile,
423            exp.Rank,
424            exp.RowNumber,
425            exp.RtrimmedLength,
426            exp.Second,
427            exp.Seq1,
428            exp.Seq2,
429            exp.Seq4,
430            exp.WidthBucket,
431        }
432    },
433    **{
434        expr_type: {"returns": exp.DType.OBJECT}
435        for expr_type in {
436            exp.ApproxPercentileAccumulate,
437            exp.ApproxPercentileCombine,
438            exp.ApproxTopKAccumulate,
439            exp.ApproxTopKCombine,
440            exp.ObjectAgg,
441            exp.ParseIp,
442            exp.ParseUrl,
443            exp.XMLGet,
444        }
445    },
446    **{
447        expr_type: {"returns": exp.DType.MAP}
448        for expr_type in {
449            exp.MapCat,
450            exp.MapDelete,
451            exp.MapInsert,
452            exp.MapPick,
453        }
454    },
455    **{
456        expr_type: {"returns": exp.DType.FILE}
457        for expr_type in {
458            exp.ToFile,
459        }
460    },
461    **{
462        expr_type: {"returns": exp.DType.TIME}
463        for expr_type in {
464            exp.TimeFromParts,
465            exp.TsOrDsToTime,
466        }
467    },
468    **{
469        expr_type: {"returns": exp.DType.TIMESTAMPLTZ}
470        for expr_type in {
471            exp.CurrentTimestamp,
472            exp.Localtimestamp,
473        }
474    },
475    **{
476        expr_type: {"returns": exp.DType.TINYINT}
477        for expr_type in {
478            exp.DayOfMonth,
479            exp.DayOfWeek,
480            exp.DayOfYear,
481            exp.Quarter,
482        }
483    },
484    **{
485        expr_type: {"returns": exp.DType.VARCHAR}
486        for expr_type in {
487            exp.AIAgg,
488            exp.AIClassify,
489            exp.AISummarizeAgg,
490            exp.Base64DecodeString,
491            exp.Base64Encode,
492            exp.CheckJson,
493            exp.CheckXml,
494            exp.Collate,
495            exp.Collation,
496            exp.CurrentAccount,
497            exp.CurrentAccountName,
498            exp.CurrentAvailableRoles,
499            exp.CurrentClient,
500            exp.CurrentDatabase,
501            exp.CurrentIpAddress,
502            exp.CurrentSchemas,
503            exp.CurrentSecondaryRoles,
504            exp.CurrentSession,
505            exp.CurrentStatement,
506            exp.CurrentTransaction,
507            exp.CurrentWarehouse,
508            exp.CurrentOrganizationUser,
509            exp.CurrentRegion,
510            exp.CurrentRoleType,
511            exp.CurrentOrganizationName,
512            exp.DecompressString,
513            exp.HexDecodeString,
514            exp.Hex,
515            exp.Randstr,
516            exp.RegexpExtract,
517            exp.RegexpReplace,
518            exp.Replace,
519            exp.Soundex,
520            exp.SoundexP123,
521            exp.SplitPart,
522            exp.Strtok,
523            exp.TryBase64DecodeString,
524            exp.TryHexDecodeString,
525            exp.Uuid,
526        }
527    },
528    **{
529        expr_type: {"returns": exp.DType.VARIANT}
530        for expr_type in {
531            exp.Minhash,
532            exp.MinhashCombine,
533        }
534    },
535    **{
536        expr_type: {"annotator": _annotate_variance}
537        for expr_type in (
538            exp.Variance,
539            exp.VariancePop,
540        )
541    },
542    exp.ArgMax: {"annotator": _annotate_arg_max_min},
543    exp.ArgMin: {"annotator": _annotate_arg_max_min},
544    exp.ConcatWs: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
545    exp.ConvertTimezone: {
546        "annotator": lambda self, e: self._set_type(
547            e,
548            exp.DType.TIMESTAMPNTZ if e.args.get("source_tz") else exp.DType.TIMESTAMPTZ,
549        )
550    },
551    exp.DateAdd: {"annotator": _annotate_date_or_time_add},
552    exp.DecodeCase: {"annotator": _annotate_decode_case},
553    exp.HashAgg: {
554        "annotator": lambda self, e: self._set_type(
555            e, exp.DataType.from_str("NUMBER(19, 0)", dialect="snowflake")
556        )
557    },
558    exp.Median: {"annotator": _annotate_median},
559    exp.Reverse: {"annotator": _annotate_reverse},
560    exp.StrToTime: {"annotator": _annotate_str_to_time},
561    exp.TimeAdd: {"annotator": _annotate_date_or_time_add},
562    exp.TimestampFromParts: {"annotator": _annotate_timestamp_from_parts},
563    exp.WithinGroup: {"annotator": _annotate_within_group},
564}