From 28cc22419e32a65fea2d1678400265b8cabc3aff Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Thu, 15 Sep 2022 18:46:17 +0200 Subject: Adding upstream version 6.0.4. Signed-off-by: Daniel Baumann --- sqlglot/dialects/presto.py | 216 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 sqlglot/dialects/presto.py (limited to 'sqlglot/dialects/presto.py') diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py new file mode 100644 index 0000000..ca913e4 --- /dev/null +++ b/sqlglot/dialects/presto.py @@ -0,0 +1,216 @@ +from sqlglot import exp, transforms +from sqlglot.dialects.dialect import ( + Dialect, + format_time_lambda, + if_sql, + no_ilike_sql, + no_safe_divide_sql, + rename_func, + str_position_sql, + struct_extract_sql, +) +from sqlglot.dialects.mysql import MySQL +from sqlglot.generator import Generator +from sqlglot.helper import csv, list_get +from sqlglot.parser import Parser +from sqlglot.tokens import Tokenizer, TokenType + + +def _approx_distinct_sql(self, expression): + accuracy = expression.args.get("accuracy") + accuracy = ", " + self.sql(accuracy) if accuracy else "" + return f"APPROX_DISTINCT({self.sql(expression, 'this')}{accuracy})" + + +def _concat_ws_sql(self, expression): + sep, *args = expression.expressions + sep = self.sql(sep) + if len(args) > 1: + return f"ARRAY_JOIN(ARRAY[{csv(*(self.sql(e) for e in args))}], {sep})" + return f"ARRAY_JOIN({self.sql(args[0])}, {sep})" + + +def _datatype_sql(self, expression): + sql = self.datatype_sql(expression) + if expression.this == exp.DataType.Type.TIMESTAMPTZ: + sql = f"{sql} WITH TIME ZONE" + return sql + + +def _date_parse_sql(self, expression): + return f"DATE_PARSE({self.sql(expression, 'this')}, '%Y-%m-%d %H:%i:%s')" + + +def _explode_to_unnest_sql(self, expression): + if isinstance(expression.this, (exp.Explode, exp.Posexplode)): + return self.sql( + exp.Join( + this=exp.Unnest( + expressions=[expression.this.this], + alias=expression.args.get("alias"), + ordinality=isinstance(expression.this, exp.Posexplode), + ), + kind="cross", + ) + ) + return self.lateral_sql(expression) + + +def _initcap_sql(self, expression): + regex = r"(\w)(\w*)" + return f"REGEXP_REPLACE({self.sql(expression, 'this')}, '{regex}', x -> UPPER(x[1]) || LOWER(x[2]))" + + +def _no_sort_array(self, expression): + if expression.args.get("asc") == exp.FALSE: + comparator = "(a, b) -> CASE WHEN a < b THEN 1 WHEN a > b THEN -1 ELSE 0 END" + else: + comparator = None + args = csv(self.sql(expression, "this"), comparator) + return f"ARRAY_SORT({args})" + + +def _schema_sql(self, expression): + if isinstance(expression.parent, exp.Property): + columns = ", ".join(f"'{c.text('this')}'" for c in expression.expressions) + return f"ARRAY[{columns}]" + + for schema in expression.parent.find_all(exp.Schema): + if isinstance(schema.parent, exp.Property): + expression = expression.copy() + expression.expressions.extend(schema.expressions) + + return self.schema_sql(expression) + + +def _quantile_sql(self, expression): + self.unsupported("Presto does not support exact quantiles") + return f"APPROX_PERCENTILE({self.sql(expression, 'this')}, {self.sql(expression, 'quantile')})" + + +def _str_to_time_sql(self, expression): + return f"DATE_PARSE({self.sql(expression, 'this')}, {self.format_time(expression)})" + + +def _ts_or_ds_to_date_sql(self, expression): + time_format = self.format_time(expression) + if time_format and time_format not in (Presto.time_format, Presto.date_format): + return f"CAST({_str_to_time_sql(self, expression)} AS DATE)" + return ( + f"CAST(SUBSTR(CAST({self.sql(expression, 'this')} AS VARCHAR), 1, 10) AS DATE)" + ) + + +def _ts_or_ds_add_sql(self, expression): + this = self.sql(expression, "this") + e = self.sql(expression, "expression") + unit = self.sql(expression, "unit") or "'day'" + return f"DATE_ADD({unit}, {e}, DATE_PARSE(SUBSTR({this}, 1, 10), {Presto.date_format}))" + + +class Presto(Dialect): + index_offset = 1 + null_ordering = "nulls_are_last" + time_format = "'%Y-%m-%d %H:%i:%S'" + time_mapping = MySQL.time_mapping + + class Tokenizer(Tokenizer): + KEYWORDS = { + **Tokenizer.KEYWORDS, + "ROW": TokenType.STRUCT, + } + + class Parser(Parser): + FUNCTIONS = { + **Parser.FUNCTIONS, + "APPROX_DISTINCT": exp.ApproxDistinct.from_arg_list, + "CARDINALITY": exp.ArraySize.from_arg_list, + "CONTAINS": exp.ArrayContains.from_arg_list, + "DATE_ADD": lambda args: exp.DateAdd( + this=list_get(args, 2), + expression=list_get(args, 1), + unit=list_get(args, 0), + ), + "DATE_DIFF": lambda args: exp.DateDiff( + this=list_get(args, 2), + expression=list_get(args, 1), + unit=list_get(args, 0), + ), + "DATE_FORMAT": format_time_lambda(exp.TimeToStr, "presto"), + "DATE_PARSE": format_time_lambda(exp.StrToTime, "presto"), + "FROM_UNIXTIME": exp.UnixToTime.from_arg_list, + "STRPOS": exp.StrPosition.from_arg_list, + "TO_UNIXTIME": exp.TimeToUnix.from_arg_list, + } + + class Generator(Generator): + + STRUCT_DELIMITER = ("(", ")") + + WITH_PROPERTIES = [ + exp.PartitionedByProperty, + exp.FileFormatProperty, + exp.SchemaCommentProperty, + exp.AnonymousProperty, + exp.TableFormatProperty, + ] + + TYPE_MAPPING = { + **Generator.TYPE_MAPPING, + exp.DataType.Type.INT: "INTEGER", + exp.DataType.Type.FLOAT: "REAL", + exp.DataType.Type.BINARY: "VARBINARY", + exp.DataType.Type.TEXT: "VARCHAR", + exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMP", + exp.DataType.Type.STRUCT: "ROW", + } + + TRANSFORMS = { + **Generator.TRANSFORMS, + **transforms.UNALIAS_GROUP, + exp.ApproxDistinct: _approx_distinct_sql, + exp.Array: lambda self, e: f"ARRAY[{self.expressions(e, flat=True)}]", + exp.ArrayContains: rename_func("CONTAINS"), + exp.ArraySize: rename_func("CARDINALITY"), + exp.BitwiseAnd: lambda self, e: f"BITWISE_AND({self.sql(e, 'this')}, {self.sql(e, 'expression')})", + exp.BitwiseLeftShift: lambda self, e: f"BITWISE_ARITHMETIC_SHIFT_LEFT({self.sql(e, 'this')}, {self.sql(e, 'expression')})", + exp.BitwiseNot: lambda self, e: f"BITWISE_NOT({self.sql(e, 'this')})", + exp.BitwiseOr: lambda self, e: f"BITWISE_OR({self.sql(e, 'this')}, {self.sql(e, 'expression')})", + exp.BitwiseRightShift: lambda self, e: f"BITWISE_ARITHMETIC_SHIFT_RIGHT({self.sql(e, 'this')}, {self.sql(e, 'expression')})", + exp.BitwiseXor: lambda self, e: f"BITWISE_XOR({self.sql(e, 'this')}, {self.sql(e, 'expression')})", + exp.ConcatWs: _concat_ws_sql, + exp.DataType: _datatype_sql, + exp.DateAdd: lambda self, e: f"""DATE_ADD({self.sql(e, 'unit') or "'day'"}, {self.sql(e, 'expression')}, {self.sql(e, 'this')})""", + exp.DateDiff: lambda self, e: f"""DATE_DIFF({self.sql(e, 'unit') or "'day'"}, {self.sql(e, 'expression')}, {self.sql(e, 'this')})""", + exp.DateStrToDate: lambda self, e: f"CAST(DATE_PARSE({self.sql(e, 'this')}, {Presto.date_format}) AS DATE)", + exp.DateToDi: lambda self, e: f"CAST(DATE_FORMAT({self.sql(e, 'this')}, {Presto.dateint_format}) AS INT)", + exp.DiToDate: lambda self, e: f"CAST(DATE_PARSE(CAST({self.sql(e, 'this')} AS VARCHAR), {Presto.dateint_format}) AS DATE)", + exp.FileFormatProperty: lambda self, e: self.property_sql(e), + exp.If: if_sql, + exp.ILike: no_ilike_sql, + exp.Initcap: _initcap_sql, + exp.Lateral: _explode_to_unnest_sql, + exp.Levenshtein: rename_func("LEVENSHTEIN_DISTANCE"), + exp.PartitionedByProperty: lambda self, e: f"PARTITIONED_BY = {self.sql(e.args['value'])}", + exp.Quantile: _quantile_sql, + exp.SafeDivide: no_safe_divide_sql, + exp.Schema: _schema_sql, + exp.SortArray: _no_sort_array, + exp.StrPosition: str_position_sql, + exp.StrToDate: lambda self, e: f"CAST({_str_to_time_sql(self, e)} AS DATE)", + exp.StrToTime: _str_to_time_sql, + exp.StrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {self.format_time(e)}))", + exp.StructExtract: struct_extract_sql, + exp.TableFormatProperty: lambda self, e: f"TABLE_FORMAT = '{e.text('value').upper()}'", + exp.TimeStrToDate: _date_parse_sql, + exp.TimeStrToTime: _date_parse_sql, + exp.TimeStrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {Presto.time_format}))", + exp.TimeToStr: lambda self, e: f"DATE_FORMAT({self.sql(e, 'this')}, {self.format_time(e)})", + exp.TimeToUnix: rename_func("TO_UNIXTIME"), + exp.TsOrDiToDi: lambda self, e: f"CAST(SUBSTR(REPLACE(CAST({self.sql(e, 'this')} AS VARCHAR), '-', ''), 1, 8) AS INT)", + exp.TsOrDsAdd: _ts_or_ds_add_sql, + exp.TsOrDsToDate: _ts_or_ds_to_date_sql, + exp.UnixToStr: lambda self, e: f"DATE_FORMAT(FROM_UNIXTIME({self.sql(e, 'this')}), {self.format_time(e)})", + exp.UnixToTime: rename_func("FROM_UNIXTIME"), + exp.UnixToTimeStr: lambda self, e: f"CAST(FROM_UNIXTIME({self.sql(e, 'this')}) AS VARCHAR)", + } -- cgit v1.2.3