sqlglot.dialects.spark2

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, transforms
  6from sqlglot.dialects.dialect import (
  7    binary_from_function,
  8    build_formatted_time,
  9    is_parse_json,
 10    pivot_column_names,
 11    rename_func,
 12    trim_sql,
 13    unit_to_str,
 14)
 15from sqlglot.dialects.hive import Hive
 16from sqlglot.helper import seq_get, ensure_list
 17from sqlglot.tokens import TokenType
 18from sqlglot.transforms import (
 19    preprocess,
 20    remove_unique_constraints,
 21    ctas_with_tmp_tables_to_create_tmp_view,
 22    move_schema_columns_to_partitioned_by,
 23)
 24
 25if t.TYPE_CHECKING:
 26    from sqlglot._typing import E
 27
 28    from sqlglot.optimizer.annotate_types import TypeAnnotator
 29
 30
 31def _map_sql(self: Spark2.Generator, expression: exp.Map) -> str:
 32    keys = expression.args.get("keys")
 33    values = expression.args.get("values")
 34
 35    if not keys or not values:
 36        return self.func("MAP")
 37
 38    return self.func("MAP_FROM_ARRAYS", keys, values)
 39
 40
 41def _build_as_cast(to_type: str) -> t.Callable[[t.List], exp.Expression]:
 42    return lambda args: exp.Cast(this=seq_get(args, 0), to=exp.DataType.build(to_type))
 43
 44
 45def _str_to_date(self: Spark2.Generator, expression: exp.StrToDate) -> str:
 46    time_format = self.format_time(expression)
 47    if time_format == Hive.DATE_FORMAT:
 48        return self.func("TO_DATE", expression.this)
 49    return self.func("TO_DATE", expression.this, time_format)
 50
 51
 52def _unix_to_time_sql(self: Spark2.Generator, expression: exp.UnixToTime) -> str:
 53    scale = expression.args.get("scale")
 54    timestamp = expression.this
 55
 56    if scale is None:
 57        return self.sql(exp.cast(exp.func("from_unixtime", timestamp), exp.DataType.Type.TIMESTAMP))
 58    if scale == exp.UnixToTime.SECONDS:
 59        return self.func("TIMESTAMP_SECONDS", timestamp)
 60    if scale == exp.UnixToTime.MILLIS:
 61        return self.func("TIMESTAMP_MILLIS", timestamp)
 62    if scale == exp.UnixToTime.MICROS:
 63        return self.func("TIMESTAMP_MICROS", timestamp)
 64
 65    unix_seconds = exp.Div(this=timestamp, expression=exp.func("POW", 10, scale))
 66    return self.func("TIMESTAMP_SECONDS", unix_seconds)
 67
 68
 69def _unalias_pivot(expression: exp.Expression) -> exp.Expression:
 70    """
 71    Spark doesn't allow PIVOT aliases, so we need to remove them and possibly wrap a
 72    pivoted source in a subquery with the same alias to preserve the query's semantics.
 73
 74    Example:
 75        >>> from sqlglot import parse_one
 76        >>> expr = parse_one("SELECT piv.x FROM tbl PIVOT (SUM(a) FOR b IN ('x')) piv")
 77        >>> print(_unalias_pivot(expr).sql(dialect="spark"))
 78        SELECT piv.x FROM (SELECT * FROM tbl PIVOT(SUM(a) FOR b IN ('x'))) AS piv
 79    """
 80    if isinstance(expression, exp.From) and expression.this.args.get("pivots"):
 81        pivot = expression.this.args["pivots"][0]
 82        if pivot.alias:
 83            alias = pivot.args["alias"].pop()
 84            return exp.From(
 85                this=expression.this.replace(
 86                    exp.select("*")
 87                    .from_(expression.this.copy(), copy=False)
 88                    .subquery(alias=alias, copy=False)
 89                )
 90            )
 91
 92    return expression
 93
 94
 95def _unqualify_pivot_columns(expression: exp.Expression) -> exp.Expression:
 96    """
 97    Spark doesn't allow the column referenced in the PIVOT's field to be qualified,
 98    so we need to unqualify it.
 99
100    Example:
101        >>> from sqlglot import parse_one
102        >>> expr = parse_one("SELECT * FROM tbl PIVOT (SUM(tbl.sales) FOR tbl.quarter IN ('Q1', 'Q2'))")
103        >>> print(_unqualify_pivot_columns(expr).sql(dialect="spark"))
104        SELECT * FROM tbl PIVOT(SUM(tbl.sales) FOR quarter IN ('Q1', 'Q1'))
105    """
106    if isinstance(expression, exp.Pivot):
107        expression.set(
108            "fields", [transforms.unqualify_columns(field) for field in expression.fields]
109        )
110
111    return expression
112
113
114def temporary_storage_provider(expression: exp.Expression) -> exp.Expression:
115    # spark2, spark, Databricks require a storage provider for temporary tables
116    provider = exp.FileFormatProperty(this=exp.Literal.string("parquet"))
117    expression.args["properties"].append("expressions", provider)
118    return expression
119
120
121def _annotate_by_similar_args(
122    self: TypeAnnotator, expression: E, *args: str, target_type: exp.DataType | exp.DataType.Type
123) -> E:
124    """
125    Infers the type of the expression according to the following rules:
126    - If all args are of the same type OR any arg is of target_type, the expr is inferred as such
127    - If any arg is of UNKNOWN type and none of target_type, the expr is inferred as UNKNOWN
128    """
129    self._annotate_args(expression)
130
131    expressions: t.List[exp.Expression] = []
132    for arg in args:
133        arg_expr = expression.args.get(arg)
134        expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
135
136    last_datatype = None
137
138    has_unknown = False
139    for expr in expressions:
140        if expr.is_type(exp.DataType.Type.UNKNOWN):
141            has_unknown = True
142        elif expr.is_type(target_type):
143            has_unknown = False
144            last_datatype = target_type
145            break
146        else:
147            last_datatype = expr.type
148
149    self._set_type(expression, exp.DataType.Type.UNKNOWN if has_unknown else last_datatype)
150    return expression
151
152
153class Spark2(Hive):
154    ALTER_TABLE_SUPPORTS_CASCADE = False
155
156    ANNOTATORS = {
157        **Hive.ANNOTATORS,
158        exp.Substring: lambda self, e: self._annotate_by_args(e, "this"),
159        exp.Concat: lambda self, e: _annotate_by_similar_args(
160            self, e, "expressions", target_type=exp.DataType.Type.TEXT
161        ),
162        exp.Pad: lambda self, e: _annotate_by_similar_args(
163            self, e, "this", "fill_pattern", target_type=exp.DataType.Type.TEXT
164        ),
165    }
166
167    class Tokenizer(Hive.Tokenizer):
168        HEX_STRINGS = [("X'", "'"), ("x'", "'")]
169
170        KEYWORDS = {
171            **Hive.Tokenizer.KEYWORDS,
172            "TIMESTAMP": TokenType.TIMESTAMPTZ,
173        }
174
175    class Parser(Hive.Parser):
176        TRIM_PATTERN_FIRST = True
177        CHANGE_COLUMN_ALTER_SYNTAX = True
178
179        FUNCTIONS = {
180            **Hive.Parser.FUNCTIONS,
181            "AGGREGATE": exp.Reduce.from_arg_list,
182            "APPROX_PERCENTILE": exp.ApproxQuantile.from_arg_list,
183            "BOOLEAN": _build_as_cast("boolean"),
184            "DATE": _build_as_cast("date"),
185            "DATE_TRUNC": lambda args: exp.TimestampTrunc(
186                this=seq_get(args, 1), unit=exp.var(seq_get(args, 0))
187            ),
188            "DAYOFMONTH": lambda args: exp.DayOfMonth(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
189            "DAYOFWEEK": lambda args: exp.DayOfWeek(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
190            "DAYOFYEAR": lambda args: exp.DayOfYear(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
191            "DOUBLE": _build_as_cast("double"),
192            "FLOAT": _build_as_cast("float"),
193            "FORMAT_STRING": exp.Format.from_arg_list,
194            "FROM_UTC_TIMESTAMP": lambda args, dialect: exp.AtTimeZone(
195                this=exp.cast(
196                    seq_get(args, 0) or exp.Var(this=""),
197                    exp.DataType.Type.TIMESTAMP,
198                    dialect=dialect,
199                ),
200                zone=seq_get(args, 1),
201            ),
202            "INT": _build_as_cast("int"),
203            "MAP_FROM_ARRAYS": exp.Map.from_arg_list,
204            "RLIKE": exp.RegexpLike.from_arg_list,
205            "SHIFTLEFT": binary_from_function(exp.BitwiseLeftShift),
206            "SHIFTRIGHT": binary_from_function(exp.BitwiseRightShift),
207            "STRING": _build_as_cast("string"),
208            "SLICE": exp.ArraySlice.from_arg_list,
209            "TIMESTAMP": _build_as_cast("timestamp"),
210            "TO_TIMESTAMP": lambda args: (
211                _build_as_cast("timestamp")(args)
212                if len(args) == 1
213                else build_formatted_time(exp.StrToTime, "spark")(args)
214            ),
215            "TO_UNIX_TIMESTAMP": exp.StrToUnix.from_arg_list,
216            "TO_UTC_TIMESTAMP": lambda args, dialect: exp.FromTimeZone(
217                this=exp.cast(
218                    seq_get(args, 0) or exp.Var(this=""),
219                    exp.DataType.Type.TIMESTAMP,
220                    dialect=dialect,
221                ),
222                zone=seq_get(args, 1),
223            ),
224            "TRUNC": lambda args: exp.DateTrunc(unit=seq_get(args, 1), this=seq_get(args, 0)),
225            "WEEKOFYEAR": lambda args: exp.WeekOfYear(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
226        }
227
228        FUNCTION_PARSERS = {
229            **Hive.Parser.FUNCTION_PARSERS,
230            "BROADCAST": lambda self: self._parse_join_hint("BROADCAST"),
231            "BROADCASTJOIN": lambda self: self._parse_join_hint("BROADCASTJOIN"),
232            "MAPJOIN": lambda self: self._parse_join_hint("MAPJOIN"),
233            "MERGE": lambda self: self._parse_join_hint("MERGE"),
234            "SHUFFLEMERGE": lambda self: self._parse_join_hint("SHUFFLEMERGE"),
235            "MERGEJOIN": lambda self: self._parse_join_hint("MERGEJOIN"),
236            "SHUFFLE_HASH": lambda self: self._parse_join_hint("SHUFFLE_HASH"),
237            "SHUFFLE_REPLICATE_NL": lambda self: self._parse_join_hint("SHUFFLE_REPLICATE_NL"),
238        }
239
240        def _parse_drop_column(self) -> t.Optional[exp.Drop | exp.Command]:
241            return self._match_text_seq("DROP", "COLUMNS") and self.expression(
242                exp.Drop, this=self._parse_schema(), kind="COLUMNS"
243            )
244
245        def _pivot_column_names(self, aggregations: t.List[exp.Expression]) -> t.List[str]:
246            if len(aggregations) == 1:
247                return []
248            return pivot_column_names(aggregations, dialect="spark")
249
250    class Generator(Hive.Generator):
251        QUERY_HINTS = True
252        NVL2_SUPPORTED = True
253        CAN_IMPLEMENT_ARRAY_ANY = True
254        ALTER_SET_TYPE = "TYPE"
255
256        PROPERTIES_LOCATION = {
257            **Hive.Generator.PROPERTIES_LOCATION,
258            exp.EngineProperty: exp.Properties.Location.UNSUPPORTED,
259            exp.AutoIncrementProperty: exp.Properties.Location.UNSUPPORTED,
260            exp.CharacterSetProperty: exp.Properties.Location.UNSUPPORTED,
261            exp.CollateProperty: exp.Properties.Location.UNSUPPORTED,
262        }
263
264        TS_OR_DS_EXPRESSIONS = (
265            *Hive.Generator.TS_OR_DS_EXPRESSIONS,
266            exp.DayOfMonth,
267            exp.DayOfWeek,
268            exp.DayOfYear,
269            exp.WeekOfYear,
270        )
271
272        TRANSFORMS = {
273            **Hive.Generator.TRANSFORMS,
274            exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
275            exp.ArraySum: lambda self,
276            e: f"AGGREGATE({self.sql(e, 'this')}, 0, (acc, x) -> acc + x, acc -> acc)",
277            exp.ArrayToString: rename_func("ARRAY_JOIN"),
278            exp.ArraySlice: rename_func("SLICE"),
279            exp.AtTimeZone: lambda self, e: self.func(
280                "FROM_UTC_TIMESTAMP", e.this, e.args.get("zone")
281            ),
282            exp.BitwiseLeftShift: rename_func("SHIFTLEFT"),
283            exp.BitwiseRightShift: rename_func("SHIFTRIGHT"),
284            exp.Create: preprocess(
285                [
286                    remove_unique_constraints,
287                    lambda e: ctas_with_tmp_tables_to_create_tmp_view(
288                        e, temporary_storage_provider
289                    ),
290                    move_schema_columns_to_partitioned_by,
291                ]
292            ),
293            exp.DateFromParts: rename_func("MAKE_DATE"),
294            exp.DateTrunc: lambda self, e: self.func("TRUNC", e.this, unit_to_str(e)),
295            exp.DayOfMonth: rename_func("DAYOFMONTH"),
296            exp.DayOfWeek: rename_func("DAYOFWEEK"),
297            # (DAY_OF_WEEK(datetime) % 7) + 1 is equivalent to DAYOFWEEK_ISO(datetime)
298            exp.DayOfWeekIso: lambda self, e: f"(({self.func('DAYOFWEEK', e.this)} % 7) + 1)",
299            exp.DayOfYear: rename_func("DAYOFYEAR"),
300            exp.Format: rename_func("FORMAT_STRING"),
301            exp.From: transforms.preprocess([_unalias_pivot]),
302            exp.FromTimeZone: lambda self, e: self.func(
303                "TO_UTC_TIMESTAMP", e.this, e.args.get("zone")
304            ),
305            exp.LogicalAnd: rename_func("BOOL_AND"),
306            exp.LogicalOr: rename_func("BOOL_OR"),
307            exp.Map: _map_sql,
308            exp.Pivot: transforms.preprocess([_unqualify_pivot_columns]),
309            exp.Reduce: rename_func("AGGREGATE"),
310            exp.RegexpReplace: lambda self, e: self.func(
311                "REGEXP_REPLACE",
312                e.this,
313                e.expression,
314                e.args["replacement"],
315                e.args.get("position"),
316            ),
317            exp.Select: transforms.preprocess(
318                [
319                    transforms.eliminate_qualify,
320                    transforms.eliminate_distinct_on,
321                    transforms.unnest_to_explode,
322                    transforms.any_to_exists,
323                ]
324            ),
325            exp.StrToDate: _str_to_date,
326            exp.StrToTime: lambda self, e: self.func("TO_TIMESTAMP", e.this, self.format_time(e)),
327            exp.TimestampTrunc: lambda self, e: self.func("DATE_TRUNC", unit_to_str(e), e.this),
328            exp.Trim: trim_sql,
329            exp.UnixToTime: _unix_to_time_sql,
330            exp.VariancePop: rename_func("VAR_POP"),
331            exp.WeekOfYear: rename_func("WEEKOFYEAR"),
332            exp.WithinGroup: transforms.preprocess(
333                [transforms.remove_within_group_for_percentiles]
334            ),
335        }
336        TRANSFORMS.pop(exp.ArraySort)
337        TRANSFORMS.pop(exp.ILike)
338        TRANSFORMS.pop(exp.Left)
339        TRANSFORMS.pop(exp.MonthsBetween)
340        TRANSFORMS.pop(exp.Right)
341
342        WRAP_DERIVED_VALUES = False
343        CREATE_FUNCTION_RETURN_AS = False
344
345        def struct_sql(self, expression: exp.Struct) -> str:
346            from sqlglot.generator import Generator
347
348            return Generator.struct_sql(self, expression)
349
350        def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) -> str:
351            arg = expression.this
352            is_json_extract = isinstance(
353                arg, (exp.JSONExtract, exp.JSONExtractScalar)
354            ) and not arg.args.get("variant_extract")
355
356            # We can't use a non-nested type (eg. STRING) as a schema
357            if expression.to.args.get("nested") and (is_parse_json(arg) or is_json_extract):
358                schema = f"'{self.sql(expression, 'to')}'"
359                return self.func("FROM_JSON", arg if is_json_extract else arg.this, schema)
360
361            if is_parse_json(expression):
362                return self.func("TO_JSON", arg)
363
364            return super(Hive.Generator, self).cast_sql(expression, safe_prefix=safe_prefix)
365
366        def fileformatproperty_sql(self, expression: exp.FileFormatProperty) -> str:
367            if expression.args.get("hive_format"):
368                return super().fileformatproperty_sql(expression)
369
370            return f"USING {expression.name.upper()}"
371
372        def altercolumn_sql(self, expression: exp.AlterColumn) -> str:
373            this = self.sql(expression, "this")
374            new_name = self.sql(expression, "rename_to") or this
375            comment = self.sql(expression, "comment")
376            if new_name == this:
377                if comment:
378                    return f"ALTER COLUMN {this} COMMENT {comment}"
379                return super(Hive.Generator, self).altercolumn_sql(expression)
380            return f"RENAME COLUMN {this} TO {new_name}"
381
382        def renamecolumn_sql(self, expression: exp.RenameColumn) -> str:
383            return super(Hive.Generator, self).renamecolumn_sql(expression)