sqlglot.dialects.doris

  1from __future__ import annotations
  2
  3from sqlglot import exp
  4from sqlglot.dialects.dialect import (
  5    approx_count_distinct_sql,
  6    build_timestamp_trunc,
  7    property_sql,
  8    rename_func,
  9    time_format,
 10    unit_to_str,
 11)
 12from sqlglot.dialects.mysql import MySQL
 13from sqlglot.tokens import TokenType
 14
 15
 16def _lag_lead_sql(self, expression: exp.Lag | exp.Lead) -> str:
 17    return self.func(
 18        "LAG" if isinstance(expression, exp.Lag) else "LEAD",
 19        expression.this,
 20        expression.args.get("offset") or exp.Literal.number(1),
 21        expression.args.get("default") or exp.null(),
 22    )
 23
 24
 25class Doris(MySQL):
 26    DATE_FORMAT = "'yyyy-MM-dd'"
 27    DATEINT_FORMAT = "'yyyyMMdd'"
 28    TIME_FORMAT = "'yyyy-MM-dd HH:mm:ss'"
 29
 30    class Parser(MySQL.Parser):
 31        FUNCTIONS = {
 32            **MySQL.Parser.FUNCTIONS,
 33            "COLLECT_SET": exp.ArrayUniqueAgg.from_arg_list,
 34            "DATE_TRUNC": build_timestamp_trunc,
 35            "MONTHS_ADD": exp.AddMonths.from_arg_list,
 36            "REGEXP": exp.RegexpLike.from_arg_list,
 37            "TO_DATE": exp.TsOrDsToDate.from_arg_list,
 38        }
 39
 40        FUNCTION_PARSERS = MySQL.Parser.FUNCTION_PARSERS.copy()
 41        FUNCTION_PARSERS.pop("GROUP_CONCAT")
 42
 43        PROPERTY_PARSERS = {
 44            **MySQL.Parser.PROPERTY_PARSERS,
 45            "PROPERTIES": lambda self: self._parse_wrapped_properties(),
 46            "UNIQUE": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 47            "PARTITION BY": lambda self: self._parse_partition_by_opt_range(),
 48        }
 49
 50        def _parse_partitioning_granularity_dynamic(self) -> exp.PartitionByRangePropertyDynamic:
 51            self._match_text_seq("FROM")
 52            start = self._parse_wrapped(self._parse_string)
 53            self._match_text_seq("TO")
 54            end = self._parse_wrapped(self._parse_string)
 55            self._match_text_seq("INTERVAL")
 56            number = self._parse_number()
 57            unit = self._parse_var(any_token=True)
 58            every = self.expression(exp.Interval, this=number, unit=unit)
 59            return self.expression(
 60                exp.PartitionByRangePropertyDynamic, start=start, end=end, every=every
 61            )
 62
 63        def _parse_partition_definition(self) -> exp.Partition:
 64            self._match_text_seq("PARTITION")
 65
 66            name = self._parse_id_var()
 67            self._match_text_seq("VALUES")
 68
 69            if self._match_text_seq("LESS", "THAN"):
 70                values = self._parse_wrapped_csv(self._parse_expression)
 71                if len(values) == 1 and values[0].name.upper() == "MAXVALUE":
 72                    values = [exp.var("MAXVALUE")]
 73
 74                part_range = self.expression(exp.PartitionRange, this=name, expressions=values)
 75                return self.expression(exp.Partition, expressions=[part_range])
 76
 77            self._match(TokenType.L_BRACKET)
 78            values = self._parse_csv(lambda: self._parse_wrapped_csv(self._parse_expression))
 79
 80            self._match(TokenType.R_BRACKET)
 81            self._match(TokenType.R_PAREN)
 82
 83            part_range = self.expression(exp.PartitionRange, this=name, expressions=values)
 84            return self.expression(exp.Partition, expressions=[part_range])
 85
 86        def _parse_partition_by_opt_range(
 87            self,
 88        ) -> exp.PartitionedByProperty | exp.PartitionByRangeProperty:
 89            if not self._match_text_seq("RANGE"):
 90                return super()._parse_partitioned_by()
 91
 92            partition_expressions = self._parse_wrapped_id_vars()
 93            self._match_l_paren()
 94
 95            if self._match_text_seq("FROM", advance=False):
 96                create_expressions = self._parse_csv(self._parse_partitioning_granularity_dynamic)
 97            elif self._match_text_seq("PARTITION", advance=False):
 98                create_expressions = self._parse_csv(self._parse_partition_definition)
 99            else:
100                create_expressions = None
101
102            self._match_r_paren()
103
104            return self.expression(
105                exp.PartitionByRangeProperty,
106                partition_expressions=partition_expressions,
107                create_expressions=create_expressions,
108            )
109
110    class Generator(MySQL.Generator):
111        LAST_DAY_SUPPORTS_DATE_PART = False
112        VARCHAR_REQUIRES_SIZE = False
113        WITH_PROPERTIES_PREFIX = "PROPERTIES"
114
115        TYPE_MAPPING = {
116            **MySQL.Generator.TYPE_MAPPING,
117            exp.DataType.Type.TEXT: "STRING",
118            exp.DataType.Type.TIMESTAMP: "DATETIME",
119            exp.DataType.Type.TIMESTAMPTZ: "DATETIME",
120        }
121
122        PROPERTIES_LOCATION = {
123            **MySQL.Generator.PROPERTIES_LOCATION,
124            exp.UniqueKeyProperty: exp.Properties.Location.POST_SCHEMA,
125            exp.PartitionByRangeProperty: exp.Properties.Location.POST_SCHEMA,
126        }
127
128        CAST_MAPPING = {}
129        TIMESTAMP_FUNC_TYPES = set()
130
131        TRANSFORMS = {
132            **MySQL.Generator.TRANSFORMS,
133            exp.AddMonths: rename_func("MONTHS_ADD"),
134            exp.ApproxDistinct: approx_count_distinct_sql,
135            exp.ArgMax: rename_func("MAX_BY"),
136            exp.ArgMin: rename_func("MIN_BY"),
137            exp.ArrayAgg: rename_func("COLLECT_LIST"),
138            exp.ArrayToString: rename_func("ARRAY_JOIN"),
139            exp.ArrayUniqueAgg: rename_func("COLLECT_SET"),
140            exp.CurrentTimestamp: lambda self, _: self.func("NOW"),
141            exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, unit_to_str(e)),
142            exp.GroupConcat: lambda self, e: self.func(
143                "GROUP_CONCAT", e.this, e.args.get("separator") or exp.Literal.string(",")
144            ),
145            exp.JSONExtractScalar: lambda self, e: self.func("JSON_EXTRACT", e.this, e.expression),
146            exp.Lag: _lag_lead_sql,
147            exp.Lead: _lag_lead_sql,
148            exp.Map: rename_func("ARRAY_MAP"),
149            exp.Property: property_sql,
150            exp.RegexpLike: rename_func("REGEXP"),
151            exp.RegexpSplit: rename_func("SPLIT_BY_STRING"),
152            exp.SchemaCommentProperty: lambda self, e: self.naked_property(e),
153            exp.Split: rename_func("SPLIT_BY_STRING"),
154            exp.StringToArray: rename_func("SPLIT_BY_STRING"),
155            exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
156            exp.TimeStrToDate: rename_func("TO_DATE"),
157            exp.TsOrDsAdd: lambda self, e: self.func("DATE_ADD", e.this, e.expression),
158            exp.TsOrDsToDate: lambda self, e: self.func("TO_DATE", e.this),
159            exp.TimeToUnix: rename_func("UNIX_TIMESTAMP"),
160            exp.TimestampTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, unit_to_str(e)),
161            exp.UnixToStr: lambda self, e: self.func(
162                "FROM_UNIXTIME", e.this, time_format("doris")(self, e)
163            ),
164            exp.UnixToTime: rename_func("FROM_UNIXTIME"),
165        }
166
167        # https://github.com/apache/doris/blob/e4f41dbf1ec03f5937fdeba2ee1454a20254015b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4#L93
168        RESERVED_KEYWORDS = {
169            "account_lock",
170            "account_unlock",
171            "add",
172            "adddate",
173            "admin",
174            "after",
175            "agg_state",
176            "aggregate",
177            "alias",
178            "all",
179            "alter",
180            "analyze",
181            "analyzed",
182            "and",
183            "anti",
184            "append",
185            "array",
186            "array_range",
187            "as",
188            "asc",
189            "at",
190            "authors",
191            "auto",
192            "auto_increment",
193            "backend",
194            "backends",
195            "backup",
196            "begin",
197            "belong",
198            "between",
199            "bigint",
200            "bin",
201            "binary",
202            "binlog",
203            "bitand",
204            "bitmap",
205            "bitmap_union",
206            "bitor",
207            "bitxor",
208            "blob",
209            "boolean",
210            "brief",
211            "broker",
212            "buckets",
213            "build",
214            "builtin",
215            "bulk",
216            "by",
217            "cached",
218            "call",
219            "cancel",
220            "case",
221            "cast",
222            "catalog",
223            "catalogs",
224            "chain",
225            "char",
226            "character",
227            "charset",
228            "check",
229            "clean",
230            "cluster",
231            "clusters",
232            "collate",
233            "collation",
234            "collect",
235            "column",
236            "columns",
237            "comment",
238            "commit",
239            "committed",
240            "compact",
241            "complete",
242            "config",
243            "connection",
244            "connection_id",
245            "consistent",
246            "constraint",
247            "constraints",
248            "convert",
249            "copy",
250            "count",
251            "create",
252            "creation",
253            "cron",
254            "cross",
255            "cube",
256            "current",
257            "current_catalog",
258            "current_date",
259            "current_time",
260            "current_timestamp",
261            "current_user",
262            "data",
263            "database",
264            "databases",
265            "date",
266            "date_add",
267            "date_ceil",
268            "date_diff",
269            "date_floor",
270            "date_sub",
271            "dateadd",
272            "datediff",
273            "datetime",
274            "datetimev2",
275            "datev2",
276            "datetimev1",
277            "datev1",
278            "day",
279            "days_add",
280            "days_sub",
281            "decimal",
282            "decimalv2",
283            "decimalv3",
284            "decommission",
285            "default",
286            "deferred",
287            "delete",
288            "demand",
289            "desc",
290            "describe",
291            "diagnose",
292            "disk",
293            "distinct",
294            "distinctpc",
295            "distinctpcsa",
296            "distributed",
297            "distribution",
298            "div",
299            "do",
300            "doris_internal_table_id",
301            "double",
302            "drop",
303            "dropp",
304            "dual",
305            "duplicate",
306            "dynamic",
307            "else",
308            "enable",
309            "encryptkey",
310            "encryptkeys",
311            "end",
312            "ends",
313            "engine",
314            "engines",
315            "enter",
316            "errors",
317            "events",
318            "every",
319            "except",
320            "exclude",
321            "execute",
322            "exists",
323            "expired",
324            "explain",
325            "export",
326            "extended",
327            "external",
328            "extract",
329            "failed_login_attempts",
330            "false",
331            "fast",
332            "feature",
333            "fields",
334            "file",
335            "filter",
336            "first",
337            "float",
338            "follower",
339            "following",
340            "for",
341            "foreign",
342            "force",
343            "format",
344            "free",
345            "from",
346            "frontend",
347            "frontends",
348            "full",
349            "function",
350            "functions",
351            "generic",
352            "global",
353            "grant",
354            "grants",
355            "graph",
356            "group",
357            "grouping",
358            "groups",
359            "hash",
360            "having",
361            "hdfs",
362            "help",
363            "histogram",
364            "hll",
365            "hll_union",
366            "hostname",
367            "hour",
368            "hub",
369            "identified",
370            "if",
371            "ignore",
372            "immediate",
373            "in",
374            "incremental",
375            "index",
376            "indexes",
377            "infile",
378            "inner",
379            "insert",
380            "install",
381            "int",
382            "integer",
383            "intermediate",
384            "intersect",
385            "interval",
386            "into",
387            "inverted",
388            "ipv4",
389            "ipv6",
390            "is",
391            "is_not_null_pred",
392            "is_null_pred",
393            "isnull",
394            "isolation",
395            "job",
396            "jobs",
397            "join",
398            "json",
399            "jsonb",
400            "key",
401            "keys",
402            "kill",
403            "label",
404            "largeint",
405            "last",
406            "lateral",
407            "ldap",
408            "ldap_admin_password",
409            "left",
410            "less",
411            "level",
412            "like",
413            "limit",
414            "lines",
415            "link",
416            "list",
417            "load",
418            "local",
419            "localtime",
420            "localtimestamp",
421            "location",
422            "lock",
423            "logical",
424            "low_priority",
425            "manual",
426            "map",
427            "match",
428            "match_all",
429            "match_any",
430            "match_phrase",
431            "match_phrase_edge",
432            "match_phrase_prefix",
433            "match_regexp",
434            "materialized",
435            "max",
436            "maxvalue",
437            "memo",
438            "merge",
439            "migrate",
440            "migrations",
441            "min",
442            "minus",
443            "minute",
444            "modify",
445            "month",
446            "mtmv",
447            "name",
448            "names",
449            "natural",
450            "negative",
451            "never",
452            "next",
453            "ngram_bf",
454            "no",
455            "non_nullable",
456            "not",
457            "null",
458            "nulls",
459            "observer",
460            "of",
461            "offset",
462            "on",
463            "only",
464            "open",
465            "optimized",
466            "or",
467            "order",
468            "outer",
469            "outfile",
470            "over",
471            "overwrite",
472            "parameter",
473            "parsed",
474            "partition",
475            "partitions",
476            "password",
477            "password_expire",
478            "password_history",
479            "password_lock_time",
480            "password_reuse",
481            "path",
482            "pause",
483            "percent",
484            "period",
485            "permissive",
486            "physical",
487            "plan",
488            "process",
489            "plugin",
490            "plugins",
491            "policy",
492            "preceding",
493            "prepare",
494            "primary",
495            "proc",
496            "procedure",
497            "processlist",
498            "profile",
499            "properties",
500            "property",
501            "quantile_state",
502            "quantile_union",
503            "query",
504            "quota",
505            "random",
506            "range",
507            "read",
508            "real",
509            "rebalance",
510            "recover",
511            "recycle",
512            "refresh",
513            "references",
514            "regexp",
515            "release",
516            "rename",
517            "repair",
518            "repeatable",
519            "replace",
520            "replace_if_not_null",
521            "replica",
522            "repositories",
523            "repository",
524            "resource",
525            "resources",
526            "restore",
527            "restrictive",
528            "resume",
529            "returns",
530            "revoke",
531            "rewritten",
532            "right",
533            "rlike",
534            "role",
535            "roles",
536            "rollback",
537            "rollup",
538            "routine",
539            "row",
540            "rows",
541            "s3",
542            "sample",
543            "schedule",
544            "scheduler",
545            "schema",
546            "schemas",
547            "second",
548            "select",
549            "semi",
550            "sequence",
551            "serializable",
552            "session",
553            "set",
554            "sets",
555            "shape",
556            "show",
557            "signed",
558            "skew",
559            "smallint",
560            "snapshot",
561            "soname",
562            "split",
563            "sql_block_rule",
564            "start",
565            "starts",
566            "stats",
567            "status",
568            "stop",
569            "storage",
570            "stream",
571            "streaming",
572            "string",
573            "struct",
574            "subdate",
575            "sum",
576            "superuser",
577            "switch",
578            "sync",
579            "system",
580            "table",
581            "tables",
582            "tablesample",
583            "tablet",
584            "tablets",
585            "task",
586            "tasks",
587            "temporary",
588            "terminated",
589            "text",
590            "than",
591            "then",
592            "time",
593            "timestamp",
594            "timestampadd",
595            "timestampdiff",
596            "tinyint",
597            "to",
598            "transaction",
599            "trash",
600            "tree",
601            "triggers",
602            "trim",
603            "true",
604            "truncate",
605            "type",
606            "type_cast",
607            "types",
608            "unbounded",
609            "uncommitted",
610            "uninstall",
611            "union",
612            "unique",
613            "unlock",
614            "unsigned",
615            "update",
616            "use",
617            "user",
618            "using",
619            "value",
620            "values",
621            "varchar",
622            "variables",
623            "variant",
624            "vault",
625            "verbose",
626            "version",
627            "view",
628            "warnings",
629            "week",
630            "when",
631            "where",
632            "whitelist",
633            "with",
634            "work",
635            "workload",
636            "write",
637            "xor",
638            "year",
639        }
640
641        def partition_sql(self, expression: exp.Partition) -> str:
642            parent = expression.parent
643            if isinstance(parent, exp.PartitionByRangeProperty):
644                return ", ".join(self.sql(e) for e in expression.expressions)
645            return super().partition_sql(expression)
646
647        def partitionrange_sql(self, expression: exp.PartitionRange) -> str:
648            name = self.sql(expression, "this")
649            values = expression.expressions
650
651            if len(values) != 1:
652                # Multiple values: use VALUES [ ... )
653                if values and isinstance(values[0], list):
654                    values_sql = ", ".join(
655                        f"({', '.join(self.sql(v) for v in inner)})" for inner in values
656                    )
657                else:
658                    values_sql = ", ".join(f"({self.sql(v)})" for v in values)
659
660                return f"PARTITION {name} VALUES [{values_sql})"
661
662            return f"PARTITION {name} VALUES LESS THAN ({self.sql(values[0])})"
663
664        def partitionbyrangepropertydynamic_sql(self, expression):
665            # Generates: FROM ("start") TO ("end") INTERVAL N UNIT
666            start = self.sql(expression, "start")
667            end = self.sql(expression, "end")
668            every = expression.args.get("every")
669
670            if every:
671                number = self.sql(every, "this")
672                interval = f"INTERVAL {number} {self.sql(every, 'unit')}"
673            else:
674                interval = ""
675
676            return f"FROM ({start}) TO ({end}) {interval}"
677
678        def partitionbyrangeproperty_sql(self, expression):
679            partition_expressions = ", ".join(
680                self.sql(e) for e in expression.args.get("partition_expressions") or []
681            )
682            create_expressions = expression.args.get("create_expressions") or []
683            # Handle both static and dynamic partition definitions
684            create_sql = ", ".join(self.sql(e) for e in create_expressions)
685            return f"PARTITION BY RANGE ({partition_expressions}) ({create_sql})"