sqlglot.typing.hive

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.typing import EXPRESSION_METADATA
  7
  8if t.TYPE_CHECKING:
  9    from sqlglot.optimizer.annotate_types import TypeAnnotator
 10
 11
 12def _annotate_floor_ceil(self: TypeAnnotator, expression: exp.Expr) -> exp.Expr:
 13    # Decimal results and numeric literals need precision inference that we don't model here.
 14    if (
 15        expression.args.get("decimals") is not None
 16        or expression.this.is_type(exp.DType.DECIMAL, exp.DType.UNKNOWN)
 17        or expression.this.unnest().is_number
 18    ):
 19        return self._set_type(expression, exp.DType.UNKNOWN)
 20    return self._set_type(expression, exp.DType.BIGINT)
 21
 22
 23EXPRESSION_METADATA = {
 24    **EXPRESSION_METADATA,
 25    **{expr_type: {"annotator": _annotate_floor_ceil} for expr_type in {exp.Ceil, exp.Floor}},
 26    **{
 27        expr_type: {"returns": exp.DType.BINARY}
 28        for expr_type in {
 29            exp.Encode,
 30            exp.Unhex,
 31        }
 32    },
 33    **{
 34        expr_type: {"returns": exp.DType.DOUBLE}
 35        for expr_type in {
 36            exp.Corr,
 37            exp.MonthsBetween,
 38            exp.Sign,
 39        }
 40    },
 41    **{
 42        expr_type: {"returns": exp.DType.VARCHAR}
 43        for expr_type in {
 44            exp.AddMonths,
 45            exp.CurrentDatabase,
 46            exp.Hex,
 47            exp.JSONExtractScalar,
 48            exp.JSONFormat,
 49            exp.NextDay,
 50            exp.RegexpExtract,
 51            exp.RegexpReplace,
 52            exp.Replace,
 53            exp.Soundex,
 54        }
 55    },
 56    **{
 57        expr_type: {"returns": exp.DType.BIGINT}
 58        for expr_type in {
 59            exp.Factorial,
 60            exp.IntDiv,
 61            exp.StrToUnix,
 62        }
 63    },
 64    **{
 65        expr_type: {"returns": exp.DType.INT}
 66        for expr_type in {
 67            exp.ArraySize,
 68            exp.DenseRank,
 69            exp.Month,
 70            exp.Ntile,
 71            exp.Rank,
 72            exp.RowNumber,
 73            exp.Second,
 74            exp.Minute,
 75        }
 76    },
 77    **{
 78        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
 79        for expr_type in {
 80            exp.ArrayDistinct,
 81            exp.ArrayExcept,
 82            exp.First,
 83            exp.Last,
 84            exp.Negative,
 85            exp.Reverse,
 86        }
 87    },
 88    exp.ArrayIntersect: {"annotator": lambda self, e: self._annotate_by_args(e, "expressions")},
 89    exp.ApproxQuantile: {"annotator": lambda self, e: self._annotate_by_args(e, "quantile")},
 90    exp.Coalesce: {
 91        "annotator": lambda self, e: self._annotate_by_args(e, "this", "expressions", promote=True)
 92    },
 93    exp.Grouping: {"returns": exp.DType.BIGINT},
 94    exp.If: {"annotator": lambda self, e: self._annotate_by_args(e, "true", "false", promote=True)},
 95    exp.PercentileDisc: {"returns": exp.DType.DOUBLE},
 96    exp.Quantile: {"annotator": lambda self, e: self._annotate_by_args(e, "quantile")},
 97    exp.RegexpSplit: {"returns": exp.DataType.from_str("ARRAY<STRING>")},
 98    exp.StrToMap: {"returns": exp.DataType.from_str("MAP<STRING, STRING>")},
 99    exp.WithinGroup: {"annotator": lambda self, e: self._annotate_by_args(e, "this")},
100}