sqlglot.dialects.presto

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, generator, parser, tokens, transforms
  6from sqlglot.dialects.dialect import (
  7    Dialect,
  8    binary_from_function,
  9    bool_xor_sql,
 10    date_trunc_to_time,
 11    encode_decode_sql,
 12    format_time_lambda,
 13    if_sql,
 14    left_to_substring_sql,
 15    no_ilike_sql,
 16    no_pivot_sql,
 17    no_safe_divide_sql,
 18    no_timestamp_sql,
 19    regexp_extract_sql,
 20    rename_func,
 21    right_to_substring_sql,
 22    struct_extract_sql,
 23    timestamptrunc_sql,
 24    timestrtotime_sql,
 25)
 26from sqlglot.dialects.mysql import MySQL
 27from sqlglot.helper import apply_index_offset, seq_get
 28from sqlglot.tokens import TokenType
 29
 30
 31def _approx_distinct_sql(self: Presto.Generator, expression: exp.ApproxDistinct) -> str:
 32    accuracy = expression.args.get("accuracy")
 33    accuracy = ", " + self.sql(accuracy) if accuracy else ""
 34    return f"APPROX_DISTINCT({self.sql(expression, 'this')}{accuracy})"
 35
 36
 37def _explode_to_unnest_sql(self: Presto.Generator, expression: exp.Lateral) -> str:
 38    if isinstance(expression.this, exp.Explode):
 39        expression = expression.copy()
 40        return self.sql(
 41            exp.Join(
 42                this=exp.Unnest(
 43                    expressions=[expression.this.this],
 44                    alias=expression.args.get("alias"),
 45                    offset=isinstance(expression.this, exp.Posexplode),
 46                ),
 47                kind="cross",
 48            )
 49        )
 50    return self.lateral_sql(expression)
 51
 52
 53def _initcap_sql(self: Presto.Generator, expression: exp.Initcap) -> str:
 54    regex = r"(\w)(\w*)"
 55    return f"REGEXP_REPLACE({self.sql(expression, 'this')}, '{regex}', x -> UPPER(x[1]) || LOWER(x[2]))"
 56
 57
 58def _no_sort_array(self: Presto.Generator, expression: exp.SortArray) -> str:
 59    if expression.args.get("asc") == exp.false():
 60        comparator = "(a, b) -> CASE WHEN a < b THEN 1 WHEN a > b THEN -1 ELSE 0 END"
 61    else:
 62        comparator = None
 63    return self.func("ARRAY_SORT", expression.this, comparator)
 64
 65
 66def _schema_sql(self: Presto.Generator, expression: exp.Schema) -> str:
 67    if isinstance(expression.parent, exp.Property):
 68        columns = ", ".join(f"'{c.name}'" for c in expression.expressions)
 69        return f"ARRAY[{columns}]"
 70
 71    if expression.parent:
 72        for schema in expression.parent.find_all(exp.Schema):
 73            column_defs = schema.find_all(exp.ColumnDef)
 74            if column_defs and isinstance(schema.parent, exp.Property):
 75                expression = expression.copy()
 76                expression.expressions.extend(column_defs)
 77
 78    return self.schema_sql(expression)
 79
 80
 81def _quantile_sql(self: Presto.Generator, expression: exp.Quantile) -> str:
 82    self.unsupported("Presto does not support exact quantiles")
 83    return f"APPROX_PERCENTILE({self.sql(expression, 'this')}, {self.sql(expression, 'quantile')})"
 84
 85
 86def _str_to_time_sql(
 87    self: Presto.Generator, expression: exp.StrToDate | exp.StrToTime | exp.TsOrDsToDate
 88) -> str:
 89    return f"DATE_PARSE({self.sql(expression, 'this')}, {self.format_time(expression)})"
 90
 91
 92def _ts_or_ds_to_date_sql(self: Presto.Generator, expression: exp.TsOrDsToDate) -> str:
 93    time_format = self.format_time(expression)
 94    if time_format and time_format not in (Presto.TIME_FORMAT, Presto.DATE_FORMAT):
 95        return exp.cast(_str_to_time_sql(self, expression), "DATE").sql(dialect="presto")
 96    return exp.cast(exp.cast(expression.this, "TIMESTAMP", copy=True), "DATE").sql(dialect="presto")
 97
 98
 99def _ts_or_ds_add_sql(self: Presto.Generator, expression: exp.TsOrDsAdd) -> str:
100    this = expression.this
101
102    if not isinstance(this, exp.CurrentDate):
103        this = exp.cast(exp.cast(expression.this, "TIMESTAMP", copy=True), "DATE")
104
105    return self.func(
106        "DATE_ADD",
107        exp.Literal.string(expression.text("unit") or "day"),
108        expression.expression,
109        this,
110    )
111
112
113def _approx_percentile(args: t.List) -> exp.Expression:
114    if len(args) == 4:
115        return exp.ApproxQuantile(
116            this=seq_get(args, 0),
117            weight=seq_get(args, 1),
118            quantile=seq_get(args, 2),
119            accuracy=seq_get(args, 3),
120        )
121    if len(args) == 3:
122        return exp.ApproxQuantile(
123            this=seq_get(args, 0), quantile=seq_get(args, 1), accuracy=seq_get(args, 2)
124        )
125    return exp.ApproxQuantile.from_arg_list(args)
126
127
128def _from_unixtime(args: t.List) -> exp.Expression:
129    if len(args) == 3:
130        return exp.UnixToTime(
131            this=seq_get(args, 0),
132            hours=seq_get(args, 1),
133            minutes=seq_get(args, 2),
134        )
135    if len(args) == 2:
136        return exp.UnixToTime(this=seq_get(args, 0), zone=seq_get(args, 1))
137
138    return exp.UnixToTime.from_arg_list(args)
139
140
141def _parse_element_at(args: t.List) -> exp.SafeBracket:
142    this = seq_get(args, 0)
143    index = seq_get(args, 1)
144    assert isinstance(this, exp.Expression) and isinstance(index, exp.Expression)
145    return exp.SafeBracket(this=this, expressions=apply_index_offset(this, [index], -1))
146
147
148def _unnest_sequence(expression: exp.Expression) -> exp.Expression:
149    if isinstance(expression, exp.Table):
150        if isinstance(expression.this, exp.GenerateSeries):
151            unnest = exp.Unnest(expressions=[expression.this])
152
153            if expression.alias:
154                return exp.alias_(unnest, alias="_u", table=[expression.alias], copy=False)
155            return unnest
156    return expression
157
158
159def _first_last_sql(self: Presto.Generator, expression: exp.First | exp.Last) -> str:
160    """
161    Trino doesn't support FIRST / LAST as functions, but they're valid in the context
162    of MATCH_RECOGNIZE, so we need to preserve them in that case. In all other cases
163    they're converted into an ARBITRARY call.
164
165    Reference: https://trino.io/docs/current/sql/match-recognize.html#logical-navigation-functions
166    """
167    if isinstance(expression.find_ancestor(exp.MatchRecognize, exp.Select), exp.MatchRecognize):
168        return self.function_fallback_sql(expression)
169
170    return rename_func("ARBITRARY")(self, expression)
171
172
173class Presto(Dialect):
174    INDEX_OFFSET = 1
175    NULL_ORDERING = "nulls_are_last"
176    TIME_FORMAT = MySQL.TIME_FORMAT
177    TIME_MAPPING = MySQL.TIME_MAPPING
178    STRICT_STRING_CONCAT = True
179    SUPPORTS_SEMI_ANTI_JOIN = False
180
181    # https://github.com/trinodb/trino/issues/17
182    # https://github.com/trinodb/trino/issues/12289
183    # https://github.com/prestodb/presto/issues/2863
184    RESOLVES_IDENTIFIERS_AS_UPPERCASE = None
185
186    class Tokenizer(tokens.Tokenizer):
187        KEYWORDS = {
188            **tokens.Tokenizer.KEYWORDS,
189            "START": TokenType.BEGIN,
190            "MATCH_RECOGNIZE": TokenType.MATCH_RECOGNIZE,
191            "ROW": TokenType.STRUCT,
192            "IPADDRESS": TokenType.IPADDRESS,
193            "IPPREFIX": TokenType.IPPREFIX,
194        }
195
196    class Parser(parser.Parser):
197        FUNCTIONS = {
198            **parser.Parser.FUNCTIONS,
199            "ARBITRARY": exp.AnyValue.from_arg_list,
200            "APPROX_DISTINCT": exp.ApproxDistinct.from_arg_list,
201            "APPROX_PERCENTILE": _approx_percentile,
202            "BITWISE_AND": binary_from_function(exp.BitwiseAnd),
203            "BITWISE_NOT": lambda args: exp.BitwiseNot(this=seq_get(args, 0)),
204            "BITWISE_OR": binary_from_function(exp.BitwiseOr),
205            "BITWISE_XOR": binary_from_function(exp.BitwiseXor),
206            "CARDINALITY": exp.ArraySize.from_arg_list,
207            "CONTAINS": exp.ArrayContains.from_arg_list,
208            "DATE_ADD": lambda args: exp.DateAdd(
209                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
210            ),
211            "DATE_DIFF": lambda args: exp.DateDiff(
212                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
213            ),
214            "DATE_FORMAT": format_time_lambda(exp.TimeToStr, "presto"),
215            "DATE_PARSE": format_time_lambda(exp.StrToTime, "presto"),
216            "DATE_TRUNC": date_trunc_to_time,
217            "ELEMENT_AT": _parse_element_at,
218            "FROM_HEX": exp.Unhex.from_arg_list,
219            "FROM_UNIXTIME": _from_unixtime,
220            "FROM_UTF8": lambda args: exp.Decode(
221                this=seq_get(args, 0), replace=seq_get(args, 1), charset=exp.Literal.string("utf-8")
222            ),
223            "NOW": exp.CurrentTimestamp.from_arg_list,
224            "REGEXP_EXTRACT": lambda args: exp.RegexpExtract(
225                this=seq_get(args, 0), expression=seq_get(args, 1), group=seq_get(args, 2)
226            ),
227            "REGEXP_REPLACE": lambda args: exp.RegexpReplace(
228                this=seq_get(args, 0),
229                expression=seq_get(args, 1),
230                replacement=seq_get(args, 2) or exp.Literal.string(""),
231            ),
232            "ROW": exp.Struct.from_arg_list,
233            "SEQUENCE": exp.GenerateSeries.from_arg_list,
234            "SPLIT_TO_MAP": exp.StrToMap.from_arg_list,
235            "STRPOS": lambda args: exp.StrPosition(
236                this=seq_get(args, 0), substr=seq_get(args, 1), instance=seq_get(args, 2)
237            ),
238            "TO_UNIXTIME": exp.TimeToUnix.from_arg_list,
239            "TO_HEX": exp.Hex.from_arg_list,
240            "TO_UTF8": lambda args: exp.Encode(
241                this=seq_get(args, 0), charset=exp.Literal.string("utf-8")
242            ),
243        }
244
245        FUNCTION_PARSERS = parser.Parser.FUNCTION_PARSERS.copy()
246        FUNCTION_PARSERS.pop("TRIM")
247
248    class Generator(generator.Generator):
249        INTERVAL_ALLOWS_PLURAL_FORM = False
250        JOIN_HINTS = False
251        TABLE_HINTS = False
252        QUERY_HINTS = False
253        IS_BOOL_ALLOWED = False
254        TZ_TO_WITH_TIME_ZONE = True
255        NVL2_SUPPORTED = False
256        STRUCT_DELIMITER = ("(", ")")
257        LIMIT_ONLY_LITERALS = True
258
259        PROPERTIES_LOCATION = {
260            **generator.Generator.PROPERTIES_LOCATION,
261            exp.LocationProperty: exp.Properties.Location.UNSUPPORTED,
262            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
263        }
264
265        TYPE_MAPPING = {
266            **generator.Generator.TYPE_MAPPING,
267            exp.DataType.Type.INT: "INTEGER",
268            exp.DataType.Type.FLOAT: "REAL",
269            exp.DataType.Type.BINARY: "VARBINARY",
270            exp.DataType.Type.TEXT: "VARCHAR",
271            exp.DataType.Type.TIMETZ: "TIME",
272            exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMP",
273            exp.DataType.Type.STRUCT: "ROW",
274            exp.DataType.Type.DATETIME: "TIMESTAMP",
275            exp.DataType.Type.DATETIME64: "TIMESTAMP",
276        }
277
278        TRANSFORMS = {
279            **generator.Generator.TRANSFORMS,
280            exp.AnyValue: rename_func("ARBITRARY"),
281            exp.ApproxDistinct: _approx_distinct_sql,
282            exp.ApproxQuantile: rename_func("APPROX_PERCENTILE"),
283            exp.Array: lambda self, e: f"ARRAY[{self.expressions(e, flat=True)}]",
284            exp.ArrayConcat: rename_func("CONCAT"),
285            exp.ArrayContains: rename_func("CONTAINS"),
286            exp.ArraySize: rename_func("CARDINALITY"),
287            exp.BitwiseAnd: lambda self, e: f"BITWISE_AND({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
288            exp.BitwiseLeftShift: lambda self, e: f"BITWISE_ARITHMETIC_SHIFT_LEFT({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
289            exp.BitwiseNot: lambda self, e: f"BITWISE_NOT({self.sql(e, 'this')})",
290            exp.BitwiseOr: lambda self, e: f"BITWISE_OR({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
291            exp.BitwiseRightShift: lambda self, e: f"BITWISE_ARITHMETIC_SHIFT_RIGHT({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
292            exp.BitwiseXor: lambda self, e: f"BITWISE_XOR({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
293            exp.Cast: transforms.preprocess([transforms.epoch_cast_to_ts]),
294            exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
295            exp.DateAdd: lambda self, e: self.func(
296                "DATE_ADD", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
297            ),
298            exp.DateDiff: lambda self, e: self.func(
299                "DATE_DIFF", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
300            ),
301            exp.DateStrToDate: lambda self, e: f"CAST(DATE_PARSE({self.sql(e, 'this')}, {Presto.DATE_FORMAT}) AS DATE)",
302            exp.DateToDi: lambda self, e: f"CAST(DATE_FORMAT({self.sql(e, 'this')}, {Presto.DATEINT_FORMAT}) AS INT)",
303            exp.DateSub: lambda self, e: self.func(
304                "DATE_ADD",
305                exp.Literal.string(e.text("unit") or "day"),
306                e.expression * -1,
307                e.this,
308            ),
309            exp.Decode: lambda self, e: encode_decode_sql(self, e, "FROM_UTF8"),
310            exp.DiToDate: lambda self, e: f"CAST(DATE_PARSE(CAST({self.sql(e, 'this')} AS VARCHAR), {Presto.DATEINT_FORMAT}) AS DATE)",
311            exp.Encode: lambda self, e: encode_decode_sql(self, e, "TO_UTF8"),
312            exp.FileFormatProperty: lambda self, e: f"FORMAT='{e.name.upper()}'",
313            exp.First: _first_last_sql,
314            exp.Group: transforms.preprocess([transforms.unalias_group]),
315            exp.GroupConcat: lambda self, e: self.func(
316                "ARRAY_JOIN", self.func("ARRAY_AGG", e.this), e.args.get("separator")
317            ),
318            exp.Hex: rename_func("TO_HEX"),
319            exp.If: if_sql(),
320            exp.ILike: no_ilike_sql,
321            exp.Initcap: _initcap_sql,
322            exp.ParseJSON: rename_func("JSON_PARSE"),
323            exp.Last: _first_last_sql,
324            exp.Lateral: _explode_to_unnest_sql,
325            exp.Left: left_to_substring_sql,
326            exp.Levenshtein: rename_func("LEVENSHTEIN_DISTANCE"),
327            exp.LogicalAnd: rename_func("BOOL_AND"),
328            exp.LogicalOr: rename_func("BOOL_OR"),
329            exp.Pivot: no_pivot_sql,
330            exp.Quantile: _quantile_sql,
331            exp.RegexpExtract: regexp_extract_sql,
332            exp.Right: right_to_substring_sql,
333            exp.SafeBracket: lambda self, e: self.func(
334                "ELEMENT_AT", e.this, seq_get(apply_index_offset(e.this, e.expressions, 1), 0)
335            ),
336            exp.SafeDivide: no_safe_divide_sql,
337            exp.Schema: _schema_sql,
338            exp.Select: transforms.preprocess(
339                [
340                    transforms.eliminate_qualify,
341                    transforms.eliminate_distinct_on,
342                    transforms.explode_to_unnest(1),
343                    transforms.eliminate_semi_and_anti_joins,
344                ]
345            ),
346            exp.SortArray: _no_sort_array,
347            exp.StrPosition: rename_func("STRPOS"),
348            exp.StrToDate: lambda self, e: f"CAST({_str_to_time_sql(self, e)} AS DATE)",
349            exp.StrToMap: rename_func("SPLIT_TO_MAP"),
350            exp.StrToTime: _str_to_time_sql,
351            exp.StrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {self.format_time(e)}))",
352            exp.StructExtract: struct_extract_sql,
353            exp.Table: transforms.preprocess([_unnest_sequence]),
354            exp.Timestamp: no_timestamp_sql,
355            exp.TimestampTrunc: timestamptrunc_sql,
356            exp.TimeStrToDate: timestrtotime_sql,
357            exp.TimeStrToTime: timestrtotime_sql,
358            exp.TimeStrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {Presto.TIME_FORMAT}))",
359            exp.TimeToStr: lambda self, e: f"DATE_FORMAT({self.sql(e, 'this')}, {self.format_time(e)})",
360            exp.TimeToUnix: rename_func("TO_UNIXTIME"),
361            exp.TryCast: transforms.preprocess([transforms.epoch_cast_to_ts]),
362            exp.TsOrDiToDi: lambda self, e: f"CAST(SUBSTR(REPLACE(CAST({self.sql(e, 'this')} AS VARCHAR), '-', ''), 1, 8) AS INT)",
363            exp.TsOrDsAdd: _ts_or_ds_add_sql,
364            exp.TsOrDsToDate: _ts_or_ds_to_date_sql,
365            exp.Unhex: rename_func("FROM_HEX"),
366            exp.UnixToStr: lambda self, e: f"DATE_FORMAT(FROM_UNIXTIME({self.sql(e, 'this')}), {self.format_time(e)})",
367            exp.UnixToTime: rename_func("FROM_UNIXTIME"),
368            exp.UnixToTimeStr: lambda self, e: f"CAST(FROM_UNIXTIME({self.sql(e, 'this')}) AS VARCHAR)",
369            exp.VariancePop: rename_func("VAR_POP"),
370            exp.With: transforms.preprocess([transforms.add_recursive_cte_column_names]),
371            exp.WithinGroup: transforms.preprocess(
372                [transforms.remove_within_group_for_percentiles]
373            ),
374            exp.Xor: bool_xor_sql,
375        }
376
377        def struct_sql(self, expression: exp.Struct) -> str:
378            if any(isinstance(arg, (exp.EQ, exp.Slice)) for arg in expression.expressions):
379                self.unsupported("Struct with key-value definitions is unsupported.")
380                return self.function_fallback_sql(expression)
381
382            return rename_func("ROW")(self, expression)
383
384        def interval_sql(self, expression: exp.Interval) -> str:
385            unit = self.sql(expression, "unit")
386            if expression.this and unit.lower().startswith("week"):
387                return f"({expression.this.name} * INTERVAL '7' day)"
388            return super().interval_sql(expression)
389
390        def transaction_sql(self, expression: exp.Transaction) -> str:
391            modes = expression.args.get("modes")
392            modes = f" {', '.join(modes)}" if modes else ""
393            return f"START TRANSACTION{modes}"
394
395        def generateseries_sql(self, expression: exp.GenerateSeries) -> str:
396            start = expression.args["start"]
397            end = expression.args["end"]
398            step = expression.args.get("step")
399
400            if isinstance(start, exp.Cast):
401                target_type = start.to
402            elif isinstance(end, exp.Cast):
403                target_type = end.to
404            else:
405                target_type = None
406
407            if target_type and target_type.is_type("timestamp"):
408                to = target_type.copy()
409
410                if target_type is start.to:
411                    end = exp.cast(end, to)
412                else:
413                    start = exp.cast(start, to)
414
415            return self.func("SEQUENCE", start, end, step)
416
417        def offset_limit_modifiers(
418            self, expression: exp.Expression, fetch: bool, limit: t.Optional[exp.Fetch | exp.Limit]
419        ) -> t.List[str]:
420            return [
421                self.sql(expression, "offset"),
422                self.sql(limit),
423            ]
424
425        def create_sql(self, expression: exp.Create) -> str:
426            """
427            Presto doesn't support CREATE VIEW with expressions (ex: `CREATE VIEW x (cola)` then `(cola)` is the expression),
428            so we need to remove them
429            """
430            kind = expression.args["kind"]
431            schema = expression.this
432            if kind == "VIEW" and schema.expressions:
433                expression = expression.copy()
434                expression.this.set("expressions", None)
435            return super().create_sql(expression)