sqlglot.parsers.spark

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.trie import new_trie
  7from sqlglot.dialects.dialect import build_date_delta, build_like
  8from sqlglot.helper import ensure_list, seq_get
  9from sqlglot.parsers.hive import build_with_ignore_nulls
 10from sqlglot.parsers.spark2 import Spark2Parser, build_as_cast
 11from sqlglot.tokens import TokenType
 12
 13
 14def _build_datediff(args: list) -> exp.Expr:
 15    """
 16    Although Spark docs don't mention the "unit" argument, Spark3 added support for
 17    it at some point. Databricks also supports this variant (see below).
 18
 19    For example, in spark-sql (v3.3.1):
 20    - SELECT DATEDIFF('2020-01-01', '2020-01-05') results in -4
 21    - SELECT DATEDIFF(day, '2020-01-01', '2020-01-05') results in 4
 22
 23    See also:
 24    - https://docs.databricks.com/sql/language-manual/functions/datediff3.html
 25    - https://docs.databricks.com/sql/language-manual/functions/datediff.html
 26    """
 27    unit = None
 28    this = seq_get(args, 0)
 29    expression = seq_get(args, 1)
 30
 31    if len(args) == 3:
 32        unit = exp.var(t.cast(exp.Expr, this).name)
 33        this = args[2]
 34
 35    return exp.DateDiff(
 36        this=exp.TsOrDsToDate(this=this), expression=exp.TsOrDsToDate(this=expression), unit=unit
 37    )
 38
 39
 40def _build_dateadd(args: list) -> exp.Expr:
 41    expression = seq_get(args, 1)
 42
 43    if len(args) == 2:
 44        # DATE_ADD(startDate, numDays INTEGER)
 45        # https://docs.databricks.com/en/sql/language-manual/functions/date_add.html
 46        return exp.TsOrDsAdd(
 47            this=seq_get(args, 0), expression=expression, unit=exp.Literal.string("DAY")
 48        )
 49
 50    # DATE_ADD / DATEADD / TIMESTAMPADD(unit, value integer, expr)
 51    # https://docs.databricks.com/en/sql/language-manual/functions/date_add3.html
 52    return exp.TimestampAdd(this=seq_get(args, 2), expression=expression, unit=seq_get(args, 0))
 53
 54
 55class SparkParser(Spark2Parser):
 56    NO_PAREN_FUNCTIONS = {
 57        **Spark2Parser.NO_PAREN_FUNCTIONS,
 58        TokenType.SESSION_USER: exp.SessionUser,
 59    }
 60
 61    SET_PARSERS = {
 62        **Spark2Parser.SET_PARSERS,
 63        "VAR": lambda self: self._parse_set_item_assignment("VARIABLE"),
 64        "VARIABLE": lambda self: self._parse_set_item_assignment("VARIABLE"),
 65    }
 66
 67    SET_TRIE = new_trie(key.split(" ") for key in SET_PARSERS)
 68
 69    FUNCTIONS = {
 70        **Spark2Parser.FUNCTIONS,
 71        "ANY_VALUE": build_with_ignore_nulls(exp.AnyValue),
 72        "ARRAY_INSERT": lambda args: exp.ArrayInsert(
 73            this=seq_get(args, 0),
 74            position=seq_get(args, 1),
 75            expression=seq_get(args, 2),
 76            offset=1,
 77        ),
 78        "BIT_AND": exp.BitwiseAndAgg.from_arg_list,
 79        "BIT_GET": exp.Getbit.from_arg_list,
 80        "BIT_OR": exp.BitwiseOrAgg.from_arg_list,
 81        "BIT_XOR": exp.BitwiseXorAgg.from_arg_list,
 82        "BIT_COUNT": exp.BitwiseCount.from_arg_list,
 83        "CURDATE": exp.CurrentDate.from_arg_list,
 84        "DATE_ADD": _build_dateadd,
 85        "DATEADD": _build_dateadd,
 86        "MAKE_TIMESTAMP": exp.TimestampFromParts.from_arg_list,
 87        "TIMESTAMPADD": _build_dateadd,
 88        "TIMESTAMPDIFF": build_date_delta(exp.TimestampDiff),
 89        "TRY_ADD": exp.SafeAdd.from_arg_list,
 90        "TRY_DIVIDE": exp.SafeDivide.from_arg_list,
 91        "TRY_MULTIPLY": exp.SafeMultiply.from_arg_list,
 92        "TRY_SUBTRACT": exp.SafeSubtract.from_arg_list,
 93        "DATEDIFF": _build_datediff,
 94        "DATE_DIFF": _build_datediff,
 95        "JSON_OBJECT_KEYS": exp.JSONKeys.from_arg_list,
 96        "LISTAGG": exp.GroupConcat.from_arg_list,
 97        "TIMESTAMP_LTZ": build_as_cast("TIMESTAMP_LTZ"),
 98        "TIMESTAMP_NTZ": build_as_cast("TIMESTAMP_NTZ"),
 99        "TRY_ELEMENT_AT": lambda args: exp.Bracket(
100            this=seq_get(args, 0),
101            expressions=ensure_list(seq_get(args, 1)),
102            offset=1,
103            safe=True,
104        ),
105        "LIKE": build_like(exp.Like),
106        "ILIKE": build_like(exp.ILike),
107    }
108
109    PLACEHOLDER_PARSERS = {
110        **Spark2Parser.PLACEHOLDER_PARSERS,
111        TokenType.L_BRACE: lambda self: self._parse_query_parameter(),
112    }
113
114    def _parse_query_parameter(self) -> exp.Expr | None:
115        this = self._parse_id_var()
116        self._match(TokenType.R_BRACE)
117        return self.expression(exp.Placeholder(this=this, widget=True))
118
119    FUNCTION_PARSERS = {
120        **Spark2Parser.FUNCTION_PARSERS,
121        "SUBSTR": lambda self: self._parse_substring(),
122    }
123
124    STATEMENT_PARSERS = {
125        **Spark2Parser.STATEMENT_PARSERS,
126        TokenType.DECLARE: lambda self: self._parse_declare(),
127    }
128
129    def _parse_generated_as_identity(
130        self,
131    ) -> (
132        exp.GeneratedAsIdentityColumnConstraint
133        | exp.ComputedColumnConstraint
134        | exp.GeneratedAsRowColumnConstraint
135    ):
136        this = super()._parse_generated_as_identity()
137        if this.expression:
138            return self.expression(exp.ComputedColumnConstraint(this=this.expression))
139        return this
140
141    def _parse_pivot_aggregation(self) -> exp.Expr | None:
142        # Spark 3+ and Databricks support non aggregate functions in PIVOT too, e.g
143        # PIVOT (..., 'foo' AS bar FOR col_to_pivot IN (...))
144        aggregate_expr = self._parse_function() or self._parse_disjunction()
145        return self._parse_alias(aggregate_expr)