sqlglot.parsers.duckdb

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, parser
  6from sqlglot.trie import new_trie
  7from sqlglot.dialects.dialect import (
  8    binary_from_function,
  9    build_default_decimal_type,
 10    build_formatted_time,
 11    build_regexp_extract,
 12    date_trunc_to_time,
 13    pivot_column_names,
 14)
 15from sqlglot.helper import seq_get
 16from sqlglot.parser import IsPredicatePrecedence, binary_range_parser
 17from sqlglot.tokens import TokenType
 18from collections.abc import Collection
 19
 20
 21def _build_list_sort(name: str, desc: bool = False) -> t.Callable[[list], exp.Expr]:
 22    # Orders are case-insensitive strings, e.g. LIST_SORT(x, 'desc', 'nulls first'), and NULLs default to last
 23    def _builder(args: list) -> exp.Expr:
 24        orders = [exp.Literal.string("DESC"), *args[1:]] if desc else args[1:]
 25        order_name = _list_sort_order_name(seq_get(orders, 0), "ASC")
 26        null_order_name = _list_sort_order_name(seq_get(orders, 1), "NULLS LAST")
 27        valid = order_name in ("ASC", "DESC") and null_order_name in ("NULLS FIRST", "NULLS LAST")
 28        if not valid:
 29            return exp.Anonymous(this=name, expressions=args)
 30
 31        return exp.ArraySort(
 32            this=seq_get(args, 0),
 33            asc=exp.false() if order_name == "DESC" else None,
 34            nulls_first=exp.Boolean(this=null_order_name == "NULLS FIRST"),
 35        )
 36
 37    return _builder
 38
 39
 40def _list_sort_order_name(order: exp.Expr | None, default: str) -> str | None:
 41    if order is None:
 42        return default
 43    return order.name.upper() if order.is_string else None
 44
 45
 46def _build_array_prepend(args: list) -> exp.Expr:
 47    return exp.ArrayPrepend(this=seq_get(args, 1), expression=seq_get(args, 0))
 48
 49
 50def _build_date_diff(args: list) -> exp.Expr:
 51    return exp.DateDiff(this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0))
 52
 53
 54def _build_generate_series(end_exclusive: bool = False) -> t.Callable[[list], exp.GenerateSeries]:
 55    def _builder(args: list) -> exp.GenerateSeries:
 56        # Check https://duckdb.org/docs/sql/functions/nested.html#range-functions
 57        if len(args) == 1:
 58            # DuckDB uses 0 as a default for the series' start when it's omitted
 59            args.insert(0, exp.Literal.number("0"))
 60
 61        gen_series = exp.GenerateSeries.from_arg_list(args)
 62        gen_series.set("is_end_exclusive", end_exclusive)
 63
 64        return gen_series
 65
 66    return _builder
 67
 68
 69def _build_make_timestamp(args: list) -> exp.Expr:
 70    if len(args) == 1:
 71        return exp.UnixToTime(this=seq_get(args, 0), scale=exp.UnixToTime.MICROS)
 72
 73    return exp.TimestampFromParts(
 74        year=seq_get(args, 0),
 75        month=seq_get(args, 1),
 76        day=seq_get(args, 2),
 77        hour=seq_get(args, 3),
 78        min=seq_get(args, 4),
 79        sec=seq_get(args, 5),
 80    )
 81
 82
 83def _show_parser(*args: t.Any, **kwargs: t.Any) -> t.Callable[[DuckDBParser], exp.Show]:
 84    def _parse(self: DuckDBParser) -> exp.Show:
 85        return self._parse_show_duckdb(*args, **kwargs)
 86
 87    return _parse
 88
 89
 90def _convert_text_type(dtype: exp.DataType) -> exp.DataType:
 91    dtype.set("expressions", None)
 92    return dtype
 93
 94
 95class DuckDBParser(parser.Parser):
 96    IS_PREDICATE_PRECEDENCE = IsPredicatePrecedence.LOW
 97
 98    MAP_KEYS_ARE_ARBITRARY_EXPRESSIONS = True
 99    PIVOT_COLUMN_NAMING = "agg_name_if_aliased_or_multiple"
100
101    NO_PAREN_FUNCTIONS = {
102        **parser.Parser.NO_PAREN_FUNCTIONS,
103        TokenType.LOCALTIME: exp.Localtime,
104        TokenType.LOCALTIMESTAMP: exp.Localtimestamp,
105        TokenType.CURRENT_CATALOG: exp.CurrentCatalog,
106        TokenType.SESSION_USER: exp.SessionUser,
107    }
108
109    BITWISE = {k: v for k, v in parser.Parser.BITWISE.items() if k != TokenType.CARET}
110
111    COLUMN_OPERATORS = {
112        k: v
113        for k, v in parser.Parser.COLUMN_OPERATORS.items()
114        if k not in (TokenType.ARROW, TokenType.DARROW)
115    }
116
117    JSON_OPERATORS = {
118        TokenType.ARROW: parser.build_json_extract,
119        TokenType.DARROW: parser.build_json_extract_scalar,
120    }
121
122    RANGE_PARSERS = {
123        **parser.Parser.RANGE_PARSERS,
124        TokenType.DAMP: binary_range_parser(exp.ArrayOverlaps),
125        TokenType.TILDE: binary_range_parser(exp.RegexpFullMatch),
126    }
127
128    EXPONENT = {
129        **parser.Parser.EXPONENT,
130        TokenType.CARET: exp.Pow,
131        TokenType.DSTAR: exp.Pow,
132    }
133
134    FUNCTIONS_WITH_ALIASED_ARGS = {*parser.Parser.FUNCTIONS_WITH_ALIASED_ARGS, "STRUCT_PACK"}
135
136    SHOW_PARSERS = {
137        "TABLES": _show_parser("TABLES"),
138        "ALL TABLES": _show_parser("ALL TABLES"),
139    }
140
141    FUNCTIONS = {
142        **{k: v for k, v in parser.Parser.FUNCTIONS.items() if k not in ("DATE_SUB", "GLOB")},
143        **dict.fromkeys(
144            ("ARRAY_DISTINCT", "LIST_DISTINCT"),
145            lambda args: exp.ArrayDistinct(this=seq_get(args, 0), check_null=False),
146        ),
147        "ANY_VALUE": lambda args: exp.IgnoreNulls(this=exp.AnyValue.from_arg_list(args)),
148        "ARRAY_PREPEND": _build_array_prepend,
149        "ARRAY_REVERSE_SORT": _build_list_sort("ARRAY_REVERSE_SORT", desc=True),
150        "ARRAY_INTERSECT": lambda args: exp.ArrayIntersect(expressions=args),
151        "ARRAY_SORT": _build_list_sort("ARRAY_SORT"),
152        "ARRAY_SLICE": lambda args: exp.ArraySlice(
153            this=seq_get(args, 0),
154            start=seq_get(args, 1),
155            end=seq_get(args, 2),
156            step=seq_get(args, 3),
157            is_length=False,
158        ),
159        "BIT_AND": exp.BitwiseAndAgg.from_arg_list,
160        "BIT_OR": exp.BitwiseOrAgg.from_arg_list,
161        "BIT_XOR": exp.BitwiseXorAgg.from_arg_list,
162        "CURRENT_LOCALTIMESTAMP": exp.Localtimestamp.from_arg_list,
163        "DATEDIFF": _build_date_diff,
164        "DATE_DIFF": _build_date_diff,
165        "DATE_TRUNC": date_trunc_to_time,
166        "DATETRUNC": date_trunc_to_time,
167        "DECODE": lambda args: exp.Decode(
168            this=seq_get(args, 0), charset=exp.Literal.string("utf-8")
169        ),
170        "EDITDIST3": exp.Levenshtein.from_arg_list,
171        "ENCODE": lambda args: exp.Encode(
172            this=seq_get(args, 0), charset=exp.Literal.string("utf-8")
173        ),
174        "EPOCH": exp.TimeToUnix.from_arg_list,
175        "EPOCH_MS": lambda args: exp.UnixToTime(this=seq_get(args, 0), scale=exp.UnixToTime.MILLIS),
176        "FROM_HEX": exp.Unhex.from_arg_list,
177        "GENERATE_SERIES": _build_generate_series(),
178        "GET_CURRENT_TIME": exp.CurrentTime.from_arg_list,
179        "GET_BIT": lambda args: exp.Getbit(
180            this=seq_get(args, 0), expression=seq_get(args, 1), zero_is_msb=True
181        ),
182        "JARO_WINKLER_SIMILARITY": exp.JarowinklerSimilarity.from_arg_list,
183        "JSON": exp.ParseJSON.from_arg_list,
184        "JSON_ARRAY": lambda args: exp.JSONArray(expressions=args),
185        "JSON_EXTRACT_PATH": parser.build_extract_json_with_path(exp.JSONExtract),
186        "JSON_EXTRACT_STRING": parser.build_extract_json_with_path(exp.JSONExtractScalar),
187        "LIST": exp.ArrayAgg.from_arg_list,
188        "LIST_APPEND": exp.ArrayAppend.from_arg_list,
189        "LIST_CONCAT": parser.build_array_concat,
190        "LIST_CONTAINS": exp.ArrayContains.from_arg_list,
191        "LIST_COSINE_DISTANCE": exp.CosineDistance.from_arg_list,
192        "LIST_DISTANCE": exp.EuclideanDistance.from_arg_list,
193        "LIST_FILTER": exp.ArrayFilter.from_arg_list,
194        "LIST_HAS": exp.ArrayContains.from_arg_list,
195        "LIST_HAS_ANY": exp.ArrayOverlaps.from_arg_list,
196        "LIST_MAX": exp.ArrayMax.from_arg_list,
197        "LIST_MIN": exp.ArrayMin.from_arg_list,
198        "LIST_PREPEND": _build_array_prepend,
199        "LIST_REVERSE_SORT": _build_list_sort("LIST_REVERSE_SORT", desc=True),
200        "LIST_SORT": _build_list_sort("LIST_SORT"),
201        "LIST_TRANSFORM": exp.Transform.from_arg_list,
202        "LIST_VALUE": lambda args: exp.Array(
203            expressions=args, value_constructor=len(args) == 1 and isinstance(args[0], exp.Query)
204        ),
205        "MAKE_DATE": exp.DateFromParts.from_arg_list,
206        "MAKE_TIME": exp.TimeFromParts.from_arg_list,
207        "MAKE_TIMESTAMP": _build_make_timestamp,
208        "QUANTILE_CONT": exp.PercentileCont.from_arg_list,
209        "QUANTILE_DISC": exp.PercentileDisc.from_arg_list,
210        "RANGE": _build_generate_series(end_exclusive=True),
211        "REGEXP_EXTRACT": build_regexp_extract(exp.RegexpExtract),
212        "REGEXP_EXTRACT_ALL": build_regexp_extract(exp.RegexpExtractAll),
213        "REGEXP_MATCHES": exp.RegexpLike.from_arg_list,
214        "REGEXP_REPLACE": lambda args: exp.RegexpReplace(
215            this=seq_get(args, 0),
216            expression=seq_get(args, 1),
217            replacement=seq_get(args, 2),
218            modifiers=seq_get(args, 3),
219            single_replace=True,
220        ),
221        "SHA256": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(256)),
222        "STRFTIME": build_formatted_time(exp.TimeToStr),
223        "STRING_SPLIT": exp.Split.from_arg_list,
224        "STRING_SPLIT_REGEX": exp.RegexpSplit.from_arg_list,
225        "STRING_TO_ARRAY": exp.Split.from_arg_list,
226        "STRPTIME": build_formatted_time(exp.StrToTime),
227        "STRUCT_PACK": exp.Struct.from_arg_list,
228        "STR_SPLIT": exp.Split.from_arg_list,
229        "STR_SPLIT_REGEX": exp.RegexpSplit.from_arg_list,
230        "TODAY": exp.CurrentDate.from_arg_list,
231        "TIME_BUCKET": exp.DateBin.from_arg_list,
232        "TO_TIMESTAMP": exp.UnixToTime.from_arg_list,
233        "UNNEST": exp.Explode.from_arg_list,
234        "VERSION": exp.CurrentVersion.from_arg_list,
235        "XOR": binary_from_function(exp.BitwiseXor),
236    }
237
238    FUNCTION_PARSERS = {
239        **{k: v for k, v in parser.Parser.FUNCTION_PARSERS.items() if k != "DECODE"},
240        **dict.fromkeys(
241            ("GROUP_CONCAT", "LISTAGG", "STRINGAGG"), lambda self: self._parse_string_agg()
242        ),
243        "APPROX_QUANTILE": lambda self: self._parse_distinct_arg_function(exp.ApproxQuantile),
244        "QUANTILE": lambda self: self._parse_distinct_arg_function(exp.Quantile),
245        "QUANTILE_CONT": lambda self: self._parse_distinct_arg_function(exp.PercentileCont),
246        "QUANTILE_DISC": lambda self: self._parse_distinct_arg_function(exp.PercentileDisc),
247    }
248
249    NO_PAREN_FUNCTION_PARSERS = {
250        **parser.Parser.NO_PAREN_FUNCTION_PARSERS,
251        "MAP": lambda self: self._parse_map(),
252        "@": lambda self: exp.Abs(this=self._parse_bitwise()),
253    }
254
255    PLACEHOLDER_PARSERS = {
256        **parser.Parser.PLACEHOLDER_PARSERS,
257        TokenType.PARAMETER: lambda self: (
258            self.expression(exp.Placeholder(this=self._prev.text))
259            if self._match(TokenType.NUMBER) or self._match_set(self.ID_VAR_TOKENS)
260            else None
261        ),
262    }
263
264    TYPE_CONVERTERS = {
265        # https://duckdb.org/docs/sql/data_types/numeric
266        exp.DType.DECIMAL: build_default_decimal_type(precision=18, scale=3),
267        # https://duckdb.org/docs/sql/data_types/text
268        exp.DType.TEXT: _convert_text_type,
269    }
270
271    STATEMENT_PARSERS = {
272        **parser.Parser.STATEMENT_PARSERS,
273        TokenType.ATTACH: lambda self: self._parse_attach_detach(),
274        TokenType.DETACH: lambda self: self._parse_attach_detach(is_attach=False),
275        TokenType.FORCE: lambda self: self._parse_force(),
276        TokenType.INSTALL: lambda self: self._parse_install(),
277        TokenType.SHOW: lambda self: self._parse_show(),
278    }
279
280    SET_PARSERS = {
281        **parser.Parser.SET_PARSERS,
282        "VARIABLE": lambda self: self._parse_set_item_assignment("VARIABLE"),
283    }
284
285    SHOW_TRIE = new_trie(key.split(" ") for key in SHOW_PARSERS)
286    SET_TRIE = new_trie(key.split(" ") for key in SET_PARSERS)
287
288    def _parse_function_properties(self) -> exp.Properties | None:
289        if self._match(TokenType.TABLE):
290            return exp.Properties(
291                expressions=[
292                    exp.ReturnsProperty(
293                        this=exp.Schema(this=exp.var("TABLE")),
294                        is_table=True,
295                    ),
296                ]
297            )
298        return super()._parse_function_properties()
299
300    def _parse_lambda(self, alias: bool = False) -> exp.Expr | None:
301        index = self._index
302        if not self._match_text_seq("LAMBDA"):
303            return super()._parse_lambda(alias=alias)
304
305        expressions = self._parse_csv(self._parse_lambda_arg)
306        if not self._match(TokenType.COLON):
307            self._retreat(index)
308            return None
309
310        this = self._replace_lambda(self._parse_assignment(), expressions)
311        return self.expression(exp.Lambda(this=this, expressions=expressions, colon=True))
312
313    def _parse_expression(self) -> exp.Expr | None:
314        # DuckDB supports prefix aliases, e.g. foo: 1
315        if self._next.token_type == TokenType.COLON:
316            alias = self._parse_id_var(tokens=self.ALIAS_TOKENS)
317            self._match(TokenType.COLON)
318            comments = self._prev_comments
319
320            this = self._parse_assignment()
321            if isinstance(this, exp.Expr):
322                # Moves the comment next to the alias in `alias: expr /* comment */`
323                comments += this.pop_comments() or []
324
325            return self.expression(exp.Alias(this=this, alias=alias), comments=comments)
326
327        return super()._parse_expression()
328
329    def _parse_table(
330        self,
331        schema: bool = False,
332        joins: bool = False,
333        alias_tokens: Collection[TokenType] | None = None,
334        parse_bracket: bool = False,
335        is_db_reference: bool = False,
336        parse_partition: bool = False,
337        consume_pipe: bool = False,
338    ) -> exp.Expr | None:
339        # DuckDB supports prefix aliases, e.g. FROM foo: bar
340        if self._next.token_type == TokenType.COLON:
341            alias = self._parse_table_alias(alias_tokens=alias_tokens or self.TABLE_ALIAS_TOKENS)
342            self._match(TokenType.COLON)
343            comments = self._prev_comments
344        else:
345            alias = None
346            comments = []
347
348        table = super()._parse_table(
349            schema=schema,
350            joins=joins,
351            alias_tokens=alias_tokens,
352            parse_bracket=parse_bracket,
353            is_db_reference=is_db_reference,
354            parse_partition=parse_partition,
355        )
356        if isinstance(table, exp.Expr) and isinstance(alias, exp.TableAlias):
357            # Moves the comment next to the alias in `alias: table /* comment */`
358            comments += table.pop_comments() or []
359            alias.comments = alias.pop_comments() + comments
360            table.set("alias", alias)
361
362        return table
363
364    def _parse_table_sample(self, as_modifier: bool = False) -> exp.TableSample | None:
365        # https://duckdb.org/docs/sql/samples.html
366        sample = super()._parse_table_sample(as_modifier=as_modifier)
367        if sample and not sample.args.get("method"):
368            if sample.args.get("size"):
369                sample.set("method", exp.var("RESERVOIR"))
370            else:
371                sample.set("method", exp.var("SYSTEM"))
372
373        return sample
374
375    def _parse_bracket(self, this: exp.Expr | None = None) -> exp.Expr | None:
376        bracket = super()._parse_bracket(this)
377
378        if self.dialect.version < (1, 2) and isinstance(bracket, exp.Bracket):
379            # https://duckdb.org/2025/02/05/announcing-duckdb-120.html#breaking-changes
380            bracket.set("returns_list_for_maps", True)
381
382        return bracket
383
384    def _parse_map(self) -> exp.ToMap | exp.Map:
385        if self._match(TokenType.L_BRACE, advance=False):
386            return self.expression(exp.ToMap(this=self._parse_bracket()))
387
388        args = self._parse_wrapped_csv(self._parse_assignment)
389        return self.expression(exp.Map(keys=seq_get(args, 0), values=seq_get(args, 1)))
390
391    def _parse_struct_types(self, type_required: bool = False) -> exp.Expr | None:
392        return self._parse_field_def()
393
394    def _pivot_column_names(self, aggregations: list[exp.Expr]) -> list[str]:
395        if len(aggregations) == 1:
396            return super()._pivot_column_names(aggregations)
397        return pivot_column_names(aggregations, dialect="duckdb")
398
399    def _parse_attach_detach(self, is_attach: bool = True) -> exp.Attach | exp.Detach:
400        def _parse_attach_option() -> exp.AttachOption:
401            return self.expression(
402                exp.AttachOption(
403                    this=self._parse_var(any_token=True),
404                    expression=self._parse_field(any_token=True),
405                )
406            )
407
408        self._match(TokenType.DATABASE)
409        exists = self._parse_exists(not_=is_attach)
410        this = self._parse_alias(self._parse_primary_or_var(), explicit=True)
411
412        if self._match(TokenType.L_PAREN, advance=False):
413            expressions = self._parse_wrapped_csv(_parse_attach_option)
414        else:
415            expressions = None
416
417        return (
418            self.expression(exp.Attach(this=this, exists=exists, expressions=expressions))
419            if is_attach
420            else self.expression(exp.Detach(this=this, exists=exists))
421        )
422
423    def _parse_show_duckdb(self, this: str) -> exp.Show:
424        from_ = self._parse_table(schema=True) if self._match(TokenType.FROM) else None
425        return self.expression(exp.Show(this=this, from_=from_))
426
427    def _parse_force(self) -> exp.Install | exp.Command:
428        # FORCE can only be followed by INSTALL or CHECKPOINT
429        # In the case of CHECKPOINT, we fallback
430        if not self._match(TokenType.INSTALL):
431            return self._parse_as_command(self._prev)
432
433        return self._parse_install(force=True)
434
435    def _parse_install(self, force: bool = False) -> exp.Install:
436        return self.expression(
437            exp.Install(
438                this=self._parse_id_var(),
439                from_=self._parse_var_or_string() if self._match(TokenType.FROM) else None,
440                force=force,
441            )
442        )
443
444    def _parse_primary(self) -> exp.Expr | None:
445        if self._match_pair(TokenType.HASH, TokenType.NUMBER):
446            return exp.PositionalColumn(this=exp.Literal.number(self._prev.text))
447
448        return super()._parse_primary()