sqlglot.generators.spark2

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, transforms
  6from sqlglot.dialects.dialect import (
  7    array_sort_order,
  8    array_slice_to_slice_sql,
  9    bracket_to_element_at_sql,
 10    is_parse_json,
 11    rename_func,
 12    timestamptrunc_sql,
 13    weekstart_unit_to_str,
 14)
 15from sqlglot.generators.hive import HIVE_DATE_FORMAT, HiveGenerator, HIVE_TS_OR_DS_EXPRESSIONS
 16from sqlglot.transforms import (
 17    preprocess,
 18    remove_unique_constraints,
 19    ctas_with_tmp_tables_to_create_tmp_view,
 20    move_schema_columns_to_partitioned_by,
 21)
 22
 23
 24def _array_sort_sql(self: Spark2Generator, expression: exp.ArraySort) -> str:
 25    if expression.expression:
 26        return self.func("ARRAY_SORT", expression.this, expression.expression)
 27
 28    # ARRAY_SORT puts NULLs last, while SORT_ARRAY puts them first exactly when ascending
 29    order = array_sort_order(expression)
 30    if order is None:
 31        asc = expression.args.get("asc")
 32        if expression.args.get("nulls_first") != asc:
 33            self.unsupported("SORT_ARRAY puts NULLs first exactly when ascending")
 34        return self.func("SORT_ARRAY", expression.this, asc)
 35
 36    asc, nulls_first = order
 37    if asc:
 38        return self.func("SORT_ARRAY" if nulls_first else "ARRAY_SORT", expression.this)
 39    if nulls_first:
 40        return self.func("REVERSE", self.func("ARRAY_SORT", expression.this))
 41    return self.func("SORT_ARRAY", expression.this, exp.false())
 42
 43
 44def _json_format_sql(self: Spark2Generator, expression: exp.JSONFormat) -> str:
 45    this = expression.this
 46
 47    if is_parse_json(this):
 48        if this.this.is_string:
 49            # Since FROM_JSON requires a nested type, we always wrap the json string with
 50            # an array to ensure that "naked" strings like "'a'" will be handled correctly
 51            wrapped_json = exp.Literal.string(f"[{this.this.name}]")
 52
 53            from_json = self.func(
 54                "FROM_JSON", wrapped_json, self.func("SCHEMA_OF_JSON", wrapped_json)
 55            )
 56            to_json = self.func("TO_JSON", from_json)
 57
 58            # This strips the [, ] delimiters of the dummy array printed by TO_JSON
 59            return self.func("REGEXP_EXTRACT", to_json, "'^.(.*).$'", "1")
 60        return self.sql(this)
 61
 62    return self.func("TO_JSON", this, expression.args.get("options"))
 63
 64
 65def _map_sql(self: Spark2Generator, expression: exp.Map) -> str:
 66    keys = expression.args.get("keys")
 67    values = expression.args.get("values")
 68
 69    if not keys or not values:
 70        return self.func("MAP")
 71
 72    return self.func("MAP_FROM_ARRAYS", keys, values)
 73
 74
 75def _str_to_date(self: Spark2Generator, expression: exp.StrToDate) -> str:
 76    time_format = self.format_time(expression)
 77    if time_format == HIVE_DATE_FORMAT:
 78        return self.func("TO_DATE", expression.this)
 79    return self.func("TO_DATE", expression.this, time_format)
 80
 81
 82def _unix_to_time_sql(self: Spark2Generator, expression: exp.UnixToTime) -> str:
 83    scale = expression.args.get("scale")
 84    timestamp = expression.this
 85
 86    if scale is None:
 87        return self.sql(exp.cast(exp.func("from_unixtime", timestamp), exp.DType.TIMESTAMP))
 88    if scale == exp.UnixToTime.SECONDS:
 89        return self.func("TIMESTAMP_SECONDS", timestamp)
 90    if scale == exp.UnixToTime.MILLIS:
 91        return self.func("TIMESTAMP_MILLIS", timestamp)
 92    if scale == exp.UnixToTime.MICROS:
 93        return self.func("TIMESTAMP_MICROS", timestamp)
 94
 95    unix_seconds = exp.Div(this=timestamp, expression=exp.func("POW", 10, scale))
 96    return self.func("TIMESTAMP_SECONDS", unix_seconds)
 97
 98
 99def _unalias_pivot(expression: exp.Expr) -> exp.Expr:
100    """
101    Spark doesn't allow PIVOT aliases, so we need to remove them and possibly wrap a
102    pivoted source in a subquery with the same alias to preserve the query's semantics.
103
104    Example:
105        >>> from sqlglot import parse_one
106        >>> expr = parse_one("SELECT piv.x FROM tbl PIVOT (SUM(a) FOR b IN ('x')) piv")
107        >>> print(_unalias_pivot(expr).sql(dialect="spark"))
108        SELECT piv.x FROM (SELECT * FROM tbl PIVOT(SUM(a) FOR b IN ('x'))) AS piv
109    """
110    if isinstance(expression, exp.From) and expression.this.args.get("pivots"):
111        pivot = expression.this.args["pivots"][0]
112        if pivot.alias:
113            alias = pivot.args["alias"].pop()
114            return exp.From(
115                this=expression.this.replace(
116                    exp.select("*")
117                    .from_(expression.this.copy(), copy=False)
118                    .subquery(alias=alias, copy=False)
119                )
120            )
121
122    return expression
123
124
125def temporary_storage_provider(expression: exp.Expr) -> exp.Expr:
126    # spark2, spark, Databricks require a storage provider for temporary tables
127    provider = exp.FileFormatProperty(this=exp.Literal.string("parquet"))
128    expression.args["properties"].append("expressions", provider)
129    return expression
130
131
132class Spark2Generator(HiveGenerator):
133    QUERY_HINTS = True
134    NVL2_SUPPORTED = True
135    CAN_IMPLEMENT_ARRAY_ANY = True
136    ALTER_SET_TYPE = "TYPE"
137    PARSE_JSON_NAME: str | None = None
138
139    PROPERTIES_LOCATION = {
140        **HiveGenerator.PROPERTIES_LOCATION,
141        exp.EngineProperty: exp.Properties.Location.UNSUPPORTED,
142        exp.AutoIncrementProperty: exp.Properties.Location.UNSUPPORTED,
143        exp.CharacterSetProperty: exp.Properties.Location.UNSUPPORTED,
144        exp.CollateProperty: exp.Properties.Location.UNSUPPORTED,
145    }
146
147    TS_OR_DS_EXPRESSIONS: t.ClassVar = (
148        *HIVE_TS_OR_DS_EXPRESSIONS,
149        exp.DayOfMonth,
150        exp.DayOfWeek,
151        exp.DayOfYear,
152        exp.WeekOfYear,
153    )
154
155    TRANSFORMS = {
156        k: v
157        for k, v in {
158            **HiveGenerator.TRANSFORMS,
159            exp.ArraySort: _array_sort_sql,
160            exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
161            exp.ArraySum: lambda self, e: (
162                f"AGGREGATE({self.sql(e, 'this')}, 0, (acc, x) -> acc + x, acc -> acc)"
163            ),
164            exp.ArrayToString: rename_func("ARRAY_JOIN"),
165            exp.ArraySlice: array_slice_to_slice_sql,
166            exp.AtTimeZone: lambda self, e: self.func(
167                "FROM_UTC_TIMESTAMP", e.this, e.args.get("zone")
168            ),
169            exp.BitwiseLeftShift: rename_func("SHIFTLEFT"),
170            exp.BitwiseRightShift: rename_func("SHIFTRIGHT"),
171            exp.Create: preprocess(
172                [
173                    remove_unique_constraints,
174                    lambda e: ctas_with_tmp_tables_to_create_tmp_view(
175                        e, temporary_storage_provider
176                    ),
177                    move_schema_columns_to_partitioned_by,
178                ]
179            ),
180            exp.DateFromParts: rename_func("MAKE_DATE"),
181            exp.DateTrunc: lambda self, e: self.func(
182                "TRUNC", e.this, weekstart_unit_to_str(self, e)
183            ),
184            exp.DayOfMonth: rename_func("DAYOFMONTH"),
185            exp.DayOfWeek: rename_func("DAYOFWEEK"),
186            # (DAY_OF_WEEK(datetime) % 7) + 1 is equivalent to DAYOFWEEK_ISO(datetime)
187            exp.DayOfWeekIso: lambda self, e: f"(({self.func('DAYOFWEEK', e.this)} % 7) + 1)",
188            exp.DayOfYear: rename_func("DAYOFYEAR"),
189            exp.Format: rename_func("FORMAT_STRING"),
190            exp.From: transforms.preprocess([_unalias_pivot]),
191            exp.FromTimeZone: lambda self, e: self.func(
192                "TO_UTC_TIMESTAMP", e.this, e.args.get("zone")
193            ),
194            exp.JSONFormat: _json_format_sql,
195            exp.LogicalAnd: rename_func("BOOL_AND"),
196            exp.LogicalOr: rename_func("BOOL_OR"),
197            exp.Map: _map_sql,
198            exp.Pivot: transforms.preprocess([transforms.unqualify_pivot_fields]),
199            exp.Reduce: rename_func("AGGREGATE"),
200            exp.RegexpReplace: lambda self, e: self.func(
201                "REGEXP_REPLACE",
202                e.this,
203                e.expression,
204                e.args["replacement"],
205                e.args.get("position"),
206            ),
207            exp.Select: transforms.preprocess(
208                [
209                    transforms.eliminate_qualify,
210                    transforms.eliminate_distinct_on,
211                    transforms.unnest_to_explode,
212                    transforms.any_to_exists,
213                ]
214            ),
215            exp.SHA2Digest: lambda self, e: self.func(
216                "SHA2", e.this, e.args.get("length") or exp.Literal.number(256)
217            ),
218            exp.StrToDate: _str_to_date,
219            exp.StrToTime: lambda self, e: self.func("TO_TIMESTAMP", e.this, self.format_time(e)),
220            exp.TimestampTrunc: timestamptrunc_sql(),
221            exp.UnixToTime: _unix_to_time_sql,
222            exp.VariancePop: rename_func("VAR_POP"),
223            exp.WeekOfYear: rename_func("WEEKOFYEAR"),
224            exp.WithinGroup: transforms.preprocess(
225                [transforms.remove_within_group_for_percentiles]
226            ),
227            exp.ILike: None,
228            exp.Left: None,
229            exp.MonthsBetween: None,
230            exp.Right: None,
231        }.items()
232        if v is not None
233    }
234
235    WRAP_DERIVED_VALUES = False
236    CREATE_FUNCTION_RETURN_AS = False
237
238    def struct_sql(self, expression: exp.Struct) -> str:
239        from sqlglot.generator import Generator
240
241        return Generator.struct_sql(self, expression)
242
243    def cast_sql(self, expression: exp.Cast, safe_prefix: str | None = None) -> str:
244        arg = expression.this
245        if (
246            isinstance(arg, exp.Literal)
247            and arg.is_number
248            and expression.to.is_type(exp.DType.DECIMAL)
249            and len(expression.to.expressions) == 2
250            and isinstance(expression.to.expressions[1].this, exp.Neg)
251        ):
252            # Spark 2 supports negative scales in literals, but not in DECIMAL type syntax.
253            return f"{self.sql(arg)}BD"
254
255        is_json_extract = isinstance(
256            arg, (exp.JSONExtract, exp.JSONExtractScalar)
257        ) and not arg.args.get("variant_extract")
258
259        # We can't use a non-nested type (eg. STRING) as a schema
260        if expression.to.args.get("nested") and (is_parse_json(arg) or is_json_extract):
261            schema = f"'{self.sql(expression, 'to')}'"
262            return self.func("FROM_JSON", arg if is_json_extract else arg.this, schema)
263
264        if is_parse_json(expression):
265            return self.func("TO_JSON", arg)
266
267        return super(HiveGenerator, self).cast_sql(expression, safe_prefix=safe_prefix)
268
269    def fileformatproperty_sql(self, expression: exp.FileFormatProperty) -> str:
270        if expression.args.get("hive_format"):
271            return super().fileformatproperty_sql(expression)
272
273        return f"USING {expression.name.upper()}"
274
275    def altercolumn_sql(self, expression: exp.AlterColumn) -> str:
276        if expression.args.get("exists"):
277            self.unsupported("ALTER COLUMN IF EXISTS is not supported by this dialect")
278
279        this = self.sql(expression, "this")
280        new_name = self.sql(expression, "rename_to") or this
281        comment = self.sql(expression, "comment")
282        if new_name == this:
283            if comment:
284                return f"ALTER COLUMN {this} COMMENT {comment}"
285            return super(HiveGenerator, self).altercolumn_sql(expression)
286        return f"RENAME COLUMN {this} TO {new_name}"
287
288    def renamecolumn_sql(self, expression: exp.RenameColumn) -> str:
289        return super(HiveGenerator, self).renamecolumn_sql(expression)
290
291    def bracket_sql(self, expression: exp.Bracket) -> str:
292        if expression.args.get("safe") is False:
293            return bracket_to_element_at_sql(self, expression)
294
295        return super().bracket_sql(expression)