sqlglot.generators.spark2

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