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    date_trunc_to_time,
 10    format_time_lambda,
 11    if_sql,
 12    left_to_substring_sql,
 13    no_ilike_sql,
 14    no_pivot_sql,
 15    no_safe_divide_sql,
 16    regexp_extract_sql,
 17    rename_func,
 18    right_to_substring_sql,
 19    struct_extract_sql,
 20    timestamptrunc_sql,
 21    timestrtotime_sql,
 22)
 23from sqlglot.dialects.mysql import MySQL
 24from sqlglot.errors import UnsupportedError
 25from sqlglot.helper import apply_index_offset, seq_get
 26from sqlglot.tokens import TokenType
 27
 28
 29def _approx_distinct_sql(self: generator.Generator, expression: exp.ApproxDistinct) -> str:
 30    accuracy = expression.args.get("accuracy")
 31    accuracy = ", " + self.sql(accuracy) if accuracy else ""
 32    return f"APPROX_DISTINCT({self.sql(expression, 'this')}{accuracy})"
 33
 34
 35def _datatype_sql(self: generator.Generator, expression: exp.DataType) -> str:
 36    sql = self.datatype_sql(expression)
 37    if expression.is_type("timestamptz"):
 38        sql = f"{sql} WITH TIME ZONE"
 39    return sql
 40
 41
 42def _explode_to_unnest_sql(self: generator.Generator, expression: exp.Lateral) -> str:
 43    if isinstance(expression.this, (exp.Explode, exp.Posexplode)):
 44        return self.sql(
 45            exp.Join(
 46                this=exp.Unnest(
 47                    expressions=[expression.this.this],
 48                    alias=expression.args.get("alias"),
 49                    ordinality=isinstance(expression.this, exp.Posexplode),
 50                ),
 51                kind="cross",
 52            )
 53        )
 54    return self.lateral_sql(expression)
 55
 56
 57def _initcap_sql(self: generator.Generator, expression: exp.Initcap) -> str:
 58    regex = r"(\w)(\w*)"
 59    return f"REGEXP_REPLACE({self.sql(expression, 'this')}, '{regex}', x -> UPPER(x[1]) || LOWER(x[2]))"
 60
 61
 62def _decode_sql(self: generator.Generator, expression: exp.Decode) -> str:
 63    _ensure_utf8(expression.args["charset"])
 64    return self.func("FROM_UTF8", expression.this, expression.args.get("replace"))
 65
 66
 67def _encode_sql(self: generator.Generator, expression: exp.Encode) -> str:
 68    _ensure_utf8(expression.args["charset"])
 69    return f"TO_UTF8({self.sql(expression, 'this')})"
 70
 71
 72def _no_sort_array(self: generator.Generator, expression: exp.SortArray) -> str:
 73    if expression.args.get("asc") == exp.false():
 74        comparator = "(a, b) -> CASE WHEN a < b THEN 1 WHEN a > b THEN -1 ELSE 0 END"
 75    else:
 76        comparator = None
 77    return self.func("ARRAY_SORT", expression.this, comparator)
 78
 79
 80def _schema_sql(self: generator.Generator, expression: exp.Schema) -> str:
 81    if isinstance(expression.parent, exp.Property):
 82        columns = ", ".join(f"'{c.name}'" for c in expression.expressions)
 83        return f"ARRAY[{columns}]"
 84
 85    if expression.parent:
 86        for schema in expression.parent.find_all(exp.Schema):
 87            if isinstance(schema.parent, exp.Property):
 88                expression = expression.copy()
 89                expression.expressions.extend(schema.expressions)
 90
 91    return self.schema_sql(expression)
 92
 93
 94def _quantile_sql(self: generator.Generator, expression: exp.Quantile) -> str:
 95    self.unsupported("Presto does not support exact quantiles")
 96    return f"APPROX_PERCENTILE({self.sql(expression, 'this')}, {self.sql(expression, 'quantile')})"
 97
 98
 99def _str_to_time_sql(
100    self: generator.Generator, expression: exp.StrToDate | exp.StrToTime | exp.TsOrDsToDate
101) -> str:
102    return f"DATE_PARSE({self.sql(expression, 'this')}, {self.format_time(expression)})"
103
104
105def _ts_or_ds_to_date_sql(self: generator.Generator, expression: exp.TsOrDsToDate) -> str:
106    time_format = self.format_time(expression)
107    if time_format and time_format not in (Presto.TIME_FORMAT, Presto.DATE_FORMAT):
108        return exp.cast(_str_to_time_sql(self, expression), "DATE").sql(dialect="presto")
109    return exp.cast(exp.cast(expression.this, "TIMESTAMP"), "DATE").sql(dialect="presto")
110
111
112def _ts_or_ds_add_sql(self: generator.Generator, expression: exp.TsOrDsAdd) -> str:
113    this = expression.this
114
115    if not isinstance(this, exp.CurrentDate):
116        this = exp.cast(exp.cast(expression.this, "TIMESTAMP"), "DATE")
117
118    return self.func(
119        "DATE_ADD",
120        exp.Literal.string(expression.text("unit") or "day"),
121        expression.expression,
122        this,
123    )
124
125
126def _ensure_utf8(charset: exp.Literal) -> None:
127    if charset.name.lower() != "utf-8":
128        raise UnsupportedError(f"Unsupported charset {charset}")
129
130
131def _approx_percentile(args: t.List) -> exp.Expression:
132    if len(args) == 4:
133        return exp.ApproxQuantile(
134            this=seq_get(args, 0),
135            weight=seq_get(args, 1),
136            quantile=seq_get(args, 2),
137            accuracy=seq_get(args, 3),
138        )
139    if len(args) == 3:
140        return exp.ApproxQuantile(
141            this=seq_get(args, 0), quantile=seq_get(args, 1), accuracy=seq_get(args, 2)
142        )
143    return exp.ApproxQuantile.from_arg_list(args)
144
145
146def _from_unixtime(args: t.List) -> exp.Expression:
147    if len(args) == 3:
148        return exp.UnixToTime(
149            this=seq_get(args, 0),
150            hours=seq_get(args, 1),
151            minutes=seq_get(args, 2),
152        )
153    if len(args) == 2:
154        return exp.UnixToTime(this=seq_get(args, 0), zone=seq_get(args, 1))
155
156    return exp.UnixToTime.from_arg_list(args)
157
158
159def _parse_element_at(args: t.List) -> exp.SafeBracket:
160    this = seq_get(args, 0)
161    index = seq_get(args, 1)
162    assert isinstance(this, exp.Expression) and isinstance(index, exp.Expression)
163    return exp.SafeBracket(this=this, expressions=apply_index_offset(this, [index], -1))
164
165
166def _unnest_sequence(expression: exp.Expression) -> exp.Expression:
167    if isinstance(expression, exp.Table):
168        if isinstance(expression.this, exp.GenerateSeries):
169            unnest = exp.Unnest(expressions=[expression.this])
170
171            if expression.alias:
172                return exp.alias_(unnest, alias="_u", table=[expression.alias], copy=False)
173            return unnest
174    return expression
175
176
177class Presto(Dialect):
178    INDEX_OFFSET = 1
179    NULL_ORDERING = "nulls_are_last"
180    TIME_FORMAT = MySQL.TIME_FORMAT
181    TIME_MAPPING = MySQL.TIME_MAPPING
182    STRICT_STRING_CONCAT = True
183
184    # https://github.com/trinodb/trino/issues/17
185    # https://github.com/trinodb/trino/issues/12289
186    # https://github.com/prestodb/presto/issues/2863
187    RESOLVES_IDENTIFIERS_AS_UPPERCASE = None
188
189    class Tokenizer(tokens.Tokenizer):
190        KEYWORDS = {
191            **tokens.Tokenizer.KEYWORDS,
192            "START": TokenType.BEGIN,
193            "MATCH_RECOGNIZE": TokenType.MATCH_RECOGNIZE,
194            "ROW": TokenType.STRUCT,
195        }
196
197    class Parser(parser.Parser):
198        FUNCTIONS = {
199            **parser.Parser.FUNCTIONS,
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            "SEQUENCE": exp.GenerateSeries.from_arg_list,
228            "STRPOS": lambda args: exp.StrPosition(
229                this=seq_get(args, 0), substr=seq_get(args, 1), instance=seq_get(args, 2)
230            ),
231            "TO_UNIXTIME": exp.TimeToUnix.from_arg_list,
232            "TO_HEX": exp.Hex.from_arg_list,
233            "TO_UTF8": lambda args: exp.Encode(
234                this=seq_get(args, 0), charset=exp.Literal.string("utf-8")
235            ),
236        }
237        FUNCTION_PARSERS = parser.Parser.FUNCTION_PARSERS.copy()
238        FUNCTION_PARSERS.pop("TRIM")
239
240    class Generator(generator.Generator):
241        INTERVAL_ALLOWS_PLURAL_FORM = False
242        JOIN_HINTS = False
243        TABLE_HINTS = False
244        QUERY_HINTS = False
245        IS_BOOL_ALLOWED = False
246        STRUCT_DELIMITER = ("(", ")")
247
248        PROPERTIES_LOCATION = {
249            **generator.Generator.PROPERTIES_LOCATION,
250            exp.LocationProperty: exp.Properties.Location.UNSUPPORTED,
251            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
252        }
253
254        TYPE_MAPPING = {
255            **generator.Generator.TYPE_MAPPING,
256            exp.DataType.Type.INT: "INTEGER",
257            exp.DataType.Type.FLOAT: "REAL",
258            exp.DataType.Type.BINARY: "VARBINARY",
259            exp.DataType.Type.TEXT: "VARCHAR",
260            exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMP",
261            exp.DataType.Type.STRUCT: "ROW",
262        }
263
264        TRANSFORMS = {
265            **generator.Generator.TRANSFORMS,
266            exp.ApproxDistinct: _approx_distinct_sql,
267            exp.ApproxQuantile: rename_func("APPROX_PERCENTILE"),
268            exp.Array: lambda self, e: f"ARRAY[{self.expressions(e, flat=True)}]",
269            exp.ArrayConcat: rename_func("CONCAT"),
270            exp.ArrayContains: rename_func("CONTAINS"),
271            exp.ArraySize: rename_func("CARDINALITY"),
272            exp.BitwiseAnd: lambda self, e: f"BITWISE_AND({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
273            exp.BitwiseLeftShift: lambda self, e: f"BITWISE_ARITHMETIC_SHIFT_LEFT({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
274            exp.BitwiseNot: lambda self, e: f"BITWISE_NOT({self.sql(e, 'this')})",
275            exp.BitwiseOr: lambda self, e: f"BITWISE_OR({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
276            exp.BitwiseRightShift: lambda self, e: f"BITWISE_ARITHMETIC_SHIFT_RIGHT({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
277            exp.BitwiseXor: lambda self, e: f"BITWISE_XOR({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
278            exp.Cast: transforms.preprocess([transforms.epoch_cast_to_ts]),
279            exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
280            exp.DataType: _datatype_sql,
281            exp.DateAdd: lambda self, e: self.func(
282                "DATE_ADD", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
283            ),
284            exp.DateDiff: lambda self, e: self.func(
285                "DATE_DIFF", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
286            ),
287            exp.DateStrToDate: lambda self, e: f"CAST(DATE_PARSE({self.sql(e, 'this')}, {Presto.DATE_FORMAT}) AS DATE)",
288            exp.DateToDi: lambda self, e: f"CAST(DATE_FORMAT({self.sql(e, 'this')}, {Presto.DATEINT_FORMAT}) AS INT)",
289            exp.Decode: _decode_sql,
290            exp.DiToDate: lambda self, e: f"CAST(DATE_PARSE(CAST({self.sql(e, 'this')} AS VARCHAR), {Presto.DATEINT_FORMAT}) AS DATE)",
291            exp.Encode: _encode_sql,
292            exp.FileFormatProperty: lambda self, e: f"FORMAT='{e.name.upper()}'",
293            exp.Group: transforms.preprocess([transforms.unalias_group]),
294            exp.Hex: rename_func("TO_HEX"),
295            exp.If: if_sql,
296            exp.ILike: no_ilike_sql,
297            exp.Initcap: _initcap_sql,
298            exp.Lateral: _explode_to_unnest_sql,
299            exp.Left: left_to_substring_sql,
300            exp.Levenshtein: rename_func("LEVENSHTEIN_DISTANCE"),
301            exp.LogicalAnd: rename_func("BOOL_AND"),
302            exp.LogicalOr: rename_func("BOOL_OR"),
303            exp.Pivot: no_pivot_sql,
304            exp.Quantile: _quantile_sql,
305            exp.RegexpExtract: regexp_extract_sql,
306            exp.Right: right_to_substring_sql,
307            exp.SafeBracket: lambda self, e: self.func(
308                "ELEMENT_AT", e.this, seq_get(apply_index_offset(e.this, e.expressions, 1), 0)
309            ),
310            exp.SafeDivide: no_safe_divide_sql,
311            exp.Schema: _schema_sql,
312            exp.Select: transforms.preprocess(
313                [
314                    transforms.eliminate_qualify,
315                    transforms.eliminate_distinct_on,
316                    transforms.explode_to_unnest,
317                ]
318            ),
319            exp.SortArray: _no_sort_array,
320            exp.StrPosition: rename_func("STRPOS"),
321            exp.StrToDate: lambda self, e: f"CAST({_str_to_time_sql(self, e)} AS DATE)",
322            exp.StrToTime: _str_to_time_sql,
323            exp.StrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {self.format_time(e)}))",
324            exp.StructExtract: struct_extract_sql,
325            exp.Table: transforms.preprocess([_unnest_sequence]),
326            exp.TimestampTrunc: timestamptrunc_sql,
327            exp.TimeStrToDate: timestrtotime_sql,
328            exp.TimeStrToTime: timestrtotime_sql,
329            exp.TimeStrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {Presto.TIME_FORMAT}))",
330            exp.TimeToStr: lambda self, e: f"DATE_FORMAT({self.sql(e, 'this')}, {self.format_time(e)})",
331            exp.TimeToUnix: rename_func("TO_UNIXTIME"),
332            exp.TryCast: transforms.preprocess([transforms.epoch_cast_to_ts]),
333            exp.TsOrDiToDi: lambda self, e: f"CAST(SUBSTR(REPLACE(CAST({self.sql(e, 'this')} AS VARCHAR), '-', ''), 1, 8) AS INT)",
334            exp.TsOrDsAdd: _ts_or_ds_add_sql,
335            exp.TsOrDsToDate: _ts_or_ds_to_date_sql,
336            exp.Unhex: rename_func("FROM_HEX"),
337            exp.UnixToStr: lambda self, e: f"DATE_FORMAT(FROM_UNIXTIME({self.sql(e, 'this')}), {self.format_time(e)})",
338            exp.UnixToTime: rename_func("FROM_UNIXTIME"),
339            exp.UnixToTimeStr: lambda self, e: f"CAST(FROM_UNIXTIME({self.sql(e, 'this')}) AS VARCHAR)",
340            exp.VariancePop: rename_func("VAR_POP"),
341            exp.With: transforms.preprocess([transforms.add_recursive_cte_column_names]),
342            exp.WithinGroup: transforms.preprocess(
343                [transforms.remove_within_group_for_percentiles]
344            ),
345        }
346
347        def interval_sql(self, expression: exp.Interval) -> str:
348            unit = self.sql(expression, "unit")
349            if expression.this and unit.lower().startswith("week"):
350                return f"({expression.this.name} * INTERVAL '7' day)"
351            return super().interval_sql(expression)
352
353        def transaction_sql(self, expression: exp.Transaction) -> str:
354            modes = expression.args.get("modes")
355            modes = f" {', '.join(modes)}" if modes else ""
356            return f"START TRANSACTION{modes}"
357
358        def generateseries_sql(self, expression: exp.GenerateSeries) -> str:
359            start = expression.args["start"]
360            end = expression.args["end"]
361            step = expression.args.get("step")
362
363            if isinstance(start, exp.Cast):
364                target_type = start.to
365            elif isinstance(end, exp.Cast):
366                target_type = end.to
367            else:
368                target_type = None
369
370            if target_type and target_type.is_type("timestamp"):
371                to = target_type.copy()
372
373                if target_type is start.to:
374                    end = exp.cast(end, to)
375                else:
376                    start = exp.cast(start, to)
377
378            return self.func("SEQUENCE", start, end, step)
379
380        def offset_limit_modifiers(
381            self, expression: exp.Expression, fetch: bool, limit: t.Optional[exp.Fetch | exp.Limit]
382        ) -> t.List[str]:
383            return [
384                self.sql(expression, "offset"),
385                self.sql(limit),
386            ]